Refactor Interaction and Better Testing

openqdc/datasets/base.py

@@ -243,6 +243,14 @@ def data_keys(self):
             keys.remove("forces")
         return keys
 
+    @property
+    def pkl_data_keys(self):
+        return list(self.pkl_data_types.keys())
+
+    @property
+    def pkl_data_types(self):
+        return {"name": str, "subset": str, "n_atoms": np.int32}
+
     @property
     def data_types(self):
         return {
@@ -257,8 +265,8 @@ def data_shapes(self):
         return {
             "atomic_inputs": (-1, NB_ATOMIC_FEATURES),
             "position_idx_range": (-1, 2),
-            "energies": (-1, len(self.energy_target_names)),
-            "forces": (-1, 3, len(self.force_target_names)),
+            "energies": (-1, len(self.energy_methods)),
+            "forces": (-1, 3, len(self.force_methods)),
         }
 
     def _set_units(self, en, ds):
@@ -332,8 +340,14 @@ def save_preprocess(self, data_dict):
 
         # save smiles and subset
         local_path = p_join(self.preprocess_path, "props.pkl")
-        for key in ["name", "subset"]:
-            data_dict[key] = np.unique(data_dict[key], return_inverse=True)
+
+        # assert that (required) pkl keys are present in data_dict
+        assert all([key in data_dict.keys() for key in self.pkl_data_keys])
+
+        # store unique and inverse indices for str-based pkl keys
+        for key in self.pkl_data_keys:
+            if self.pkl_data_types[key] == str:
+                data_dict[key] = np.unique(data_dict[key], return_inverse=True)
 
         with open(local_path, "wb") as f:
             pkl.dump(data_dict, f)
@@ -369,7 +383,10 @@ def read_preprocess(self, overwrite_local_cache=False):
         pull_locally(filename, overwrite=overwrite_local_cache)
         with open(filename, "rb") as f:
             tmp = pkl.load(f)
-            for key in ["name", "subset", "n_atoms"]:
+            all_pkl_keys = set(tmp.keys()) - set(self.data_keys)
+            # assert required pkl_keys are present in all_pkl_keys
+            assert all([key in all_pkl_keys for key in self.pkl_data_keys])
+            for key in all_pkl_keys:
                 x = tmp.pop(key)
                 if len(x) == 2:
                     self.data[key] = x[0][x[1]]

openqdc/datasets/interaction/L7.py

@@ -1,5 +1,7 @@
 import os
-from typing import Dict, List
+from dataclasses import dataclass
+from functools import partial
+from typing import Dict, List, Optional
 
 import numpy as np
 import yaml
@@ -10,42 +12,49 @@
 from openqdc.utils.constants import ATOM_TABLE
 
 
-class DataItemYAMLObj:
-    def __init__(self, name, shortname, geometry, reference_value, setup, group, tags):
-        self.name = name
-        self.shortname = shortname
-        self.geometry = geometry
-        self.reference_value = reference_value
-        self.setup = setup
-        self.group = group
-        self.tags = tags
-
-
-class DataSetYAMLObj:
-    def __init__(self, name, references, text, method_energy, groups_by, groups, global_setup, method_geometry=None):
-        self.name = name
-        self.references = references
-        self.text = text
-        self.method_energy = method_energy
-        self.method_geometry = method_geometry
-        self.groups_by = groups_by
-        self.groups = groups
-        self.global_setup = global_setup
-
-
-def data_item_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode):
-    return DataItemYAMLObj(**loader.construct_mapping(node))
+@dataclass
+class DataSet:
+    description: Dict
+    items: List[Dict]
+    alternative_reference: Dict
 
 
-def dataset_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode):
-    return DataSetYAMLObj(**loader.construct_mapping(node))
+@dataclass
+class DataItemYAMLObj:
+    name: str
+    shortname: str
+    geometry: str
+    reference_value: float
+    setup: Dict
+    group: str
+    tags: str
+
+
+@dataclass
+class DataSetDescription:
+    name: Dict
+    references: str
+    text: str
+    groups_by: str
+    groups: List[str]
+    global_setup: Dict
+    method_energy: str
+    method_geometry: Optional[str] = None
 
 
 def get_loader():
     """Add constructors to PyYAML loader."""
+
+    def constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode, cls):
+        return cls(**loader.construct_mapping(node))
+
     loader = yaml.SafeLoader
