Interaction Datasets

README.md

@@ -53,7 +53,7 @@ openqdc download Spice QMugs
 6. QM Level of Theory
  -->
 
-We provide support for the following publicly available QM Datasets.
+We provide support for the following publicly available QM Potential Energy Datasets.
 
 # Potential Energy
 
@@ -82,3 +82,11 @@ We provide support for the following publicly available QM Datasets.
 | --- | --- | --- | --- | --- | --- | --- | --- |
 | [DES370K](https://www.nature.com/articles/s41597-021-00833-x) | 3,700 | 370,000 | 100 | No | 20 | CCSD(T) | Yes |
 | [DES5M](https://www.nature.com/articles/s41597-021-00833-x) | 3,700 | 5,000,000 | 1351 | No | 20 | SNS-MP2 | Yes |
+
+We also provide support for the following publicly available QM Noncovalent Interaction Energy Datasets.
+
+| Dataset |
+| --- |
+| [DES370K](https://www.nature.com/articles/s41597-021-00833-x) |
+| [DES5M](https://www.nature.com/articles/s41597-021-00833-x)   |
+| [Metcalf](https://pubs.aip.org/aip/jcp/article/152/7/074103/1059677/Approaches-for-machine-learning-intermolecular) |

openqdc/datasets/__init__.py

@@ -24,7 +24,6 @@
     "ani1ccx": ANI1CCX,
     "ani1x": ANI1X,
     "comp6": COMP6,
-    "des": DES,
     "gdml": GDML,
     "geom": GEOM,
     "iso17": ISO17,

openqdc/datasets/interaction/L7.py

@@ -0,0 +1,70 @@
+import os
+from typing import Dict, List
+
+import numpy as np
+from loguru import logger
+from ruamel.yaml import YAML
+
+from openqdc.datasets.interaction import BaseInteractionDataset
+from openqdc.utils.molecule import atom_table
+
+
+class L7(BaseInteractionDataset):
+    __name__ = "L7"
+    __energy_unit__ = "hartree"
+    __distance_unit__ = "ang"
+    __forces_unit__ = "hartree/ang"
+    __energy_methods__ = [
+        "CSD(T) | QCISD(T)",
+        "DLPNO-CCSD(T)",
+        "MP2/CBS",
+        "MP2C/CBS",
+        "fixed",
+        "DLPNO-CCSD(T0)",
+        "LNO-CCSD(T)",
+        "FN-DMC",
+    ]
+
+    energy_target_names = []
+
+    def read_raw_entries(self) -> List[Dict]:
+        yaml_fpath = os.path.join(self.root, "l7.yaml")
+        logger.info(f"Reading L7 interaction data from {self.root}")
+        yaml_file = open(yaml_fpath, "r")
+        yaml = YAML()
+        data = []
+        data_dict = yaml.load(yaml_file)
+        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"]):
+            energies = []
+            name = np.array([item["shortname"]])
+            fname = item["geometry"].split(":")[1]
+            energies.append(item["reference_value"])
+            xyz_file = open(os.path.join(self.root, f"{fname}.xyz"), "r")
+            lines = list(map(lambda x: x.strip().split(), xyz_file.readlines()))
+            lines.pop(1)
+            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 = np.array([energies], dtype=np.float32)
+            pos = np.array(lines[1:])[:, 1:].astype(np.float32)
+            elems = np.array(lines[1:])[:, 0]
+            atomic_nums = np.expand_dims(np.array([atom_table.GetAtomicNumber(x) for x in elems]), axis=1)
+            natoms0 = n_atoms_first[0]
+            natoms1 = n_atoms[0] - natoms0
+            charges = np.expand_dims(np.array([charge0] * natoms0 + [charge1] * natoms1), axis=1)
+            atomic_inputs = np.concatenate((atomic_nums, charges, pos), axis=-1, dtype=np.float32)
+
+            item = dict(
+                energies=energies,
+                subset=subset,
+                n_atoms=n_atoms,
+                n_atoms_first=n_atoms_first,
+                atomic_inputs=atomic_inputs,
+                name=name,
+            )
+            data.append(item)
+        return data

openqdc/datasets/interaction/X40.py

@@ -0,0 +1,67 @@
+import os
+from typing import Dict, List
+
+import numpy as np
+from loguru import logger
+from ruamel.yaml import YAML
+
+from openqdc.datasets.interaction import BaseInteractionDataset
+from openqdc.utils.molecule import atom_table
+
+
+class X40(BaseInteractionDataset):
+    __name__ = "X40"
+    __energy_unit__ = "hartree"
+    __distance_unit__ = "ang"
+    __forces_unit__ = "hartree/ang"
+    __energy_methods__ = [
+        "default",
+        "MP2/CBS",
+        "dCCSD(T)/haDZ",
+        "dCCSD(T)/haTZ",
+        "MP2.5/CBS(aDZ)",
+    ]
+
+    energy_target_names = []
+
+    def read_raw_entries(self) -> List[Dict]:
+        yaml_fpath = os.path.join(self.root, "x40.yaml")
+        logger.info(f"Reading X40 interaction data from {self.root}")
+        yaml_file = open(yaml_fpath, "r")
+        yaml = YAML()
+        data = []
+        data_dict = yaml.load(yaml_file)
+        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"]):
+            energies = []
+            name = np.array([item["shortname"]])
+            energies.append(float(item["reference_value"]))
+            xyz_file = open(os.path.join(self.root, f"{item['shortname']}.xyz"), "r")
+            lines = list(map(lambda x: x.strip().split(), xyz_file.readlines()))
+            setup = lines.pop(1)
+            n_atoms = np.array([int(lines[0][0])], dtype=np.int32)
+            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 = np.array([energies], dtype=np.float32)
+            pos = np.array(lines[1:])[:, 1:].astype(np.float32)
+            elems = np.array(lines[1:])[:, 0]
+            atomic_nums = np.expand_dims(np.array([atom_table.GetAtomicNumber(x) for x in elems]), axis=1)
+            natoms0 = n_atoms_first[0]
+            natoms1 = n_atoms[0] - natoms0
+            charges = np.expand_dims(np.array([charge0] * natoms0 + [charge1] * natoms1), axis=1)
+            atomic_inputs = np.concatenate((atomic_nums, charges, pos), axis=-1, dtype=np.float32)
+
+            item = dict(
+                energies=energies,
+                subset=subset,
+                n_atoms=n_atoms,
+                n_atoms_first=n_atoms_first,
+                atomic_inputs=atomic_inputs,
+                name=name,
+            )
+            data.append(item)
+        return data

openqdc/datasets/interaction/__init__.py

@@ -1,3 +1,69 @@
-from .des import DES
+import importlib
+import os
+from typing import TYPE_CHECKING  # noqa F401
 
-AVAILABLE_DATASETS = {"des": DES}
+# The below lazy import logic is coming from openff-toolkit:
+# https://github.com/openforcefield/openff-toolkit/blob/b52879569a0344878c40248ceb3bd0f90348076a/openff/toolkit/__init__.py#L44
+
+# Dictionary of objects to lazily import; maps the object's name to its module path
+
+_lazy_imports_obj = {
+    "BaseInteractionDataset": "openqdc.datasets.interaction.base",
+    "DES370K": "openqdc.datasets.interaction.des370k",
+    "DES5M": "openqdc.datasets.interaction.des5m",
+    "Metcalf": "openqdc.datasets.interaction.metcalf",
+    "DESS66": "openqdc.datasets.interaction.dess66",
+    "DESS66x8": "openqdc.datasets.interaction.dess66x8",
+    "L7": "openqdc.datasets.interaction.L7",
+    "X40": "openqdc.datasets.interaction.X40",
+    "Splinter": "openqdc.datasets.interaction.splinter",
+}
+
+_lazy_imports_mod = {}
+
+
+def __getattr__(name):
+    """Lazily import objects from _lazy_imports_obj or _lazy_imports_mod
+
+    Note that this method is only called by Python if the name cannot be found
+    in the current module."""
+    obj_mod = _lazy_imports_obj.get(name)
+    if obj_mod is not None:
+        mod = importlib.import_module(obj_mod)
+        return mod.__dict__[name]
+
+    lazy_mod = _lazy_imports_mod.get(name)
+    if lazy_mod is not None:
+        return importlib.import_module(lazy_mod)
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    """Add _lazy_imports_obj and _lazy_imports_mod to dir(<module>)"""
+    keys = (*globals().keys(), *_lazy_imports_obj.keys(), *_lazy_imports_mod.keys())
+    return sorted(keys)
+
+
+if TYPE_CHECKING or os.environ.get("OPENQDC_DISABLE_LAZY_LOADING", "0") == "1":
+    from .base import BaseInteractionDataset
+    from .des5m import DES5M
+    from .des370k import DES370K
+    from .dess66 import DESS66
+    from .dess66x8 import DESS66x8
+    from .L7 import L7
+    from .metcalf import Metcalf
+    from .splinter import Splinter
+    from .X40 import X40
+
+    __all__ = [
+        "BaseInteractionDataset",
+        "DES370K",
+        "DES5M",
+        "Metcalf",
+        "DESS66",
+        "DESS66x8",
+        "L7",
+        "X40",
+        "Splinter",
+    ]

openqdc/datasets/interaction/base.py

@@ -0,0 +1,48 @@
+from typing import Dict, List, Optional
+
+import numpy as np
+
+from openqdc.datasets.base import BaseDataset
+from openqdc.utils.constants import NB_ATOMIC_FEATURES
+
+
+class BaseInteractionDataset(BaseDataset):
+    def __init__(
+        self,
+        energy_unit: Optional[str] = None,
+        distance_unit: Optional[str] = None,
+        overwrite_local_cache: bool = False,
+        cache_dir: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            energy_unit=energy_unit,
+            distance_unit=distance_unit,
+            overwrite_local_cache=overwrite_local_cache,
+            cache_dir=cache_dir,
+        )
+
+    def collate_list(self, list_entries: List[Dict]):
+        # concatenate entries
+        print(list_entries[0])
+        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"))
+        print(csum)
+        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)),
+        }