xpcsjax — JAX-native XPCS NLSQ fitting

xpcsjax logo JAX-native Python 3.12+ License

xpcsjax is a unified JAX-native package for non-linear least-squares (NLSQ) fitting of X-ray Photon Correlation Spectroscopy (XPCS) data. v0.1 consolidates the homodyne and heterodyne NLSQ pipelines from the upstream homodyne and heterodyne packages into a single JAX-first codebase with a small, lazy-loaded public API.

Scope of v0.1

xpcsjax is NLSQ-only by design. Bayesian sampling — NumPyro, BlackJAX, ArviZ, CMC (Consensus Monte Carlo), NUTS, HMC, parallel tempering — is out of scope for this package and will not be added. Users needing Bayesian XPCS analysis should use the upstream homodyne or heterodyne packages.

A 30-second tour

from xpcsjax import load_xpcs_data, fit_nlsq

data   = load_xpcs_data("config.yaml")
result = fit_nlsq(data, "config.yaml")

print(result.parameters)
print(result.reduced_chi_squared)

The same two-function workflow drives both homodyne and heterodyne modes; the ConfigManager decides which physics model and which NLSQ strategy to use, based on the analysis_mode field in your YAML.

Where to go next

Install & first fit

Set up a uv-managed environment and run the homodyne quickstart against the bundled example config.

Quickstart
User guide

Data loading, analysis modes, NLSQ strategies, and how to read an OptimizationResult.

User guide
Theory

Homodyne and heterodyne models, anti-degeneracy, transport-coefficient formalism, and the rationale behind the 5-layer defence system.

Theory and physics
API reference

The seven lazy public symbols, plus the submodule surface that backs them. Auto-generated from the live source.

API reference

Public API at a glance

The package exposes seven symbols, all lazy-loaded via a module-level __getattr__:

Symbol

Purpose

xpcsjax.data.xpcs_loader.load_xpcs_data()

Read XPCS HDF5 + YAML config into a homodyne/heterodyne-ready dict.

xpcsjax.optimization.nlsq.fit_nlsq()

Run the NLSQ fit. Dispatches to homodyne or heterodyne path based on analysis_mode.

xpcsjax.config.ConfigManager

Load, validate, and query a YAML config; canonical source of bounds and active-parameter lists.

xpcsjax.core.HomodyneModel

Hybrid stateful + JIT model for static / laminar-flow homodyne XPCS.

xpcsjax.core.HeterodyneModel

Two-component (sample + reference) heterodyne model with 14 physics parameters.

xpcsjax.optimization.nlsq.results.OptimizationResult

Dataclass holding fitted parameters, covariance, χ², diagnostics, and quality flags.

xpcsjax.viz.nlsq_plots.generate_nlsq_plots()

Generate diagnostic plots and serialize NPZ + JSON artifacts for every φ angle after fitting. Supports homodyne and heterodyne (individual per-angle mode). Optional Datashader fast path + parallel rendering.

See Public API for the full public surface and API reference for the per-submodule autodoc reference.

Documentation map

Project

Indices