NMR Analysis Plans

from warp_md import NmrIredPlan

plan = NmrIredPlan(
    vector_pairs,
    order=2,
    length_scale=0.1,
    pbc="none",
    corr_mode="tensor",
    return_corr=True
)

import warp_md as wmd
import numpy as np

system = wmd.System("protein.pdb")
traj = wmd.open_trajectory("md.nc", system)

# Define N-H bond vectors
atom_table = system.atom_table()
nh_pairs = []

for i, atom in enumerate(atom_table["name"]):
    if atom == "N":
        resid = atom_table["resid"][i]
        # Find corresponding H
        for j, h_atom in enumerate(atom_table["name"]):
            if h_atom == "H" and atom_table["resid"][j] == resid:
                nh_pairs.append((i, j))
                break

print(f"Found {len(nh_pairs)} N-H vectors")

# Calculate iRED vectors and correlation matrix
plan = wmd.NmrIredPlan(
    nh_pairs,
    order=2,
    length_scale=0.1,
    corr_mode="tensor",
    return_corr=True
)

vectors, corr_matrix = plan.run(traj, system)

# Calculate order parameters S2
n_frames = vectors.shape[0]
m = np.einsum('fpi,fpj->pij', vectors, vectors) / n_frames
tr_m2 = np.einsum('pij,pij->p', m, m)
s2 = 1.5 * tr_m2 - 0.5

print("Order parameters S²:")
for i, (n_idx, h_idx) in enumerate(nh_pairs[:10]):
    resid = atom_table["resid"][n_idx]
    print(f"  Residue {resid}: S² = {s2[i]:.3f}")

from warp_md.analysis import nh_order_parameters

# Automatic detection using selection
s2 = nh_order_parameters(
    traj,
    system,
    vector_pairs=None,
    selection="protein and (name N or name H)",
    order=2,
    method="tensor"
)

print(f"Order parameters shape: {s2.shape}")
print(f"Mean S²: {s2.mean():.3f}")

import matplotlib.pyplot as plt
import numpy as np
from warp_md.analysis import nh_order_parameters

# Calculate order parameters
s2 = nh_order_parameters(
    traj,
    system,
    selection="protein and backbone",
    method="tensor"
)

# Get residue numbers
atom_table = system.atom_table()
n_indices = [i for i, name in enumerate(atom_table["name"]) if name == "N"]
resids = [atom_table["resid"][i] for i in n_indices[:len(s2)]]

# Plot
plt.figure(figsize=(12, 4))
plt.plot(resids, s2, 'o-')
plt.axhline(0.8, color='r', linestyle='--', alpha=0.5, label='Rigid')
plt.axhline(0.6, color='orange', linestyle='--', alpha=0.5, label='Flexible')
plt.xlabel("Residue Number")
plt.ylabel("Order Parameter S²")
plt.ylim(0, 1)
plt.legend()
plt.title("Backbone Dynamics")
plt.grid(alpha=0.3)
plt.show()

# Identify flexible regions
flexible_residues = [resids[i] for i, val in enumerate(s2) if val < 0.6]
print(f"Flexible regions: {flexible_residues}")

from warp_md.analysis import jcoupling

# Define dihedral angles (4 atoms each)
dihedrals = [
    (10, 11, 12, 13),  # Example: H-N-C-H dihedral
    (20, 21, 22, 23),
]

j_values = jcoupling(
    traj,
    system,
    dihedral_indices=dihedrals,
    karplus=(6.4, -1.4, 1.9),  # A, B, C parameters
    phase_deg=0.0
)

print(f"J-coupling shape: {j_values.shape}")
print(f"Average J-coupling: {j_values.mean(axis=0)} Hz")

# ³J(Hᴺ-Hᴿ) for proteins (Vuister & Bax)
karplus_hn_ha = (7.09, -1.42, 1.55)

# ³J(H-H) for nucleic acids
karplus_nucleic = (9.5, -1.1, 0.9)

# ³J(C-C) for carbohydrates
karplus_carb = (3.5, -0.9, 0.3)

import numpy as np
import matplotlib.pyplot as plt
from warp_md.analysis import jcoupling

# Experimental J-coupling data
exp_j_couplings = {
    5: 8.2,   # Residue: J value (Hz)
    6: 7.5,
    7: 9.1,
    # ...
}

# Calculate from simulation
dihedrals = []
residues = []
for resid in exp_j_couplings.keys():
    # H-N-Ca-Ha dihedral
    dih = (
        system.select(f"resid {resid} and name H").indices[0],
        system.select(f"resid {resid} and name N").indices[0],
        system.select(f"resid {resid} and name CA").indices[0],
        system.select(f"resid {resid} and name HA").indices[0],
    )
    dihedrals.append(dih)
    residues.append(resid)

calc_j = jcoupling(
    traj,
    system,
    dihedral_indices=dihedrals,
    karplus=(7.09, -1.42, 1.55)
)

# Average over trajectory
calc_j_avg = calc_j.mean(axis=0)

# Compare
plt.figure(figsize=(8, 6))
exp_values = [exp_j_couplings[r] for r in residues]
plt.scatter(exp_values, calc_j_avg)
plt.plot([6, 10], [6, 10], 'r--', label='Perfect agreement')
plt.xlabel("Experimental ³J (Hz)")
plt.ylabel("Calculated ³J (Hz)")
plt.title("J-coupling Validation")
plt.legend()
plt.grid(alpha=0.3)
plt.show()

# Calculate RMSD
rmsd = np.sqrt(np.mean((np.array(exp_values) - calc_j_avg)**2))
print(f"RMSD from experiment: {rmsd:.2f} Hz")