torchref.base.targets.triton package

Triton kernels for geometry-target math.

Each module in this package provides a Triton-backed implementation of the corresponding eager-mode function in torchref.base.targets.

Naming convention: <target>_math_triton. The eager versions remain in the parent package as reference implementations and equivalence baselines.

torchref.base.targets.triton.adp_simu_math_triton(b, pair_indices, simu_sigma)

Triton-backed ADP similarity (SIMU) Gaussian NLL.

Drop-in replacement for torchref.base.targets.adp.adp_simu_math().

torchref.base.targets.triton.angle_math_triton(xyz, idx, references_rad, sigmas_rad)

Triton-backed bond-angle Gaussian NLL with analytic backward.

Drop-in replacement for torchref.base.targets.angle.angle_math().

torchref.base.targets.triton.bhattacharyya_xray_loss_math_triton(F_obs, F_calc, sigma_d, sigma_m, mask)

Triton-backed Bhattacharyya overlap loss.

Drop-in replacement for torchref.base.targets.xray_bhattacharyya.bhattacharyya_xray_loss_math().

torchref.base.targets.triton.bond_math_triton(xyz, idx, references, sigmas)

Triton-backed bond-length Gaussian NLL.

Drop-in replacement for torchref.base.targets.bond.bond_math().

torchref.base.targets.triton.chiral_math_triton(xyz, indices, ideal_volumes, sigmas)

Triton-backed chiral-volume Gaussian NLL with analytic backward.

torchref.base.targets.triton.gaussian_xray_loss_math_triton(F_obs, F_calc, sigma, mask)

Triton-backed Gaussian NLL X-ray loss.

torchref.base.targets.triton.ls_xray_loss_math_triton(F_obs, F_calc, sigma, mask, weighting='sigma')

Triton-backed least-squares X-ray loss.

torchref.base.targets.triton.ml_xray_loss_math_triton(F_obs, F_calc, sigma, centric_flags, mask)

Triton-backed maximum-likelihood X-ray loss.

Drop-in replacement for torchref.base.targets.xray_ml.ml_xray_loss_math().

torchref.base.targets.triton.nonbonded_heavy_math_triton(xyz, indices, min_distances, symop_indices, cell_offsets, symop_matrices, symop_translations, fractional_matrix, inv_fractional_matrix, c_rep, r_exp, buffer_, sigma_vdw)

Triton-backed heavy-heavy VDW prolsq NLL with analytic backward.

torchref.base.targets.triton.planarity_math_triton(xyz, plane_groups)

Triton-backed planarity NLL with analytic backward.

Drop-in replacement for torchref.base.targets.planarity.planarity_math(). SVD-derived plane normals are computed on the host (detached, same as eager); the gather + project + NLL + sum and the gradient scatter run in Triton.

torchref.base.targets.triton.ramachandran_math_triton(xyz, phi_idx, psi_idx, nll_surfaces, surface_type)

Triton-backed Ramachandran bilinear-interp NLL with analytic backward.

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

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

Submodules

• torchref.base.targets.triton.adp_simu module
  • adp_simu_math_triton()
• torchref.base.targets.triton.angle module
  • angle_math_triton()
• torchref.base.targets.triton.bond module
  • bond_math_triton()
• torchref.base.targets.triton.chiral module
  • chiral_math_triton()
• torchref.base.targets.triton.nonbonded module
  • nonbonded_heavy_math_triton()
• torchref.base.targets.triton.place_hydrogens module
  • place_riding_hydrogens_triton()
• torchref.base.targets.triton.planarity module
  • planarity_math_triton()
• torchref.base.targets.triton.ramachandran module
  • ramachandran_math_triton()
• torchref.base.targets.triton.torsion module
  • torsion_omega_math_triton()
  • torsion_unimodal_full_math_triton()
• torchref.base.targets.triton.xray_bhattacharyya module
  • bhattacharyya_xray_loss_math_triton()
• torchref.base.targets.triton.xray_gaussian module
  • gaussian_xray_loss_math_triton()
• torchref.base.targets.triton.xray_ls module
  • ls_xray_loss_math_triton()
• torchref.base.targets.triton.xray_ml module
  • ml_xray_loss_math_triton()