DSFCoulombModel

class torch_sim.models.electrostatics.DSFCoulombModel(cutoff=10.0, *, alpha=0.2, device=None, dtype=torch.float64, compute_forces=True, compute_stress=True, neighbor_list_fn=torchsim_nl)[source]

Bases: ModelInterface

Damped Shifted Force electrostatics as a ModelInterface.

Uses the nvalchemiops DSF kernel for O(N) electrostatic energy, forces, and (optionally) stress. All user-facing quantities are in metal units (Angstrom / eV); the Coulomb constant ke is baked in.

Per-atom partial charges are read from state.partial_charges.

Parameters:

  • cutoff (float) – Real-space cutoff in Angstrom.

  • alpha (float) – DSF damping parameter. 0.0 gives shifted-force bare Coulomb.

  • device (device | None) – Compute device. Defaults to CUDA if available, else CPU.

  • dtype (dtype) – Floating-point dtype. Defaults to torch.float64.

  • compute_forces (bool) – Whether to return forces. Defaults to True.

  • compute_stress (bool) – Whether to return stress. Defaults to True.

  • neighbor_list_fn (Callable) – Neighbor-list constructor. Defaults to torchsim_nl.

forward(state, **_kwargs)[source]

Compute DSF electrostatic energy, forces, and stress.

Parameters:

  • state (SimState) – Simulation state with partial_charges set as an atom extra (shape [n_atoms]).

  • **_kwargs (object) – Unused; accepted for interface compatibility.

Returns:

dict with "energy" [n_systems], "forces" [n_atoms, 3], and (if compute_stress) "stress" [n_systems, 3, 3].

Return type:

dict[str, Tensor]