refactor: (1) rename minimization benchmark, (2) fix links to docs in UI (#111)

Author: chrbrunk

docs/source/api_reference/index.rst

@@ -30,7 +30,7 @@ Benchmark implementations
     small_molecules/noncovalent_interactions
     small_molecules/tautomers
     small_molecules/ring_planarity
-    small_molecules/minimization
+    small_molecules/reference_geometry_stability
     small_molecules/bond_length_distribution
     small_molecules/radial_distribution
     small_molecules/solvent_radial_distribution

docs/source/api_reference/small_molecules/bond_length_distribution.rst

@@ -19,4 +19,4 @@ Bond length distribution
 
 .. autoclass:: BondLengthDistributionModelOutput
 
-.. autoclass:: MoleculeSimulationOutput
+.. autoclass:: MoleculeModelOutput

docs/source/api_reference/small_molecules/minimization.rst

@@ -0,0 +1,22 @@
+.. _reference_geometry_stability_api:
+
+Reference Geometry Stability
+============================
+
+.. module:: mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability
+
+.. autoclass:: ReferenceGeometryStabilityBenchmark
+
+    .. automethod:: __init__
+
+    .. automethod:: run_model
+
+    .. automethod:: analyze
+
+.. autoclass:: ReferenceGeometryStabilityResult
+
+.. autoclass:: ReferenceGeometryStabilityDatasetResult
+
+.. autoclass:: ReferenceGeometryStabilityModelOutput
+
+.. autoclass:: MoleculeModelOutput

docs/source/api_reference/small_molecules/reference_geometry_stability.rst

@@ -15,7 +15,7 @@ sampling, stability, and interactions with other molecules.
     Noncovalent Interactions <noncovalent_interactions>
     Ring planarity <ring_planarity>
     Tautomers <tautomers>
-    Reference geometry stability <minimization>
+    Reference geometry stability <reference_geometry_stability>
     Bond length distribution <bond_length_distribution>
     Reactivity <reactivity>
     Nudged elastic band <nudged_elastic_band>

docs/source/benchmarks/small_molecules/index.rst

@@ -37,7 +37,7 @@ fingerprints of the original datasets.
 Interpretation
 --------------
 
-The energy minimization benchmark can be understood as a sanity-check benchmark to verify that the **MLIP** recognizes a true geometrical
+The reference geometry stability benchmark can be understood as a sanity-check benchmark to verify that the **MLIP** recognizes a true geometrical
 minimum of a molecule as such. Often, true ground state structures, with forces on all atoms being essentially zero, are absent from training data.
 In this case the ability of the **MLIP** to correctly interpolate between near-minimum structures is tested. In any case, the ground-truth geometry
 should be retained and not deviate from the **QM** geometry. Since the **QM** geometry is also the starting structure, the geometry should ideally

…chmarks/small_molecules/minimization.rst …lecules/reference_geometry_stability.rstdocs/source/benchmarks/small_molecules/minimization.rst renamed to docs/source/benchmarks/small_molecules/reference_geometry_stability.rst docs/source/benchmarks/small_molecules/minimization.rst renamed to docs/source/benchmarks/small_molecules/reference_geometry_stability.rst

@@ -49,6 +49,11 @@
     ReactivityModelOutput,
     ReactivityResult,
 )
+from mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability import (  # noqa: E501
+    ReferenceGeometryStabilityBenchmark,
+    ReferenceGeometryStabilityModelOutput,
+    ReferenceGeometryStabilityResult,
+)
 from mlipaudit.benchmarks.ring_planarity.ring_planarity import (
     RingPlanarityBenchmark,
     RingPlanarityModelOutput,
@@ -64,11 +69,6 @@
     ScalingModelOutput,
     ScalingResult,
 )
