heterodyne 2.0.1

CPU-optimized JAX-based heterodyne scattering analysis for XPCS under nonequilibrium conditions

Heterodyne

[!IMPORTANT] Deprecation Notice: The homodyne and heterodyne packages have been merged into a single consolidated ecosystem: xpcsjax. This repository is now deprecated and maintained primarily for archival and reproducibility of the original PNAS 2024 and 2025 papers. All active development, performance optimizations, and future updates have moved to xpcsjax. We strongly recommend migrating your workflows to the new package.

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}$

Parameter	Description	Default	Units
D0_ref	Reference diffusion prefactor	1e4	Ų/sα+1
alpha_ref	Reference transport exponent	0.0	—
D_offset_ref	Reference transport rate offset	0.0	Ų/s

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

Parameter	Description	Default	Units
D0_sample	Sample diffusion prefactor	1e4	Ų/sα+1
alpha_sample	Sample transport exponent	0.0	—
D_offset_sample	Sample transport rate offset	0.0	Ų/s

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

Parameter	Description	Default	Units
v0	Velocity prefactor	1e3	Å/sβ+1
beta	Velocity exponent (0 = constant velocity)	0.0	—
v_offset	Velocity offset (negative for reversal)	0.0	Å/s

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

Parameter	Description	Default	Units
f0	Fraction amplitude	0.5	—
f1	Exponential rate (0 = constant fraction)	0.0	s⁻¹
f2	Time shift	0.0	s
f3	Baseline offset	0.0	—

Flow angle

Parameter	Description	Default	Units
phi0	Flow angle offset relative to q-vector	0.0	degrees

Per-angle scaling (2 parameters per detector angle)

Parameter	Description	Default
contrast	Optical contrast β (speckle contrast)	0.5
offset	Baseline offset	1.0

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:

Backend	Use case
CPUBackend	Single-node sequential or vectorized chains
MultiprocessingBackend	Parallel shards across CPU cores (default for chain_method="parallel")
WorkerPoolBackend / PersistentWorkerPool	Long analyses — amortises JIT cost across shards
PBSBackend	HPC cluster job-array submission (ALCF, NERSC)
PjitBackend	Experimental multi-device via pjit

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

Command	Short alias	Purpose
heterodyne	ht	Run XPCS analysis (NLSQ/CMC)
heterodyne-config	ht-config	Generate and validate config files
heterodyne-config-xla	ht-config-xla	Configure XLA device settings
heterodyne-validate	ht-validate	System validation
heterodyne-post-install	ht-post-install	Install shell completion (bash/zsh/fish)
heterodyne-cleanup	ht-cleanup	Remove shell completion files
hexp	—	Plot experimental data (skip optimization)
hsim	—	Plot simulated C2 heatmaps from config

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 ([e-mail]) -- Argonne National Laboratory
• Hongrui He ([e-mail]) -- Argonne National Laboratory