LennardJonesModel

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

Bases: Module, ModelInterface

Lennard-Jones potential energy and force calculator.

Implements the Lennard-Jones 12-6 potential for molecular dynamics simulations. This model calculates pairwise interactions between atoms and supports either full pairwise calculation or neighbor list-based optimization for efficiency.

Variables:

  • sigma (Tensor) – Length parameter controlling particle size/repulsion distance.

  • epsilon (Tensor) – Energy parameter controlling interaction strength.

  • 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:

Example:

# Basic usage with default parameters
lj_model = LennardJonesModel(device=torch.device("cuda"))
results = lj_model(sim_state)

# Custom parameterization for Argon
ar_model = LennardJonesModel(
    sigma=3.405,  # Å
    epsilon=0.0104,  # eV
    cutoff=8.5,  # Å
    compute_stress=True,
)
unbatched_forward(state)[source]

Compute Lennard-Jones 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, neighbor lists, and property calculations.

Parameters:

state (SimState) – Single, non-batched simulation state 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 handles two different approaches: 1. Neighbor list approach: Efficient for larger systems 2. Full pairwise calculation: Better for small systems

The implementation applies cutoff distance to both approaches for consistency.

forward(state)[source]

Compute Lennard-Jones energies, forces, and stresses for a system.

Main entry point for Lennard-Jones 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)

  • ”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]

Raises:

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

Example:

# Compute properties for a simulation state
model = LennardJonesModel(compute_stress=True)
results = model(sim_state)

energy = results["energy"]  # Shape: [n_batches]
forces = results["forces"]  # Shape: [n_atoms, 3]
stress = results["stress"]  # Shape: [n_batches, 3, 3]
energies = results["energies"]  # Shape: [n_atoms]
stresses = results["stresses"]  # Shape: [n_atoms, 3, 3]