feat: implement dataset that only needs a dataframe

docs/source/api/data.rst

@@ -15,6 +15,14 @@ Structure datasets
 .. automodule:: mofdscribe.datasets.bw_dataset
     :members:
 
+
 .. automodule:: mofdscribe.datasets.arabg_dataset
     :members:
 
+
+.. automodule:: mofdscribe.datasets.arcmof_dataset
+    :members:
+
+
+.. automodule:: mofdscribe.datasets.structuredataset
+    :members:

src/mofdscribe/datasets/structuredataset.py

@@ -10,8 +10,9 @@
 from mofdscribe.datasets.dataset import AbstractStructureDataset
 from mofdscribe.datasets.utils import compress_dataset
 from mofdscribe.types import PathType
+from loguru import logger
 
-__all__ = ["StructureDataset"]
+__all__ = ["StructureDataset", "FrameDataset"]
 
 
 class StructureDataset(AbstractStructureDataset):
@@ -173,3 +174,92 @@ def from_folder_and_dataframe(
             undecorated_scaffold_hash_column,
             density_column,
         )
+
+
+class FrameDataset(AbstractStructureDataset):
+    """Dataset containing structure information read from a dataframe."""
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        structure_name_column: str,
+        year_column: Optional[str] = None,
+        label_columns: Optional[List[str]] = None,
+        decorated_graph_hash_column: Optional[str] = None,
+        undecorated_graph_hash_column: Optional[str] = None,
+        decorated_scaffold_hash_column: Optional[str] = None,
+        undecorated_scaffold_hash_column: Optional[str] = None,
+        density_column: Optional[str] = None,
+    ):
+        """Initialize the dataset.
+
+        Args:
+            df (pd.DataFrame): Dataframe containing the structures.
+            structure_name_column (str): Name of the column containing the structure names.
+            year_column (str, optional): Name of the column containing the year of the structure.
+                Defaults to None.
+            label_columns (Optional[List[str]], optional): List of columns containing the labels.
+                Defaults to None.
+            decorated_graph_hash_column (str, optional): Name of the column containing the decorated graph hash.
+                Defaults to None.
+            undecorated_graph_hash_column (str, optional): Name of the column containing the undecorated graph hash.
+                Defaults to None.
+            decorated_scaffold_hash_column (str, optional): Name of the column containing the decorated scaffold hash.
+                Defaults to None.
+            undecorated_scaffold_hash_column (str, optional): Name of the column containing the undecorated scaffold
+                hash. Defaults to None.
+            density_column (str, optional): Name of the column containing the density of the structure.
+                Defaults to None.
+        """
+        super().__init__()
+        logger.warning("FrameDataset support is experimental. Some splitter integrations may not work.")
+        self._df = df
+        compress_dataset(self._df)
+        self._structure_name_column = structure_name_column
+        self._year_column = year_column
+        self._label_columns = list(label_columns) if label_columns is not None else tuple()
+        self._decorated_graph_hash_column = decorated_graph_hash_column
+        self._undecorated_graph_hash_column = undecorated_graph_hash_column
+        self._decorated_scaffold_hash_column = decorated_scaffold_hash_column
+        self._undecorated_scaffold_hash_column = undecorated_scaffold_hash_column
+        self._density_column = density_column
+
+        self._years = None if year_column is None else self._df[year_column]
+        self._labels = None if label_columns is None else self._df[label_columns].values
+        self._decorated_graph_hashes = (
+            None
+            if decorated_graph_hash_column is None
+            else self._df[decorated_graph_hash_column].values
+        )
+        self._undecorated_graph_hashes = (
+            None
+            if undecorated_graph_hash_column is None
+            else self._df[undecorated_graph_hash_column].values
+        )
+        self._decorated_scaffold_hashes = (
+            None
+            if decorated_scaffold_hash_column is None
+            else self._df[decorated_scaffold_hash_column].values
+        )
+        self._undecorated_scaffold_hashes = (
+            None
+            if undecorated_scaffold_hash_column is None
+            else self._df[undecorated_scaffold_hash_column].values
+        )
+        self._densities = None if density_column is None else self._df[density_column].values
+
+    def __len__(self):
+        """Return number of structures in the dataset."""
+        return len(self._df)
+
+    @property
+    def available_features(self) -> List[str]:
+        return self._featurenames
+
+    @property
+    def available_labels(self) -> List[str]:
+        return self._labelnames
+
+    def get_labels(self, idx: Iterable[int], labelnames: Iterable[str] = None) -> np.ndarray:
+        labelnames = labelnames if labelnames is not None else self._labelnames
+        return self._df.iloc[idx][list(labelnames)].values

tests/datasets/test_structuredataset.py

@@ -1,6 +1,6 @@
 import pandas as pd
 
-from mofdscribe.datasets.structuredataset import StructureDataset
+from mofdscribe.datasets.structuredataset import StructureDataset, FrameDataset
 
 
 def test_structuredataset(dataset_files, dataset_folder):
@@ -37,3 +37,14 @@ def test_structuredataset(dataset_files, dataset_folder):
     assert len(ds) == 2
     hashes = ds.get_decorated_graph_hashes([0, 1])
     assert len(hashes) == 2
+
+
+def test_framedataset(dataset_files):
+    _, frame = dataset_files
+    frame = pd.read_json(frame[0])
+    ds = FrameDataset(frame,  structure_name_column="info.basename",
+        decorated_graph_hash_column="info.decorated_graph_hash")
+    # only two of them are in the dataframe
+    assert len(ds) == 2
+    hashes = ds.get_decorated_graph_hashes([0, 1])
+    assert len(hashes) == 2