Welcome to RheoJAX’s documentation!¶

RheoJAX is a unified, JAX-accelerated rheological analysis package that provides a modern, high-performance framework for analyzing experimental rheology data. Built on JAX for automatic differentiation and GPU acceleration, RheoJAX combines powerful numerical capabilities with an intuitive API for seamless end-to-end analysis.

At a Glance¶

Models	53 models across 22 families (classical, fractional, flow, constitutive ODE, transient network, vitrimer, and more)
Bayesian	Full MCMC inference (NumPyro NUTS) with NLSQ warm-start for all models
Transforms	11 transforms (FFT, mastercurve/TTS, SRFS, SPP, OWChirp, derivatives, mutation number, Prony conversion, spectrum inversion, LVE envelope, Cox-Merz)
Performance	10-100x speedups via JAX; 5-270x optimization via NLSQ
DMTA/DMA	Automatic E* ↔ G* conversion for 45 oscillation-capable models with tensile deformation mode
Notebooks	249 tutorial notebooks across 21 categories covering all model families and protocols

Key Features¶

JAX-First Architecture: All numerical operations use JAX for automatic CPU/GPU dispatch, providing 10-100x speedups with automatic differentiation for optimization
Comprehensive Data Support: Automatic test mode detection (relaxation, creep, oscillation, rotation) with support for multiple file formats (TRIOS, CSV, Excel, Anton Paar)
Flexible Parameter System: Type-safe parameter management with bounds, constraints, and optimization support
Publication-Quality Visualization: Matplotlib-based plots with three built-in styles (default, publication, presentation)
Extensible Design: Plugin system for custom models and transforms with registry-based discovery

Quick Start¶

Installation¶

pip install rheojax

Basic Usage¶

from rheojax.io.readers import auto_read
from rheojax.visualization import plot_rheo_data
import matplotlib.pyplot as plt

# Load data (auto-detect format)
data = auto_read("stress_relaxation.txt")

# Check detected test mode
print(f"Test mode: {data.test_mode}")  # relaxation

# Visualize
fig, ax = plot_rheo_data(data, style='publication')
plt.show()

Working with Parameters¶

from rheojax.core import ParameterSet

# Create parameter set
params = ParameterSet()
params.add("E", value=1000.0, bounds=(100, 10000), units="Pa")
params.add("tau", value=1.0, bounds=(0.1, 100), units="s")

# Get/set values
E = params.get_value("E")
params.set_value("tau", 2.5)

Contents¶

User Guide

User Guide: Graduate Student Learning Pathway

API Reference

API Reference

Tutorial Notebooks

Tutorial Notebooks

Model Verification

Model Verification Reports

Additional Resources

Community and Support¶

Getting Help

Read the User Guide: Graduate Student Learning Pathway
Report issues on GitHub
Contact: wchen@anl.gov

Contributing

See Contributing to RheoJAX for guidelines
Check Architecture Overview for design principles

Citation

If you use RheoJAX in your research, please cite:

@software{rheojax2024,
  title = {RheoJAX: JAX-Powered Rheological Analysis with Bayesian Inference},
  year = {2024-2026},
  author = {Wei Chen},
  url = {https://github.com/imewei/rheojax},
  version = {0.6.1}
}

Indices and Tables¶

License¶

RheoJAX is released under the MIT License. See LICENSE file for details.