Molecular Dynamics#

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

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][source]#

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'])[source]#

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.

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

Properties calculator can handle (energy, forces, …)

class apax.md.function_transformations.FunctionTransformation[source]#
class apax.md.function_transformations.GaussianAcceleratedMolecularDynamics[source]#

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.UncertaintyDrivenDynamics[source]#

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.nvt.md_setup(model_config: Config, md_config: MDConfig)[source]#

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.nvt.run_md(model_config: Config, md_config: MDConfig, log_level='error')[source]#

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.nvt.run_nvt(system: ~apax.md.sim_utils.System, sim_fns, ensemble, sim_dir: ~pathlib.Path, n_steps: int, n_inner: int, extra_capacity: int, rng_key: int, load_momenta: bool = False, restart: bool = True, checkpoint_interval: int = 50000, traj_handler: ~apax.md.io.TrajHandler = <apax.md.io.TrajHandler object>, disable_pbar: bool = False)[source]#

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.