torchff

Contents

torchff#

Bond#

Bond energy and force terms (harmonic and Amoeba-style) for force field calculations.

torchff.bond.compute_harmonic_bond_energy(coords: Tensor, bonds: Tensor, b0: Tensor, k: Tensor) → Tensor[source]#

Compute harmonic bond energies via custom CUDA/CPU ops.

Energy per bond: \(E = (1/2) k (r - r_0)^2\) with \(r\) the bond length.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • bonds (torch.Tensor) – Shape (M, 2), integer indices; each row is (i, j).

  • b0 (torch.Tensor) – Shape (M,), equilibrium bond lengths.

  • k (torch.Tensor) – Shape (M,), force constants.

Returns:

Scalar total harmonic bond energy.

Return type:

torch.Tensor

torchff.bond.compute_harmonic_bond_energy_ref(coords, bonds, b0, k)[source]#

Reference implementation of harmonic bond energy (PyTorch only).

Same formula as compute_harmonic_bond_energy(); used when custom ops are disabled.

class torchff.bond.HarmonicBond(use_customized_ops: bool = False)[source]#

Bases: Module

Harmonic bond energy module: \(E = (1/2) k (r - r_0)^2\).

Dispatches to custom ops or a PyTorch reference implementation based on use_customized_ops.

__init__(use_customized_ops: bool = False)[source]#
Parameters:

use_customized_ops (bool, optional) – If True, use custom CUDA/CPU kernels; otherwise use reference implementation.

forward(coords, bonds, b0, k)[source]#

Compute total harmonic bond energy.

Parameters:
Returns:

Scalar total energy.

Return type:

torch.Tensor

torchff.bond.compute_harmonic_bond_forces(coords: Tensor, bonds: Tensor, b0: Tensor, k: Tensor, forces: Tensor)[source]#

Compute harmonic bond forces in-place (no autograd).

Adds bond force contributions into forces. Intended for fast MD; backward is not supported.

Parameters:
torchff.bond.compute_amoeba_bond_energy(coords: Tensor, bonds: Tensor, b0: Tensor, k: Tensor, cubic: float = -2.55, quartic: float = 3.793125) → Tensor[source]#

Compute Amoeba-style bond energies via custom ops.

\(E = k (b - b_0)^2 [1 + c_3 (b - b_0) + c_4 (b - b_0)^2]\) with \(b\) the bond length and \(c_3, c_4\) the cubic and quartic coefficients.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • bonds (torch.Tensor) – Shape (M, 2), bond indices.

  • b0 (torch.Tensor) – Shape (M,), equilibrium lengths.

  • k (torch.Tensor) – Shape (M,), force constants.

  • cubic (float, optional) – Cubic coefficient in the polynomial (default -2.55).

  • quartic (float, optional) – Quartic coefficient (default 3.793125).

Returns:

Scalar total Amoeba bond energy.

Return type:

torch.Tensor

torchff.bond.compute_amoeba_bond_energy_ref(coords, bonds, b0, k, cubic: float = -2.55, quartic: float = 3.793125)[source]#

Reference implementation of Amoeba bond energy (PyTorch only).

Same formula as compute_amoeba_bond_energy(); used when custom ops are disabled.

class torchff.bond.AmoebaBond(use_customized_ops: bool = False)[source]#

Bases: Module

Amoeba-style bond energy with cubic and quartic correction terms.

Dispatches to custom ops or reference based on use_customized_ops. See HarmonicBond for constructor and dispatch behavior.

__init__(use_customized_ops: bool = False)[source]#

Same as HarmonicBond.__init__.

forward(coords, bonds, b0, k, cubic: float = -2.55, quartic: float = 3.793125)[source]#

Compute total Amoeba bond energy.

Parameters:
Returns:

Scalar total energy.

Return type:

torch.Tensor

torchff.bond.compute_morse_bond_energy(coords: Tensor, bonds: Tensor, b0: Tensor, k: Tensor, d: Tensor) → Tensor[source]#

Compute Morse bond energies via custom CUDA/CPU ops.

Energy per bond: \(E = D [1 - e^{-\beta (r - r_0)}]^2\) with \(\beta = \sqrt{k/(2D)}\), \(r\) the bond length, \(r_0\) the equilibrium length, \(k\) the stiffness, and \(D\) the well depth.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • bonds (torch.Tensor) – Shape (M, 2), integer indices; each row is (i, j).

  • b0 (torch.Tensor) – Shape (M,), equilibrium bond lengths \(r_0\).

  • k (torch.Tensor) – Shape (M,), stiffness parameters \(k\).

  • d (torch.Tensor) – Shape (M,), well depth \(D\) (must be positive).

Returns:

Scalar total Morse bond energy.

Return type:

torch.Tensor

torchff.bond.compute_morse_bond_energy_ref(coords, bonds, b0, k, d)[source]#

Reference implementation of Morse bond energy (PyTorch only).

Same formula as compute_morse_bond_energy(); used when custom ops are disabled.

class torchff.bond.MorseBond(use_customized_ops: bool = False)[source]#

Bases: Module

Morse bond energy module.

Dispatches to custom ops or a PyTorch reference implementation based on use_customized_ops.

__init__(use_customized_ops: bool = False)[source]#
forward(coords, bonds, b0, k, d)[source]#

Compute total Morse bond energy.

Parameters:
Returns:

Scalar total energy.

Return type:

torch.Tensor

torchff.bond.compute_morse_bond_forces(coords: Tensor, bonds: Tensor, b0: Tensor, k: Tensor, d: Tensor, forces: Tensor)[source]#

Compute Morse bond forces in-place (no autograd).

Adds bond force contributions into forces.

Parameters:

Angle#

Angle energy and force terms (harmonic and Amoeba-style) for force field calculations.

torchff.angle.compute_angles_ref(coords: Tensor, angles: Tensor) → Tensor[source]#

Reference implementation: angle (theta) at the central atom for each triple (PyTorch only).

For each angle (i, j, k), computes the angle at j between vectors (coords[i] - coords[j]) and (coords[k] - coords[j]).

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), integer indices; each row is (i, j, k).

Returns:

Shape (M,), angle in radians for each triple.

Return type:

torch.Tensor

