Ase dataset updates

ocpmodels/datasets/ase_datasets.py

@@ -1,13 +1,15 @@
+from __future__ import annotations
+
 import bisect
 import copy
-import functools
-import glob
 import logging
 import os
 import warnings
 from abc import ABC, abstractmethod
+from functools import cache, reduce
+from glob import glob
 from pathlib import Path
-from typing import List
+from typing import Any, Callable, Optional
 
 import ase
 import numpy as np
@@ -18,6 +20,7 @@
 from ocpmodels.common.registry import registry
 from ocpmodels.datasets.lmdb_database import LMDBDatabase
 from ocpmodels.datasets.target_metadata_guesser import guess_property_metadata
+from ocpmodels.modules.transforms import DataTransforms
 from ocpmodels.preprocessing import AtomsToGraphs
 
 
@@ -65,33 +68,38 @@ class AseAtomsDataset(Dataset, ABC):
     """
 
     def __init__(
-        self, config, transform=None, atoms_transform=apply_one_tags
+        self,
+        config: dict,
+        atoms_transform: Callable[
+            [ase.Atoms, Any, ...], ase.Atoms
+        ] = apply_one_tags,
+        transform=None,  # is this deprecated?
     ) -> None:
         self.config = config
 
         a2g_args = config.get("a2g_args", {})
-        if a2g_args is None:
-            a2g_args = {}
 
         # Make sure we always include PBC info in the resulting atoms objects
         a2g_args["r_pbc"] = True
         self.a2g = AtomsToGraphs(**a2g_args)
 
-        self.transform = transform
+        self.key_mapping = self.config.get("key_mapping", None)
+        self.transforms = DataTransforms(self.config.get("transforms", {}))
+
         self.atoms_transform = atoms_transform
 
         if self.config.get("keep_in_memory", False):
-            self.__getitem__ = functools.cache(self.__getitem__)
+            self.__getitem__ = cache(self.__getitem__)
 
-        self.ids = self.load_dataset_get_ids(config)
+        self.ids = self._load_dataset_get_ids(config)
 
     def __len__(self) -> int:
         return len(self.ids)
 
     def __getitem__(self, idx):
         # Handle slicing
         if isinstance(idx, slice):
-            return [self[i] for i in range(*idx.indices(len(self.ids)))]
+            return [self[i] for i in range(*idx.indices(len(self)))]
 
         # Get atoms object via derived class method
         atoms = self.get_atoms_object(self.ids[idx])
@@ -105,10 +113,10 @@ def __getitem__(self, idx):
         sid = atoms.info.get("sid", self.ids[idx])
         try:
             sid = tensor([sid])
+        except (RuntimeError, ValueError, TypeError):
             warnings.warn(
                 "Supplied sid is not numeric (or missing). Using dataset indices instead."
             )
-        except:
             sid = tensor([idx])
 
         fid = atoms.info.get("fid", tensor([0]))
@@ -118,11 +126,17 @@ def __getitem__(self, idx):
         data_object.fid = fid
         data_object.natoms = len(atoms)
 
+        if self.key_mapping is not None:
+            for _property in self.key_mapping:
+                # catch for test data not containing labels
+                if _property in data_object:
+                    new_property = self.key_mapping[_property]
+                    if new_property not in data_object:
+                        data_object[new_property] = data_object[_property]
+                        del data_object[_property]
+
         # Transform data object
-        if self.transform is not None:
-            data_object = self.transform(
-                data_object, **self.config.get("transform_args", {})
-            )
+        data_object = self.transforms(data_object)
 
         if self.config.get("include_relaxed_energy", False):
             data_object.y_relaxed = self.get_relaxed_energy(self.ids[idx])
@@ -137,7 +151,7 @@ def get_atoms_object(self, identifier):
         )
 
     @abstractmethod
-    def load_dataset_get_ids(self, config):
+    def _load_dataset_get_ids(self, config):
         # This function should return a list of ids that can be used to index into the database
         raise NotImplementedError(
             "Every ASE dataset needs to declare a function to load the dataset and return a list of ids."
@@ -147,7 +161,7 @@ def close_db(self) -> None:
         # This method is sometimes called by a trainer
         pass
 
-    def guess_target_metadata(self, num_samples: int = 100):
+    def get_metadata(self, num_samples: int = 100):
         metadata = {}
 
         if num_samples < len(self):
@@ -169,9 +183,6 @@ def guess_target_metadata(self, num_samples: int = 100):
 
         return metadata
 
-    def get_metadata(self):
-        return self.guess_target_metadata()
-
 
 @registry.register_dataset("ase_read")
 class AseReadDataset(AseAtomsDataset):
@@ -196,7 +207,7 @@ class AseReadDataset(AseAtomsDataset):
                     default options will work for most users
 
                     If you are using this for a training dataset, set
-                    "r_energy":True and/or "r_forces":True as appropriate
+                    "r_energy":True, "r_forces":True, and/or "r_stress":True as appropriate
                     In that case, energy/forces must be in the files you read (ex. OUTCAR)
 
             ase_read_args (dict): Keyword arguments for ase.io.read()
@@ -213,14 +224,18 @@ class AseReadDataset(AseAtomsDataset):
 
             transform_args (dict): Additional keyword arguments for the transform callable
 
+            key_mapping (dict[str, str]): Dictionary specifying a mapping between the name of a property used
+                in the model with the corresponding property as it was named in the dataset. Only need to use if
+                the name is different.
+
         atoms_transform (callable, optional): Additional preprocessing function applied to the Atoms
                     object. Useful for applying tags, for example.
 
         transform (callable, optional): Additional preprocessing function for the Data object
 
     """
 
