Molecular Dynamics

class apax.md.ase_calc.ASECalculator(model_dir: Path | list[Path], dr_threshold: float = 0.5, transformations: list[Callable] = [], padding_factor: float = 1.5, disable_cell_list: bool = False, **kwargs)

ASE Calculator for apax models. Always implements energy and force predictions. Stress predictions and corresponding uncertainties are added to implemented_properties based on whether the stress flag is set in the model config and whether a model ensemble is loaded.

batch_eval(atoms_list: list[Atoms], batch_size: int = 64, silent: bool = False) → list[Atoms]

Evaluate the model on a list of Atoms. This is preferable to assigning the calculator to each Atoms instance for 2 reasons: 1. Processing can be abtched, which is advantageous for larger datasets. 2. Inputs are padded so no recompilation is triggered when evaluating differently sized systems.

Parameters:

• atoms_list – List of Atoms to be evaluated.

• batch_size – Processing batch size. Does not affect results, only speed and memory requirements.

• silent – Whether or not to suppress progress bars.

Returns:

List of Atoms with labels predicted by the model.

Return type:

evaluated_atoms_list

calculate(atoms, properties=['energy'], system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])

Do the calculation.

properties: list of str

List of what needs to be calculated. Can be any combination of ‘energy’, ‘forces’, ‘stress’, ‘dipole’, ‘charges’, ‘magmom’ and ‘magmoms’.

system_changes: list of str

List of what has changed since last calculation. Can be any combination of these six: ‘positions’, ‘numbers’, ‘cell’, ‘pbc’, ‘initial_charges’ and ‘initial_magmoms’.

Subclasses need to implement this, but can ignore properties and system_changes if they want. Calculated properties should be inserted into results dictionary like shown in this dummy example:

self.results = {'energy': 0.0,
                'forces': np.zeros((len(atoms), 3)),
                'stress': np.zeros(6),
                'dipole': np.zeros(3),
                'charges': np.zeros(len(atoms)),
                'magmom': 0.0,
                'magmoms': np.zeros(len(atoms))}

The subclass implementation should first call this implementation to set the atoms attribute and create any missing directories.

get_descriptors(frames: list[Atoms], processing_batch_size: int = 1, only_use_n_layers: int | None = None, should_average: bool = True) → ndarray

Compute the descriptors for a list of Atoms.

Parameters:

• frames (list[ase.Atoms]) – List of Atoms to compute descriptors for.

• processing_batch_size (int, default = 1) – Batch size for processing the frames. This does not affect the results, only the speed and memory requirements.

• only_use_n_layers (int | None, default = None) – If specified, only the first only_use_n_layers layers of the feature model will be used to compute the descriptors. If None, all layers will be used.

• should_average (bool, default = True) – Whether to average the descriptors over the atomic species.

Returns:

Array of computed descriptors for each frame with shape (n_frames, n_descriptors) or (n_frames, n_atoms, n_descriptors) if should_average=False.

Return type:

np.ndarray

implemented_properties: List[str] = ['energy', 'forces']

Properties calculator can handle (energy, forces, …)

class apax.md.function_transformations.GaussianAcceleratedMolecularDynamics(energy_target: float, spring_constant: float)

Applies a boost potential to the system that pulls it towards a target energy. https://pubs.acs.org/doi/10.1021/acs.jctc.5b00436

Parameters:

• energy_target (float) – Target potential energy below which to apply the boost potential.

• spring_constant (float) – Spring constant of the boost potential.

class apax.md.function_transformations.GlobalCalibration(energy_factor: float, forces_factor: float)

Applies a global calibration to energy and force uncertainties. Energy ensemble predictions are rescaled according to EQ 7 in https://doi.org/10.1063/5.0036522

Parameters:

• energy_factor (float) – Global calibration factor by which to scale the energy uncertainty.

• forces_factor (float) – Global calibration factor by which to scale the force uncertainties.

class apax.md.function_transformations.ProcessStress

Remove Volume factor from stress predictions.

class apax.md.function_transformations.UncertaintyDrivenDynamics(height: float, width: float)

