torchref.base.targets.ramachandran module

Ramachandran restraint via bilinear interpolation on a precomputed NLL surface.

torchref.base.targets.ramachandran.ramachandran_math(xyz, phi_idx, psi_idx, nll_surfaces, surface_type)[source]

Ramachandran bilinear-interpolated NLL.

Mirrors RamachandranTarget.forward.

Dispatches to torchref.base.targets.triton.ramachandran_math_triton() on CUDA float32 (~10× faster fwd+bw on A100). Falls back to eager otherwise.

Parameters:

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

  • phi_idx (torch.Tensor) – (N, 4) atom indices for the two backbone dihedrals.

  • psi_idx (torch.Tensor) – (N, 4) atom indices for the two backbone dihedrals.

  • nll_surfaces (torch.Tensor) – (n_surface_types, 360, 360) precomputed NLL = -log P(φ, ψ | type).

  • surface_type (torch.Tensor) – (N,) integer type per residue.