torchff.angle.compute_angles(coords: Tensor, angles: Tensor) → Tensor[source]#

Compute angle values via custom CUDA/CPU ops.

Same geometry as compute_angles_ref().

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), integer indices; each row is (i, j, k).

Returns:

Shape (M,), angle in radians for each triple.

Return type:

torch.Tensor

class torchff.angle.Angles(use_customized_ops: bool = False)[source]#

Bases: Module

Plane angle values at central atoms (radians), matching compute_angles_ref() geometry.

Dispatches to custom ops or reference based on use_customized_ops.

__init__(use_customized_ops: bool = False)[source]#
forward(coords: Tensor, angles: Tensor) → Tensor[source]#
torchff.angle.compute_harmonic_angle_energy(coords: Tensor, angles: Tensor, theta0: Tensor, k: Tensor) → Tensor[source]#

Compute harmonic angle energies via custom CUDA/CPU ops.

Energy per angle: \(E = (1/2) k (\theta - \theta_0)^2\) with \(\theta\) in radians.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), integer indices; each row is (i, j, k), j the central atom.

  • theta0 (torch.Tensor) – Shape (M,), equilibrium angles in radians.

  • k (torch.Tensor) – Shape (M,), force constants.

Returns:

Scalar total harmonic angle energy.

Return type:

torch.Tensor

torchff.angle.compute_harmonic_angle_energy_ref(coords, angles, theta0, k)[source]#

Reference implementation of harmonic angle energy (PyTorch only).

Uses compute_angles_ref() then same formula as compute_harmonic_angle_energy(); used when custom ops are disabled.

class torchff.angle.HarmonicAngle(use_customized_ops: bool = False)[source]#

Bases: Module

Harmonic angle energy module: \(E = (1/2) k (\theta - \theta_0)^2\).

Dispatches to custom ops or a PyTorch reference implementation based on use_customized_ops.

__init__(use_customized_ops: bool = False)[source]#
Parameters:

use_customized_ops (bool, optional) – If True, use custom CUDA/CPU kernels; otherwise use reference implementation.

forward(coords, angles, theta0, k)[source]#

Compute total harmonic angle energy.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), angle indices (i, j, k).

  • theta0 (torch.Tensor) – Shape (M,), equilibrium angles in radians.

  • k (torch.Tensor) – Shape (M,), force constants.

Returns:

Scalar total energy.

Return type:

torch.Tensor

torchff.angle.compute_harmonic_angle_forces(coords: Tensor, angles: Tensor, theta0: Tensor, k: Tensor, forces: Tensor) → Tensor[source]#

Compute harmonic angle forces in-place (no autograd).

Adds angle force contributions into forces. Intended for fast MD; backward is not supported.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), angle indices.

  • theta0 (torch.Tensor) – Shape (M,), equilibrium angles in radians.

  • k (torch.Tensor) – Shape (M,), force constants.

  • forces (torch.Tensor) – Shape (N, 3), modified in-place with angle forces.

Returns:

Reference to forces.

Return type:

torch.Tensor

torchff.angle.compute_amoeba_angle_energy(coords: Tensor, angles: Tensor, theta0: Tensor, k: Tensor, cubic: float = -0.014, quartic: float = 5.6e-05, pentic: float = -7e-07, sextic: float = 2.2e-08) → Tensor[source]#

Compute Amoeba-style angle energies via custom ops.

\(E = k (\theta - \theta_0)^2 [1 + c_3 \Delta\theta + c_4 \Delta\theta^2 + c_5 \Delta\theta^3 + c_6 \Delta\theta^4]\) with \(\Delta\theta = \theta - \theta_0\) in radians.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), angle indices (i, j, k).

  • theta0 (torch.Tensor) – Shape (M,), equilibrium angles in radians.

  • k (torch.Tensor) – Shape (M,), force constants.

  • cubic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

  • quartic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

  • pentic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

  • sextic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

Returns:

Scalar total Amoeba angle energy.

Return type:

torch.Tensor

torchff.angle.compute_amoeba_angle_energy_ref(coords, angles, theta0, k, cubic: float = -0.014, quartic: float = 5.6e-05, pentic: float = -7e-07, sextic: float = 2.2e-08)[source]#

Reference implementation of Amoeba angle energy (PyTorch only).

Uses compute_angles_ref(); same formula as compute_amoeba_angle_energy(); used when custom ops are disabled.

class torchff.angle.AmoebaAngle(use_customized_ops: bool = False)[source]#

Bases: Module

Amoeba-style angle energy with cubic through sextic correction terms.

Dispatches to custom ops or reference based on use_customized_ops. See HarmonicAngle for constructor and dispatch behavior.

__init__(use_customized_ops: bool = False)[source]#

Same as HarmonicAngle.__init__.

forward(coords, angles, theta0, k, cubic: float = -0.014, quartic: float = 5.6e-05, pentic: float = -7e-07, sextic: float = 2.2e-08)[source]#

Compute total Amoeba angle energy.

Parameters:

  • coords – Same as HarmonicAngle.forward().

  • angles – Same as HarmonicAngle.forward().

  • theta0 – Same as HarmonicAngle.forward().

  • k – Same as HarmonicAngle.forward().

  • cubic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

  • quartic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

  • pentic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

  • sextic (float, optional) – Polynomial coefficients (defaults: -0.014, 5.6e-5, -7.0e-7, 2.2e-8).

Returns:

Scalar total energy.

Return type:

torch.Tensor

torchff.angle.compute_cosine_angle_energy(coords: Tensor, angles: Tensor, theta0: Tensor, k: Tensor) → Tensor[source]#

Cosine-style angle energy via custom ops.

Energy per angle: \(E = (1/2) k (\cos\theta - \cos\theta_0)^2\).

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • angles (torch.Tensor) – Shape (M, 3), angle indices (i, j, k).

  • theta0 (torch.Tensor) – Shape (M,), equilibrium angles in radians.

  • k (torch.Tensor) – Shape (M,), force constants.

Returns:

Scalar total cosine angle energy.

Return type:

torch.Tensor

torchff.angle.compute_cosine_angle_energy_ref(coords, angles, theta0, k)[source]#

Reference implementation of cosine angle energy (PyTorch only).

Same formula as compute_cosine_angle_energy().