-    def load_dataset_get_ids(self, config) -> List[Path]:
+    def _load_dataset_get_ids(self, config) -> list[Path]:
         self.ase_read_args = config.get("ase_read_args", {})
 
         if ":" in self.ase_read_args.get("index", ""):
@@ -286,7 +301,7 @@ class AseReadMultiStructureDataset(AseAtomsDataset):
                     default options will work for most users
 
                     If you are using this for a training dataset, set
-                    "r_energy":True and/or "r_forces":True as appropriate
+                    "r_energy":True, "r_forces":True, and/or "r_stress":True as appropriate
                     In that case, energy/forces must be in the files you read (ex. OUTCAR)
 
             ase_read_args (dict): Keyword arguments for ase.io.read()
@@ -305,13 +320,17 @@ class AseReadMultiStructureDataset(AseAtomsDataset):
 
             transform_args (dict): Additional keyword arguments for the transform callable
 
+            key_mapping (dict[str, str]): Dictionary specifying a mapping between the name of a property used
+                in the model with the corresponding property as it was named in the dataset. Only need to use if
+                the name is different.
+
         atoms_transform (callable, optional): Additional preprocessing function applied to the Atoms
             object. Useful for applying tags, for example.
 
         transform (callable, optional): Additional preprocessing function for the Data object
     """
 
-    def load_dataset_get_ids(self, config):
+    def _load_dataset_get_ids(self, config):
         self.ase_read_args = config.get("ase_read_args", {})
         if not hasattr(self.ase_read_args, "index"):
             self.ase_read_args["index"] = ":"
@@ -374,32 +393,6 @@ def get_relaxed_energy(self, identifier):
         return relaxed_atoms.get_potential_energy(apply_constraint=False)
 
 
-class dummy_list(list):
-    def __init__(self, max) -> None:
-        self.max = max
-        return
-
-    def __len__(self):
-        return self.max
-
-    def __getitem__(self, idx):
-        # Handle slicing
-        if isinstance(idx, slice):
-            return [self[i] for i in range(*idx.indices(self.max))]
-
-        # Cast idx as int since it could be a tensor index
-        idx = int(idx)
-
-        # Handle negative indices (referenced from end)
-        if idx < 0:
-            idx += self.max
-
-        if 0 <= idx < self.max:
-            return idx
-        else:
-            raise IndexError
-
-
 @registry.register_dataset("ase_db")
 class AseDBDataset(AseAtomsDataset):
     """
@@ -435,7 +428,7 @@ class AseDBDataset(AseAtomsDataset):
                     default options will work for most users
 
                     If you are using this for a training dataset, set
-                    "r_energy":True and/or "r_forces":True as appropriate
+                    "r_energy":True, "r_forces":True, and/or "r_stress":True as appropriate
                     In that case, energy/forces must be in the database
 
             keep_in_memory (bool): Store data in memory. This helps avoid random reads if you need
@@ -444,23 +437,34 @@ class AseDBDataset(AseAtomsDataset):
 
             atoms_transform_args (dict): Additional keyword arguments for the atoms_transform callable
 
