Reference Geometry Stability¶

class mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.ReferenceGeometryStabilityBenchmark(force_field: ForceField | Calculator, data_input_dir: str | PathLike = './data', run_mode: RunMode | Literal['dev', 'fast', 'standard'] = RunMode.STANDARD)¶

Benchmark for reference geometry stability.

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 reference_geometry_stability.

Type:: str

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”.

Type:: str

result_class¶

A reference to the type of BenchmarkResult that will determine the return type of self.analyze(). The result class type is ReferenceGeometryStabilityResult.

Type:: type[mlipaudit.benchmark.BenchmarkResult] | None

model_output_class¶

A reference to the ReferenceGeometryStabilityModelOutput class.

Type:: type[mlipaudit.benchmark.ModelOutput] | None

required_elements¶

The set of atomic element types that are present in the benchmark’s input files.

Type:: set[str] | None

skip_if_elements_missing¶

Whether the benchmark should be skipped entirely if there are some atomic element types that the model cannot handle. If False, the benchmark must have its own custom logic to handle missing atomic element types. For this benchmark, the attribute is set to True.

Type:: bool

__init__(force_field: ForceField | Calculator, data_input_dir: str | PathLike = './data', run_mode: RunMode | Literal['dev', 'fast', 'standard'] = RunMode.STANDARD) → None¶

Initializes the benchmark.

Parameters:

force_field – The force field model to be benchmarked.
data_input_dir – The local input data directory. Defaults to “./data”. If the subdirectory “{data_input_dir}/{benchmark_name}” exists, the benchmark expects the relevant data to be in there, otherwise it will download it from HuggingFace.
run_mode – Whether to run the standard benchmark length, a faster version, or a very fast development version. Subclasses should ensure that when RunMode.DEV, their benchmark runs in a much shorter timeframe, by running on a reduced number of test cases, for instance. Implementing RunMode.FAST being different from RunMode.STANDARD is optional and only recommended for very long-running benchmarks. This argument can also be passed as a string “dev”, “fast”, or “standard”.

Raises:

ChemicalElementsMissingError – If initialization is attempted with a force field that cannot perform inference on the required elements.
ValueError – If force field type is not compatible.

run_model() → None¶

Run an MD simulation for each structure.

The MD simulation is performed using the JAX MD engine and starts from the reference structure. The model output is saved in the model_output attribute.

analyze() → ReferenceGeometryStabilityResult¶

Calculates the RMSD between the MLIP and reference structures.

The RMSD is calculated for each structure in the inference_results attribute. The results are stored in the analysis_results attribute. For every structure, the results contain the heavy atom RMSD of the last simulation frame with respect to the reference structure.

Returns:: A ReferenceGeometryStabilityResult object with the benchmark results.
Raises:: RuntimeError – If called before run_model().

class mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.ReferenceGeometryStabilityResult(*, failed: bool = False, score: Annotated[float | None, Ge(ge=0), Le(le=1)] = None, openff_neutral: ReferenceGeometryStabilityDatasetResult, openff_charged: ReferenceGeometryStabilityDatasetResult, avg_rmsd: Annotated[float, Ge(ge=0)] | None = None)¶

Results object for reference geometry stability benchmark.

openff_neutral¶

The results for the openff neutral dataset.

Type:: mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.ReferenceGeometryStabilityDatasetResult

openff_charged¶

The results for the openff charged dataset.

Type:: mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.ReferenceGeometryStabilityDatasetResult

avg_rmsd¶

The average rmsd across all datasets.

Type:: float | None

failed¶

Whether all the simulations failed and no analysis could be performed. Defaults to False.

Type:: bool

score¶

The final score for the benchmark between 0 and 1.

Type:: float | None

class mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.ReferenceGeometryStabilityDatasetResult(*, rmsd_values: list[Annotated[float, Ge(ge=0)] | None], avg_rmsd: Annotated[float, Ge(ge=0)] | None = None, num_exploded: Annotated[int, Ge(ge=0)] = 0, num_bad_rmsds: Annotated[int, Ge(ge=0)] = 0, failed: bool = False)¶

Result for a single dataset.

rmsd_values¶

The list of rmsd values for each molecule.

Type:: list[float | None]

avg_rmsd¶

The average rmsd across all molecules in the dataset.

Type:: float | None

num_exploded¶

The number of molecules that exploded or failed during minimization or that failed the simulation. Defaults to 0.

Type:: int

num_bad_rmsds¶

The number of molecules that we consider to have a poor rmsd score. Defaults to 0.

Type:: int

failed¶

Whether all the simulations or inferences failed and no analysis could be performed. Defaults to False.

Type:: bool

class mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.ReferenceGeometryStabilityModelOutput(*, openff_neutral: list[MoleculeModelOutput], openff_charged: list[MoleculeModelOutput])¶

ModelOutput object for reference geometry stability benchmark.

openff_neutral¶

A list of simulation states for each molecule in the dataset, including those that failed.

Type:: list[mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.MoleculeModelOutput]

openff_charged¶

A list of simulation states for each molecule in the dataset, including those that failed.

Type:: list[mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.MoleculeModelOutput]

class mlipaudit.benchmarks.reference_geometry_stability.reference_geometry_stability.MoleculeModelOutput(*, molecule_name: str, simulation_state: SimulationState | None = None, failed: bool = False)¶

Stores the simulation state for a molecule.

molecule_name¶

The name of the molecule.

Type:: str

simulation_state¶

The simulation state. None if the simulation failed.

Type:: mlip.simulation.state.SimulationState | None

failed¶

Whether the simulation failed on the molecule. Defaults to False.

Type:: bool