torchref.base.targets.nonbonded module

Non-bonded (VDW) heavy-heavy repulsion NLL — prolsq mode with symmetry mates.

torchref.base.targets.nonbonded.nonbonded_heavy_math(xyz, indices, min_distances, symop_indices, cell_offsets, symop_matrices, symop_translations, fractional_matrix, inv_fractional_matrix, c_rep, r_exp, buffer, sigma_vdw)[source]

Heavy-heavy VDW prolsq repulsion NLL.

Matches the prolsq branch of NonBondedTarget.forward and the symmetry-aware pair-position gather from NonBondedTarget._compute_positions. The H-VDW contribution added by NonBondedHTarget is excluded — see nonbonded_h (TBD) for that.

Dispatches to torchref.base.targets.triton.nonbonded_heavy_math_triton() on CUDA float32 (~1.4× faster fwd+bw on A100 — forward kernel is memory-bound but the analytic backward saves most of the win). Falls back to eager otherwise.

Parameters:

  • xyz (torch.Tensor) – (N_atoms, 3) Cartesian coordinates of the ASU.

  • indices (torch.Tensor) – (N, 2) per-pair atom indices.

  • min_distances (torch.Tensor) – (N,) VDW threshold per pair.

  • symop_indices (torch.Tensor, optional) – (N,) symmetry-operator index per pair; 0 = identity.

  • cell_offsets (torch.Tensor, optional) – (N, 3) fractional cell offsets per pair.

  • symop_matrices (torch.Tensor, optional) – (n_symops, 3, 3) and (n_symops, 3) — the symmetry operator table.

  • symop_translations (torch.Tensor, optional) – (n_symops, 3, 3) and (n_symops, 3) — the symmetry operator table.

  • fractional_matrix (torch.Tensor) – cell.fractional_matrix and its inverse (3, 3).

  • inv_fractional_matrix (torch.Tensor) – cell.fractional_matrix and its inverse (3, 3).

  • c_rep (torch.Tensor) – Scalar repulsion coefficient, exponent, and effective tolerance.

  • r_exp (torch.Tensor) – Scalar repulsion coefficient, exponent, and effective tolerance.

  • sigma_vdw (torch.Tensor) – Scalar repulsion coefficient, exponent, and effective tolerance.

  • buffer (float) – Distance buffer in Å.