SoftSphereModel

class torch_sim.models.soft_sphere.SoftSphereModel(sigma=1.0, epsilon=1.0, alpha=2.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

Calculator for soft sphere potential energies and forces.

Implements a model for computing properties based on the soft sphere potential, which describes purely repulsive interactions between particles. This potential is useful for modeling systems where particles should not overlap but don’t have attractive interactions, such as granular materials and some colloidal systems.

The potential energy between particles i and j is:

V_ij(r) = epsilon * (sigma/r)^alpha

Variables:

  • sigma (Tensor) – Effective particle diameter in distance units.

  • epsilon (Tensor) – Energy scale parameter in energy units.

  • alpha (Tensor) – Exponent controlling repulsion steepness (typically ≥ 2).

  • cutoff (Tensor) – Cutoff distance for interactions.

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

  • _device (device) – Computation device (CPU/GPU).

  • _dtype (dtype) – Data type for tensor calculations.

  • _compute_forces (bool) – Whether to compute 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.

Parameters:

Examples

```py # Basic usage with default parameters model = SoftSphereModel() results = model(sim_state)

# Custom parameters for colloidal system colloid_model = SoftSphereModel(

sigma=2.0, # particle diameter in nm epsilon=10.0, # energy scale in kJ/mol alpha=12.0, # steep repulsion for hard colloids compute_stress=True,

)

# Get forces for a system with periodic boundary conditions results = colloid_model(

SimState(

positions=positions, cell=box_vectors, pbc=torch.tensor([True, True, True]),

)

) forces = results[“forces”] # shape: [n_particles, 3] ```

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 for pair interactions, including neighbor list construction, distance calculations, and property computation.

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 can work with both neighbor list and full pairwise calculations. The soft sphere potential is purely repulsive, and forces are truncated at the cutoff distance.

forward(state)[source]

Compute soft sphere potential energies, forces, and stresses for a system.

Main entry point for soft sphere 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 = SoftSphereModel(compute_forces=True) results = model(sim_state)

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