-    loader.add_constructor("!ruby/object:ProtocolDataset::DataSetItem", data_item_constructor)
-    loader.add_constructor("!ruby/object:ProtocolDataset::DataSetDescription", dataset_constructor)
+
+    loader.add_constructor("!ruby/object:ProtocolDataset::DataSet", partial(constructor, cls=DataSet))
+    loader.add_constructor("!ruby/object:ProtocolDataset::DataSetItem", partial(constructor, cls=DataItemYAMLObj))
+    loader.add_constructor(
+        "!ruby/object:ProtocolDataset::DataSetDescription", partial(constructor, cls=DataSetDescription)
+    )
     return loader
 
 
@@ -62,7 +71,7 @@ class L7(BaseInteractionDataset):
     http://cuby4.molecular.cz/dataset_l7.html
     """
 
-    __name__ = "L7"
+    __name__ = "l7"
     __energy_unit__ = "kcal/mol"
     __distance_unit__ = "ang"
     __forces_unit__ = "kcal/mol/ang"
@@ -87,10 +96,10 @@ def read_raw_entries(self) -> List[Dict]:
         yaml_file = open(yaml_fpath, "r")
         data = []
         data_dict = yaml.load(yaml_file, Loader=get_loader())
-        charge0 = int(data_dict["description"].global_setup["molecule_a"]["charge"])
-        charge1 = int(data_dict["description"].global_setup["molecule_b"]["charge"])
+        charge0 = int(data_dict.description.global_setup["molecule_a"]["charge"])
+        charge1 = int(data_dict.description.global_setup["molecule_b"]["charge"])
 
-        for idx, item in enumerate(data_dict["items"]):
+        for idx, item in enumerate(data_dict.items):
             energies = []
             name = np.array([item.shortname])
             fname = item.geometry.split(":")[1]
@@ -101,7 +110,7 @@ def read_raw_entries(self) -> List[Dict]:
             n_atoms = np.array([int(lines[0][0])], dtype=np.int32)
             n_atoms_first = np.array([int(item.setup["molecule_a"]["selection"].split("-")[1])], dtype=np.int32)
             subset = np.array([item.group])
-            energies += [float(val[idx]) for val in list(data_dict["alternative_reference"].values())]
+            energies += [float(val[idx]) for val in list(data_dict.alternative_reference.values())]
             energies = np.array([energies], dtype=np.float32)
             pos = np.array(lines[1:])[:, 1:].astype(np.float32)
             elems = np.array(lines[1:])[:, 0]

openqdc/datasets/interaction/X40.py

@@ -25,7 +25,7 @@ class X40(BaseInteractionDataset):
     http://cuby4.molecular.cz/dataset_x40.html
     """
 
-    __name__ = "X40"
+    __name__ = "x40"
     __energy_unit__ = "hartree"
     __distance_unit__ = "ang"
     __forces_unit__ = "hartree/ang"
@@ -48,10 +48,10 @@ def read_raw_entries(self) -> List[Dict]:
         yaml_file = open(yaml_fpath, "r")
         data = []
         data_dict = yaml.load(yaml_file, Loader=get_loader())
-        charge0 = int(data_dict["description"].global_setup["molecule_a"]["charge"])
-        charge1 = int(data_dict["description"].global_setup["molecule_b"]["charge"])
+        charge0 = int(data_dict.description.global_setup["molecule_a"]["charge"])
+        charge1 = int(data_dict.description.global_setup["molecule_b"]["charge"])
 
-        for idx, item in enumerate(data_dict["items"]):
+        for idx, item in enumerate(data_dict.items):
             energies = []
             name = np.array([item.shortname])
             energies.append(float(item.reference_value))
@@ -62,7 +62,7 @@ def read_raw_entries(self) -> List[Dict]:
             n_atoms_first = setup[0].split("-")[1]
             n_atoms_first = np.array([int(n_atoms_first)], dtype=np.int32)
             subset = np.array([item.group])
-            energies += [float(val[idx]) for val in list(data_dict["alternative_reference"].values())]
+            energies += [float(val[idx]) for val in list(data_dict.alternative_reference.values())]
             energies = np.array([energies], dtype=np.float32)
             pos = np.array(lines[1:])[:, 1:].astype(np.float32)
             elems = np.array(lines[1:])[:, 0]

openqdc/datasets/interaction/base.py

@@ -1,52 +1,26 @@
 import os
-import pickle as pkl
 from os.path import join as p_join
-from typing import Dict, List, Optional
+from typing import Optional
 
 import numpy as np
 from ase.io.extxyz import write_extxyz
-from loguru import logger
 from sklearn.utils import Bunch
 
 from openqdc.datasets.base import BaseDataset
