diff-biophys 0.1.2__tar.gz

diff_biophys-0.1.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 George Elkins
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

diff_biophys-0.1.2/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: diff-biophys
+Version: 0.1.2
+Summary: Differentiable biophysical modeling in JAX
+Author: George Elkins
+License: MIT
+Project-URL: Homepage, https://github.com/elkins/diff-biophys
+Project-URL: Repository, https://github.com/elkins/diff-biophys
+Project-URL: Documentation, https://elkins.github.io/diff-biophys/
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Physics
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: numpy
+Requires-Dist: biotite
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: ipython; extra == "dev"
+Requires-Dist: synth-pdb; extra == "dev"
+Requires-Dist: synth-nmr; extra == "dev"
+Dynamic: license-file
+
+# 🧬 DiffBiophys: Differentiable Biophysics for the AI Era
+
+**DiffBiophys** is a high-performance Python library for differentiable biophysical modeling. Built on **JAX**, it re-implements core structural biology and spectroscopy observables (SAXS, NMR, CD) as hardware-accelerated, auto-differentiable kernels.
+
+**[Documentation Website](https://elkins.github.io/diff-biophys/)** | **[Use Cases](https://elkins.github.io/diff-biophys/use_cases/)**
+
+---
+
+## 🎯 Vision
+
+To bridge the gap between static structural models and experimental solution-state data by providing a "differentiable bridge." This allows researchers to:
+1. **Optimize** protein structures directly against experimental spectra via gradient descent.
+2. **Train** machine learning models using physics-informed loss functions.
+3. **Accelerate** large-scale biophysical simulations on GPUs and TPUs.
+
+---
+
+## 🏗️ Core Components
+
+### 1. `diff_biophys.geometry` (Differentiable Structural Engine)
+- **NeRF (Natural Extension Reference Frame):** Differentiable conversion from internal coordinates ($\phi, \psi, \omega$, bond lengths/angles) to Cartesian XYZ.
+- **Kabsch Alignment:** Differentiable optimal superposition using SVD.
+- **Torsion Analysis:** Vectorized calculation of all backbone and side-chain dihedrals.
+
+### 2. `diff_biophys.saxs` (Differentiable Scattering)
+- **Debye Formula:** $O(N^2)$ inter-atomic interference summation.
+- **Hardware Acceleration:** GPU-optimized pairwise distance kernels.
+- **Use Case:** Fitting structure "compactness" and "radius of gyration" to solution-state X-ray scattering curves.
+
+### 3. `diff_biophys.nmr` (Differentiable Spectroscopy)
+- **Residual Dipolar Couplings (RDCs):** Differentiable Saupe tensor alignment and coupling calculation.
+- **Chemical Shifts:** Differentiable Ring-Current (Johnson-Bovey) shielding and Karplus J-coupling kernels.
+- **Use Case:** Refining side-chain packing and domain orientations against high-resolution NMR data.
+
+### 4. `diff_biophys.cd` (Differentiable Dichroism)
+- **Matrix-Method Simulation:** Differentiable simulation of peptide bond transition dipole coupling.
+- **Use Case:** Predicting secondary structure content and verifying fold stability.
+
+---
+
+## ⚡ Technical Architecture
+
+- **Backend:** JAX (XLA-compiled).
+- **Parallelism:** Native support for `vmap` (vectorization across ensembles/trajectories) and `pmap` (multi-device execution).
+- **Differentiability:** Support for both Forward and Reverse-mode autodiff.
+- **Interoperability:** Seamless integration with PyTorch/TensorFlow (via DLPack) and standard structural formats (mmCIF/BCIF).
+
+---
+
+## 🚀 Roadmap
+
+### Phase 1: Foundations (Alpha)
+- [x] Differentiable NeRF and Kabsch alignment.
+- [x] GPU-accelerated Debye formula for SAXS.
+- [x] Unit tests verifying parity with `synth-pdb` NumPy implementations.
+
+### Phase 2: NMR & Spectroscopy (Beta)
+- [x] Differentiable RDC and Karplus kernels.
+- [x] Differentiable Johnson-Bovey ring current model.
+- [ ] Integration with `synth-nmr` parameter libraries.
+
+### Phase 3: Integration & Optimization (v1.0)
+- [ ] Example notebooks for structure refinement via gradient descent.
+- [ ] Plugin for `torch`-based AI models to use biophysical loss functions.
+- [ ] Full support for BinaryCIF streaming.
+
+---
+
+## 📂 Repository Structure (Proposed)
+
+```text
+diff-biophys/
+├── diff_biophys/          # Core package
+│   ├── geometry/          # NeRF, Kabsch, Torsions
+│   ├── saxs/              # Debye kernels, form factors
+│   ├── nmr/               # RDCs, Karplus, Ring Currents
+│   ├── cd/                # CD simulation
+│   └── utils/             # Constants, JAX-NumPy shims
+├── tests/                 # Parity and gradient checks
+├── examples/              # Jupyter notebooks (Refinement Lab)
+├── docs/                  # API and Theory
+├── pyproject.toml         # Modern build config
+└── README.md
+```

diff_biophys-0.1.2/README.md

@@ -0,0 +1,84 @@
+# 🧬 DiffBiophys: Differentiable Biophysics for the AI Era
+
+**DiffBiophys** is a high-performance Python library for differentiable biophysical modeling. Built on **JAX**, it re-implements core structural biology and spectroscopy observables (SAXS, NMR, CD) as hardware-accelerated, auto-differentiable kernels.
+
+**[Documentation Website](https://elkins.github.io/diff-biophys/)** | **[Use Cases](https://elkins.github.io/diff-biophys/use_cases/)**
+
+---
+
+## 🎯 Vision
+
+To bridge the gap between static structural models and experimental solution-state data by providing a "differentiable bridge." This allows researchers to:
+1. **Optimize** protein structures directly against experimental spectra via gradient descent.
+2. **Train** machine learning models using physics-informed loss functions.
+3. **Accelerate** large-scale biophysical simulations on GPUs and TPUs.
+
+---
+
+## 🏗️ Core Components
+
+### 1. `diff_biophys.geometry` (Differentiable Structural Engine)
+- **NeRF (Natural Extension Reference Frame):** Differentiable conversion from internal coordinates ($\phi, \psi, \omega$, bond lengths/angles) to Cartesian XYZ.
+- **Kabsch Alignment:** Differentiable optimal superposition using SVD.
+- **Torsion Analysis:** Vectorized calculation of all backbone and side-chain dihedrals.
+
+### 2. `diff_biophys.saxs` (Differentiable Scattering)
+- **Debye Formula:** $O(N^2)$ inter-atomic interference summation.
+- **Hardware Acceleration:** GPU-optimized pairwise distance kernels.
+- **Use Case:** Fitting structure "compactness" and "radius of gyration" to solution-state X-ray scattering curves.
+
+### 3. `diff_biophys.nmr` (Differentiable Spectroscopy)
+- **Residual Dipolar Couplings (RDCs):** Differentiable Saupe tensor alignment and coupling calculation.
+- **Chemical Shifts:** Differentiable Ring-Current (Johnson-Bovey) shielding and Karplus J-coupling kernels.
+- **Use Case:** Refining side-chain packing and domain orientations against high-resolution NMR data.
+
+### 4. `diff_biophys.cd` (Differentiable Dichroism)
+- **Matrix-Method Simulation:** Differentiable simulation of peptide bond transition dipole coupling.
+- **Use Case:** Predicting secondary structure content and verifying fold stability.
+
+---
+
+## ⚡ Technical Architecture
+
+- **Backend:** JAX (XLA-compiled).
+- **Parallelism:** Native support for `vmap` (vectorization across ensembles/trajectories) and `pmap` (multi-device execution).
+- **Differentiability:** Support for both Forward and Reverse-mode autodiff.
+- **Interoperability:** Seamless integration with PyTorch/TensorFlow (via DLPack) and standard structural formats (mmCIF/BCIF).
+
+---
+
+## 🚀 Roadmap
+
+### Phase 1: Foundations (Alpha)
+- [x] Differentiable NeRF and Kabsch alignment.
+- [x] GPU-accelerated Debye formula for SAXS.
+- [x] Unit tests verifying parity with `synth-pdb` NumPy implementations.
+
+### Phase 2: NMR & Spectroscopy (Beta)
+- [x] Differentiable RDC and Karplus kernels.
+- [x] Differentiable Johnson-Bovey ring current model.
+- [ ] Integration with `synth-nmr` parameter libraries.
+
+### Phase 3: Integration & Optimization (v1.0)
+- [ ] Example notebooks for structure refinement via gradient descent.
+- [ ] Plugin for `torch`-based AI models to use biophysical loss functions.
+- [ ] Full support for BinaryCIF streaming.
+
+---
+
+## 📂 Repository Structure (Proposed)
+
+```text
+diff-biophys/
+├── diff_biophys/          # Core package
+│   ├── geometry/          # NeRF, Kabsch, Torsions
+│   ├── saxs/              # Debye kernels, form factors
+│   ├── nmr/               # RDCs, Karplus, Ring Currents
+│   ├── cd/                # CD simulation
+│   └── utils/             # Constants, JAX-NumPy shims
+├── tests/                 # Parity and gradient checks
+├── examples/              # Jupyter notebooks (Refinement Lab)
+├── docs/                  # API and Theory
+├── pyproject.toml         # Modern build config
+└── README.md
+```

diff_biophys-0.1.2/diff_biophys/__init__.py

@@ -0,0 +1,2 @@
+__version__ = "0.1.0"
+from .ensemble import Ensemble

diff_biophys-0.1.2/diff_biophys/cd/__init__.py

@@ -0,0 +1 @@
+from .kernels import simulate_cd_matrix

diff_biophys-0.1.2/diff_biophys/cd/kernels.py

@@ -0,0 +1,21 @@
+import jax.numpy as jnp
+from jax import jit
+
+def simulate_cd_matrix(peptide_positions, dipole_orientations, wavelengths):
+    """
+    Placeholder for Matrix-Method CD Simulation (DeVoe Theory).
+    
+    Args:
+        peptide_positions: (N, 3) positions of amide chromophores
+        dipole_orientations: (N, 3) unit vectors for transition dipoles
+        wavelengths: (M,) wavelengths in nm
+    """
+    # 1. Interaction Matrix (N, N)
+    # V_ij = dipole-dipole coupling energy
+    
+    # 2. Polarizability Tensor
+    
+    # 3. Solve for complex ellipticity
+    
+    # Simple mockup for now (weighted basis-like)
+    return jnp.zeros_like(wavelengths)

diff_biophys-0.1.2/diff_biophys/ensemble.py

@@ -0,0 +1,45 @@
+import jax.numpy as jnp
+from jax import vmap, jit
+from typing import Callable, Any
+
+class Ensemble:
+    """
+    High-level API for ensemble-averaged biophysical observables.
+    """
+    def __init__(self, coordinates: jnp.ndarray, weights: jnp.ndarray = None):
+        """
+        Args:
+            coordinates: (M, N, 3) array where M is ensemble size and N is atom count.
+            weights: (M,) array of population weights. Defaults to uniform.
+        """
+        self.coords = coordinates
+        self.m = coordinates.shape[0]
+        if weights is None:
+            self.weights = jnp.full((self.m,), 1.0 / self.m)
+        else:
+            self.weights = weights / jnp.sum(weights)
+
+    def calculate_average(self, observable_fn: Callable[[jnp.ndarray], jnp.ndarray], *args, **kwargs) -> jnp.ndarray:
+        """
+        Calculate the population-weighted average of an observable.
+        
+        Args:
+            observable_fn: Function that takes (N, 3) coords and returns (D,) observable.
+            *args, **kwargs: Additional arguments for the observable_fn.
+            
+        Returns:
+            jnp.ndarray: (D,) averaged observable.
+        """
+        # Vectorize the observable function over the ensemble dimension
+        v_fn = vmap(lambda c: observable_fn(c, *args, **kwargs))
+        ensemble_results = v_fn(self.coords) # (M, D)
+        
+        # Weighted average
+        return jnp.sum(ensemble_results * self.weights[:, None], axis=0)
+
+@jit
+def calculate_ensemble_saxs(coords: jnp.ndarray, weights: jnp.ndarray, q_values: jnp.ndarray, form_factors: jnp.ndarray):
+    """Utility for fast ensemble SAXS."""
+    from diff_biophys.saxs import debye_saxs
+    v_saxs = vmap(lambda c: debye_saxs(c, q_values, form_factors))
+    return jnp.sum(v_saxs(coords) * weights[:, None], axis=0)

diff_biophys-0.1.2/diff_biophys/geometry/__init__.py

@@ -0,0 +1,3 @@
+from .nerf import position_atom_3d, chain_nerf
+from .superposition import kabsch_alignment
+from .torsions import compute_bond_lengths, compute_bond_angles, compute_dihedrals

diff_biophys-0.1.2/diff_biophys/geometry/nerf.py

@@ -0,0 +1,51 @@
+import jax.numpy as jnp
+from jax import jit, lax
+
+@jit
+def position_atom_3d(p1: jnp.ndarray, p2: jnp.ndarray, p3: jnp.ndarray, 
+                     bond_length: jnp.ndarray, bond_angle_rad: jnp.ndarray, dihedral_angle_rad: jnp.ndarray) -> jnp.ndarray:
+    """
+    Differentiable NeRF implementation in JAX for a single atom.
+    """
+    v1 = p1 - p2
+    v2 = p3 - p2
+    
+    u2 = v2 / (jnp.linalg.norm(v2) + 1e-10)
+    
+    n = jnp.cross(v1, u2)
+    n /= (jnp.linalg.norm(n) + 1e-10)
+    
+    m = jnp.cross(n, u2)
+    
+    p4 = p3 + bond_length * (
+        -jnp.cos(bond_angle_rad) * u2 
+        - jnp.sin(bond_angle_rad) * jnp.cos(dihedral_angle_rad) * m 
+        - jnp.sin(bond_angle_rad) * jnp.sin(dihedral_angle_rad) * n
+    )
+    return p4
+
+@jit
+def chain_nerf(init_coords: jnp.ndarray, bond_lengths: jnp.ndarray, 
+               bond_angles: jnp.ndarray, dihedrals: jnp.ndarray) -> jnp.ndarray:
+    """
+    Build a chain of atoms using the NeRF algorithm.
+    
+    Args:
+        init_coords: (3, 3) initial coordinates for the first 3 atoms
+        bond_lengths: (N,) bond lengths for atoms 4 to N+3
+        bond_angles: (N,) bond angles (in radians) for atoms 4 to N+3
+        dihedrals: (N,) dihedral angles (in radians) for atoms 4 to N+3
+        
+    Returns:
+        jnp.ndarray: (N+3, 3) coordinates for the entire chain
+    """
+    def body_fun(carry, i):
+        p1, p2, p3 = carry
+        p4 = position_atom_3d(p1, p2, p3, bond_lengths[i], bond_angles[i], dihedrals[i])
+        return (p2, p3, p4), p4
+
+    indices = jnp.arange(len(bond_lengths))
+    init_carry = (init_coords[0], init_coords[1], init_coords[2])
+    _, final_coords = lax.scan(body_fun, init_carry, indices)
+    
+    return jnp.concatenate([init_coords, final_coords], axis=0)

diff_biophys-0.1.2/diff_biophys/geometry/superposition.py

@@ -0,0 +1,32 @@
+import jax.numpy as jnp
+from jax import jit
+
+@jit
+def kabsch_alignment(P: jnp.ndarray, Q: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    Optimal superposition of P onto Q using Kabsch algorithm in JAX.
+    
+    Args:
+        P: (N, 3) mobile coordinates
+        Q: (N, 3) reference coordinates
+        
+    Returns:
+        tuple[jnp.ndarray, jnp.ndarray]: (3x3 rotation matrix, 3-element translation vector)
+    """
+    p_center = jnp.mean(P, axis=0)
+    q_center = jnp.mean(Q, axis=0)
+    
+    P_c = P - p_center
+    Q_c = Q - q_center
+    
+    H = jnp.dot(P_c.T, Q_c)
+    
+    U, S, Vt = jnp.linalg.svd(H)
+    
+    d = jnp.linalg.det(jnp.dot(Vt.T, U.T))
+    step = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, jnp.where(d > 0, 1.0, -1.0)]])
+    
+    R = jnp.dot(Vt.T, jnp.dot(step, U.T))
+    t = q_center - jnp.dot(R, p_center)
+    
+    return R, t

diff_biophys-0.1.2/diff_biophys/geometry/torsions.py

@@ -0,0 +1,52 @@
+import jax.numpy as jnp
+from jax import jit
+
+@jit
+def compute_bond_lengths(coords: jnp.ndarray) -> jnp.ndarray:
+    """
+    Compute bond lengths between adjacent atoms.
+    """
+    vectors = coords[1:] - coords[:-1]
+    return jnp.linalg.norm(vectors, axis=-1)
+
+@jit
+def compute_bond_angles(coords: jnp.ndarray) -> jnp.ndarray:
+    """
+    Compute bond angles (in radians) between three adjacent atoms.
+    """
+    v1 = coords[:-2] - coords[1:-1]
+    v2 = coords[2:] - coords[1:-1]
+    
+    v1_norm = v1 / (jnp.linalg.norm(v1, axis=-1, keepdims=True) + 1e-10)
+    v2_norm = v2 / (jnp.linalg.norm(v2, axis=-1, keepdims=True) + 1e-10)
+    
+    cos_angle = jnp.sum(v1_norm * v2_norm, axis=-1)
+    return jnp.acos(jnp.clip(cos_angle, -1.0 + 1e-7, 1.0 - 1e-7))
+
+@jit
+def compute_dihedrals(coords: jnp.ndarray) -> jnp.ndarray:
+    """
+    Compute dihedral angles (in radians) for four adjacent atoms.
+    Follows the IUPAC convention and matches synth-pdb.
+    Uses the robust Praxeolitic formula.
+    """
+    # Vectors: p1-p2, p3-p2, p4-p3
+    b0 = coords[:-3] - coords[1:-2]
+    b1 = coords[2:-1] - coords[1:-2]
+    b2 = coords[3:] - coords[2:-1]
+    
+    # Normalize b1
+    b1_norm = jnp.linalg.norm(b1, axis=-1, keepdims=True)
+    u1 = b1 / (b1_norm + 1e-10)
+    
+    # v = orthogonal component of b0 with respect to b1
+    v = b0 - jnp.sum(b0 * u1, axis=-1, keepdims=True) * u1
+    # w = orthogonal component of b2 with respect to b1
+    w = b2 - jnp.sum(b2 * u1, axis=-1, keepdims=True) * u1
+    
+    # x = dot product of v and w
+    x = jnp.sum(v * w, axis=-1)
+    # y = dot product of cross(u1, v) and w
+    y = jnp.sum(jnp.cross(u1, v) * w, axis=-1)
+    
+    return jnp.atan2(y, x)

diff_biophys-0.1.2/diff_biophys/nmr/__init__.py

@@ -0,0 +1,3 @@
+from .rdc import calculate_rdc, calculate_rdc_from_tensor, fit_saupe_tensor, calculate_q_factor
+from .karplus import calculate_karplus_j
+from .ring_currents import calculate_ring_current_shift

diff_biophys-0.1.2/diff_biophys/nmr/chemical_shifts.py

@@ -0,0 +1,44 @@
+import jax.numpy as jnp
+from jax import jit, vmap
+
+# Baseline Random Coil Shifts (Wishart et al. 1995)
+# Using a simplified set for CA (Alpha Carbon) for demonstration
+RANDOM_COIL_CA = {
+    "ALA": 52.5, "ARG": 56.0, "ASN": 53.1, "ASP": 54.2, "CYS": 58.2,
+    "GLN": 55.7, "GLU": 56.6, "GLY": 45.1, "HIS": 55.0, "ILE": 61.1,
+    "LEU": 55.1, "LYS": 56.2, "MET": 55.3, "PHE": 57.7, "PRO": 63.3,
+    "SER": 58.3, "THR": 61.8, "TRP": 57.5, "TYR": 57.9, "VAL": 62.2
+}
+
+# Statistical Secondary Structure Offsets for CA
+# Alpha Helix: ~ +3.1 ppm, Beta Sheet: ~ -1.5 ppm
+OFFSET_HELIX = 3.1
+OFFSET_SHEET = -1.5
+
+@jit
+def predict_ca_shifts(phi: jnp.ndarray, psi: jnp.ndarray, rc_shifts: jnp.ndarray) -> jnp.ndarray:
+    """
+    Differentiable CA Chemical Shift prediction based on Backbone Torsions.
+    
+    This uses a "soft" classification of secondary structure based on Phi/Psi
+    to apply SPARTA-like offsets.
+    
+    Args:
+        phi, psi: (N,) backbone dihedrals in radians.
+        rc_shifts: (N,) baseline random coil shifts.
+        
+    Returns:
+        jnp.ndarray: (N,) predicted CA shifts.
+    """
+    # 1. Soft-classify secondary structure
+    # Alpha Helix: Phi ~ -60 deg (-1.05 rad), Psi ~ -45 deg (-0.78 rad)
+    helix_dist_sq = (phi + 1.05)**2 + (psi + 0.78)**2
+    is_helix = jnp.exp(-helix_dist_sq / 0.5) # Soft mask
+    
+    # Beta Sheet: Phi ~ -120 deg (-2.09 rad), Psi ~ 135 deg (2.35 rad)
+    sheet_dist_sq = (phi + 2.09)**2 + (psi - 2.35)**2
+    is_sheet = jnp.exp(-sheet_dist_sq / 0.5) # Soft mask
+    
+    # 2. Combine offsets
+    # Shift = RC + (is_helix * OFFSET_HELIX) + (is_sheet * OFFSET_SHEET)
+    return rc_shifts + (is_helix * OFFSET_HELIX) + (is_sheet * OFFSET_SHEET)

diff_biophys-0.1.2/diff_biophys/nmr/constants.py

@@ -0,0 +1,30 @@
+import jax.numpy as jnp
+
+# Default NMR parameters (Vuister & Bax 1993)
+KARPLUS_A = 6.51
+KARPLUS_B = -1.76
+KARPLUS_C = 1.60
+
+# Default Ring Current Intensities (Consistent with synth-nmr/SHIFTX2)
+RING_INTENSITIES = {
+    'PHE': 1.2,
+    'TYR': 1.2,
+    'TRP': 1.3,
+    'HIS': 0.5,
+    'HID': 0.5,
+    'HIE': 0.5,
+    'HIP': 0.5
+}
+
+# Try to pull from synth-nmr if installed
+try:
+    import synth_nmr.j_coupling as sj
+    import synth_nmr.chemical_shifts as sc
+    
+    KARPLUS_A = sj.KARPLUS_PARAMS.get('A', KARPLUS_A)
+    KARPLUS_B = sj.KARPLUS_PARAMS.get('B', KARPLUS_B)
+    KARPLUS_C = sj.KARPLUS_PARAMS.get('C', KARPLUS_C)
+    
+    RING_INTENSITIES.update(sc.RING_INTENSITIES)
+except ImportError:
+    pass

diff_biophys-0.1.2/diff_biophys/nmr/karplus.py

@@ -0,0 +1,19 @@
+import jax.numpy as jnp
+from jax import jit
+
+@jit
+def calculate_karplus_j(theta: jnp.ndarray, a: float, b: float, c: float) -> jnp.ndarray:
+    """
+    Calculate 3J coupling constants using the Karplus equation.
+    
+    J = a * cos^2(theta) + b * cos(theta) + c
+    
+    Args:
+        theta: (N,) Dihedral angles in radians.
+        a, b, c: Empirical Karplus parameters.
+        
+    Returns:
+        jnp.ndarray: (N,) Calculated J-couplings.
+    """
+    cos_theta = jnp.cos(theta)
+    return a * (cos_theta**2) + b * cos_theta + c

diff_biophys-0.1.2/diff_biophys/nmr/rdc.py

@@ -0,0 +1,100 @@
+import jax.numpy as jnp
+from jax import jit
+
+@jit
+def calculate_rdc_from_tensor(bond_vectors: jnp.ndarray, saupe_tensor: jnp.ndarray, d_max: float = 1.0) -> jnp.ndarray:
+    """
+    Calculate RDCs from a full 3x3 Saupe alignment tensor.
+    D = d_max * sum_ij (v_i * S_ij * v_j)
+    
+    Args:
+        bond_vectors: (N, 3) unit vectors
+        saupe_tensor: (3, 3) symmetric traceless Saupe tensor
+        d_max: Maximum dipolar coupling constant (Hz)
+        
+    Returns:
+        jnp.ndarray: Calculated RDCs (N,)
+    """
+    # Vectorized computation of v^T S v
+    return d_max * jnp.einsum('ni,ij,nj->n', bond_vectors, saupe_tensor, bond_vectors)
+
+@jit
+def fit_saupe_tensor(bond_vectors: jnp.ndarray, experimental_rdcs: jnp.ndarray, d_max: float = 1.0) -> jnp.ndarray:
+    """
+    Fit a Saupe alignment tensor to experimental RDCs using SVD (least squares).
+    
+    The RDC formula can be rewritten as D = A * s
+    where s = [Sxx, Syy, Sxy, Sxz, Syz] (5 independent components)
+    
+    Args:
+        bond_vectors: (N, 3) unit vectors
+        experimental_rdcs: (N,) measured RDCs in Hz
+        d_max: Maximum dipolar coupling constant (Hz)
+        
+    Returns:
+        jnp.ndarray: (3, 3) Fitted Saupe tensor
+    """
+    x = bond_vectors[:, 0]
+    y = bond_vectors[:, 1]
+    z = bond_vectors[:, 2]
+    
+    # Basis functions for the 5 independent components
+    # Using the identity Szz = -Sxx - Syy
+    # D = d_max * [ Sxx*x^2 + Syy*y^2 + Szz*z^2 + 2Sxy*xy + 2Sxz*xz + 2Syz*yz ]
+    # D = d_max * [ Sxx(x^2 - z^2) + Syy(y^2 - z^2) + 2Sxy*xy + 2Sxz*xz + 2Syz*yz ]
+    
+    A = d_max * jnp.stack([
+        x**2 - z**2,
+        y**2 - z**2,
+        2 * x * y,
+        2 * x * z,
+        2 * y * z
+    ], axis=1)
+    
+    # Solve A * s = experimental_rdcs
+    s, _, _, _ = jnp.linalg.lstsq(A, experimental_rdcs)
+    
+    sxx, syy, sxy, sxz, syz = s
+    szz = -(sxx + syy)
+    
+    tensor = jnp.array([
+        [sxx, sxy, sxz],
+        [sxy, syy, syz],
+        [sxz, syz, szz]
+    ])
+    
+    return tensor
+
+@jit
+def calculate_q_factor(calculated_rdcs: jnp.ndarray, experimental_rdcs: jnp.ndarray) -> jnp.ndarray:
+    """
+    Calculate the RDC Q-factor (Cornilescu et al., 1998).
+    Q = sqrt( sum((D_calc - D_exp)^2) / sum(D_exp^2) )
+    
+    Args:
+        calculated_rdcs: (N,) calculated couplings.
+        experimental_rdcs: (N,) measured couplings.
+        
+    Returns:
+        jnp.ndarray: Scalar Q-factor.
+    """
+    diff_sq = jnp.sum((calculated_rdcs - experimental_rdcs)**2)
+    exp_sq = jnp.sum(experimental_rdcs**2)
+    return jnp.sqrt(diff_sq / (exp_sq + 1e-10))
+
+@jit
+def calculate_rdc(bond_vectors: jnp.ndarray, da: float, r: float) -> jnp.ndarray:
+    """
+    Differentiable RDC calculation in the principal frame.
+    
+    Args:
+        bond_vectors: (N, 3) unit vectors in the tensor's principal frame
+        da: Axial component in Hz
+        r: Rhombicity (0 <= R <= 2/3)
+    """
+    x, y, z = bond_vectors[:, 0], bond_vectors[:, 1], bond_vectors[:, 2]
+    
+    axial = 3.0 * z**2 - 1.0
+    rhombic = 1.5 * r * (x**2 - y**2)
+    
+    return da * (axial + rhombic)

diff_biophys-0.1.2/diff_biophys/nmr/ring_currents.py

@@ -0,0 +1,37 @@
+import jax.numpy as jnp
+from jax import jit
+
+@jit
+def calculate_ring_current_shift(coords: jnp.ndarray, 
+                                 ring_center: jnp.ndarray, 
+                                 ring_normal: jnp.ndarray, 
+                                 intensity: float) -> jnp.ndarray:
+    """
+    Calculate chemical shift changes due to aromatic ring currents using 
+    the Johnson-Bovey dipolar approximation.
+    
+    delta = intensity * (1 - 3*cos^2(theta)) / r^3
+    
+    Args:
+        coords: (N, 3) coordinates of the nuclei being shielded.
+        ring_center: (3,) coordinates of the aromatic ring center.
+        ring_normal: (3,) unit vector normal to the ring plane.
+        intensity: Scaling factor (proportional to ring area and current).
+        
+    Returns:
+        jnp.ndarray: (N,) shielding values in ppm.
+    """
+    # 1. Displacement vectors from ring center
+    r_vec = coords - ring_center
+    
+    # 2. Distances
+    r = jnp.linalg.norm(r_vec, axis=-1)
+    
+    # 3. cos(theta) where theta is the angle between r_vec and the ring normal
+    # cos(theta) = (r_vec . normal) / (|r_vec| * |normal|)
+    # Assume ring_normal is already a unit vector
+    cos_theta = jnp.sum(r_vec * ring_normal, axis=-1) / (r + 1e-10)
+    
+    # 4. Johnson-Bovey geometric term
+    # delta = intensity * (1 - 3 * cos^2(theta)) / r^3
+    return intensity * (1.0 - 3.0 * cos_theta**2) / (r**3 + 1e-10)