chore: code changes for test-case updates (#62)

Author: leonwehrhan

src/mlipaudit/benchmarks/folding_stability/folding_stability.py

@@ -44,6 +44,13 @@
     "orexin_beta_1cq0_nmr",
 ]
 
+BOX_SIZES = {
+    "chignolin_1uao_xray": [23.98, 22.45, 20.68],
+    "trp_cage_2jof_xray": [29.33, 29.74, 23.59],
+    "amyloid_beta_1ba6_nmr": [51.90, 33.74, 39.50],
+    "orexin_beta_1cq0_nmr": [40.30, 29.56, 33.97],
+}
+
 SIMULATION_CONFIG = {
     "num_steps": 100_000,
     "snapshot_interval": 100,
@@ -182,9 +189,13 @@ def run_model(self) -> None:
             simulation_states=[],
         )
 
-        structure_names = (
-            STRUCTURE_NAMES[:1] if self.run_mode == RunMode.DEV else STRUCTURE_NAMES
-        )
+        if self.run_mode == RunMode.DEV:
+            structure_names = STRUCTURE_NAMES[:1]
+        elif self.run_mode == RunMode.FAST:
+            structure_names = STRUCTURE_NAMES[:2]
+        else:
+            structure_names = STRUCTURE_NAMES
+
         if self.run_mode == RunMode.DEV:
             md_kwargs = SIMULATION_CONFIG_FAST
         else:
@@ -195,7 +206,9 @@ def run_model(self) -> None:
             xyz_filename = structure_name + ".xyz"
             atoms = ase_read(self.data_input_dir / self.name / xyz_filename)
 
-            md_engine = get_simulation_engine(atoms, self.force_field, **md_kwargs)
+            md_engine = get_simulation_engine(
+                atoms, self.force_field, box=BOX_SIZES[structure_name], **md_kwargs
+            )
             md_engine.run()
 
             final_state = md_engine.state
@@ -239,11 +252,18 @@ def analyze(self) -> FoldingStabilityResult:
             topology_filename = structure_name + ".pdb"
             ref_filename = structure_name + "_ref.pdb"
 
