Heterodyne

CPU-optimized JAX package for heterodyne X-ray Photon Correlation Spectroscopy (XPCS) analysis under nonequilibrium conditions, implementing the theoretical framework from He et al. PNAS 2024 and He et al. PNAS 2025 for characterizing transport properties in flowing soft matter systems.

Heterodyne Model

The package implements a two-component scattering model (PNAS 2025 SI Eqs. S-77–S-98) where scattered light from a moving sample interferes with that from a static reference, producing oscillations in the cross-correlation whose frequency encodes the sample velocity.

Two-Time Correlation (Eq. S-95)

$$c_2(\vec{q}, t_1, t_2) = 1 + \frac{\beta}{f^2} \left[ C_{\text{ref}} + C_{\text{sample}} + C_{\text{cross}} \right]$$

with three physical contributions:

$$C_{\text{ref}} = [x_r(t_1) \cdot x_r(t_2)]^2 \exp\left(-q^2 \int_{t_1}^{t_2} J_r(t') dt'\right)$$

$$C_{\text{sample}} = [x_s(t_1) \cdot x_s(t_2)]^2 \exp\left(-q^2 \int_{t_1}^{t_2} J_s(t') dt'\right)$$

$$C_{\text{cross}} = 2 x_r(t_1) x_r(t_2) x_s(t_1) x_s(t_2) \exp\left(-\frac{1}{2} q^2 \int_{t_1}^{t_2} [J_s(t') + J_r(t')] dt'\right) \cos\left[q \cos(\varphi) \int_{t_1}^{t_2} \mathbb{E}[v] dt'\right]$$

$$f^2 = [x_s(t_1)^2 + x_r(t_1)^2][x_s(t_2)^2 + x_r(t_2)^2]$$

where $x_s(t)$ is the sample fraction, $x_r(t) = 1 - x_s(t)$ the reference fraction, $\beta$ the optical contrast, $\varphi$ the angle between velocity and $\vec{q}$, and $f^2$ a normalization ensuring $c_2(t, t) = 1 + \beta$ on the diagonal.

One-Time Equilibrium Form (Eq. S-98)

When all parameters are time-independent and $J_n(t) = 6D_n$ (Wiener process), the two-time correlation reduces to a function of lag time $\tau = t_2 - t_1$ only:

$$g_2(\vec{q}, \tau) = 1 + \beta \left[ (1 - x)^2 e^{-6q^2 D_r \tau} + x^2 e^{-6q^2 D_s \tau} + 2x(1 - x) e^{-3q^2(D_r + D_s)\tau} \cos[q \cos(\varphi) \mathbb{E}[v] \tau] \right]$$

where $x = x_s^2 / (x_s^2 + x_r^2)$ is the equilibrium sample intensity fraction.

Fitting Model

The implementation wraps the correlation with per-angle scaling parameters:

$$c_2^{\text{model}} = \text{offset} + \text{contrast} \times \frac{C_{\text{ref}} + C_{\text{sample}} + C_{\text{cross}}}{f^2}$$

Rate Functions

Each component has its own power-law transport coefficient, and the sample has an additional velocity rate:

$$J(t) = D_0 t^\alpha + D_{\text{offset}} \qquad v(t) = v_0 t^\beta + v_{\text{offset}}$$

All time integrals are evaluated numerically via cumulative trapezoid — no analytical antiderivatives are ever used, ensuring correctness for the general power-law form.

Parameters

The model has 14 physics parameters organized into five groups, plus 2 per-angle scaling parameters:

Reference transport — $J_r(t) = D_{0,r} t^{\alpha_r} + D_{\text{offset},r}$

Sample transport — $J_s(t) = D_{0,s} t^{\alpha_s} + D_{\text{offset},s}$

Velocity — $v(t) = v_0 t^\beta + v_{\text{offset}}$

Sample fraction — $f_s(t) = f_0 \exp(f_1 (t - f_2)) + f_3$, where $f_s(t) \in [0, 1]$

Flow angle

Per-angle scaling (2 parameters per detector angle)

Total: 14 physics + 2 scaling parameters per angle.

Installation

pip install heterodyne

For development:

git clone https://github.com/imewei/heterodyne.git
cd heterodyne
make dev    # or: uv sync

Requirements: Python 3.12+, CPU-only (no GPU). Runs on Linux, macOS, and Windows.

Quick Start

CLI

# Generate a config template
heterodyne-config --output config.yaml

# Run NLSQ optimization
heterodyne --method nlsq --config config.yaml

# Run Consensus Monte Carlo for uncertainty quantification
heterodyne --method cmc --config config.yaml

Python API

from heterodyne.data import load_xpcs_data
from heterodyne.config import ConfigManager
from heterodyne.optimization import fit_nlsq_jax
from heterodyne.optimization.cmc import fit_cmc_jax

# Load data and config
data = load_xpcs_data("experiment.hdf5")
config = ConfigManager("config.yaml")

# NLSQ trust-region optimization
nlsq_result = fit_nlsq_jax(data, config)

# CMC with NLSQ warm-start for Bayesian uncertainty
cmc_result = fit_cmc_jax(data, config, nlsq_result=nlsq_result)

Data Flow

YAML config --> XPCSDataLoader(HDF5) --> HeterodyneModel --> NLSQ or CMC --> Results (JSON + NPZ)

Optimization Methods

NLSQ (primary) -- JAX-native trust-region Levenberg-Marquardt with automatic anti-degeneracy defense, CMA-ES global search (including BIPOP bi-population restarts) for multi-scale problems, and memory-aware routing for large datasets. Post-fit quality assessment is available via validate_fit_quality in heterodyne.optimization.nlsq.validation (configurable chi-squared, residual, and parameter thresholds via FitQualityConfig).

CMC (secondary) -- Consensus Monte Carlo using NumPyro NUTS sampling with automatic sharding, NLSQ warm-start priors, and pluggable execution backends. Backends:

Produces publication-quality posterior distributions with ArviZ diagnostics. Configurable chain_method (sequential, vectorized, parallel) and combination_method (consensus, mixture) control sampling and aggregation. Built-in convergence checking (R-hat, ESS, BFMI) with actionable recommendations is available via check_convergence in heterodyne.optimization.cmc.diagnostics.

Configuration

Heterodyne uses YAML configuration files. Generate a template:

heterodyne-config --output config.yaml

Key sections:

experimental_data:
  file_path: "data.h5"
optimization:
  method: "nlsq"
  nlsq:
    anti_degeneracy:
      per_angle_mode: "auto"   # auto, constant, individual, fourier
  cmc:
    chain_method: "vectorized"      # sequential | vectorized | parallel
    combination_method: "consensus" # consensus | mixture
    sharding:
      max_points_per_shard: "auto"

CLI Commands

Plotting Shortcuts

# Inspect experimental data before fitting
hexp --config config.yaml

# Preview simulated C2 heatmaps from current parameters
hsim --config config.yaml --contrast 0.3 --offset-sim 1.0

Shell completion and additional aliases (e.g. het, het-nlsq, het-cmc) are available after running heterodyne-post-install --interactive.

Development

make test       # All tests
make test-fast  # Exclude slow tests
make quality    # Format + lint + type-check + shellcheck
make verify     # Full local CI verification before pushing

Documentation

• Full documentation — API reference, theory, user guide, architecture deep-dives
• Changelog

Citation

If you use Heterodyne in your research, please cite:

@article{He2024,
  author  = {He, Hongrui and Liang, Hao and Chu, Miaoqi and Jiang, Zhang and
             de Pablo, Juan J and Tirrell, Matthew V and Narayanan, Suresh
             and Chen, Wei},
  title   = {Transport coefficient approach for characterizing nonequilibrium
             dynamics in soft matter},
  journal = {Proceedings of the National Academy of Sciences},
  volume  = {121},
  number  = {31},
  year    = {2024},
  doi     = {10.1073/pnas.2401162121}
}

@article{He2025,
  author  = {He, Hongrui and Liang, Heyi and Chu, Miaoqi and Jiang, Zhang and
             de Pablo, Juan J and Tirrell, Matthew V and Narayanan, Suresh
             and Chen, Wei},
  title   = {Bridging microscopic dynamics and rheology in the yielding
             of charged colloidal suspensions},
  journal = {Proceedings of the National Academy of Sciences},
  volume  = {122},
  number  = {42},
  year    = {2025},
  doi     = {10.1073/pnas.2514216122}
}

License

MIT License -- see LICENSE for details.

Authors

• Wei Chen (weichen@anl.gov) -- Argonne National Laboratory
• Hongrui He (hhe@anl.gov) -- Argonne National Laboratory