bug fix and simplifying interaction dataset

Author: Nikhil Shenoy

openqdc/datasets/base.py

@@ -341,8 +341,8 @@ def save_preprocess(self, data_dict):
         # save smiles and subset
         local_path = p_join(self.preprocess_path, "props.pkl")
 
-        # assert that required keys are present in data_dict
-        assert all([key in self.pkl_data_keys for key in data_dict.keys()])
+        # 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:

openqdc/datasets/interaction/base.py

@@ -1,6 +1,6 @@
 import os
 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
@@ -23,21 +23,6 @@ def pkl_data_types(self):
             "n_atoms_first": np.int32,
         }
 
-    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
-
     def __getitem__(self, idx: int):
         shift = MAX_CHARGE
         p_start, p_end = self.data["position_idx_range"][idx]