MorseModel

class torch_sim.models.morse.MorseModel(sigma=1.0, epsilon=5.0, alpha=5.0, device=None, dtype=torch.float32, *, compute_forces=False, compute_stress=False, per_atom_energies=False, per_atom_stresses=False, use_neighbor_list=True, cutoff=None)[source]

Bases: Module, ModelInterface

Morse potential energy and force calculator.

Implements the Morse potential for molecular dynamics simulations. This model is particularly useful for modeling covalent bonds as it can accurately describe bond stretching, breaking, and anharmonic behavior. Unlike the Lennard-Jones potential, Morse is often better for cases where accurate dissociation energy and bond dynamics are important.

Variables:

  • sigma (Tensor) – Equilibrium bond length (r_e) in distance units.

  • epsilon (Tensor) – Dissociation energy (D_e) in energy units.

  • alpha (Tensor) – Parameter controlling the width/steepness of the potential.

  • cutoff (Tensor) – Distance cutoff for truncating potential calculation.

  • device (device) – Device where calculations are performed.

  • dtype (dtype) – Data type used for calculations.

  • compute_forces (bool) – Whether to compute atomic forces.

  • compute_stress (bool) – Whether to compute stress tensor.

  • per_atom_energies (bool) – Whether to compute per-atom energy decomposition.

  • per_atom_stresses (bool) – Whether to compute per-atom stress decomposition.

  • use_neighbor_list (bool) – Whether to use neighbor list optimization.

Parameters:

Examples

```py # Basic usage with default parameters morse_model = MorseModel(device=torch.device(“cuda”)) results = morse_model(sim_state)

# Model parameterized for O-H bonds in water, atomic units oh_model = MorseModel(

sigma=0.96, epsilon=4.52, alpha=2.0, compute_forces=True, compute_stress=True,

)

unbatched_forward(state)[source]

Compute Morse potential properties for a single unbatched system.

Internal implementation that processes a single, non-batched simulation state. This method handles the core computations of pair interactions, including neighbor list construction, distance calculations, and property computation.

Parameters:

state (SimState | StateDict) – Single, non-batched simulation state or equivalent dictionary containing atomic positions, cell vectors, and other system information.

Returns:

Dictionary of computed properties:

  • ”energy”: Total potential energy (scalar)

  • ”forces”: Atomic forces with shape [n_atoms, 3] (if

    compute_forces=True)

  • ”stress”: Stress tensor with shape [3, 3] (if compute_stress=True)

  • ”energies”: Per-atom energies with shape [n_atoms] (if

    per_atom_energies=True)

  • ”stresses”: Per-atom stresses with shape [n_atoms, 3, 3] (if

    per_atom_stresses=True)

Return type:

dict[str, Tensor]

Notes

This method can work with both neighbor list and full pairwise calculations. In both cases, interactions are truncated at the cutoff distance.

forward(state)[source]

Compute Morse potential energies, forces, and stresses for a system.

Main entry point for Morse potential calculations that handles batched states by dispatching each batch to the unbatched implementation and combining results.

Parameters:

state (SimState | StateDict) – Input state containing atomic positions, cell vectors, and other system information. Can be a SimState object or a dictionary with the same keys.

Returns:

Dictionary of computed properties:

  • ”energy”: Potential energy with shape [n_batches]

  • ”forces”: Atomic forces with shape [n_atoms, 3]

    (if compute_forces=True)

  • ”stress”: Stress tensor with shape [n_batches, 3, 3]

    (if compute_stress=True)

  • May include additional outputs based on configuration

Return type:

dict[str, Tensor]

Raises:

ValueError – If batch cannot be inferred for multi-cell systems.

Examples

```py # Compute properties for a simulation state model = MorseModel(compute_forces=True) results = model(sim_state)

energy = results[“energy”] # Shape: [n_batches] forces = results[“forces”] # Shape: [n_atoms, 3] ```