class torchff.angle.CosineAngle(use_customized_ops: bool = False)[source]#

Bases: Module

Cosine angle energy: \(E = (1/2) k (\cos\theta - \cos\theta_0)^2\).

Dispatches to custom ops or reference based on use_customized_ops.

__init__(use_customized_ops: bool = False)[source]#
forward(coords, angles, theta0, k)[source]#
torchff.angle.compute_cosine_angle_forces(coords: Tensor, angles: Tensor, theta0: Tensor, k: Tensor, forces: Tensor) → Tensor[source]#

Cosine angle forces in-place (no autograd).

Parameters:

Torsion#

Torsion (dihedral) angles, harmonic torsion energy, and periodic torsion energies.

torchff.torsion.compute_torsion_ref(coords: Tensor, torsions: Tensor) → Tensor[source]#

Reference implementation of torsion (dihedral) angles using PyTorch ops.

For each torsion (i, j, k, l), computes the dihedral angle between the planes formed by (i, j, k) and (j, k, l). The convention matches the CUDA implementation used by the custom operator.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • torsions (torch.Tensor) – Shape (M, 4), integer indices; each row is (i, j, k, l).

Returns:

Shape (M,), torsion angles in radians.

Return type:

torch.Tensor

torchff.torsion.compute_torsion(coords: Tensor, torsions: Tensor) → Tensor[source]#

Compute torsion (dihedral) angles via custom CUDA/CPU ops.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • torsions (torch.Tensor) – Shape (M, 4), integer indices; each row is (i, j, k, l).

Returns:

Shape (M,), torsion angles in radians.

Return type:

torch.Tensor

class torchff.torsion.Torsion(use_customized_ops: bool = False)[source]#

Bases: Module

Torsion (dihedral) angle module returning raw angle values in radians.

Dispatches to custom ops or a PyTorch reference implementation based on use_customized_ops.

__init__(use_customized_ops: bool = False)[source]#
Parameters:

use_customized_ops (bool, optional) – If True, use custom CUDA/CPU kernels; otherwise use reference implementation.

forward(coords: Tensor, torsions: Tensor) → Tensor[source]#

Compute torsion (dihedral) angles.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • torsions (torch.Tensor) – Shape (M, 4), torsion indices (i, j, k, l).

Returns:

Shape (M,), torsion angles in radians.

Return type:

torch.Tensor

torchff.torsion.compute_periodic_torsion_energy(coords: Tensor, torsions: Tensor, fc: Tensor, periodicity: Tensor, phase: Tensor) → Tensor[source]#

Compute periodic torsion energies

torchff.torsion.compute_periodic_torsion_energy_ref(coords, torsions, fc, periodicity, phase)[source]#
class torchff.torsion.PeriodicTorsion(use_customized_ops: bool = False)[source]#

Bases: Module

__init__(use_customized_ops: bool = False)[source]#
forward(coords, torsions, fc, periodicity, phase)[source]#
torchff.torsion.compute_harmonic_torsion_energy_ref(coords: Tensor, torsions: Tensor, k: Tensor, torsion_ref: Tensor | None = None) → Tensor[source]#

Reference harmonic torsion energy using PyTorch dihedral angles.

Total energy is sum_i (k_i / 2) * (phi_i - phi0_i)^2 when torsion_ref (phi0) is given; if torsion_ref is None, uses sum_i (k_i / 2) * phi_i^2.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • torsions (torch.Tensor) – Shape (M, 4), integer indices (i, j, k, l) per torsion.

  • k (torch.Tensor) – Shape (M,), force constants.

  • torsion_ref (torch.Tensor, optional) – Shape (M,), equilibrium torsion angles in radians. If None, the equilibrium angle is taken as zero for every term.

Returns:

Scalar total energy.

Return type:

torch.Tensor

class torchff.torsion.HarmonicTorsion(use_customized_ops: bool = False)[source]#

Bases: Module

Harmonic torsion energy: sum_i (k_i/2) (phi_i - phi0_i)^2.

Dihedral angles phi_i come from Torsion (custom or reference ops). The energy reduction is pure PyTorch; there is no separate fused CUDA kernel.

__init__(use_customized_ops: bool = False)[source]#
Parameters:

use_customized_ops (bool, optional) – If True, dihedral angles use compute_torsion(); otherwise compute_torsion_ref().

forward(coords: Tensor, torsions: Tensor, k: Tensor, torsion_ref: Tensor | None = None) → Tensor[source]#

Compute total harmonic torsion energy.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • torsions (torch.Tensor) – Shape (M, 4), torsion indices (i, j, k, l).

  • k (torch.Tensor) – Shape (M,), force constants.

  • torsion_ref (torch.Tensor, optional) – Shape (M,), equilibrium angles in radians. If None, uses sum (k/2) phi^2.

Returns:

Scalar total energy.

Return type:

torch.Tensor

torchff.torsion.compute_periodic_torsion_forces(coords: Tensor, torsions: Tensor, fc: Tensor, periodicity: Tensor, phase: Tensor, forces: Tensor) → None[source]#

Compute periodic torsion forces in-place, backward calculation does not supported, used for fast MD

Coulomb#

torchff.coulomb.compute_coulomb_energy(coords: Tensor, pairs: Tensor, box: Tensor, charges: Tensor, coulomb_constant: float, cutoff: float, do_shift: bool = True) → Tensor[source]#

Compute Coulomb energies using customized CUDA/C++ ops

torchff.coulomb.compute_coulomb_energy_ref(coords: Tensor, pairs: Tensor, box: Tensor, charges: Tensor, coulomb_constant: float, cutoff: float, do_shift: bool = True) → Tensor[source]#

Reference Coulomb implementation using native PyTorch ops

class torchff.coulomb.Coulomb(use_customized_ops: bool = False)[source]#

Bases: Module

__init__(use_customized_ops: bool = False)[source]#
forward(coords: Tensor, pairs: Tensor, box: Tensor, charges: Tensor, coulomb_constant: float, cutoff: float, do_shift: bool = True) → Tensor[source]#

Van der Waals#

Van der Waals (vdW) pair energies: Lennard-Jones 12-6 and AMOEBA buffered 14-7.