-            mdtraj_traj = create_mdtraj_trajectory_from_simulation_state(
+            mdtraj_traj_solv = create_mdtraj_trajectory_from_simulation_state(
                 simulation_state,
                 topology_path=self.data_input_dir / self.name / topology_filename,
             )
-            ase_traj = create_ase_trajectory_from_simulation_state(simulation_state)
+            ase_traj_solv = create_ase_trajectory_from_simulation_state(
+                simulation_state
+            )
+
+            non_solvent_idx = mdtraj_traj_solv.top.select("not resname HOH")
+
+            mdtraj_traj = mdtraj_traj_solv.atom_slice(non_solvent_idx)
+            ase_traj = [atoms[non_solvent_idx] for atoms in ase_traj_solv]
 
             # 1. Radius of gyration
             rg_values = [

src/mlipaudit/benchmarks/small_molecule_minimization/small_molecule_minimization.py

@@ -37,13 +37,9 @@
 
 logger = logging.getLogger("mlipaudit")
 
-QM9_NEUTRAL_FILENAME = "qm9_n100_neutral.json"
-QM9_CHARGED_FILENAME = "qm9_n10_charged.json"
-OPENFF_NEUTRAL_FILENAME = "openff_n100_neutral.json"
-OPENFF_CHARGED_FILENAME = "openff_n10_charged.json"
+OPENFF_NEUTRAL_FILENAME = "openff_n200_neutral.json"
+OPENFF_CHARGED_FILENAME = "openff_n20_charged.json"
 DATASET_PREFIXES = [
-    "qm9_neutral",
-    "qm9_charged",
     "openff_neutral",
     "openff_charged",
 ]
@@ -107,14 +103,10 @@ class SmallMoleculeMinimizationModelOutput(ModelOutput):
     """ModelOutput object for small molecule conformer minimization benchmark.
 
     Attributes:
-        qm9_neutral: A list of simulation states for each molecule in the dataset.
-        qm9_charged: A list of simulation states for each molecule in the dataset.
         openff_neutral: A list of simulation states for each molecule in the dataset.
         openff_charged: A list of simulation states for each molecule in the dataset.
     """
 
-    qm9_neutral: list[MoleculeSimulationOutput]
-    qm9_charged: list[MoleculeSimulationOutput]
     openff_neutral: list[MoleculeSimulationOutput]
     openff_charged: list[MoleculeSimulationOutput]
 
@@ -141,17 +133,13 @@ class SmallMoleculeMinimizationResult(BenchmarkResult):
     """Results object for small molecule minimization benchmark.
 
     Attributes:
-        qm9_neutral: The results for the qm9 neutral dataset.
-        qm9_charged: The results for the qm9 charged dataset.
         openff_neutral: The results for the openff neutral dataset.
         openff_charged: The results for the openff charged dataset.
         avg_rmsd: The average rmsd across all datasets.
         score: The final score for the benchmark between
             0 and 1.
     """
 
-    qm9_neutral: SmallMoleculeMinimizationDatasetResult
-    qm9_charged: SmallMoleculeMinimizationDatasetResult
     openff_neutral: SmallMoleculeMinimizationDatasetResult
     openff_charged: SmallMoleculeMinimizationDatasetResult
     avg_rmsd: NonNegativeFloat | None = None
@@ -200,8 +188,6 @@ def run_model(self) -> None:
             md_kwargs = SIMULATION_CONFIG
 
         self.model_output = SmallMoleculeMinimizationModelOutput(
-            qm9_neutral=[],
-            qm9_charged=[],
             openff_neutral=[],
             openff_charged=[],
         )
@@ -354,14 +340,6 @@ def _load_dataset_from_file(self, filename: str) -> dict[str, Molecule]:
 
         return dataset
 
-    @functools.cached_property
-    def _qm9_neutral_dataset(self) -> dict[str, Molecule]:
-        return self._load_dataset_from_file(QM9_NEUTRAL_FILENAME)
-
-    @functools.cached_property
-    def _qm9_charged_dataset(self) -> dict[str, Molecule]:
-        return self._load_dataset_from_file(QM9_CHARGED_FILENAME)
-
     @functools.cached_property
     def _openff_neutral_dataset(self) -> dict[str, Molecule]:
         return self._load_dataset_from_file(OPENFF_NEUTRAL_FILENAME)

src/mlipaudit/benchmarks/stability/stability.py

@@ -58,50 +58,57 @@ class StructureMetadata(TypedDict):
 
 
 STRUCTURES: dict[str, StructureMetadata] = {
-    # Small proteins/peptides
-    "1JRS_Leupeptin": {
-        "xyz": "71_1jrs_leupeptin.xyz",
-        "pdb": "71_1jrs_leupeptin.pdb",
-        "description": "Leupeptin inhibitor (71 atoms)",
+    # Small molecules in vacuum
+    "Small_molecule_HCNO": {
+        "xyz": "small_molecule_HCNO.xyz",
+        "pdb": "small_molecule_HCNO.pdb",
+        "description": "SMall molecule (HCNO)",
     },
-    "Chignolin": {
-        "xyz": "138_1uao_chignolin.xyz",
-        "pdb": "138_1uao_chignolin.pdb",
-        "description": "Chignolin peptide (138 atoms)",
+    "Small_molecule_Sulfur": {
+        "xyz": "small_molecule_S.xyz",
+        "pdb": "small_molecule_S.pdb",
+        "description": "Small molecule (contains Sulfur)",
     },
-    "RNA fragment": {
-        "xyz": "168_1p79_RNA.xyz",
-        "pdb": "168_1p79_RNA.pdb",
-        "description": "RNA fragment (168 atoms)",
+    "Small_molecule_Halogen": {
+        "xyz": "small_molecule_Hal.xyz",
+        "pdb": "small_molecule_Hal.pdb",
+        "description": "Small molecule (contains Halogens)",
     },
-    # Medium proteins
-    "5KGZ": {
-        "xyz": "634_5kgz.xyz",
-        "pdb": "634_5kgz.pdb",
-        "description": "Protein structure (634 atoms)",
+    # peptides in vacuum
+    "Peptide_HCNO": {
+        "xyz": "peptide_HCNO.xyz",
+        "pdb": "peptide_HCNO.pdb",
+        "description": "Neurotensin in vacuum (PDB: 2LNF)",
     },
-    "1AB7": {
-        "xyz": "1432_1ab7.xyz",
-        "pdb": "1432_1ab7.pdb",
-        "description": "Protein structure (1,432 atoms)",
+    "Peptide_cys": {
+        "xyz": "peptide_cys.xyz",
+        "pdb": "peptide_cys.pdb",
+        "description": "Cyclic peptide with cysteines in vacuum (Oxytocin; PDB: 7OFG)",
     },
-    "1BIP": {
-        "xyz": "1818_1bip.xyz",
-        "pdb": "1818_1bip.pdb",
-        "description": "Protein structure (1,818 atoms)",
+    # Medium protein in vacuum
+    "Protein": {
+        "xyz": "protein_1a7m.xyz",
+        "pdb": "protein_1a7m.pdb",
+        "description": "Protein structure in vacuum (PDB: 1A7M)",
     },
-    "1A5E": {
-        "xyz": "2301_1a5e.xyz",
-        "pdb": "2301_1a5e.pdb",
-        "description": "Protein structure (2,301 atoms)",
+    # solvated systems
+    "Peptide_solvated": {
+        "xyz": "peptide_solv.xyz",
+        "pdb": "peptide_solv.pdb",
+        "description": "Solvated Oxytocin (PDB: 7OFG)",
     },
-    "1A7M": {
-        "xyz": "2803_1a7m.xyz",
-        "pdb": "2803_1a7m.pdb",
-        "description": "Protein structure (2,803 atoms)",
+    "Peptide_solvated_ions": {
+        "xyz": "peptide_solv_ion.xyz",
+        "pdb": "peptide_solv_ion.pdb",
+        "description": "Solvated Neurotensin with counter-ions (PDB: 2LNF)",
     },
 }
 
+BOX_SIZES = {
+    "Peptide_solvated": [23.43, 28.96, 20.90],
+    "Peptide_solvated_ions": [25.62, 27.89, 37.36],
+}
+
 STRUCTURE_NAMES = list(STRUCTURES.keys())
 
 
@@ -376,7 +383,7 @@ class StabilityBenchmark(Benchmark):
     result_class = StabilityResult
     model_output_class = StabilityModelOutput
 
-    required_elements = {"N", "H", "O", "S", "P", "C"}
+    required_elements = {"N", "H", "O", "S", "P", "C", "Cl", "F"}
 
     def run_model(self) -> None:
         """Run MD for each structure.
@@ -392,16 +399,24 @@ def run_model(self) -> None:
         if self.run_mode == RunMode.DEV:
             structure_names = STRUCTURE_NAMES[:2]
         elif self.run_mode == RunMode.FAST:
-            structure_names = STRUCTURE_NAMES[:4]
+            structure_names = STRUCTURE_NAMES[:5]
 
         for structure_name in structure_names:
             logger.info("Running MD for %s", structure_name)
             xyz_filename = STRUCTURES[structure_name]["xyz"]
             atoms = ase_read(self.data_input_dir / self.name / xyz_filename)
 
-            md_engine = get_simulation_engine(
-                atoms, self.force_field, **self._md_kwargs
-            )
+            if structure_name in BOX_SIZES:
+                md_engine = get_simulation_engine(
+                    atoms,
+                    self.force_field,
+                    box=BOX_SIZES[structure_name],
+                    **self._md_kwargs,
+                )
+            else:
+                md_engine = get_simulation_engine(
+                    atoms, self.force_field, **self._md_kwargs
+                )
             md_engine.run()
 
             final_state = md_engine.state

src/mlipaudit/ui/small_molecule_minimization.py

@@ -84,7 +84,7 @@ def small_molecule_minimization_page(
         "Small molecule energy minimization benchmark. We run energy"
         " minimizations with "
         "our MLIP starting from reference structures extracted from the"
-        " QM9 dataset and "
+        " OpenFF dataset and "
         "calculate after the minimization, how much the atomic positions"
         " of the "
         "heavy atoms deviate from the reference structure. The key metric"
@@ -96,12 +96,12 @@ def small_molecule_minimization_page(
 
     st.markdown(
         "Here, we test this ability on two datasets of organic small molecules: "
-        "the QM9 dataset and the OpenFF dataset. To be able to verfify the MLIP's"
+        "the  OpenFF dataset. To be able to verfify the MLIP's "
         "ability to represent charged systems, we split the two datasets into neutral "
         " and charged subsets. "
         "To ensure that the benchmark can be run within an acceptable time, we "
-        "reduce the number of test structures to 100 for the neutral datasets"
-        " and 10 for "
+        "reduce the number of test structures to 200 for the neutral datasets"
+        " and 20 for "
         "the charged datasets. The subsets are constructed so that the chemical "
         "diversity, "
         "as represented by Morgan fingerprints, is maximized. For each of these"