Interaction

openqdc/datasets/interaction/base.py

@@ -105,7 +105,7 @@ def save_preprocess(self, data_dict):
         for key in data_dict:
             if key not in self.data_keys:
                 x = data_dict[key]
-                x[np.where(x is None)] = -1
+                x[x == None] = -1
                 data_dict[key] = np.unique(x, return_inverse=True)
 
         with open(local_path, "wb") as f:

openqdc/datasets/interaction/splinter.py

@@ -111,13 +111,13 @@ def read_raw_entries(self) -> List[Dict]:
                         protein_interaction_site_type=np.array([protein_interaction_site_type]),
                         ligand_monomer_name=np.array([ligand_monomer_name]),
                         ligand_interaction_site_type=np.array([ligand_interaction_site_type]),
-                        index=np.array([index]),
-                        r=np.array([r]),
-                        theta_P=np.array([theta_P]),
-                        tau_P=np.array([tau_P]),
-                        theta_L=np.array([theta_L]),
-                        tau_L=np.array([tau_L]),
-                        tau_PL=np.array([tau_PL]),
+                        index=np.array([index], dtype=np.float32),
+                        r=np.array([r], dtype=np.float32),
+                        theta_P=np.array([theta_P], dtype=np.float32),
+                        tau_P=np.array([tau_P], dtype=np.float32),
+                        theta_L=np.array([theta_L], dtype=np.float32),
+                        tau_L=np.array([tau_L], dtype=np.float32),
+                        tau_PL=np.array([tau_PL], dtype=np.float32),
                         name=np.array([protein_monomer_name + "." + ligand_monomer_name]),
                     )
                     data.append(item)