UDD requires an uncertainty aware model. It drives the dynamics towards higher uncertainty regions up to some maximum bias energy. https://doi.org/10.1038/s43588-023-00406-5

Parameters:

• height (float) – Maximum bias potential that can be applied

• width (float) – Width of the Gaussian bias.

apax.md.simulate.md_setup(model_config: Config, md_config: MDConfig)

Sets up the energy and neighborlist functions for an MD simulation, loads the initial structure.

Parameters:

• model_config (Config) – Configuration of the model used as an interatomic potential.

• md_config (MDConfig) – configuration of the MD simulation.

Returns:

• R – Initial positions in Angstrom.

• atomic_numbers – Atomic numbers of the system.

• masses – Atomic masses in ASE units.

• box – Side length of the cubic box.

• energy_fn – Interatomic potential.

• neighbor_fn – Neighborlist function.

• shift_fn – Shift function for the integrator.

apax.md.simulate.run_md(model_config: Config, md_config: MDConfig, log_level='error')

Utiliy function to start NVT molecualr dynamics simulations from a previously trained model.

Parameters:

• model_config (Config) – Configuration of the model used as an interatomic potential.

• md_config (MDConfig) – configuration of the MD simulation.

apax.md.simulate.run_sim(system: System, sim_fns: SimulationFunctions, ensemble, sim_dir: Path, n_steps: int, n_inner: int, extra_capacity: int, rng_key: int, traj_handler: TrajHandler, sampling_rate: int = 10, load_momenta: bool = False, restart: bool = True, checkpoint_interval: int = 50000, dynamics_checks: list[DynamicsCheckBase] = [], constraints: list[ConstraintBase] = [], disable_pbar: bool = False)

Performs NVT MD.

Parameters:

• ensemble – Thermodynamic ensemble.

• n_steps (int) – Total time steps.

• n_inner (int) – JIT compiled inner loop. Also determines atoms buffer size.

• extra_capacity (int) – Extra capacity for the neighborlist.

• rng_key (int) – RNG key used to initialize the simulation.

• restart (bool, default = True) – Whether a checkpoint should be loaded. No implemented yet.

• checkpoint_interval (int, default = 50_000) – Number of time steps between saving full simulation state checkpoints.

• sim_dir (Path) – Directory where the trajectory and simulation checkpoints will be saved.

class apax.md.dynamics_checks.DynamicsCheckBase

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.dynamics_checks.EnergyUncertaintyCheck(*, name: Literal['energy_uncertainty'] = 'energy_uncertainty', threshold: float, per_atom: bool = True)

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.dynamics_checks.ForceUncertaintyCheck(*, name: Literal['forces_uncertainty'] = 'forces_uncertainty', threshold: float)

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.dynamics_checks.RadiusCheck(*, name: Literal['radius'] = 'radius', cutoff_radius: float)

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.dynamics_checks.ReflectionCheck(*, name: Literal['reflection'] = 'reflection', cutoff_plane_height: float)

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.constraints.ConstraintBase

Base class for constraints. Constraints work by implementing a create method. This method accepts a reference state which to compare to subsequent ones and returns a callable which applies the constraint during simulations.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.constraints.FixAtoms(*, name: Literal['fixatoms'] = 'fixatoms', indices: list[int])

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.constraints.FixCenterOfMass(*, name: Literal['fixcenterofmass'] = 'fixcenterofmass', position: Literal['initial', 'origin'] | list[float] = 'initial')

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.constraints.FixRotation(*, name: Literal['fixrotation'] = 'fixrotation')

model_config = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.bias.BiasEnergyBase

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.bias.SphericalWall(*, radius: float, spring_constant: float)

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class apax.md.sim_utils.SimulationFunctions(energy_fn: Callable, auxiliary_fn: Callable, shift_fn: Callable, neighbor_fn: Callable)

class apax.md.sim_utils.System(atomic_numbers: <function array at 0x74a0558b0f40>, masses: <function array at 0x74a0558b0f40>, positions: <function array at 0x74a0558b0f40>, box: <function array at 0x74a0558b0f40>, momenta: Optional[array])