Relaxation API

relaxation

Functions

spectral_density(omega, tau_m, s2, tau_f=0.0)

Calculate Spectral Density J(w) using Lipari-Szabo Model-Free formalism.

EDUCATIONAL NOTE - BPP Theory & Spectral Density:

In NMR, relaxation is caused by local magnetic fields that "flicker" due to molecular tumbling. Spectral Density J(w) is essentially a measure of the POWER of these fluctuations at a specific frequency (w).

• BPP Theory (Bloembergen-Purcell-Pound) shows that relaxation is most efficient when the frequency of motion (1/tau_m) matches the Larmor frequency of the nuclei.
• If a protein tumbles too fast (small tau_m), the fluctuations are too high-frequency to cause efficient relaxation (Extreme Narrowing Limit).
• If it tumbles too slow (large tau_m), the energy is concentrated at low frequencies, leading to fast T2 relaxation and broad lines.

Model-Free Analysis (Lipari-Szabo) separates this into Global Tumbling (tau_m) and fast Local Motion (tau_f), weighted by the Order Parameter (S^2).

Formula: J(w) = (2/5) * [ S^2 * tm / (1 + (wtm)^2) + (1-S^2) * te / (1 + (wte)^2) ]

where te (tau_e) is the effective internal correlation time: 1/te = 1/tm + 1/tf Usually for simple MF, we assume fast motion tf << tm.

Parameters:

Name	Type	Description	Default
omega	float	Frequency (rad/s). Must be numeric.	required
tau_m	float	Global rotational correlation time (seconds). Must be numeric and positive.	required
s2	float	Generalized order parameter (0.0 to 1.0). Must be numeric and within this range.	required
tau_f	float	Fast internal correlation time (seconds). Must be numeric and non-negative. If 0 (default), internal motion is considered infinitely fast/negligible, and the function simplifies to a one-time scale model-free.	0.0

Source code in synth_nmr/relaxation.py

@njit
def spectral_density(omega: float, tau_m: float, s2: float, tau_f: float = 0.0) -> float:
    """
    Calculate Spectral Density J(w) using Lipari-Szabo Model-Free formalism.

    EDUCATIONAL NOTE - BPP Theory & Spectral Density:
    -------------------------------------------------
    In NMR, relaxation is caused by local magnetic fields that "flicker" due to
    molecular tumbling. Spectral Density J(w) is essentially a measure of the
    POWER of these fluctuations at a specific frequency (w).

    - BPP Theory (Bloembergen-Purcell-Pound) shows that relaxation is most
      efficient when the frequency of motion (1/tau_m) matches the Larmor
      frequency of the nuclei.
    - If a protein tumbles too fast (small tau_m), the fluctuations are too
      high-frequency to cause efficient relaxation (Extreme Narrowing Limit).
    - If it tumbles too slow (large tau_m), the energy is concentrated at
      low frequencies, leading to fast T2 relaxation and broad lines.

    Model-Free Analysis (Lipari-Szabo) separates this into Global Tumbling (tau_m)
    and fast Local Motion (tau_f), weighted by the Order Parameter (S^2).

    Formula:
    J(w) = (2/5) * [ S^2 * tm / (1 + (w*tm)^2) + (1-S^2) * te / (1 + (w*te)^2) ]

    where te (tau_e) is the effective internal correlation time: 1/te = 1/tm + 1/tf
    Usually for simple MF, we assume fast motion tf << tm.

    Args:
        omega: Frequency (rad/s). Must be numeric.
        tau_m: Global rotational correlation time (seconds). Must be numeric and positive.
        s2: Generalized order parameter (0.0 to 1.0). Must be numeric and within this range.
        tau_f: Fast internal correlation time (seconds). Must be numeric and non-negative.
               If 0 (default), internal motion is considered infinitely fast/negligible,
               and the function simplifies to a one-time scale model-free.
    """
    # Simple Model Free (assuming tf is very small/negligible or incorporated)
    # If tau_f is provided, calculate effective time tau_e

    # Term 1: Global tumbling
    j_global = (s2 * tau_m) / (1 + (omega * tau_m) ** 2)

    # Term 2: Fast internal motion
    # Effective correlation time 1/tau_e = 1/tau_m + 1/tau_f
    # If tau_f is 0, this term vanishes in standard simplified approximation
    # or acts as a very fast motion limit.
    j_fast = 0.0
    if tau_f > 0:
        tau_e = (tau_m * tau_f) / (tau_m + tau_f)
        j_fast = ((1 - s2) * tau_e) / (1 + (omega * tau_e) ** 2)

    return 0.4 * (j_global + j_fast)

predict_order_parameters(structure)

Predict Generalized Order Parameters (S2) based on secondary structure, termini effects, and solvent accessible surface area (SASA).

EDUCATIONAL NOTE - Lipari-Szabo Model Free:

The Order Parameter (S2) describes the amplitude of internal motion: - S2 = 1.0: Completely rigid (no internal motion relative to tumbling). - S2 = 0.0: Completely disordered (isotropic internal motion).

Typical values in proteins: - Alpha Helices / Beta Sheets: S2 ~ 0.85 (Very rigid H-bond network) - Loops / Turns: S2 ~ 0.60 - 0.70 (Flexible) - Termini (N/C): S2 ~ 0.40 - 0.50 (Fraying)

Parameters:

Name	Type	Description	Default
structure	AtomArray	The biotite.structure.AtomArray containing the protein.	required

Returns:

Type	Description
Dict[int, float]	A dictionary mapping each residue ID to its predicted S2 value.

Raises:

Type	Description
TypeError	If the input is not a biotite.structure.AtomArray.

Source code in synth_nmr/relaxation.py

def predict_order_parameters(structure: struc.AtomArray) -> Dict[int, float]:
    """
    Predict Generalized Order Parameters (S2) based on secondary structure,
    termini effects, and solvent accessible surface area (SASA).

    EDUCATIONAL NOTE - Lipari-Szabo Model Free:
    ===========================================
    The Order Parameter (S2) describes the amplitude of internal motion:
    - S2 = 1.0: Completely rigid (no internal motion relative to tumbling).
    - S2 = 0.0: Completely disordered (isotropic internal motion).

    Typical values in proteins:
    - Alpha Helices / Beta Sheets: S2 ~ 0.85 (Very rigid H-bond network)
    - Loops / Turns: S2 ~ 0.60 - 0.70 (Flexible)
    - Termini (N/C): S2 ~ 0.40 - 0.50 (Fraying)

    Args:
        structure: The biotite.structure.AtomArray containing the protein.

    Returns:
        A dictionary mapping each residue ID to its predicted S2 value.

    Raises:
        TypeError: If the input is not a biotite.structure.AtomArray.
    """
    from synth_nmr.structure_utils import get_residue_info

    logger.info("Predicting Generalized Order Parameters (S2)...")

    if not isinstance(structure, struc.AtomArray):
        raise TypeError("Input 'structure' must be a biotite.structure.AtomArray.")
    if structure.array_length() == 0:
        return {}

    try:
        # ── Physical Basis of Order Parameters (S2) ─────────────────────────
        # In the Lipari-Szabo Model-Free formalism, the spectral density is
        # decomposed into global tumbling (tau_m) and internal motions.
        # S2 is the spatial part of the internal correlation function.
        #
        # A value of 0.85 indicates that the N-H bond vector is highly
        # constrained, sampling only a small portion of the unit sphere.
        # This is typical for well-defined secondary structural elements
        # like alpha-helices and beta-sheets, where hydrogen bonding
        # stabilizes the backbone.
        # ─────────────────────────────────────────────────────────────────────

        # Check unique residues for test compatibility (some tests mock np.unique)
        if len(np.unique(structure.res_id)) == 0:
            return {}

        chain_ids, res_ids, res_names, res_starts = get_residue_info(structure)
        if len(res_ids) == 0:
            return {}

        # ── Correlation Times and Molecular Motion ──────────────────────────
        # Protein relaxation is driven by the rotational diffusion (tumbling).
        # We model this using the Global Correlation Time (tau_m).
        # However, residues also have internal motion (fast fluctuations).
        # We use a heuristic 'internal' correlation time (tau_f) which is
        # inversely proportional to the Order Parameter (S2).
        #
        # - Rigid residues (High S2): Slower internal motion, dominant global tumbling.
        # - Flexible residues (Low S2): Significant fast internal motion (ps-ns range).
        # ─────────────────────────────────────────────────────────────────────

        ss_list = get_secondary_structure(structure)

        # Heuristic Max SASA per residue (Angstrom^2) for normalization
        MAX_SASA = 150.0

        # Calculate SASA for "Packing Awareness"
        sasa_per_residue_array = np.full(len(res_ids), MAX_SASA)
        try:
            temp_struc = structure.copy()
            temp_struc.res_name[np.isin(temp_struc.res_name, ["HIE", "HID", "HIP"])] = "HIS"
            temp_struc.res_name[temp_struc.res_name == "SEP"] = "SER"
            temp_struc.res_name[temp_struc.res_name == "TPO"] = "THR"
            temp_struc.res_name[temp_struc.res_name == "PTR"] = "TYR"

            ptm_atom_names = ["P", "O1P", "O2P", "O3P"]
            ptm_mask = np.isin(temp_struc.atom_name, ptm_atom_names)
            if np.any(ptm_mask):
                temp_struc = temp_struc[~ptm_mask]

            ion_res_names = ["ZN", "MG", "CA", "NA", "CL", "K", "FE", "CU", "MN"]
            ion_mask = np.isin(temp_struc.res_name, ion_res_names)
            if np.any(ion_mask):
                temp_struc = temp_struc[~ion_mask]

            atom_sasa = struc.sasa(temp_struc, probe_radius=1.4)
            atom_sasa = np.nan_to_num(atom_sasa, nan=50.0)

            # Vectorized aggregation of SASA per residue
            temp_res_starts = struc.get_residue_starts(temp_struc)
            for i in range(len(temp_res_starts)):
                start = temp_res_starts[i]
                end = temp_res_starts[i + 1] if i + 1 < len(temp_res_starts) else len(temp_struc)
                rid = temp_struc.res_id[start]
                idx = np.where(res_ids == rid)[0]
                if len(idx) > 0:
                    sasa_per_residue_array[idx[0]] = np.sum(atom_sasa[start:end])

        except Exception as e:
            logger.warning(f"SASA calculation failed ({e}). Using default exposed values.")

        start_res = res_ids[0]
        end_res = res_ids[-1]

        results = {}
        for i in range(len(res_ids)):
            rid = int(res_ids[i])
            ss = ss_list[i] if i < len(ss_list) else "coil"
            res_sasa = sasa_per_residue_array[i]

            # Relative SASA (0.0 = Buried, 1.0 = Exposed)
            rel_sasa = min(res_sasa / MAX_SASA, 1.0)

            # Base S2 from Secondary Structure
            base_s2 = 0.85 if ss in ["alpha", "beta"] else 0.70

            # Termini effects (Must call helper for test mocking)
            base_s2 = _apply_termini_effects(rid, start_res, end_res, base_s2)

            # Modulate by SASA (Must call helper for test mocking)
            s2 = _predict_s2_from_sasa(rel_sasa, base_s2)

            # Add realistic noise
            s2 += np.random.normal(0, 0.02)
            s2 = np.clip(s2, 0.01, 0.98)
            results[rid] = float(s2)

        logger.info(f"Successfully predicted S2 for {len(results)} residues.")
        return results

    except Exception:
        logger.error("An unexpected error occurred during S2 prediction", exc_info=True)
        raise

