core.common.relaxation.ase_utils

core.common.relaxation.ase_utils#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Utilities to interface OCP models/trainers with the Atomic Simulation Environment (ASE)

Attributes#

ASE_PROP_RESHAPE

Classes#

OCPCalculator

ASE based calculator using an OCP model

Functions#

batch_to_atoms(→ list[ase.Atoms])

Convert a data batch to ase Atoms

Module Contents#

core.common.relaxation.ase_utils.ASE_PROP_RESHAPE#
core.common.relaxation.ase_utils.batch_to_atoms(batch: torch_geometric.data.Batch, results: dict[str, torch.Tensor] | None = None, wrap_pos: bool = True, eps: float = 1e-07) → list[ase.Atoms]#

Convert a data batch to ase Atoms

Parameters:

  • batch – data batch

  • results – dictionary with predicted result tensors that will be added to a SinglePointCalculator. If no results are given no calculator will be added to the atoms objects.

  • wrap_pos – wrap positions back into the cell.

  • eps – Small number to prevent slightly negative coordinates from being wrapped.

Returns:

list of Atoms

class core.common.relaxation.ase_utils.OCPCalculator(config_yml: str | None = None, checkpoint_path: str | pathlib.Path | None = None, model_name: str | None = None, local_cache: str | None = None, trainer: str | None = None, cpu: bool = True, seed: int | None = None, only_output: list[str] | None = None)#

Bases: ase.calculators.calculator.Calculator

ASE based calculator using an OCP model

_reshaped_props#
config#
trainer#
a2g#
implemented_properties#

Properties calculator can handle (energy, forces, …)

load_checkpoint(checkpoint_path: str, checkpoint: dict | None = None) → None#

Load existing trained model

Parameters:

  • checkpoint_path – string Path to trained model

  • checkpoint – dict A pretrained checkpoint dict

calculate(atoms: ase.Atoms | torch_geometric.data.Batch, properties, system_changes) → None#

Calculate implemented properties for a single Atoms object or a Batch of them.