schrodinger.application.matsci.mlip.wrappers module

schrodinger.application.matsci.mlip.wrappers.get_uma_task(forcefield)

Return the FairChem task for an official UMA force field.

Parameters:

forcefield (str) – Validated MLFF force-field value

Return type:

str or None

Returns:

FairChem task name, or None for a non-UMA force field

class schrodinger.application.matsci.mlip.wrappers.NullWrapper(device='cpu')

Bases: object

Null potential wrapper that returns zero energies and forces.

Intended for debugging and testing the MD engine without an actual ML potential.

__init__(device='cpu')
Parameters:

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

to(device)

Move the wrapper to the specified device.

Parameters:

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

Return type:

NullWrapper

Returns:

Self, for method chaining

class schrodinger.application.matsci.mlip.wrappers.ANIWrapper(model=None, device='cpu')

Bases: object

ANI neural network potential wrapper using TorchANI.

Lazily imports torchani on first use and defaults to the ANI1x model. Returns energies in eV and forces in eV/Angstrom.

__init__(model=None, device='cpu')
Parameters:

  • model (model or None) – Pre-built TorchANI model. If None, ANI1x is loaded on first use.

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

to(device)

Move the model to the specified device.

Parameters:

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

Return type:

ANIWrapper

Returns:

Self, for method chaining

class schrodinger.application.matsci.mlip.wrappers.MlffIterativeNonPbcWrapper(forcefield, solvent, device='cpu')

Bases: object

Wrapper for iterative models registered with the MLFF infrastructure.

The model is configured for the fixed batched system on the first call. Later calls reuse that configuration and update only atomic coordinates.

__init__(forcefield, solvent, device='cpu')
Parameters:

  • forcefield (str) – Validated MLFF force-field value

  • solvent (str) – Nanoreactor implicit solvent

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

static getModel(model_type, solvent, device)

Create an iterative MLFF model.

Parameters:

  • model_type (MlffModelType) – MLFF model type

  • solvent (str) – Nanoreactor implicit solvent

  • device (torch.device) – PyTorch device

Returns:

Iterative MLFF model

static getPredictor(model, z_tensor, charge, spin_multiplicity)

Configure the iterative predictor for a fixed batch.

Parameters:

  • model (IterativeMlffModel) – Iterative MLFF model

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

  • charge (torch.Tensor) – Charge per system (M,)

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

Returns:

Iterative MLFF predictor configured for the batch

static updateCoordinates(predictor, r_tensor)

Copy the current BatchMD coordinates into the MLFF input.

Parameters:

  • predictor (MlffSystemPredictor) – Iterative MLFF predictor

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

static getResults(predictor, num_replicas, num_atoms, device, dtype)

Convert MLFF output to BatchMD energy and force tensors.

Parameters:

  • predictor (MlffSystemPredictor) – Iterative MLFF predictor

  • num_replicas (int) – Number of systems in the batch

  • num_atoms (int) – Number of atoms in each system

  • device (torch.device) – Output tensor device

  • dtype (torch.dtype) – Output tensor data type

Return type:

tuple[torch.Tensor, torch.Tensor]

Returns:

Energies in eV and forces in eV/Angstrom

class schrodinger.application.matsci.mlip.wrappers.UMAWrapperBuffered(forcefield, ckpt_path='/scr/buildbot/builds/build/internal/data/mlff_models/UMA/uma-s-1p1.pt', device='cpu', cutoff=6.0, max_neigh=50)

Bases: object

UMA wrapper for nanoreactor (M replicates of the same molecule, no padding).

Behavior:

  • Single-shot AtomicData construction.

  • Empty connectivity; then build neighbors with torch_cluster.radius_graph, similar to FAIRChemCalculator

  • Only positions are updated each step (single tensor copy).

NOTE:

  • the default value of max_neigh in the UMA repo is vague, I set it to 50 here.

__init__(forcefield, ckpt_path='/scr/buildbot/builds/build/internal/data/mlff_models/UMA/uma-s-1p1.pt', device='cpu', cutoff=6.0, max_neigh=50)
Parameters:

  • forcefield (str) – Validated UMA MLFF force-field value

  • ckpt_path (str) – Path to UMA checkpoint file

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

  • cutoff (float) – Neighbor cutoff radius in Angstrom

  • max_neigh (int) – Maximum number of neighbors per atom

buildBuffer(z_tensor, r_tensor, *, charge=None, spin=None, cell=None, pbc=None)

Build the AtomicData batch buffer from atomic numbers and positions.

Parameters:

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

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

  • charge (torch.Tensor or None) – Charge per system (M,)

  • spin (torch.Tensor or None) – Spin multiplicity per system (M,)

  • cell (torch.Tensor or None) – Unit cell (M, 3, 3) in Angstrom

  • pbc (torch.Tensor or None) – Periodic boundary conditions (M, 3)