core.components.calculate.recipes.omol#
Copyright (c) Meta Platforms, Inc. and affiliates.
This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.
OMol Evaluation Recipes#
This module provides evaluation recipes for various molecular property tasks proposed in the OMol25 paper.
The module includes functions for: - Geometry optimizations and conformer generation - Protonation state energetics - Ionization energies and electron affinities - Spin gap calculations - Protein-ligand interactions - Ligand strain energies - Distance scaling behavior - Single-point energy and force calculations
Each function follows a consistent pattern of taking input data and any ASE calculator, performing the required calculations, and returning results in a standardized format suitable for downstream evaluation on the OMol leaderboard - #TODO: add link.
Functions#
|
Perform a geometry optimization job on an atomic structure. |
|
Perform a single-point energy and force calculation. |
|
Calculate conformer energies and geometries. |
|
Calculate protonation state energies and geometries. |
|
Calculate unoptimized ionization energies and electron affinities. |
|
Calculate unoptimized spin gap energies. |
|
Calculate protein-ligand interaction energies and forces. |
|
Calculate ligand strain energies in protein-bound conformations. |
|
Calculate energies and forces at different inter-molecular distances. |
|
Perform general single-point energy and force calculations. |
Module Contents#
- core.components.calculate.recipes.omol.relax_job(atoms: ase.Atoms, calculator: ase.calculators.calculator.Calculator, opt_flags: dict[str, Any]) dict[str, Any]#
Perform a geometry optimization job on an atomic structure.
This function optimizes the geometry of an atomic structure using the provided calculator and optimization parameters. It captures both the initial and final states (energy, forces, and atomic positions) for comparison.
- Parameters:
atoms – ASE Atoms object representing the initial structure
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
opt_flags (dict) – Dictionary containing optimization parameters including: - optimizer: ASE optimizer class (e.g., BFGS, FIRE) - optimizer_kwargs: Additional kwargs for the optimizer - fmax: Force convergence criterion (eV/Å) - max_steps: Maximum number of optimization steps
- Returns:
Results organized in the following form - {
- ”initial”: {
“atoms”: MSONAtoms dictionary of initial structure, “energy”: Initial total energy (eV), “forces”: Initial forces as a list (eV/Å),
}, “final”: {
”atoms”: MSONAtoms dictionary of optimized structure, “energy”: Final total energy (eV), “forces”: Final forces as a list (eV/Å),
}
}
- Return type:
dict
Note
If optimization fails, the function logs the error and returns the last valid state rather than crashing.
- core.components.calculate.recipes.omol.single_point_job(atoms: ase.Atoms, calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Perform a single-point energy and force calculation.
This function calculates the energy and forces for a given atomic structure.
- Parameters:
atoms – ASE Atoms object representing the structure
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
”atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}
- Return type:
dict
- core.components.calculate.recipes.omol.conformers(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate conformer energies and geometries.
This function performs geometry optimizations on molecular conformers.
- Parameters:
input_data (dict) – Input data organized by molecule families, where each entry contains conformer information with initial and final structures
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”molecule_family_1”: {
- “conformer_id_1”: {
- “initial”: {
“atoms”: MSONAtoms dictionary of initial structure, “energy”: Initial total energy (eV), “forces”: Initial forces as a list (eV/Å),
}, “final”: {
”atoms”: MSONAtoms dictionary of optimized structure, “energy”: Final total energy (eV), “forces”: Final forces as a list (eV/Å),
},
}, “conformer_id_2”: { … }, …
}, “molecule_family_2”: { … }, …
}
- Return type:
dict
- core.components.calculate.recipes.omol.protonation(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate protonation state energies and geometries.
This function calculates the energies and geometries of different protonation states of molecules.
- Parameters:
input_data (dict) – Input data organized by molecule families, where each entry contains different protonation states with initial and final structures
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”molecule_family_1”: {
- “protonation_state_1”: {
- “initial”: {
“atoms”: MSONAtoms dictionary of initial structure, “energy”: Initial total energy (eV), “forces”: Initial forces as a list (eV/Å),
}, “final”: {
”atoms”: MSONAtoms dictionary of optimized structure, “energy”: Final total energy (eV), “forces”: Final forces as a list (eV/Å),
},
}, “protonation_state_2”: { … }, …
}, “molecule_family_2”: { … }, …
}
- Return type:
dict
- core.components.calculate.recipes.omol.ieea(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate unoptimized ionization energies and electron affinities.
This function performs single-point calculations on structures at different charge states to evaluate ionization energies (IE) and electron affinities (EA). No geometry optimization is performed, testing the MLIP’s ability to predict energetics of charged species at fixed geometries.
- Parameters:
input_data (dict) – Input data organized by system identifier, with each entry containing structures at different charge and spin states
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”identifier_1”: {
- “charge_state_1”: {
- “spin_state_1”: {
“atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “spin_state_2”: { … }, …
}, “charge_state_2”: { … }, …
}, “identifier_2”: { … },
}
- Return type:
dict
- core.components.calculate.recipes.omol.spin_gap(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate unoptimized spin gap energies.
This function performs single-point calculations on structures at different spin states to evaluate spin gaps (energy differences between different spin multiplicities). No geometry optimization is performed.
- Parameters:
input_data (dict) – Input data organized by system identifier, with each entry containing structures at different spin states
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”identifier_1”: {
- “spin_state_1”: {
“atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “spin_state_2”: { … },
}, “identifier_2”: { … },
}
- Return type:
dict
- core.components.calculate.recipes.omol.ligand_pocket(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate protein-ligand interaction energies and forces.
This function performs single-point calculations on protein-ligand systems, calculating energies and forces for the complex and individual components (ligand, pocket, ligand_pocket). This enables evaluation of interaction energies and binding affinity predictions.
- Parameters:
input_data (dict) – Input data organized by system identifier, with each entry containing ASE Atoms objects for ligand, pocket, and complex
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”identifier_1”: {
- “ligand”: {
“atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “pocket”: {
”atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “ligand_pocket”: {
”atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
},
}, “identifier_2”: { … },
}
- Return type:
dict
- core.components.calculate.recipes.omol.ligand_strain(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate ligand strain energies in protein-bound conformations.
This function calculates strain energies by comparing the energy of a ligand in its bioactive (protein-bound) conformation with its global minimum energy conformation in the gas phase.
- Parameters:
input_data (dict) – Input data organized by system identifier, with each entry containing: - bioactive_conf: Ligand in bioactive conformation - conformers: List of (initial, final) conformer pairs for gas phase
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”identifier_1”: {
- “bioactive”: {
“atoms”: MSONAtoms dictionary of bioactive conformation, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “gas_phase”: {
- ”0”: { # conformer index
- ”initial”: {
“atoms”: MSONAtoms dictionary of initial structure, “energy”: Initial total energy (eV), “forces”: Initial forces as a list (eV/Å),
}, “final”: {
”atoms”: MSONAtoms dictionary of optimized structure, “energy”: Final total energy (eV), “forces”: Final forces as a list (eV/Å),
}
}, “1”: { … }, …
},
}, “identifier_2”: { … }
}
- Return type:
dict
- core.components.calculate.recipes.omol.distance_scaling(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Calculate energies and forces at different inter-molecular distances.
This function performs single-point calculations on molecular systems where inter-molecular distances have been systematically varied. This tests the MLIP’s ability to capture both short-range repulsion and long-range attraction in potential energy surfaces.
- Parameters:
input_data (dict) – Input data organized by domain type (vertical), then by system identifier, then by distance scale factor, containing ASE Atoms objects at different inter-molecular separations
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”vertical_1”: {
- “identifier_1”: {
- “short_range_scaled_complex_X”: {
“atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “short_range_scaled_complex_Y”: { … }, “long_range_scaled_complex_Z”: { … }, …
}, “identifier_2”: { … }, …
}, “vertical_2”: { … }, …
- Return type:
dict
- core.components.calculate.recipes.omol.singlepoint(input_data: dict[str, Any], calculator: ase.calculators.calculator.Calculator) dict[str, Any]#
Perform general single-point energy and force calculations.
This is a general-purpose function for performing single-point calculations on arbitrary molecular structures.
- Parameters:
input_data (dict) – Input data organized by system identifier, with each entry containing an ASE Atoms object
calculator – ASE calculator object (e.g., FAIRChemCalculator) to use for energy/force calculations
- Returns:
Results organized in the following form - {
- ”identifier_1”: {
“atoms”: MSONAtoms dictionary of the structure, “energy”: Total energy (eV), “forces”: Forces as a list (eV/Å),
}, “identifier_2”: { … },
}
- Return type:
dict