diff-biophys 0.1.2__tar.gz → 0.1.3__tar.gz

diff_biophys-0.1.3/PKG-INFO

@@ -0,0 +1,182 @@
+Metadata-Version: 2.4
+Name: diff-biophys
+Version: 0.1.3
+Summary: Differentiable biophysical modeling in JAX
+Author: George Elkins
+License-Expression: 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: 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
+Requires-Python: >=3.10
+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>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: ruff>=0.6.0; extra == "dev"
+Requires-Dist: mypy>=1.8.0; extra == "dev"
+Requires-Dist: pre-commit>=3.0.0; extra == "dev"
+Requires-Dist: ipython; extra == "dev"
+Requires-Dist: synth-pdb; extra == "dev"
+Requires-Dist: synth-nmr; extra == "dev"
+Requires-Dist: synth-core; extra == "dev"
+Provides-Extra: torch
+Requires-Dist: torch>=2.0.0; extra == "torch"
+Provides-Extra: examples
+Requires-Dist: optax>=0.1.0; extra == "examples"
+Requires-Dist: matplotlib>=3.5.0; extra == "examples"
+Dynamic: license-file
+
+# 🧬 DiffBiophys: Differentiable Biophysics for the AI Era
+
+[![Tests](https://github.com/elkins/diff-biophys/actions/workflows/test.yml/badge.svg)](https://github.com/elkins/diff-biophys/actions/workflows/test.yml)
+[![PyPI version](https://img.shields.io/pypi/v/diff-biophys.svg)](https://pypi.org/project/diff-biophys/)
+[![Python 3.10+](https://img.shields.io/pypi/pyversions/diff-biophys.svg)](https://pypi.org/project/diff-biophys/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![codecov](https://codecov.io/gh/elkins/diff-biophys/branch/main/graph/badge.svg)](https://codecov.io/gh/elkins/diff-biophys)
+[![JAX](https://img.shields.io/badge/backend-JAX-9cf.svg)](https://github.com/google/jax)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![Checked with mypy](https://img.shields.io/badge/type%20checked-mypy-blue)](https://mypy-lang.org/)
+
+**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.
+- **Hydration Shell Correction:** Excluded-volume solvent subtraction (Fraser et al. 1978).
+- **Hardware Acceleration:** GPU-optimized pairwise distance kernels via JAX `vmap`.
+- **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. Includes SVD-based tensor fitting.
+- **Chemical Shifts:** Differentiable ring-current (Johnson-Bovey) shielding and softmax-weighted secondary structure Cα shift predictor.
+- **Karplus J-coupling:** Parameterizable 3J coupling equation (Vuister & Bax 1993 defaults).
+- **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 via DeVoe theory.
+- **Status:** ✅ Implemented. Supports frequency-dependent coupled-oscillator response.
+
+---
+
+## ⚡ Technical Architecture
+
+- **Backend:** JAX (XLA-compiled) — supports CPU, GPU, and TPU.
+- **Parallelism:** Native support for `vmap` (vectorization across ensembles/trajectories) and `pmap` (multi-device execution).
+- **Differentiability:** Forward and reverse-mode autodiff through all kernels.
+- **Interoperability:** JAX arrays are compatible with NumPy and can be exchanged with PyTorch via `dlpack` (user-managed conversion).
+
+---
+
+## 🚀 Roadmap
+
+### Phase 1: Foundations (Alpha)
+- [x] Differentiable NeRF and Kabsch alignment.
+- [x] GPU-accelerated Debye formula for SAXS with hydration shell correction.
+- [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.
+- [x] Integration with `synth-nmr` parameter libraries (optional dependency).
+
+### Phase 3: Integration & Optimization (v1.0)
+- [x] Full CD matrix-method implementation (DeVoe theory).
+- [ ] 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
+
+```text
+diff-biophys/
+├── diff_biophys/          # Core package
+│   ├── geometry/          # NeRF, Kabsch, Torsions
+│   ├── saxs/              # Debye kernels, form factors
+│   ├── nmr/               # RDCs, Karplus, Ring Currents, Chemical Shifts
+│   ├── cd/                # CD simulation (DeVoe Matrix Method)
+│   └── ensemble.py        # Ensemble averaging API
+├── tests/                 # Parity, gradient, and scientific validation checks
+├── examples/              # Jupyter notebooks (Refinement Lab)
+├── docs/                  # API and Theory
+├── pyproject.toml         # Modern build config
+└── README.md
+```
+
+## 🚀 Installation
+
+```bash
+pip install diff-biophys
+```
+
+For GPU support (CUDA):
+```bash
+pip install "jax[cuda12]" diff-biophys
+```
+
+## 🤝 Contributing
+
+Contributions are welcome from both ML and structural biology communities! Please open an issue or pull request on [GitHub](https://github.com/elkins/diff-biophys). Run `pre-commit run --all-files` before submitting.
+
+## 🔗 Related Projects
+
+diff-biophys is the **differentiable engine** powering the higher-level tools in this ecosystem:
+
+- [synth-pdb](https://github.com/elkins/synth-pdb) — Synthetic structure generation (uses NumPy implementations)
+- [synth-nmr](https://github.com/elkins/synth-nmr) — NMR observables (optional dependency)
+- [synth-saxs](https://github.com/elkins/synth-saxs) — SAXS profile simulator
+- [diff-fret](https://github.com/elkins/diff-fret) — Differentiable FRET (new)
+- [diff-hdx](https://github.com/elkins/diff-hdx) — Differentiable HDX-MS (new)
+- [diff-epr](https://github.com/elkins/diff-epr) — Differentiable EPR/DEER (new)
+- [diff-ensemble](https://github.com/elkins/diff-ensemble) — IDP ensemble VAE (depends on diff-biophys)
+- [TorsionTuner](https://github.com/elkins/TorsionTuner) — GNN refinement (depends on diff-biophys)
+- [resonance-flow](https://github.com/elkins/resonance-flow) — NMR-guided folding (depends on diff-biophys)
+
+## ⚖️ License
+
+MIT License — see [LICENSE](LICENSE) for details.
+
+## 📖 Citation
+
+```bibtex
+@software{diff_biophys,
+  author  = {Elkins, George},
+  title   = {diff-biophys: Differentiable biophysics kernels for JAX},
+  year    = {2024},
+  url     = {https://github.com/elkins/diff-biophys},
+  version = {0.1.2}
+}
+```

diff_biophys-0.1.3/README.md

@@ -0,0 +1,141 @@
+# 🧬 DiffBiophys: Differentiable Biophysics for the AI Era
+
+[![Tests](https://github.com/elkins/diff-biophys/actions/workflows/test.yml/badge.svg)](https://github.com/elkins/diff-biophys/actions/workflows/test.yml)
+[![PyPI version](https://img.shields.io/pypi/v/diff-biophys.svg)](https://pypi.org/project/diff-biophys/)
+[![Python 3.10+](https://img.shields.io/pypi/pyversions/diff-biophys.svg)](https://pypi.org/project/diff-biophys/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![codecov](https://codecov.io/gh/elkins/diff-biophys/branch/main/graph/badge.svg)](https://codecov.io/gh/elkins/diff-biophys)
+[![JAX](https://img.shields.io/badge/backend-JAX-9cf.svg)](https://github.com/google/jax)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![Checked with mypy](https://img.shields.io/badge/type%20checked-mypy-blue)](https://mypy-lang.org/)
+
+**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.
+- **Hydration Shell Correction:** Excluded-volume solvent subtraction (Fraser et al. 1978).
+- **Hardware Acceleration:** GPU-optimized pairwise distance kernels via JAX `vmap`.
+- **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. Includes SVD-based tensor fitting.
+- **Chemical Shifts:** Differentiable ring-current (Johnson-Bovey) shielding and softmax-weighted secondary structure Cα shift predictor.
+- **Karplus J-coupling:** Parameterizable 3J coupling equation (Vuister & Bax 1993 defaults).
+- **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 via DeVoe theory.
+- **Status:** ✅ Implemented. Supports frequency-dependent coupled-oscillator response.
+
+---
+
+## ⚡ Technical Architecture
+
+- **Backend:** JAX (XLA-compiled) — supports CPU, GPU, and TPU.
+- **Parallelism:** Native support for `vmap` (vectorization across ensembles/trajectories) and `pmap` (multi-device execution).
+- **Differentiability:** Forward and reverse-mode autodiff through all kernels.
+- **Interoperability:** JAX arrays are compatible with NumPy and can be exchanged with PyTorch via `dlpack` (user-managed conversion).
+
+---
+
+## 🚀 Roadmap
+
+### Phase 1: Foundations (Alpha)
+- [x] Differentiable NeRF and Kabsch alignment.
+- [x] GPU-accelerated Debye formula for SAXS with hydration shell correction.
+- [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.
+- [x] Integration with `synth-nmr` parameter libraries (optional dependency).
+
+### Phase 3: Integration & Optimization (v1.0)
+- [x] Full CD matrix-method implementation (DeVoe theory).
+- [ ] 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
+
+```text
+diff-biophys/
+├── diff_biophys/          # Core package
+│   ├── geometry/          # NeRF, Kabsch, Torsions
+│   ├── saxs/              # Debye kernels, form factors
+│   ├── nmr/               # RDCs, Karplus, Ring Currents, Chemical Shifts
+│   ├── cd/                # CD simulation (DeVoe Matrix Method)
+│   └── ensemble.py        # Ensemble averaging API
+├── tests/                 # Parity, gradient, and scientific validation checks
+├── examples/              # Jupyter notebooks (Refinement Lab)
+├── docs/                  # API and Theory
+├── pyproject.toml         # Modern build config
+└── README.md
+```
+
+## 🚀 Installation
+
+```bash
+pip install diff-biophys
+```
+
+For GPU support (CUDA):
+```bash
+pip install "jax[cuda12]" diff-biophys
+```
+
+## 🤝 Contributing
+
+Contributions are welcome from both ML and structural biology communities! Please open an issue or pull request on [GitHub](https://github.com/elkins/diff-biophys). Run `pre-commit run --all-files` before submitting.
+
+## 🔗 Related Projects
+
+diff-biophys is the **differentiable engine** powering the higher-level tools in this ecosystem:
+
+- [synth-pdb](https://github.com/elkins/synth-pdb) — Synthetic structure generation (uses NumPy implementations)
+- [synth-nmr](https://github.com/elkins/synth-nmr) — NMR observables (optional dependency)
+- [synth-saxs](https://github.com/elkins/synth-saxs) — SAXS profile simulator
+- [diff-fret](https://github.com/elkins/diff-fret) — Differentiable FRET (new)
+- [diff-hdx](https://github.com/elkins/diff-hdx) — Differentiable HDX-MS (new)
+- [diff-epr](https://github.com/elkins/diff-epr) — Differentiable EPR/DEER (new)
+- [diff-ensemble](https://github.com/elkins/diff-ensemble) — IDP ensemble VAE (depends on diff-biophys)
+- [TorsionTuner](https://github.com/elkins/TorsionTuner) — GNN refinement (depends on diff-biophys)
+- [resonance-flow](https://github.com/elkins/resonance-flow) — NMR-guided folding (depends on diff-biophys)
+
+## ⚖️ License
+
+MIT License — see [LICENSE](LICENSE) for details.
+
+## 📖 Citation
+
+```bibtex
+@software{diff_biophys,
+  author  = {Elkins, George},
+  title   = {diff-biophys: Differentiable biophysics kernels for JAX},
+  year    = {2024},
+  url     = {https://github.com/elkins/diff-biophys},
+  version = {0.1.2}
+}
+```

diff_biophys-0.1.3/diff_biophys/__init__.py

@@ -0,0 +1,8 @@
+__version__ = "0.1.3"
+from . import cryo_em
+from .ensemble import Ensemble
+
+try:
+    from . import torch_interop
+except ImportError:
+    pass

diff_biophys-0.1.3/diff_biophys/cd/kernels.py

@@ -0,0 +1,87 @@
+import jax
+import jax.numpy as jnp
+
+
+def simulate_cd_matrix(
+    peptide_positions: jnp.ndarray,
+    dipole_orientations: jnp.ndarray,
+    wavelengths: jnp.ndarray,
+    f_osc: float = 0.2,
+    gamma: float = 10.0,
+    lambda_0: float = 190.0,
+) -> jnp.ndarray:
+    """
+    Matrix-Method CD Simulation (DeVoe Theory).
+
+    Implements the coupled-oscillator model for transition dipole coupling.
+    Calculates the interaction matrix and solves for the complex polarizability
+    response to determine molar ellipticity.
+
+    Args:
+        peptide_positions: (N, 3) positions of amide chromophores in Angstroms.
+        dipole_orientations: (N, 3) unit vectors for transition dipoles.
+        wavelengths: (M,) wavelengths in nm to simulate.
+        f_osc: Oscillator strength of the transition (default 0.2 for pi->pi*).
+        gamma: Linewidth parameter in nm (default 10.0).
+        lambda_0: Resonance wavelength in nm (default 190.0).
+
+    Returns:
+        Molar ellipticity [θ] in deg cm^2 / dmol (M,).
+    """
+    n_chromophores = peptide_positions.shape[0]
+
+    # 1. Compute dipole-dipole interaction matrix V_ij
+    # V_ij = (1/r^3) * [ mu_i . mu_j - 3(mu_i . r_ij)(mu_j . r_ij) ]
+    diff = peptide_positions[:, None, :] - peptide_positions[None, :, :]
+    dist_sq = jnp.sum(diff**2, axis=-1)
+
+    # Safe distance for gradients (avoid sqrt(0) and 1/0)
+    # 1e-9 is a safe epsilon for float32
+    mask = dist_sq > 0
+    safe_dist_sq = jnp.where(mask, dist_sq, 1.0)
+    r_ij = jnp.sqrt(safe_dist_sq)
+    r_ij_inv3 = jnp.where(mask, 1.0 / r_ij**3, 0.0)
+
+    # Unit vectors between chromophores
+    r_hat = diff * jnp.where(mask[:, :, None], 1.0 / r_ij[:, :, None], 0.0)
+
+    # Dot products
+    mu_i_mu_j = jnp.sum(dipole_orientations[:, None, :] * dipole_orientations[None, :, :], axis=-1)
+    mu_i_r = jnp.sum(dipole_orientations[:, None, :] * r_hat, axis=-1)
+    mu_j_r = jnp.sum(dipole_orientations[None, :, :] * r_hat, axis=-1)
+
+    # Interaction energy V (N, N)
+    V = r_ij_inv3 * (mu_i_mu_j - 3 * mu_i_r * mu_j_r)
+
+    # 2. Frequency-dependent response
+    def compute_at_wavelength(lmbda: jnp.ndarray) -> jnp.ndarray:
+        # Complex polarizability alpha(lambda)
+        # Lorentzian-like response
+        denom = (1.0 / lmbda**2 - 1.0 / lambda_0**2) + 1j * (gamma / (lmbda * lambda_0**2))
+        alpha = f_osc / denom
+
+        # Interaction matrix (I - alpha * V)
+        # alpha is scalar for all identical chromophores here
+        M = jnp.eye(n_chromophores) - alpha * V
+
+        # We'll use the matrix inverse to find the coupled response
+        # Note: jnp.linalg.inv is differentiable but can be sensitive
+        inv_M = jnp.linalg.inv(M)
+
+        # Geometric factor for CD (Scalar triple product mu_i x mu_j . r_ij)
+        # This represents the chiral arrangement.
+        cross_mu = jnp.cross(dipole_orientations[:, None, :], dipole_orientations[None, :, :])
+        R_ij = jnp.sum(cross_mu * diff, axis=-1)
+
+        # Total CD response at this wavelength
+        coupled_V = inv_M @ (alpha * V)
+        cd_val = jnp.imag(jnp.sum(coupled_V * R_ij))
+
+        return cd_val
+
+    # Vectorize over wavelengths
+    cd_spectrum = jax.vmap(compute_at_wavelength)(wavelengths)
+
+    # Scale to molar ellipticity (arbitrary units for this kernel,
+    # should be calibrated to exp data)
+    return cd_spectrum * 1e5

diff_biophys-0.1.3/diff_biophys/cryo_em.py

@@ -0,0 +1,88 @@
+import jax
+import jax.numpy as jnp
+
+
+@jax.jit
+def compute_fsc(
+    data1: jax.Array, data2: jax.Array, voxel_size: tuple[float, float, float]
+) -> tuple[jax.Array, jax.Array]:
+    """
+    Compute the fully differentiable Fourier Shell Correlation (FSC) between two 3D maps using JAX.
+    Returns frequencies and correlation values.
+
+    This function matches the implementation in synth-core, but uses jax.numpy
+    so that gradients can flow through the FSC calculation to the input maps.
+    """
+    # Fourier transforms in JAX
+    f1 = jnp.fft.rfftn(data1)
+    f2 = jnp.fft.rfftn(data2)
+
+    # Cross-spectral density and power spectra are computed using float arithmetic
+    # We avoid complex multiplication to parallel the numpy memory stability fix
+    cross = f1.real * f2.real + f1.imag * f2.imag
+    p1 = f1.real**2 + f1.imag**2
+    p2 = f2.real**2 + f2.imag**2
+
+    # Calculate radial bins
+    nz, ny, nx = data1.shape
+    kz = jnp.fft.fftfreq(nz, d=voxel_size[0])
+    ky = jnp.fft.fftfreq(ny, d=voxel_size[1])
+    kx = jnp.fft.rfftfreq(nx, d=voxel_size[2])
+
+    # Create 3D grid of frequencies
+    kz_grid, ky_grid, kx_grid = jnp.meshgrid(kz, ky, kx, indexing="ij")
+
+    # Calculate magnitude of frequency vector for each voxel
+    k = jnp.sqrt(kz_grid**2 + ky_grid**2 + kx_grid**2)
+
+    # Flatten everything
+    k = k.ravel()
+    cross = cross.ravel()
+    p1 = p1.ravel()
+    p2 = p2.ravel()
+
+    # Sort by frequency
+    idx = jnp.argsort(k)
+    k_sorted = k[idx]
+    cross_sorted = cross[idx]
+    p1_sorted = p1[idx]
+    p2_sorted = p2[idx]
+
+    n_bins = min(nx, ny, nz) // 2
+    k_max = k_sorted[-1]
+    k_eps = k_max / (10 * n_bins)
+    bins = jnp.linspace(k_eps, k_max, n_bins + 1)
+
+    # We use vmap to compute the bin sums to keep the function differentiable and JIT-compatible
+    # We avoid python loops with dynamic shapes.
+
+    def compute_bin(i: jax.Array) -> tuple[jax.Array, jax.Array, jax.Array]:
+        bin_start = bins[i]
+        bin_end = bins[i + 1]
+        mask = (k_sorted >= bin_start) & (k_sorted < bin_end)
+
+        sum_cross = jnp.sum(jnp.where(mask, cross_sorted, 0.0))
+        sum_p1 = jnp.sum(jnp.where(mask, p1_sorted, 0.0))
+        sum_p2 = jnp.sum(jnp.where(mask, p2_sorted, 0.0))
+
+        num = sum_cross
+        den = jnp.sqrt(sum_p1 * sum_p2)
+
+        # Avoid division by zero
+        val = jnp.where(den > 0, num / den, 0.0)
+        # Clamp to [-1, 1]
+        val = jnp.clip(val, -1.0, 1.0)
+        freq = (bin_start + bin_end) / 2.0
+
+        # We need to return valid mask too, because some bins might be empty
+        is_valid = jnp.any(mask)
+        return freq, val, is_valid
+
+    indices = jnp.arange(n_bins)
+    freqs, vals, is_valid = jax.vmap(compute_bin)(indices)
+
+    # Note: jnp.where with dynamic sizes breaks JIT if we don't pad.
+    # For a fully differentiable metric, we typically pad with 0s or NaNs, or return the full array.
+    # We will return the full array but mask out invalid frequencies with NaN or 0.
+
+    return freqs, vals

{diff_biophys-0.1.2 → diff_biophys-0.1.3}/diff_biophys/ensemble.py

@@ -1,45 +1,62 @@
+from collections.abc import Callable
+from typing import Any, cast
+
 import jax.numpy as jnp
-from jax import vmap, jit
-from typing import Callable, Any
+from jax import jit, vmap
+
 
 class Ensemble:
     """
     High-level API for ensemble-averaged biophysical observables.
     """
-    def __init__(self, coordinates: jnp.ndarray, weights: jnp.ndarray = None):
+
+    coords: jnp.ndarray
+    weights: jnp.ndarray
+    m: int
+
+    def __init__(self, coordinates: jnp.ndarray, weights: jnp.ndarray | None = 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]
+        self.m = int(coordinates.shape[0])
         if weights is None:
             self.weights = jnp.full((self.m,), 1.0 / self.m)
         else:
-            self.weights = weights / jnp.sum(weights)
+            self.weights = cast(jnp.ndarray, weights / jnp.sum(weights))
 
-    def calculate_average(self, observable_fn: Callable[[jnp.ndarray], jnp.ndarray], *args, **kwargs) -> jnp.ndarray:
+    def calculate_average(
+        self,
+        observable_fn: Callable[..., jnp.ndarray],
+        *args: Any,
+        **kwargs: Any,
+    ) -> 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)
-        
+        ensemble_results = v_fn(self.coords)  # (M, D)
+
         # Weighted average
-        return jnp.sum(ensemble_results * self.weights[:, None], axis=0)
+        return cast(jnp.ndarray, 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):
+def calculate_ensemble_saxs(
+    coords: jnp.ndarray, weights: jnp.ndarray, q_values: jnp.ndarray, form_factors: jnp.ndarray
+) -> 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)
+    return cast(jnp.ndarray, jnp.sum(v_saxs(coords) * weights[:, None], axis=0))

diff_biophys-0.1.3/diff_biophys/geometry/__init__.py

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

diff_biophys-0.1.3/diff_biophys/geometry/nerf.py

@@ -0,0 +1,84 @@
+from typing import Any, cast
+
+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.
+
+    Places atom p4 given three reference atoms (p1, p2, p3) and the internal
+    coordinates (bond length, bond angle, dihedral angle) that define its
+    position relative to p3.
+
+    Args:
+        p1, p2, p3: (3,) reference atom coordinates.
+        bond_length: Scalar distance p3→p4 in Ångströms.
+        bond_angle_rad: Scalar bond angle ∠(p2, p3, p4) in radians.
+        dihedral_angle_rad: Scalar dihedral angle ∠(p1, p2, p3, p4) in radians.
+
+    Returns:
+        jnp.ndarray: (3,) Cartesian coordinates of the new atom p4.
+    """
+    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 cast(jnp.ndarray, 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: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], i: Any
+    ) -> tuple[tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], Any]:
+        p1, p2, p3 = carry
+
+        p4 = position_atom_3d(p1, p2, p3, bond_lengths[i], bond_angles[i], dihedrals[i])
+        return (p2, p3, p4), p4
+
+    # Use .shape[0] instead of len() so this works correctly under vmap
+    # and with dynamically-shaped arrays during JAX tracing.
+    indices = jnp.arange(bond_lengths.shape[0])
+    init_carry = (init_coords[0], init_coords[1], init_coords[2])
+    _, final_coords = lax.scan(body_fun, init_carry, indices)
+
+    return cast(jnp.ndarray, jnp.concatenate([init_coords, final_coords], axis=0))