torchff.vdw.compute_vdw_14_7_energy(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, cutoff: float, atom_types: Tensor | None = None) → Tensor[source]#

Compute AMOEBA buffered 14-7 vdW pair energies via custom CUDA/C++ ops.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • pairs (torch.Tensor) – Shape (P, 2), integer indices (i, j) of interacting pairs.

  • box (torch.Tensor) – Shape (3, 3) or broadcastable, periodic box (same convention as torchff.pbc).

  • sigma (torch.Tensor) – Per-pair or type-pair \(\sigma\) (see atom_types).

  • epsilon (torch.Tensor) – Per-pair or type-pair \(\epsilon\) (see atom_types).

  • cutoff (float) – Distance cutoff; interactions beyond cutoff are excluded by the kernel.

  • atom_types (torch.Tensor, optional) – If provided, used by the backend for type-based indexing together with sigma and epsilon.

Returns:

Scalar total vdW energy for the buffered 14-7 potential.

Return type:

torch.Tensor

torchff.vdw.compute_lennard_jones_energy(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, cutoff: float, atom_types: Tensor | None = None) → Tensor[source]#

Compute Lennard-Jones 12-6 vdW pair energies via custom CUDA/C++ ops.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • pairs (torch.Tensor) – Shape (P, 2), integer indices (i, j) of interacting pairs.

  • box (torch.Tensor) – Shape (3, 3) or broadcastable, periodic box (same convention as torchff.pbc).

  • sigma (torch.Tensor) – Per-pair or type-pair \(\sigma\) (see atom_types).

  • epsilon (torch.Tensor) – Per-pair or type-pair \(\epsilon\) (see atom_types).

  • cutoff (float) – Distance cutoff; interactions beyond cutoff are excluded by the kernel.

  • atom_types (torch.Tensor, optional) – If provided, used by the backend for type-based indexing together with sigma and epsilon.

Returns:

Scalar total Lennard-Jones energy.

Return type:

torch.Tensor

torchff.vdw.compute_lennard_jones_energy_ref(r_ij, sigma_ij, epsilon_ij, sum=True)[source]#

Reference Lennard-Jones 12-6 pair energy in PyTorch.

Per pair: \(E_{ij} = 4 \epsilon_{ij} \left[ (\sigma_{ij}/r_{ij})^{12} - (\sigma_{ij}/r_{ij})^6 \right]\).

Parameters:

  • r_ij (torch.Tensor) – Pair distances, shape (P,) or broadcastable.

  • sigma_ij (torch.Tensor) – \(\sigma\) for each pair, same shape as r_ij (after broadcast).

  • epsilon_ij (torch.Tensor) – \(\epsilon\) for each pair, same shape as r_ij (after broadcast).

  • sum (bool, optional) – If True (default), return the sum over pairs; otherwise return per-pair energies.

Returns:

Scalar total energy if sum is True, else shape (P,) per-pair energies.

Return type:

torch.Tensor

torchff.vdw.compute_vdw_14_7_energy_ref(r_ij, sigma_ij, epsilon_ij, sum=True)[source]#

Reference AMOEBA buffered 14-7 vdW pair energy in PyTorch.

With \(\rho = r_{ij} / \sigma_{ij}\), \(E_{ij} = \epsilon_{ij} \left( \frac{1.07}{\rho + 0.07} \right)^7 \left( \frac{1.12}{\rho^7 + 0.12} - 2 \right)\).

Parameters:

  • r_ij (torch.Tensor) – Pair distances, shape (P,) or broadcastable.

  • sigma_ij (torch.Tensor) – \(\sigma\) for each pair, same shape as r_ij (after broadcast).

  • epsilon_ij (torch.Tensor) – \(\epsilon\) for each pair, same shape as r_ij (after broadcast).

  • sum (bool, optional) – If True (default), return the sum over pairs; otherwise return per-pair energies.

Returns:

Scalar total energy if sum is True, else shape (P,) per-pair energies.

Return type:

torch.Tensor

class torchff.vdw.Vdw(function: Literal['LennardJones', 'AmoebaVdw147'] = 'LennardJones', cutoff: float | None = None, use_customized_ops: bool = False, use_type_pairs: bool = False, sum_output: bool = True, cuda_graph_compat: bool = True)[source]#

Bases: Module

Van der Waals pair energy module (Lennard-Jones 12-6 or AMOEBA buffered 14-7).

Dispatches to compute_lennard_jones_energy() / compute_vdw_14_7_energy() when use_customized_ops is True; otherwise uses minimum-image displacements via torchff.pbc.PBC and the reference formulas compute_lennard_jones_energy_ref() / compute_vdw_14_7_energy_ref().

__init__(function: Literal['LennardJones', 'AmoebaVdw147'] = 'LennardJones', cutoff: float | None = None, use_customized_ops: bool = False, use_type_pairs: bool = False, sum_output: bool = True, cuda_graph_compat: bool = True)[source]#
Parameters:

  • function ({'LennardJones', 'AmoebaVdw147'}, optional) – Potential form: standard LJ 12-6 or AMOEBA buffered 14-7.

  • cutoff (float, optional) – Stored on the module; the active cutoff is the cutoff argument to forward().

  • use_customized_ops (bool, optional) – If True, use custom CUDA/C++ kernels; otherwise use the PyTorch reference path.

  • use_type_pairs (bool, optional) – If True, sigma and epsilon are indexed by atom_types for each pair (shape (n_types, n_types)).

  • sum_output (bool, optional) – If True (default), return a scalar sum over pairs. Must be True when use_customized_ops is True because the custom kernels only return total energy. When use_customized_ops is False, if False return per-pair energies of shape (P,).

  • cuda_graph_compat (bool, optional) – If True (default), apply the cutoff with torch.where() so tensor shapes are stable; if False, distances are filtered with boolean indexing before the energy expression.

expand_type_pairs(sigma, epsilon, pairs, atom_types)[source]#
forward(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, cutoff: float, atom_types: Tensor | None = None)[source]#

