torchref.base.targets.triton.torsion module

Triton forward + analytic backward for the omega cis/trans torsion target.

torchref.base.targets.triton.torsion.torsion_omega_math_triton(xyz, idx, sigmas_deg, is_proline, w_cis_proline=0.05, w_cis_general=0.0005)[source]

Triton-backed omega cis/trans mixture NLL with analytic backward.

torchref.base.targets.triton.torsion.torsion_unimodal_full_math_triton(xyz, idx, references_deg, sigmas_deg, periods)[source]

Triton-backed full unimodal torsion NLL (dihedral + periodic wrap + von Mises NLL + sum), with analytic backward through the dihedral formula. Inputs match Restraints.restraints['torsion']['all']:

Parameters:

  • xyz ((N_atoms, 3) float32 CUDA)

  • idx ((N, 4) int64 atom indices)

  • references_deg ((N,) float32 — target angles in degrees)

  • sigmas_deg ((N,) float32 — sigmas in degrees)

  • periods ((N,) int — n-fold periodicity (≥ 1))