HDF5TrajectoryDataset (#40)

docs/beignet.datasets.md

@@ -1,6 +1,8 @@
 # beignet.datasets
 
 ::: beignet.datasets.FASTADataset
+::: beignet.datasets.HDF5TrajectoryDataset
+::: beignet.datasets.PDBTrajectoryDataset
 ::: beignet.datasets.RandomEulerAngleDataset
 ::: beignet.datasets.RandomQuaternionDataset
 ::: beignet.datasets.RandomRotationMatrixDataset
@@ -9,6 +11,7 @@
 ::: beignet.datasets.SizedSequenceDataset
 ::: beignet.datasets.SwissProtDataset
 ::: beignet.datasets.TrEMBLDataset
+::: beignet.datasets.TrajectoryDataset
 ::: beignet.datasets.UniProtDataset
 ::: beignet.datasets.UniRef50Dataset
 ::: beignet.datasets.UniRef90Dataset

pyproject.toml

@@ -22,12 +22,15 @@ requires-python = ">=3.10"
 
 [project.optional-dependencies]
 all = [
-    "beignet[docs,test]",
+    "beignet[docs,mdtraj,test]",
 ]
 docs = [
     "mkdocs-material",
     "mkdocstrings[python]",
 ]
+mdtraj = [
+    "mdtraj",
+]
 test = [
     "hypothesis",
     "numpy==1.26.4",

src/beignet/datasets/__init__.py

@@ -1,11 +1,14 @@
 from ._fasta_dataset import FASTADataset
+from ._hdf5_trajectory_dataset import HDF5TrajectoryDataset
+from ._pdb_trajectory_dataset import PDBTrajectoryDataset
 from ._random_euler_angle_dataset import RandomEulerAngleDataset
 from ._random_quaternion_dataset import RandomQuaternionDataset
 from ._random_rotation_matrix_dataset import RandomRotationMatrixDataset
 from ._random_rotation_vector_dataset import RandomRotationVectorDataset
 from ._sequence_dataset import SequenceDataset
 from ._sized_sequence_dataset import SizedSequenceDataset
 from ._swissprot_dataset import SwissProtDataset
+from ._trajectory_dataset import TrajectoryDataset
 from ._trembl_dataset import TrEMBLDataset
 from ._uniprot_dataset import UniProtDataset
 from ._uniref50_dataset import UniRef50Dataset
@@ -14,6 +17,8 @@
 
 __all__ = [
     "FASTADataset",
+    "HDF5TrajectoryDataset",
+    "PDBTrajectoryDataset",
     "RandomEulerAngleDataset",
     "RandomQuaternionDataset",
     "RandomRotationMatrixDataset",
@@ -22,6 +27,7 @@
     "SizedSequenceDataset",
     "SwissProtDataset",
     "TrEMBLDataset",
+    "TrajectoryDataset",
     "UniProtDataset",
     "UniRef100Dataset",
     "UniRef50Dataset",

src/beignet/datasets/_hdf5_trajectory_dataset.py

@@ -0,0 +1,25 @@
+from os import PathLike
+from typing import Any, Callable
+
+import mdtraj
+from mdtraj import Trajectory
+
+from ._trajectory_dataset import TrajectoryDataset
+
+
+class HDF5TrajectoryDataset(TrajectoryDataset):
+    def __init__(
+        self,
+        root: str | PathLike,
+        transform: Callable[[Trajectory], Any] | None = None,
+        stride: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            func=mdtraj.load_hdf5,
+            extension="hdf5",
+            root=root,
+            transform=transform,
+            stride=stride,
+            **kwargs,
+        )

src/beignet/datasets/_pdb_trajectory_dataset.py

@@ -0,0 +1,25 @@
+from os import PathLike
+from typing import Any, Callable
+
+import mdtraj
+from mdtraj import Trajectory
+
+from ._trajectory_dataset import TrajectoryDataset
+
+
+class PDBTrajectoryDataset(TrajectoryDataset):
+    def __init__(
+        self,
+        root: str | PathLike,
+        transform: Callable[[Trajectory], Any] | None = None,
+        stride: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            func=mdtraj.load_pdb,
+            extension="pdb",
+            root=root,
+            transform=transform,
+            stride=stride,
+            **kwargs,
+        )

src/beignet/datasets/_trajectory_dataset.py

@@ -0,0 +1,44 @@
+import functools
+from os import PathLike
+from pathlib import Path
+from typing import Any, Callable
+
+from mdtraj import Trajectory
+from torch.utils.data import Dataset
+
+
+class TrajectoryDataset(Dataset):
+    def __init__(
+        self,
+        func: Callable,
+        extension: str,
+        root: str | PathLike,
+        transform: Callable[[Trajectory], Any] | None = None,
+        stride: int | None = None,
+        **kwargs,
+    ):
+        self.func = functools.partial(func, **kwargs)
+
+        if isinstance(root, str):
+            root = Path(root)
+
+        self.root = root.resolve()
+
+        self.transform = transform
+
+        self.stride = stride
+
+        self.paths = [*self.root.glob(f"*.{extension}")]
+
+        super().__init__()
+
+    def __getitem__(self, index: int) -> Trajectory:
+        item = self.func(self.paths[index], stride=self.stride)
+
+        if self.transform:
+            item = self.transform(item)
+
+        return item
+
+    def __len__(self) -> int:
+        return len(self.paths)

tests/beignet/conftest.py

@@ -0,0 +1,10 @@
+import os.path
+
+import pytest
+
+
+@pytest.fixture
+def data_path(request):
+    directory = os.path.dirname(request.module.__file__)
+
+    return os.path.join(directory, "../data")