Added __init__ and qmugs remove hs fix

valence-labs · Oct 11, 2023 · dae4dae · dae4dae
1 parent 3fae2e3
commit dae4dae
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/src/openqdc/__init__.py b/src/openqdc/__init__.py
diff --git a/src/openqdc/datasets/__init__.py b/src/openqdc/datasets/__init__.py
@@ -1,24 +1,25 @@
+from .ani import ANI1, ANI1CCX, ANI1X
 from .comp6 import COMP6
 from .gdml import GDML
 from .geom import GEOM
 from .iso_17 import ISO17
 from .molecule3d import Molecule3D
-from .nabladft import NablaDFT
 from .orbnet_denali import OrbnetDenali
 from .qmugs import QMugs
 from .sn2_rxn import SN2RXN
 from .spice import Spice
 
 __all__ = [
+    "ANI1",
+    "ANI1CCX",
+    "ANI1X",
     "Spice",
     "GEOM",
     "QMugs",
-    "NablaDFT",
     "ISO17",
     "COMP6",
     "GDML",
     "Molecule3D",
-    "NablaDFT",
     "OrbnetDenali",
     "QMugs",
     "SN2RXN",

diff --git a/src/openqdc/datasets/qmugs.py b/src/openqdc/datasets/qmugs.py
@@ -12,7 +12,7 @@
 
 def read_mol(mol_dir):
     filenames = glob(p_join(mol_dir, "*.sdf"))
-    mols = [dm.read_sdf(f)[0] for f in filenames]
+    mols = [dm.read_sdf(f, remove_hs=False)[0] for f in filenames]
     n_confs = len(mols)
 
     if len(mols) == 0: