Skip to content

Utils Module

get_residue_rc_shifts(res_indices)

Map residue indices to their random coil CA shifts.

Source code in torsiontuner/montelione_utils.py 
11
12
13
14
def get_residue_rc_shifts(res_indices: jnp.ndarray) -> jnp.ndarray:
    """Map residue indices to their random coil CA shifts."""
    # Use JAX indexing instead of list comprehension
    return _RC_ARRAY[res_indices]

montelione_loss(phi, psi, target_shifts, res_indices)

A loss function inspired by the Montelione group's approach of validating structures against experimental chemical shifts.

Source code in torsiontuner/montelione_utils.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def montelione_loss(
    phi: jnp.ndarray,
    psi: jnp.ndarray,
    target_shifts: jnp.ndarray,
    res_indices: jnp.ndarray,
) -> jnp.ndarray:
    """
    A loss function inspired by the Montelione group's approach of
    validating structures against experimental chemical shifts.
    """
    rc_shifts = get_residue_rc_shifts(res_indices)
    pred_shifts = predict_ca_shifts(phi, psi, rc_shifts)

    # MSE between predicted and experimental CA shifts
    return jnp.mean((pred_shifts - target_shifts) ** 2)

ramachandran_penalty(phi, psi)

A soft Ramachandran-region penalty used as a backbone geometry regularizer.

Penalizes phi/psi values that fall far from the three canonical favored regions of the Ramachandran plot (alpha-helix, beta-strand, and left-handed alpha-helix), using a smooth log-sum-exp potential.

This is analogous to the Ramachandran term used in programs such as CNS and Rosetta (e.g., Engh & Huber torsion potentials), and is similar in spirit to the G-factor computed by PROCHECK.

Note: This is NOT related to the ANSURR validation method (Fowler, Sljoka & Williamson, 2020, Nature Commun. 11:6321). ANSURR measures agreement between experimental NMR backbone flexibility (RCI) and structural rigidity computed by graph-theoretic rigidity theory (FIRST algorithm); it cannot be computed from torsion angles alone.

Parameters:

Name Type Description Default
phi ndarray

(N,) backbone phi angles in radians.

 required
psi ndarray

(N,) backbone psi angles in radians.

 required

Returns:

Type Description
ndarray

Scalar penalty (lower = more geometrically regular backbone).
Source code in torsiontuner/montelione_utils.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def ramachandran_penalty(phi: jnp.ndarray, psi: jnp.ndarray) -> jnp.ndarray:
    """
    A soft Ramachandran-region penalty used as a backbone geometry regularizer.

    Penalizes phi/psi values that fall far from the three canonical
    favored regions of the Ramachandran plot (alpha-helix, beta-strand,
    and left-handed alpha-helix), using a smooth log-sum-exp potential.

    This is analogous to the Ramachandran term used in programs such as
    CNS and Rosetta (e.g., Engh & Huber torsion potentials), and is
    similar in spirit to the G-factor computed by PROCHECK.

    Note: This is NOT related to the ANSURR validation method (Fowler,
    Sljoka & Williamson, 2020, Nature Commun. 11:6321). ANSURR measures
    agreement between experimental NMR backbone flexibility (RCI) and
    structural rigidity computed by graph-theoretic rigidity theory
    (FIRST algorithm); it cannot be computed from torsion angles alone.

    Args:
        phi: (N,) backbone phi angles in radians.
        psi: (N,) backbone psi angles in radians.

    Returns:
        Scalar penalty (lower = more geometrically regular backbone).
    """
    # Alpha-helix region: phi ~ -60 deg (-1.05 rad), psi ~ -45 deg (-0.78 rad)
    alpha_dist = jnp.sqrt((phi + 1.05) ** 2 + (psi + 0.78) ** 2)
    # Beta-strand region: phi ~ -120 deg (-2.09 rad), psi ~ +135 deg (+2.35 rad)
    beta_dist = jnp.sqrt((phi + 2.09) ** 2 + (psi - 2.35) ** 2)
    # Left-handed alpha-helix: phi ~ +60 deg (+1.05 rad), psi ~ +45 deg (+0.78 rad)
    l_alpha_dist = jnp.sqrt((phi - 1.05) ** 2 + (psi - 0.78) ** 2)

    # Soft minimum distance to any favored region via log-sum-exp
    sigma = 0.5
    dist_sq = jnp.stack([alpha_dist, beta_dist, l_alpha_dist], axis=-1) ** 2
    penalty = -sigma * jax.nn.logsumexp(-dist_sq / sigma, axis=-1)

    return jnp.mean(penalty)