SoftSphereMultiModel

class torch_sim.models.soft_sphere.SoftSphereMultiModel(species=None, sigma_matrix=None, epsilon_matrix=None, alpha_matrix=None, device=None, dtype=torch.float32, *, pbc=True, compute_forces=True, compute_stress=False, per_atom_energies=False, per_atom_stresses=False, use_neighbor_list=True, cutoff=None)[source]

Bases: Module

Calculator for systems with multiple particle types.

Extends the basic soft sphere model to support multiple particle types with different interaction parameters for each pair of particle types. This enables simulation of heterogeneous systems like mixtures, composites, or biomolecular systems with different interaction strengths between different components.

This model maintains matrices of interaction parameters (sigma, epsilon, alpha) where each element [i,j] represents the parameter for interactions between particle types i and j.

Variables:

  • species (Tensor) – Particle type indices for each particle in the system.

  • sigma_matrix (Tensor) – Matrix of distance parameters for each pair of types. Shape: [n_types, n_types].

  • epsilon_matrix (Tensor) – Matrix of energy scale parameters for each pair. Shape: [n_types, n_types].

  • alpha_matrix (Tensor) – Matrix of exponents for each pair of types. Shape: [n_types, n_types].

  • cutoff (Tensor) – Maximum interaction distance.

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

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

  • periodic (bool) – Whether to use periodic boundary conditions.

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

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

Parameters:

Examples

```py # Create a binary mixture with different interaction parameters # Define interaction matrices (size 2x2 for binary system) sigma_matrix = torch.tensor(

[

[1.0, 0.8], # Type 0-0 and 0-1 interactions [0.8, 0.6], # Type 1-0 and 1-1 interactions

]

)

epsilon_matrix = torch.tensor(
[

[1.0, 0.5], # Type 0-0 and 0-1 interactions [0.5, 2.0], # Type 1-0 and 1-1 interactions

]

)

# Particle type assignments (0 or 1 for each particle) species = torch.tensor([0, 0, 1, 1, 0, 1])

# Create the model model = SoftSphereMultiModel(

species=species, sigma_matrix=sigma_matrix, epsilon_matrix=epsilon_matrix, compute_forces=True,

)

# Compute properties results = model(simulation_state) ```

unbatched_forward(state, species=None)[source]

Compute energies and forces for a single unbatched system with multiple species.

Internal implementation that processes a single, non-batched simulation state. This method handles all pair interactions between particles of different types using the appropriate interaction parameters from the parameter matrices.

Parameters:

  • state (SimState) – Single, non-batched simulation state containing atomic positions, cell vectors, and other system information.

  • species (Tensor | None) – Optional species indices to override the ones provided during initialization. Shape: [n_particles]. If None, uses the species defined at initialization. Defaults to None.

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 supports both neighbor list optimization and full pairwise calculations based on the use_neighbor_list parameter. For each pair of particles, it looks up the appropriate parameters based on the species of the two particles.

forward(state)[source]

Compute soft sphere potential properties for multi-component systems.

Main entry point for multi-species soft sphere 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 or if species information is missing.

Examples

```py # Create model for binary mixture model = SoftSphereMultiModel(

species=particle_types, sigma_matrix=distance_matrix, epsilon_matrix=strength_matrix, compute_forces=True,

)

# Calculate properties results = model(simulation_state) energy = results[“energy”] forces = results[“forces”] ```

Notes

This method requires species information either provided during initialization or included in the state object’s metadata.