Added DES5M dataset

src/openqdc/datasets/interaction/__init__.py

@@ -10,6 +10,7 @@
 _lazy_imports_obj = {
     "BaseInteractionDataset": "openqdc.datasets.interaction.base",
     "DES370K": "openqdc.datasets.interaction.des370k",
+    "DES5M": "openqdc.datasets.interaction.des5m",
     "Metcalf": "openqdc.datasets.interaction.metcalf",
 }
 
@@ -46,5 +47,6 @@ def __dir__():
     __all__ = [
         "BaseInteractionDataset",
         "DES370K",
+        "DES5M",
         "Metcalf",
     ]

src/openqdc/datasets/interaction/des5m.py

@@ -0,0 +1,104 @@
+import os
+import numpy as np
+import pandas as pd
+
+from typing import Dict, List
+
+from tqdm import tqdm
+from rdkit import Chem
+from loguru import logger
+from openqdc.datasets.interaction import DES370K
+from openqdc.utils.molecule import atom_table, molecule_groups
+
+
+class DES5M(DES370K):
+    __name__ = "des5m_interaction"
+    __energy_methods__ = [
+        "mp2/cc-pvqz",
+        "mp2/cc-pvtz",
+        "mp2/cbs",
+        "ccsd(t)/nn",  # nn
+        "sapt0/aug-cc-pwcvxz",
+        "sapt0/aug-cc-pwcvxz_es",
+        "sapt0/aug-cc-pwcvxz_ex",
+        "sapt0/aug-cc-pwcvxz_exs2",
+        "sapt0/aug-cc-pwcvxz_ind",
+        "sapt0/aug-cc-pwcvxz_exind",
+        "sapt0/aug-cc-pwcvxz_disp",
+        "sapt0/aug-cc-pwcvxz_exdisp_os",
+        "sapt0/aug-cc-pwcvxz_exdisp_ss",
+        "sapt0/aug-cc-pwcvxz_delta_HF",
+    ]
+
+    energy_target_names = [
+        "qz_MP2_all",
+        "tz_MP2_all",
+        "cbs_MP2_all",
+        "nn_CCSD(T)_all",
+        "sapt_all",
+        "sapt_es",
+        "sapt_ex",
+        "sapt_exs2",
+        "sapt_ind",
+        "sapt_exind",
+        "sapt_disp",
+        "sapt_exdisp_os",
+        "sapt_exdisp_ss",
+        "sapt_delta_HF",
+    ]
+
+    def read_raw_entries(self) -> List[Dict]:
+        self.filepath = os.path.join(self.root, "DES5M.csv")
+        logger.info(f"Reading DES5M interaction data from {self.filepath}")
+        df = pd.read_csv(self.filepath)
+        data = []
+        for idx, row in tqdm(df.iterrows(), total=df.shape[0]):
+            smiles0, smiles1 = row["smiles0"], row["smiles1"]
+            charge0, charge1 = row["charge0"], row["charge1"]
+            natoms0, natoms1 = row["natoms0"], row["natoms1"]
+            pos = np.array(list(map(float, row["xyz"].split()))).reshape(-1, 3)
+            pos0 = pos[:natoms0]
+            pos1 = pos[natoms0:]
+
+            elements = row["elements"].split()
+            elements0 = np.array(elements[:natoms0])
+            elements1 = np.array(elements[natoms0:])
+
+            atomic_nums = np.expand_dims(np.array([atom_table.GetAtomicNumber(x) for x in elements]), axis=1)
+            atomic_nums0 = np.array(atomic_nums[:natoms0])
+            atomic_nums1 = np.array(atomic_nums[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)
+            atomic_inputs0 = atomic_inputs[:natoms0, :]
+            atomic_inputs1 = atomic_inputs[natoms0:, :]
+
+            energies = np.array(row[self.energy_target_names].values).astype(np.float32)[None, :]
+
+            name = np.array([smiles0 + "." + smiles1])
+            canon_smiles0 = Chem.MolToSmiles(Chem.MolFromSmiles(smiles0))
+            canon_smiles1 = Chem.MolToSmiles(Chem.MolFromSmiles(smiles1))
+
+            subsets = []
+            # for smiles in [canon_smiles0, canon_smiles1]:
+            for smiles in [smiles0, smiles1]:
+                found = False
+                for functional_group, smiles_set in molecule_groups.items():
+                    if smiles in smiles_set:
+                        subsets.append(functional_group)
+                        found = True
+                if not found:
+                    logger.info(f"molecule group lookup failed for {smiles}")
+
+
+            item = dict(
+                energies=energies,
+                subset=np.array([subsets]),
+                n_atoms=np.array([natoms0 + natoms1], dtype=np.int32),
+                n_atoms_first=np.array([natoms0], dtype=np.int32),
+                atomic_inputs=atomic_inputs,
+                name=name,
+            )
+            data.append(item)
+        return data

src/openqdc/utils/molecule.py

@@ -50,6 +50,8 @@
                      "CSCSCC", "CSCSC", "CSC", "CSSC", "S1CCSCC1", "S1CSCSC1"]),
     "thiols": set(["CCCSS", "CCCS", "CCSS", "CCS", "CSCCCS", "CSCCS", "CSCS", "CSS", "CS", "SCCCS", "SCCS", "SS", "S"]),
     "water": set(["O"]),
+    "flourane": set(["F"]),
+    "hydrogen chloride": set(["Cl"]),
 }
 
 def z_to_formula(z):