calculate_relaxation_rates(structure, field_mhz=600.0, tau_m_ns=10.0, s2_map=None)

Calculate R1, R2, and Heteronuclear NOE for all backbone Amides (N-H).

Parameters:

Name	Type	Description	Default
structure	AtomArray	biotite.structure.AtomArray containing the protein.	required
field_mhz	float	Proton Larmor frequency in MHz (e.g., 600.0).	600.0
tau_m_ns	float	Global rotational correlation time in nanoseconds.	10.0
s2_map	Optional[Dict[int, float]]	Optional dictionary of pre-calculated Order Parameters (S2).	None

Returns:

Type	Description
Dict[int, Dict[str, float]]	Dict keyed by Residue ID -> {R1, R2, NOE, S2}.

Source code in synth_nmr/relaxation.py

def calculate_relaxation_rates(
    structure: struc.AtomArray,
    field_mhz: float = 600.0,
    tau_m_ns: float = 10.0,
    s2_map: Optional[Dict[int, float]] = None,
) -> Dict[int, Dict[str, float]]:
    """
    Calculate R1, R2, and Heteronuclear NOE for all backbone Amides (N-H).

    Args:
        structure: biotite.structure.AtomArray containing the protein.
        field_mhz: Proton Larmor frequency in MHz (e.g., 600.0).
        tau_m_ns: Global rotational correlation time in nanoseconds.
        s2_map: Optional dictionary of pre-calculated Order Parameters (S2).

    Returns:
        Dict keyed by Residue ID -> {R1, R2, NOE, S2}.
    """
    from synth_nmr.structure_utils import get_residue_info

    logger.info(f"Starting Relaxation Rates calculation (Field={field_mhz}MHz, tm={tau_m_ns}ns)...")

    if not isinstance(structure, struc.AtomArray):
        raise TypeError("Input 'structure' must be a biotite.structure.AtomArray.")
    if structure.array_length() == 0:
        logger.warning("Input 'structure' is empty. Returning no relaxation rates.")
        return {}
    if not isinstance(field_mhz, (int, float)) or field_mhz <= 0:
        raise ValueError("Parameter 'field_mhz' must be a positive numeric value.")
    if not isinstance(tau_m_ns, (int, float)) or tau_m_ns <= 0:
        raise ValueError("Parameter 'tau_m_ns' must be a positive numeric value.")
    if s2_map is not None and not isinstance(s2_map, dict):
        raise TypeError("Parameter 's2_map' must be a dictionary or None.")

    # Calculate S2 profile if not provided
    if s2_map is None:
        s2_map = predict_order_parameters(structure)

    # Physical constants and frequencies
    tau_m = tau_m_ns * 1e-9
    omega_h = 2 * np.pi * field_mhz * 1e6
    b0 = omega_h / GAMMA_H
    omega_n = GAMMA_N * b0
    d_sq = _calculate_dipolar_constant(R_NH)
    c_sq = _calculate_csa_constant(CSA_N, omega_n)

    # Get residue info
    _, res_ids, res_names, _ = get_residue_info(structure)

    # Map S2 to array matching res_ids
    s2_arr = np.array([s2_map.get(int(rid), 0.85) for rid in res_ids])

    # Find residues with both N and H atoms (backbone amides)
    n_mask = (structure.atom_name == "N") & struc.filter_amino_acids(structure)
    h_mask = (structure.atom_name == "H") & struc.filter_amino_acids(structure)

    n_res_ids = structure.res_id[n_mask]
    h_res_ids = structure.res_id[h_mask]

    has_n = np.isin(res_ids, n_res_ids)
    has_h = np.isin(res_ids, h_res_ids)

    # Valid amides: have N, H and NOT Proline
    is_pro = res_names == "PRO"
    valid_mask = has_n & has_h & (~is_pro)

    # Filter arrays to valid amides only
    active_res_ids = res_ids[valid_mask]
    active_s2 = s2_arr[valid_mask]

    # Heuristic tau_f
    tau_f_vals = (1.0 - active_s2) * 500e-12 + 50e-12

    results = {}
    for i in range(len(active_res_ids)):
        rid = active_res_ids[i]
        s2 = active_s2[i]
        tau_f = tau_f_vals[i]

        j0 = spectral_density(0, tau_m, s2, tau_f)
        jwn = spectral_density(np.abs(omega_n), tau_m, s2, tau_f)
        jwh = spectral_density(omega_h, tau_m, s2, tau_f)
        # Since omega_n is negative, omega_h - omega_n is the SUM of magnitudes
        # and omega_h + omega_n is the DIFFERENCE of magnitudes.
        j_sum = spectral_density(omega_h - omega_n, tau_m, s2, tau_f)
        j_diff = spectral_density(omega_h + omega_n, tau_m, s2, tau_f)

        r1 = 0.25 * d_sq * (j_diff + 3 * jwn + 6 * j_sum) + c_sq * jwn
        r2 = 0.125 * d_sq * (4 * j0 + j_diff + 3 * jwn + 6 * jwh + 6 * j_sum) + (
            1.0 / 6.0
        ) * c_sq * (4 * j0 + 3 * jwn)

        if r1 != 0:
            noe = 1.0 + (GAMMA_H / GAMMA_N) * 0.25 * d_sq * (6 * j_sum - j_diff) * (1.0 / r1)
        else:
            noe = np.nan
            logger.warning(f"R1 value for residue {rid} is zero, NOE cannot be calculated.")

        results[int(rid)] = {
            "R1": float(r1),
            "R2": float(r2),
            "NOE": float(noe),
            "S2": float(s2),
        }

    logger.info(f"Successfully calculated relaxation rates for {len(results)} residues.")
    return results