Compute vdW energy for the configured potential.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • pairs (torch.Tensor) – Shape (P, 2), pair indices (i, j).

  • box (torch.Tensor) – Periodic box, same convention as torchff.pbc.PBC.

  • sigma (torch.Tensor) – Per-pair (P,) or type table (T, T) when use_type_pairs is True.

  • epsilon (torch.Tensor) – Same layout as sigma.

  • cutoff (float) – Pair distance cutoff.

  • atom_types (torch.Tensor, optional) – Shape (N,), integer atom types; required when use_type_pairs is True.

Returns:

If use_customized_ops is True, scalar total energy from the custom op. Otherwise per-pair energies of shape (P,), or a scalar if sum_output is True.

Return type:

torch.Tensor

Dispersion#

Tang–Tonnies damped C6 dispersion: U = -C6 f_6(br) / r^6.

torchff.dispersion.compute_tang_tonnies_dispersion_energy(coords: Tensor, pairs: Tensor, box: Tensor, c6: Tensor, b: Tensor, cutoff: float, atom_types: Tensor | None = None) → Tensor[source]#

Total Tang–Tonnies C6 dispersion energy (custom CUDA).

torchff.dispersion.compute_tang_tonnies_dispersion_energy_ref(r_ij: Tensor, c6_ij: Tensor, b_ij: Tensor, sum: bool = True) → Tensor[source]#

Reference: U_ij = -C6_ij f_6(b_ij r_ij) / r_ij^6.

Parameters:

  • r_ij (torch.Tensor) – Pair distances, shape (P,) or broadcastable.

  • c6_ij (torch.Tensor) – Per-pair C6 and Tang–Tonnies inverse length b, same shape as r_ij after broadcast.

  • b_ij (torch.Tensor) – Per-pair C6 and Tang–Tonnies inverse length b, same shape as r_ij after broadcast.

  • sum (bool) – If True, return scalar sum over pairs; else per-pair energies (P,).

class torchff.dispersion.Dispersion(cutoff: float | None = None, use_customized_ops: bool = False, use_type_pairs: bool = False, sum_output: bool = False, cuda_graph_compat: bool = True)[source]#

Bases: Module

Tang–Tonnies C6 dispersion pair energy.

Dispatches to compute_tang_tonnies_dispersion_energy() when use_customized_ops is True; otherwise PBC minimum-image distances and compute_tang_tonnies_dispersion_energy_ref().

__init__(cutoff: float | None = None, use_customized_ops: bool = False, use_type_pairs: bool = False, sum_output: bool = False, cuda_graph_compat: bool = True) → None[source]#
expand_type_pairs(c6: Tensor, b: Tensor, pairs: Tensor, atom_types: Tensor | None)[source]#
forward(coords: Tensor, pairs: Tensor, box: Tensor, c6: Tensor, b: Tensor, cutoff: float, atom_types: Tensor | None = None) → Tensor[source]#

Slater#

Slater-type pair energies: A * P(B r) * exp(-B r) with P = (B r)^2 / 3 + B r + 1.

torchff.slater.compute_slater_energy(coords: Tensor, pairs: Tensor, box: Tensor, A: Tensor, B: Tensor, cutoff: float, atom_types: Tensor | None = None) → Tensor[source]#

Compute Slater pair energies via custom CUDA/C++ ops.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • pairs (torch.Tensor) – Shape (P, 2), integer indices (i, j) of interacting pairs.

  • box (torch.Tensor) – Shape (3, 3) or broadcastable, periodic box (same convention as torchff.pbc).

  • A (torch.Tensor) – Per-pair or type-pair amplitude \(A_{ij}\) (see atom_types).

  • B (torch.Tensor) – Per-pair or type-pair inverse length \(B_{ij}\) (see atom_types).

  • cutoff (float) – Distance cutoff; interactions beyond cutoff are excluded by the kernel.

  • atom_types (torch.Tensor, optional) – If provided, used by the backend for type-based indexing together with A and B.

Returns:

Scalar total Slater energy.

Return type:

torch.Tensor

torchff.slater.compute_slater_energy_ref(r_ij: Tensor, A_ij: Tensor, B_ij: Tensor, sum: bool = True)[source]#

Reference Slater pair energy in PyTorch.

With \(x = B_{ij} r_{ij}\), \(P = x^2/3 + x + 1\), \(E_{ij} = A_{ij} P \exp(-x)\).

Parameters:

  • r_ij (torch.Tensor) – Pair distances, shape (P,) or broadcastable.

  • A_ij (torch.Tensor) – \(A_{ij}\) for each pair, same shape as r_ij (after broadcast).

  • B_ij (torch.Tensor) – \(B_{ij}\) for each pair, same shape as r_ij (after broadcast).

  • sum (bool, optional) – If True (default), return the sum over pairs; otherwise return per-pair energies.

Returns:

Scalar total energy if sum is True, else shape (P,) per-pair energies.

Return type:

torch.Tensor

class torchff.slater.Slater(cutoff: float | None = None, use_customized_ops: bool = False, use_type_pairs: bool = False, sum_output: bool = False, cuda_graph_compat: bool = True)[source]#

Bases: Module

Slater pair energy module.

Dispatches to compute_slater_energy() when use_customized_ops is True; otherwise uses minimum-image displacements via torchff.pbc.PBC and compute_slater_energy_ref().

__init__(cutoff: float | None = None, use_customized_ops: bool = False, use_type_pairs: bool = False, sum_output: bool = False, cuda_graph_compat: bool = True)[source]#
Parameters:

  • cutoff (float, optional) – Stored on the module; the active cutoff is the cutoff argument to forward().

  • use_customized_ops (bool, optional) – If True, use custom CUDA/C++ kernels; otherwise use the PyTorch reference path.

  • use_type_pairs (bool, optional) – If True, A and B are indexed by atom_types for each pair (shape (n_types, n_types)).

  • sum_output (bool, optional) – If True, return a scalar sum over pairs. Requires use_customized_ops False.

  • cuda_graph_compat (bool, optional) – If True (default), apply the cutoff with torch.where() so tensor shapes are stable; if False, distances are filtered with boolean indexing before the energy expression.

expand_type_pairs(A: Tensor, B: Tensor, pairs: Tensor, atom_types: Tensor)[source]#
forward(coords: Tensor, pairs: Tensor, box: Tensor, A: Tensor, B: Tensor, cutoff: float, atom_types: Tensor | None = None)[source]#

Compute Slater energy.