-from openqdc.utils.constants import MAX_CHARGE, NB_ATOMIC_FEATURES
-from openqdc.utils.io import pull_locally, push_remote, to_atoms
+from openqdc.utils.constants import MAX_CHARGE
+from openqdc.utils.io import to_atoms
 
 
 class BaseInteractionDataset(BaseDataset):
     __energy_type__ = []
 
-    def collate_list(self, list_entries: List[Dict]):
-        # concatenate entries
-        res = {
-            key: np.concatenate([r[key] for r in list_entries if r is not None], axis=0)
-            for key in list_entries[0]
-            if not isinstance(list_entries[0][key], dict)
-        }
-
-        csum = np.cumsum(res.get("n_atoms"))
-        x = np.zeros((csum.shape[0], 2), dtype=np.int32)
-        x[1:, 0], x[:, 1] = csum[:-1], csum
-        res["position_idx_range"] = x
-
-        return res
-
-    @property
-    def data_shapes(self):
-        return {
-            "atomic_inputs": (-1, NB_ATOMIC_FEATURES),
-            "position_idx_range": (-1, 2),
-            "energies": (-1, len(self.__energy_methods__)),
-            "forces": (-1, 3, len(self.force_target_names)),
-        }
-
     @property
-    def data_types(self):
+    def pkl_data_types(self):
         return {
-            "atomic_inputs": np.float32,
-            "position_idx_range": np.int32,
-            "energies": np.float32,
-            "forces": np.float32,
+            "name": str,
+            "subset": str,
+            "n_atoms": np.int32,
+            "n_atoms_first": np.int32,
         }
 
     def __getitem__(self, idx: int):
@@ -68,13 +42,16 @@ def __getitem__(self, idx: int):
             forces = self._convert_array(np.array(self.data["forces"][p_start:p_end], dtype=np.float32))
 
         e0 = self._convert_array(np.array(self.__isolated_atom_energies__[..., z, c + shift].T, dtype=np.float32))
+        formation_energies = energies - e0.sum(axis=0)
 
         bunch = Bunch(
             positions=positions,
             atomic_numbers=z,
             charges=c,
             e0=e0,
             energies=energies,
+            formation_energies=formation_energies,
+            per_atom_formation_energies=formation_energies / len(z),
             name=name,
             subset=subset,
             forces=forces,
@@ -86,56 +63,6 @@ def __getitem__(self, idx: int):
 
         return bunch
 
-    def save_preprocess(self, data_dict):
-        # save memmaps
-        logger.info("Preprocessing data and saving it to cache.")
-        for key in self.data_keys:
-            local_path = p_join(self.preprocess_path, f"{key}.mmap")
-            out = np.memmap(local_path, mode="w+", dtype=data_dict[key].dtype, shape=data_dict[key].shape)
-            out[:] = data_dict.pop(key)[:]
-            out.flush()
-            push_remote(local_path, overwrite=True)
-
-        # save all other keys in props.pkl
-        local_path = p_join(self.preprocess_path, "props.pkl")
-        for key in data_dict:
-            if key not in self.data_keys:
-                x = data_dict[key]
-                x[x == None] = -1  # noqa
-                data_dict[key] = np.unique(x, return_inverse=True)
-
-        with open(local_path, "wb") as f:
-            pkl.dump(data_dict, f)
-        push_remote(local_path, overwrite=True)
-
-    def read_preprocess(self, overwrite_local_cache=False):
-        logger.info("Reading preprocessed data.")
-        logger.info(
-            f"Dataset {self.__name__} with the following units:\n\
-                     Energy: {self.energy_unit},\n\
-                     Distance: {self.distance_unit},\n\
-                     Forces: {self.force_unit if self.__force_methods__ else 'None'}"
-        )
-        self.data = {}
-        for key in self.data_keys:
-            filename = p_join(self.preprocess_path, f"{key}.mmap")
-            pull_locally(filename, overwrite=overwrite_local_cache)
-            self.data[key] = np.memmap(filename, mode="r", dtype=self.data_types[key]).reshape(self.data_shapes[key])
-
-        filename = p_join(self.preprocess_path, "props.pkl")
-        pull_locally(filename, overwrite=overwrite_local_cache)
-        with open(filename, "rb") as f:
-            tmp = pkl.load(f)
-            for key in set(tmp.keys()) - set(self.data_keys):
-                x = tmp.pop(key)
-                if len(x) == 2:
-                    self.data[key] = x[0][x[1]]
-                else:
-                    self.data[key] = x
-
-        for key in self.data:
-            logger.info(f"Loaded {key} with shape {self.data[key].shape}, dtype {self.data[key].dtype}")
-
     def get_ase_atoms(self, idx: int):
         entry = self[idx]
         at = to_atoms(entry["positions"], entry["atomic_numbers"])