SpiceVL2

valence-labs · Apr 18, 2024 · dcc1b6b · dcc1b6b
1 parent 565dc26
commit dcc1b6b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/openqdc/__init__.py b/openqdc/__init__.py
@@ -21,6 +21,7 @@ def get_project_root():
     "ANI1X": "openqdc.datasets.potential.ani",
     "Spice": "openqdc.datasets.potential.spice",
     "SpiceV2": "openqdc.datasets.potential.spice",
+    "SpiceVL2": "openqdc.datasets.potential.spice",
     "GEOM": "openqdc.datasets.potential.geom",
     "QMugs": "openqdc.datasets.potential.qmugs",
     "QMugs_V2": "openqdc.datasets.potential.qmugs",
@@ -114,7 +115,7 @@ def __dir__():
     from .datasets.potential.revmd17 import RevMD17
     from .datasets.potential.sn2_rxn import SN2RXN
     from .datasets.potential.solvated_peptides import SolvatedPeptides
-    from .datasets.potential.spice import Spice, SpiceV2
+    from .datasets.potential.spice import Spice, SpiceV2, SpiceVL2
     from .datasets.potential.tmqm import TMQM
     from .datasets.potential.transition1x import Transition1X
     from .datasets.potential.waterclusters3_30 import WaterClusters
diff --git a/openqdc/datasets/potential/__init__.py b/openqdc/datasets/potential/__init__.py
@@ -14,7 +14,7 @@
 from .revmd17 import RevMD17
 from .sn2_rxn import SN2RXN
 from .solvated_peptides import SolvatedPeptides
-from .spice import Spice, SpiceV2
+from .spice import Spice, SpiceV2, SpiceVL2
 from .tmqm import TMQM
 from .transition1x import Transition1X
 from .waterclusters3_30 import WaterClusters
@@ -41,6 +41,7 @@
     "solvatedpeptides": SolvatedPeptides,
     "spice": Spice,
     "spicev2": SpiceV2,
+    "spicevl2": SpiceVL2,
     "tmqm": TMQM,
     "transition1x": Transition1X,
     "watercluster": WaterClusters,

diff --git a/openqdc/datasets/potential/spice.py b/openqdc/datasets/potential/spice.py
@@ -145,3 +145,10 @@ def read_raw_entries(self):
         tmp = [read_record(data[mol_name], self) for i, mol_name in enumerate(tqdm(data)) if i != 40132]
 
         return tmp
+
+
+class SpiceVL2(SpiceV2):
+    __name__ = "spice_vl2"
+
+    __energy_methods__ = SpiceV2.__energy_methods__ + [PotentialMethod.GFN2_XTB, PotentialMethod.PM6]
+    energy_target_names = SpiceV2.energy_target_names + ["GFN2," "PM6"]