-            transform_args (dict): Additional keyword arguments for the transform callable
+            transforms (dict[str, dict]): Dictionary specifying data transforms as {transform_function: config}
+                    where config is a dictionary specifying arguments to the transform_function
+
+            key_mapping (dict[str, str]): Dictionary specifying a mapping between the name of a property used
+                in the model with the corresponding property as it was named in the dataset. Only need to use if
+                the name is different.
 
         atoms_transform (callable, optional): Additional preprocessing function applied to the Atoms
                     object. Useful for applying tags, for example.
 
-        transform (callable, optional): Additional preprocessing function for the Data object
+        transform (callable, optional): deprecated?
     """
 
-    def load_dataset_get_ids(self, config) -> dummy_list:
+    def _load_dataset_get_ids(self, config: dict) -> list[int]:
         if isinstance(config["src"], list):
-            filepaths = config["src"]
+            if os.path.isdir(config["src"][0]):
+                filepaths = reduce(
+                    lambda x, y: x + y,
+                    (glob(f"{path}/*db") for path in config["src"]),
+                )
+            else:
+                filepaths = config["src"]
         elif os.path.isfile(config["src"]):
             filepaths = [config["src"]]
         elif os.path.isdir(config["src"]):
-            filepaths = glob.glob(f'{config["src"]}/*')
+            filepaths = glob(f'{config["src"]}/*db')
         else:
-            filepaths = glob.glob(config["src"])
+            filepaths = glob(config["src"])
 
         self.dbs = []
 
@@ -488,16 +492,24 @@ def load_dataset_get_ids(self, config) -> dummy_list:
             if hasattr(db, "ids") and self.select_args == {}:
                 self.db_ids.append(db.ids)
             else:
+                # this is the slow alternative
                 self.db_ids.append(
                     [row.id for row in db.select(**self.select_args)]
                 )
 
         idlens = [len(ids) for ids in self.db_ids]
         self._idlen_cumulative = np.cumsum(idlens).tolist()
 
-        return dummy_list(sum(idlens))
+        return list(range(sum(idlens)))
+
+    def get_atoms_object(self, idx: int) -> ase.Atoms:
+        """Get atoms object corresponding to datapoint idx. Useful to read other properties not in data object.
+        Args:
+            idx (int): index in dataset
 
-    def get_atoms_object(self, idx):
+        Returns:
+            atoms: ASE atoms corresponding to datapoint idx
+        """
         # Figure out which db this should be indexed from.
         db_idx = bisect.bisect(self._idlen_cumulative, idx)
 
@@ -507,20 +519,25 @@ def get_atoms_object(self, idx):
             el_idx = idx - self._idlen_cumulative[db_idx - 1]
         assert el_idx >= 0
 
-        atoms_row = self.dbs[db_idx]._get_row(self.db_ids[db_idx][el_idx])
+        atoms_row = self.dbs[db_idx].get(self.db_ids[db_idx][el_idx])
         atoms = atoms_row.toatoms()
 
+        # put data back into atoms info
         if isinstance(atoms_row.data, dict):
             atoms.info.update(atoms_row.data)
 
         return atoms
 
-    def connect_db(self, address, connect_args={}):
+    @staticmethod
+    def connect_db(
+        address: str | Path, connect_args: Optional[dict] = None
+    ) -> ase.db.core.Database:
         if connect_args is None:
             connect_args = {}
         db_type = connect_args.get("type", "extract_from_name")
-        if db_type == "lmdb" or (
-            db_type == "extract_from_name" and address.split(".")[-1] == "lmdb"
+        if db_type in ("lmdb", "aselmdb") or (
+            db_type == "extract_from_name"
+            and str(address).split(".")[-1] in ("lmdb", "aselmdb")
         ):
             return LMDBDatabase(address, readonly=True, **connect_args)
         else:
@@ -531,12 +548,12 @@ def close_db(self) -> None:
             if hasattr(db, "close"):
                 db.close()
 
-    def get_metadata(self):
+    def get_metadata(self, num_samples: int = 100) -> dict:
         logging.warning(
             "You specific a folder of ASE dbs, so it's impossible to know which metadata to use. Using the first!"
         )
         if self.dbs[0].metadata == {}:
-            return self.guess_target_metadata()
+            return super().get_metadata(num_samples)
         else:
             return copy.deepcopy(self.dbs[0].metadata)