ParticleLifeModel

class torch_sim.models.particle_life.ParticleLifeModel(sigma=1.0, epsilon=1.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

Calculator for asymmetric particle interaction.

This model implements an asymmetric interaction between particles based on distance-dependent forces. The interaction is defined by three parameters: sigma, epsilon, and beta.

Parameters:
unbatched_forward(state)[source]

Compute energies and forces 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:

A dictionary containing the energy, forces, and stresses

Return type:

dict[str, Tensor]

forward(state)[source]

Compute particle life energies and forces for a system.

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

Parameters:

state (SimState | dict[Literal['positions', 'masses', 'cell', 'pbc', 'atomic_numbers', 'batch'], ~torch.Tensor]) – Input state containing atomic positions, cell vectors, and other system information. Can be a SimState object or a dictionary with the same keys.

Returns:

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.