Parameters:

  • coords (torch.Tensor) – Shape (N, 3), atom coordinates.

  • pairs (torch.Tensor) – Shape (P, 2), pair indices (i, j).

  • box (torch.Tensor) – Periodic box, same convention as torchff.pbc.PBC.

  • A (torch.Tensor) – Per-pair (P,) or type table (T, T) when use_type_pairs is True.

  • B (torch.Tensor) – Same layout as A.

  • cutoff (float) – Pair distance cutoff.

  • atom_types (torch.Tensor, optional) – Shape (N,), integer atom types; required when use_type_pairs is True.

Returns:

If use_customized_ops is True, scalar total energy from the custom op. Otherwise per-pair energies of shape (P,), or a scalar if sum_output is True.

Return type:

torch.Tensor

Neighbor List#

class torchff.nblist.NeighborList(natoms: int, exclusions: Tensor | None = None, include_self: bool = False, use_customized_ops: bool = True, algorithm: str = 'nsquared')[source]#

Bases: Module

__init__(natoms: int, exclusions: Tensor | None = None, include_self: bool = False, use_customized_ops: bool = True, algorithm: str = 'nsquared')[source]#
classmethod convert_pairs_coo_to_csr(pairs_coo: Tensor, num_atoms: int | None = None) → Tuple[Tensor, Tensor][source]#

Convert a COO pair list into CSR format for a symmetric matrix.

Given a list of index pairs in COO (coordinate) format representing non-zero entries of a symmetric matrix, this method produces the equivalent CSR (Compressed Sparse Row) representation. Because the matrix is symmetric, both directions (i, j) and (j, i) are guaranteed in the output regardless of whether the input contains one or both directions.

The conversion proceeds in four steps:

  1. Symmetrize – concatenate pairs_coo with its column-flipped copy so that both (i, j) and (j, i) are present.

  2. Encode and deduplicate – map each pair to a scalar key row * num_atoms + col and call torch.unique() on the 1-D key tensor. This is more efficient than torch.unique(dim=0) and produces keys in sorted (row-major) order.

  3. Decode – recover row and column indices from the unique keys via integer division and modulo.

  4. Build row_ptr – count neighbors per row with torch.Tensor.scatter_add_(), then take the cumulative sum to form the CSR indptr array.

Parameters:

  • pairs_coo (torch.Tensor) – Integer tensor of shape (N, 2) where each row [i, j] denotes a non-zero entry. The input may contain only one direction of a symmetric pair, both directions, or a mix.

  • num_atoms (int, optional) – Total number of atoms (i.e. the matrix dimension). When None (default), it is inferred as pairs_coo.max() + 1. Specify this explicitly when isolated atoms with no pairs exist so that row_ptr has the correct length.

Returns:

  • row_ptr (torch.Tensor) – Shape (num_atoms + 1,), dtype torch.long. row_ptr[i] is the start index in col_indices for row i; row_ptr[i+1] - row_ptr[i] gives the number of neighbors of atom i.

  • col_indices (torch.Tensor) – Shape (M,), dtype torch.long. Column indices of all non-zero entries, sorted first by row then by column within each row. M is the total number of non-zero entries after symmetrization and deduplication.

to(*args, **kwargs)[source]#
forward(coords: Tensor, box: Tensor, cutoff: float, max_npairs: int = -1, padding: bool = False)[source]#
torchff.nblist.build_neighbor_list_nsquared(coords: Tensor, box: Tensor, cutoff: float, max_npairs: int = -1, padding: bool = False, excl_row_ptr: Tensor | None = None, excl_col_indices: Tensor | None = None, include_self: bool = False, out: Tensor | None = None)[source]#

Build a neighbor list using the O(N^2) algorithm.

Parameters:

  • coords (torch.Tensor) – Atom positions, shape (natoms, 3).

  • box (torch.Tensor) – Periodic box vectors, shape (3, 3).

  • cutoff (float) – Distance cutoff.

  • max_npairs (int, optional) – Pre-allocated capacity for pairs. -1 allocates for the worst case natoms*(natoms-1)/2.

  • padding (bool, optional) – If True, return the full (possibly padded) pairs tensor instead of trimming to the actual count.

  • excl_row_ptr (torch.Tensor, optional) – CSR row-pointer tensor of shape (natoms + 1,) for the exclusion list. Must be torch.long.

  • excl_col_indices (torch.Tensor, optional) – CSR column-indices tensor for the exclusion list. Must be torch.long. Pairs (i, j) present in this sparse structure are skipped during neighbor-list construction. Use NeighborList.convert_pairs_coo_to_csr() to convert COO-format exclusion pairs into the required CSR tensors.

  • include_self (bool, optional) – If True, include self-pairs (i, i).

  • out (torch.Tensor, optional) – Pre-allocated output tensor of shape (max_npairs, 2) (dtype torch.long). When provided, the kernel writes into this tensor and max_npairs is ignored.

Returns:

  • pairs (torch.Tensor) – Neighbor pairs, shape (npairs, 2).

  • npairs (torch.Tensor) – Scalar tensor with the number of pairs found.

torchff.nblist.build_neighbor_list_cell_list(coords: Tensor, box: Tensor, cutoff: float, max_npairs: int = -1, cell_size: float = 0.4, padding: bool = False, shared: bool = False)[source]#
torchff.nblist.build_cluster_pairs(coords: Tensor, box: Tensor, cutoff: float, exclusions: Tensor | None = None, cell_size: float = 0.4, max_num_interacting_clusters: int = -1) → Tuple[Tensor, Tensor, Tensor, Tensor][source]#
torchff.nblist.decode_cluster_pairs(coords: Tensor, box: Tensor, sorted_atom_indices, cluster_exclusions, bitmask_exclusions, interacting_clusters, interacting_atoms, cutoff: float, max_npairs: int = -1, padding: bool = False) → Tuple[Tensor, Tensor][source]#
torchff.nblist.build_neighbor_list_cluster_pairs(coords: Tensor, box: Tensor, cutoff: float, exclusions: Tensor | None, cell_size: float = 0.45, max_num_interacting_clusters: int = -1, max_npairs: int = -1, padding: bool = False)[source]#

Nonbonded#