-from mlipaudit.benchmarks.small_molecule_minimization.small_molecule_minimization import (  # noqa: E501
-    SmallMoleculeMinimizationBenchmark,
-    SmallMoleculeMinimizationModelOutput,
-    SmallMoleculeMinimizationResult,
-)
 from mlipaudit.benchmarks.solvent_radial_distribution.solvent_radial_distribution import (  # noqa: E501
     SolventRadialDistributionBenchmark,
     SolventRadialDistributionModelOutput,

src/mlipaudit/benchmarks/__init__.py

@@ -70,7 +70,7 @@ class Molecule(BaseModel):
 Molecules = TypeAdapter(dict[str, Molecule])
 
 
-class MoleculeSimulationOutput(BaseModel):
+class MoleculeModelOutput(BaseModel):
     """Stores the simulation state for a molecule.
 
     Attributes:
@@ -92,7 +92,7 @@ class BondLengthDistributionModelOutput(ModelOutput):
         molecules: A list of simulation states for every molecule.
     """
 
-    molecules: list[MoleculeSimulationOutput]
+    molecules: list[MoleculeModelOutput]
 
 
 class BondLengthDistributionMoleculeResult(BaseModel):
@@ -185,7 +185,7 @@ def run_model(self) -> None:
             md_engine = get_simulation_engine(atoms, self.force_field, **md_kwargs)
             md_engine.run()
 
-            molecule_output = MoleculeSimulationOutput(
+            molecule_output = MoleculeModelOutput(
                 molecule_name=pattern_name, simulation_state=md_engine.state
             )
             molecule_outputs.append(molecule_output)

src/mlipaudit/benchmarks/bond_length_distribution/bond_length_distribution.py

@@ -83,7 +83,7 @@ class Molecule(BaseModel):
 Molecules = TypeAdapter(dict[str, Molecule])
 
 
-class MoleculeSimulationOutput(BaseModel):
+class MoleculeModelOutput(BaseModel):
     """Stores the simulation state for a molecule.
 
     Attributes:
@@ -97,26 +97,25 @@ class MoleculeSimulationOutput(BaseModel):
     simulation_state: SimulationState
 
 
-class SmallMoleculeMinimizationModelOutput(ModelOutput):
-    """ModelOutput object for small molecule conformer minimization benchmark.
+class ReferenceGeometryStabilityModelOutput(ModelOutput):
+    """ModelOutput object for reference geometry stability benchmark.
 
     Attributes:
         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.
     """
 
-    openff_neutral: list[MoleculeSimulationOutput]
-    openff_charged: list[MoleculeSimulationOutput]
+    openff_neutral: list[MoleculeModelOutput]
+    openff_charged: list[MoleculeModelOutput]
 
 
-class SmallMoleculeMinimizationDatasetResult(BaseModel):
+class ReferenceGeometryStabilityDatasetResult(BaseModel):
     """Result for a single dataset.
 
     Attributes:
         rmsd_values: The list of rmsd values for each molecule.
         avg_rmsd: The average rmsd across all molecules in the dataset.
-        num_exploded: The number of molecules that exploded during
-            minimization.
+        num_exploded: The number of molecules that exploded during minimization.
         num_bad_rmsds: The number of molecules that we consider to
             have a poor rmsd score.
     """
@@ -127,8 +126,8 @@ class SmallMoleculeMinimizationDatasetResult(BaseModel):
     num_bad_rmsds: NonNegativeInt
 
 
-class SmallMoleculeMinimizationResult(BenchmarkResult):
-    """Results object for small molecule minimization benchmark.
+class ReferenceGeometryStabilityResult(BenchmarkResult):
+    """Results object for reference geometry stability benchmark.
 
     Attributes:
         openff_neutral: The results for the openff neutral dataset.
@@ -138,25 +137,25 @@ class SmallMoleculeMinimizationResult(BenchmarkResult):
             0 and 1.
     """
 
-    openff_neutral: SmallMoleculeMinimizationDatasetResult
-    openff_charged: SmallMoleculeMinimizationDatasetResult
+    openff_neutral: ReferenceGeometryStabilityDatasetResult
+    openff_charged: ReferenceGeometryStabilityDatasetResult
     avg_rmsd: NonNegativeFloat | None = None
 
 
-class SmallMoleculeMinimizationBenchmark(Benchmark):
-    """Benchmark for small molecule minimization.
+class ReferenceGeometryStabilityBenchmark(Benchmark):
+    """Benchmark for reference geometry stability.
 
     Attributes:
         name: The unique benchmark name that should be used to run the benchmark
             from the CLI and that will determine the output folder name for the result
-            file. The name is `small_molecule_minimization`.
+            file. The name is `reference_geometry_stability`.
         category: A string that describes the category of the benchmark, used for
             example, in the UI app for grouping. Default, if not overridden,
             is "General". This benchmark's category is "Small Molecules".
         result_class: A reference to the type of `BenchmarkResult` that will determine
             the return type of `self.analyze()`. The result class type is
-            `SmallMoleculeMinimizationResult`.
-        model_output_class: A reference to the `SmallMoleculeMinimizationModelOutput`
+            `ReferenceGeometryStabilityResult`.
+        model_output_class: A reference to the `ReferenceGeometryStabilityModelOutput`
             class.
         required_elements: The set of atomic element types that are present in the
             benchmark's input files.
@@ -166,10 +165,10 @@ class SmallMoleculeMinimizationBenchmark(Benchmark):
             element types. For this benchmark, the attribute is set to True.
     """
 
-    name = "small_molecule_minimization"
+    name = "reference_geometry_stability"
     category = "Small Molecules"
-    result_class = SmallMoleculeMinimizationResult
-    model_output_class = SmallMoleculeMinimizationModelOutput
+    result_class = ReferenceGeometryStabilityResult
+    model_output_class = ReferenceGeometryStabilityModelOutput
 
     required_elements = {"N", "Cl", "H", "O", "S", "F", "P", "C", "Br"}
 
@@ -185,7 +184,7 @@ def run_model(self) -> None:
         else:
             md_kwargs = SIMULATION_CONFIG
 
-        self.model_output = SmallMoleculeMinimizationModelOutput(
+        self.model_output = ReferenceGeometryStabilityModelOutput(
             openff_neutral=[],
             openff_charged=[],
         )
@@ -206,12 +205,12 @@ def run_model(self) -> None:
                 md_engine.run()
 
                 getattr(self.model_output, dataset_prefix).append(
-                    MoleculeSimulationOutput(
+                    MoleculeModelOutput(
                         molecule_name=molecule_name, simulation_state=md_engine.state
                     )
                 )
 
-    def analyze(self) -> SmallMoleculeMinimizationResult:
+    def analyze(self) -> ReferenceGeometryStabilityResult:
         """Calculates the RMSD between the MLIP and reference structures.
 
         The RMSD is calculated for each structure in the `inference_results` attribute.
@@ -220,7 +219,7 @@ def analyze(self) -> SmallMoleculeMinimizationResult:
         respect to the reference structure.
 
         Returns:
-            A `SmallMoleculeMinimizationResult` object with the benchmark results.
+            A `ReferenceGeometryStabilityResult` object with the benchmark results.
 
         Raises:
             RuntimeError: If called before `run_model()`.
@@ -232,7 +231,7 @@ def analyze(self) -> SmallMoleculeMinimizationResult:
 
         for dataset_prefix in DATASET_PREFIXES:
             rmsd_values: list[float | None] = []
-            dataset_model_output: list[MoleculeSimulationOutput] = getattr(
+            dataset_model_output: list[MoleculeModelOutput] = getattr(
                 self.model_output, dataset_prefix
             )
             num_exploded = 0
@@ -287,7 +286,7 @@ def analyze(self) -> SmallMoleculeMinimizationResult:
                     rmsd for rmsd in rmsd_values if rmsd is not None
                 )
 
-            dataset_result = SmallMoleculeMinimizationDatasetResult(
+            dataset_result = ReferenceGeometryStabilityDatasetResult(
                 rmsd_values=rmsd_values,
                 avg_rmsd=avg_rmsd,
                 num_exploded=num_exploded,
@@ -318,7 +317,9 @@ def analyze(self) -> SmallMoleculeMinimizationResult:
                 if dataset_result.avg_rmsd is not None
             )
 
-        return SmallMoleculeMinimizationResult(**result, score=score, avg_rmsd=avg_rmsd)
+        return ReferenceGeometryStabilityResult(
+            **result, score=score, avg_rmsd=avg_rmsd
+        )
 
     def _load_dataset_from_file(self, filename: str) -> dict[str, Molecule]:
         """Helper method to load, validate, and optionally truncate a dataset.

…imization/small_molecule_minimization.py …tability/reference_geometry_stability.pysrc/mlipaudit/benchmarks/small_molecule_minimization/small_molecule_minimization.py renamed to src/mlipaudit/benchmarks/reference_geometry_stability/reference_geometry_stability.py src/mlipaudit/benchmarks/small_molecule_minimization/small_molecule_minimization.py renamed to src/mlipaudit/benchmarks/reference_geometry_stability/reference_geometry_stability.py

@@ -127,7 +127,7 @@ class RingPlanarityResult(BenchmarkResult):
     mae_deviation: float | None = None
 
 
-class MoleculeSimulationOutput(BaseModel):
+class MoleculeModelOutput(BaseModel):
     """Stores the simulation state for a molecule.
 
     Attributes:
@@ -148,7 +148,7 @@ class RingPlanarityModelOutput(ModelOutput):
         molecules: A list of simulation states for each molecule..
     """
 
-    molecules: list[MoleculeSimulationOutput]
+    molecules: list[MoleculeModelOutput]
 
 
 class RingPlanarityBenchmark(Benchmark):
@@ -204,7 +204,7 @@ def run_model(self) -> None:
             md_engine = get_simulation_engine(atoms, self.force_field, **md_kwargs)
             md_engine.run()
 
-            molecule_output = MoleculeSimulationOutput(
+            molecule_output = MoleculeModelOutput(
                 molecule_name=molecule_name, simulation_state=md_engine.state
             )
             molecule_outputs.append(molecule_output)