schrodinger.application.matsci.mlip.md module

Molecular-dynamics state and engine orchestration for MLIP simulations.

Tensor shapes use M for the number of replicas and N for the number of atoms in each replica.

Copyright Schrodinger, LLC. All rights reserved.

class schrodinger.application.matsci.mlip.md.MDVariables(z_tensor, r_tensor, charge, spin_multiplicity)

Bases: object

Container for MD simulation variables. Assumes z_tensor and r_tensor are always provided, z_tensor in the shape of (M, N), r_tensor in the shape of (M, N, 3). M is the number of systems, N is the number of atoms in each system.

__init__(z_tensor, r_tensor, charge, spin_multiplicity)
Parameters:

  • z_tensor (torch.Tensor) – Atomic numbers tensor of shape (M, N)

  • r_tensor (torch.Tensor) – Positions tensor of shape (M, N, 3) in Angstrom

  • charge (torch.Tensor) – Charge per system

  • spin_multiplicity (torch.Tensor) – Spin multiplicity per system

to(device)

Move all tensor attributes to the given device.

Parameters:

device (str) – PyTorch device name (‘cuda’ or ‘cpu’)

load(key, value)

Store a value as an attribute on this instance.

Parameters:

  • key (str) – Attribute name

  • value – Value to store

loadCheckpoint(read_from='checkpoint.pt')

Load simulation state from a checkpoint file when resuming an interrupted simulation.

Parameters:

read_from (str) – Path to the checkpoint file

saveCheckpoint(save_to='checkpoint.pt')

Save simulation state to a checkpoint file.

Parameters:

save_to (str) – Path to save the checkpoint file

initializeVelocities(temperature, random_seed_generator=None)

Initialize Maxwell-Boltzmann velocities for all atoms.

Parameters:

  • temperature (float) – Target temperature in Kelvin

  • random_seed_generator (torch.Generator or None) – PyTorch random number generator for reproducibility

class schrodinger.application.matsci.mlip.md.MDEngine(mdv, potentials, constraints=None, log_interval=5, dump_interval=10, save_interval=-1, max_steps=50, traj_file_prefix='traj_', log_file_prefix='log_', checkpoint_prefix='checkpoint_', replace_existing_files=False, trackers=None, async_writers=True, verbose=True)

Bases: object

MD simulation engine responsible for force/energy evaluation, step orchestration, and trajectory/log I/O for batched replicas.

__init__(mdv, potentials, constraints=None, log_interval=5, dump_interval=10, save_interval=-1, max_steps=50, traj_file_prefix='traj_', log_file_prefix='log_', checkpoint_prefix='checkpoint_', replace_existing_files=False, trackers=None, async_writers=True, verbose=True)
Parameters:

  • mdv (MDVariables) – Molecular dynamics variables

  • potentials (dict) – Mapping of potential names to callable potential objects (MLIP and bias terms)

  • constraints (dict or None) – Mapping of constraint names to callable constraint objects

  • log_interval (int) – Steps between log writes

  • dump_interval (int) – Steps between trajectory frame dumps

  • save_interval (int) – Steps between checkpoint saves. Use 0 to save only at the end, -1 to never save.

  • max_steps (int) – Maximum number of simulation steps

  • traj_file_prefix (str) – Prefix for trajectory output files

  • log_file_prefix (str) – Prefix for log output files

  • checkpoint_prefix (str) – Prefix for checkpoint files

  • replace_existing_files (bool) – Whether to remove existing trajectory, log, and checkpoint files before starting

  • trackers (dict or None) – Mapping of tracker names to callable tracker objects for logging additional quantities

  • async_writers (bool) – Whether to use asynchronous writers

  • verbose (bool) – Whether to print step progress

getEnergyAndForces()

Evaluate all potentials and constraints, storing per-term energies and forces on mdv.

Constraint energies are recorded but not added to the total potential energy. Constraint forces are added to the total.

Return type:

tuple

Returns:

(total_energy, total_forces) tensors

step(to_update)

Advance the simulation by one step: update variables, write trajectory/log at configured intervals, and save checkpoints.

Parameters:

to_update (dict) – Mapping of MDVariables attribute names to their new values (e.g. {‘r_tensor’: positions, ‘velocities’: vels})

closeAsyncWriters()

Flush and close all async writers.