torchff.nonbonded.compute_nonbonded_energy_from_atom_pairs(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, charges: Tensor, coul_constant, cutoff: float, do_shift: bool = True) → Tensor[source]#

Compute fused Coulomb + Lennard-Jones nonbonded energies using custom CUDA/C++ ops.

torchff.nonbonded.compute_nonbonded_energy_from_atom_pairs_ref(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, charges: Tensor, coul_constant: float, cutoff: float, do_shift: bool = True) → Tensor[source]#

Reference fused Coulomb + Lennard-Jones implementation using native PyTorch ops.

torchff.nonbonded.compute_nonbonded_forces_from_atom_pairs(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, charges: Tensor, coul_constant: float, cutoff: float, forces: Tensor) → Tensor[source]#

Compute fused Coulomb + Lennard-Jones nonbonded forces in-place using custom CUDA/C++ ops.

class torchff.nonbonded.Nonbonded(use_customized_ops: bool = False)[source]#

Bases: Module

Fused fixed-charge nonbonded (Coulomb + Lennard-Jones) interaction.

__init__(use_customized_ops: bool = False)[source]#
forward(coords: Tensor, pairs: Tensor, box: Tensor, sigma: Tensor, epsilon: Tensor, charges: Tensor, coul_constant: float, cutoff: float, do_shift: bool = True) → Tensor[source]#

Multipoles#

torchff.multipoles.computeInteractionTensor(drVec: Tensor, dampFactors: Tensor | None = None, drInv: Tensor | None = None, rank: int = 2)[source]#
torchff.multipoles.computeDampFactorsErfc(dr: Tensor, b: float, rank: int)[source]#
class torchff.multipoles.MultipolePacker(rank: int = 2)[source]#

Bases: Module

Convert monopole, dipole, and quadrupole to (N, 10) polytensor with quadrupole entries scaled so symmetry-equivalent operations are avoided.

__init__(rank: int = 2)[source]#
forward(mono: Tensor, dipo: Tensor | None = None, quad: Tensor | None = None) → Tensor[source]#
Parameters:
Returns:

Polytensor [q, ux, uy, uz, Qxx, Qxy, Qxz, Qyy, Qyz, Qzz], shape (N, 10).

Return type:

torch.Tensor

class torchff.multipoles.MultipolarInteraction(rank: int, cutoff: float, ewald_alpha: float = -1.0, prefactor: float = 1.0, return_fields: bool = False, use_customized_ops: bool = True, cuda_graph_compat: bool = True)[source]#

Bases: Module

__init__(rank: int, cutoff: float, ewald_alpha: float = -1.0, prefactor: float = 1.0, return_fields: bool = False, use_customized_ops: bool = True, cuda_graph_compat: bool = True)[source]#
forward(coords, box, pairs, q, p=None, t=None, pairs_excl=None)[source]#
class torchff.multipoles.AxisTypes(*values)[source]#

Bases: IntEnum

ZThenX = 0#
Bisector = 1#
ZBisect = 2#
ThreeFold = 3#
ZOnly = 4#
NoAxisType = 5#
LastAxisTypeIndex = 6#
torchff.multipoles.normVec(vec: Tensor) → Tensor[source]#
class torchff.multipoles.MultipolarRotation(use_customized_ops: bool = False)[source]#

Bases: Module

Multipole local frame: build rotation matrices and rotate dipole / quadrupole tensors.

Vectorized local-to-global rotation matrices for multipole sites use the same batched construction as the historical TorchFF Python path: differences positions[neighbor] - positions[site] with no periodic minimum-image wrapping. Intramolecular (or otherwise local) geometry should already lie in one periodic image so that these raw vectors match the intended local frame.

Parameters:

use_customized_ops (bool) – If True, use _compute_rotation_matrices_torchff() (CUDA custom op); if False, use _compute_rotation_matrices_python() (Python reference). Dipole and quadrupole rotation use rotateDipoles() and rotateQuadrupoles() in either case.

Notes

Rotation matrices have shape (N, 3, 3); rows are local X, Y, Z in global coordinates (hstack of xVec, yVec, zVec). Inputs z_atoms, x_atoms, y_atoms, and axis_types have shape (N,) (see AxisTypes).

__init__(use_customized_ops: bool = False) → None[source]#
classmethod compute_matrices(coords: Tensor, z_atoms: Tensor, x_atoms: Tensor, y_atoms: Tensor, axis_types: Tensor, *, use_customized_ops: bool = False) → Tensor[source]#

Shape (N, 3, 3) rotation matrices (rows = local X, Y, Z in global coordinates).

classmethod rotate_dipoles(matrices: Tensor, dipoles: Tensor) → Tensor[source]#

Rotate dipoles of shape (N, 3); returns (N, 1, 3) (same as rotateDipoles()).

classmethod rotate_quadrupoles(matrices: Tensor, quadrupoles: Tensor) → Tensor[source]#

Rotate quadrupoles of shape (N, 3, 3).

forward(coords: Tensor, z_atoms: Tensor, x_atoms: Tensor, y_atoms: Tensor, axis_types: Tensor, dipoles: Tensor, quadrupoles: Tensor | None = None) → Tensor | Tuple[Tensor, Tensor][source]#

Apply local-to-global rotation to dipoles and optionally quadrupoles.

Returns:

Rotated dipoles (N, 3). If quadrupoles is given, returns (dipoles_rot, quadrupoles_rot).

Return type:

torch.Tensor or tuple

torchff.multipoles.scaleMultipoles(mPoles: Tensor, monoScales: Tensor, dipoScales: Tensor, quadScales: Tensor)[source]#
torchff.multipoles.rotateDipoles(dipo: Tensor, rotMatrix: Tensor)[source]#
torchff.multipoles.rotateQuadrupoles(quad: Tensor, rotMatrix: Tensor)[source]#
torchff.multipoles.rotateMultipoles(mono: Tensor, dipo: Tensor, quad: Tensor, rotMatrix: Tensor)[source]#

Rotate multipoles

Parameters:
Returns:

mPoles – Multipoles [q, ux, uy, uz, Qxx, Qxy, Qxz, Qyy, Qyz, Qzz], shape (N, 10)

Return type:

torch.Tensor

torchff.multipoles.convertMultipolesToPolytensor(mono: Tensor, dipo: Tensor, quad: Tensor)[source]#

