torchref.base.targets.adp module

ADP (B-factor) restraint NLLs: similarity, KL-divergence, locality.

torchref.base.targets.adp.adp_simu_math(b, pair_indices, simu_sigma)[source]

ADP similarity (SIMU) NLL on bonded-atom B-factor differences.

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

Parameters:

  • b (torch.Tensor) – (N_atoms,) B-factors.

  • pair_indices (torch.Tensor) – (N, 2) bonded-atom pairs to compare.

  • simu_sigma (torch.Tensor) – Scalar sigma on the difference (a buffer in the target).

torchref.base.targets.adp.adp_kl_math(log_adp, target_log_std=0.2)[source]

KL divergence regularizer on log(B).

Mirrors Model.adp_kl_divergence_loss: KL between an empirical Gaussian (mean fixed to current mean(log_adp) detached, std = std(log_adp)) and a target Gaussian with the same mean but fixed std.

torchref.base.targets.adp.adp_locality_math(b, neighbor_indices, neighbor_distances)[source]

ADP locality NLL: weighted MSE on log(B) differences with KNN.

Mirrors ADPLocalityTarget.forward. Neighbor list construction is the target’s bookkeeping and is not included here.