Takes already-rotated multipoles and flattens to (N, 10) polytensor with quadrupole entries appropriately scaled so that symmetry-equivalent operations are avoided.

Parameters:
Returns:

mPoles – Multipoles [q, ux, uy, uz, Qxx, Qxy, Qxz, Qyy, Qyz, Qzz], shape (N, 10)

Return type:

torch.Tensor

torchff.multipoles.computeCartesianQuadrupoles(quad_s: Tensor)[source]#

Compute cartesian quadrupoles from spheric-harmonics quadrupoles

Parameters:

quad_s (torch.Tensor) – Quadrupoles in spherical harmonics form (Q20, Q21c, Q21s, Q22c, Q22s), shape (N, 5).

Returns:

quad – Quadrupoles in cartesian form, shape N x 3 x 3

Return type:

torch.Tensor

torchff.multipoles.computeSphericalQuadrupoles(quad_c: Tensor)[source]#

Compute cartesian quadrupoles from spheric-harmonics quadrupoles

Parameters:

quad_c (torch.Tensor) – Quadrupoles in cartesian form (Qxx, Qxy, Qxz, Qyy, Qyz, Qzz), shape (N, 6).

Returns:

quad_s – Quadrupoles in spherical harmonics form, shape (N, 5)

Return type:

torch.Tensor

Ewald Summation#

torchff.ewald.ewald_long_range(coords: Tensor, box: Tensor, q: Tensor, p: Tensor | None, t: Tensor | None, kmax: int, alpha: float)[source]#
torchff.ewald.ewald_long_range_with_fields(coords: Tensor, box: Tensor, q: Tensor, p: Tensor | None, t: Tensor | None, kmax: int, alpha: float)[source]#
class torchff.ewald.Ewald(alpha: float, kmax: int, rank: int, use_customized_ops: bool = True, return_fields: bool = True)[source]#

Bases: Module

__init__(alpha: float, kmax: int, rank: int, use_customized_ops: bool = True, return_fields: bool = True)[source]#
forward(coords, box, q, p=None, t=None)[source]#

Particle Mesh Ewald#

PME Reference#

B-Splines#

B-spline utilities for PME, matching the logic in csrc/pme/bsplines.cuh.

torchff.bsplines.get_bspline_coeff_order6(i: int) → float[source]#

B-spline coefficient for order-6 PME at index i.

Matches get_bspline_coeff_order6() in csrc/pme/bsplines.cuh.

torchff.bsplines.get_bspline_modulus(k: Tensor, K: int, order: int = 6) → Tensor[source]#

B-spline modulus for PME reciprocal-space scaling.

For each wave vector index k, computes sum_m b(m) * cos(2*pi*m*k/K) where m runs over the order-6 B-spline range and b(m) are the B-spline coefficients. Matches get_bspline_modulus_device() in csrc/pme/bsplines.cuh.

Parameters:

  • k (torch.Tensor) – 1D tensor of integer indices (e.g. 0 .. K-1 or 0 .. K//2).

  • K (int) – Grid size along this dimension.

  • order (int) – B-spline order; only 6 is implemented.

Returns:

Modulus values, same shape and device as k, dtype float.

Return type:

torch.Tensor

torchff.bsplines.compute_bspline_moduli_1d(K: int, dtype: dtype = torch.float32, device: device = None) → Tensor[source]#

Precompute the 1D B-spline modulus array for grid size K.

Returns moduli for indices 0, 1, …, K-1. For use as xmoduli or ymoduli in PME (length K).

Parameters:

  • K (int) – Grid size along this dimension.

  • dtype (torch.dtype) – Output tensor dtype.

  • device (torch.device, optional) – Output device; default is CPU.

Returns:

Shape (K,), dtype and device as requested.

Return type:

torch.Tensor

torchff.bsplines.compute_bspline_moduli_z(K3: int, dtype: dtype = torch.float32, device: device = None) → Tensor[source]#

Precompute the z-dimension B-spline modulus array for real FFT.

Returns moduli for z indices 0, 1, …, K3//2 (length K3//2 + 1), as used in the conjugate-even (rfft) reciprocal grid.

Parameters:

  • K3 (int) – Grid size along the z dimension.

  • dtype (torch.dtype) – Output tensor dtype.

  • device (torch.device, optional) – Output device; default is CPU.

Returns:

Shape (K3//2 + 1,), dtype and device as requested.

Return type:

torch.Tensor

Periodic Boundary Conditions#

class torchff.pbc.PBC(box: Tensor | None = None, box_inv: Tensor | None = None)[source]#

Bases: Module

Apply periodic boundary conditions to vectors.

Supports both fixed box (provided at construction) and variable box (passed to forward()). If neither box nor box_inv is available, returns the input vectors unchanged.

__init__(box: Tensor | None = None, box_inv: Tensor | None = None)[source]#
forward(dr_vecs: Tensor, box: Tensor | None = None, box_inv: Tensor | None = None) → Tensor[source]#

Apply PBC to displacement vectors.

Parameters:

  • dr_vecs (torch.Tensor) – Real space vectors in Cartesian, shape (N, 3).

  • box (torch.Tensor, optional) – Simulation box with axes in rows, shape (3, 3). Overrides the buffer if provided.

  • box_inv (torch.Tensor, optional) – Inverse of the box matrix, shape (3, 3). Overrides the buffer if provided.

Returns:

PBC-folded vectors in Cartesian, shape (N, 3).

Return type:

torch.Tensor

torchff.pbc.applyPBC(dr_vecs: Tensor, box: Tensor | None = None, box_inv: Tensor | None = None) → Tensor[source]#

Apply periodic boundary conditions to a set of vectors.

Functional interface to PBC. Prefer using PBC when integrating with nn.Module-based models.

Parameters:

  • dr_vecs (torch.Tensor) – Real space vectors in Cartesian, shape (N, 3).

  • box (torch.Tensor, optional) – Simulation box with axes in rows, shape (3, 3).

  • box_inv (torch.Tensor, optional) – Inverse of the box matrix, shape (3, 3).

Returns:

PBC-folded vectors in Cartesian, shape (N, 3).

Return type:

torch.Tensor

Contents