Overview

moju is a physics supervision and audit library for SciML. The core is JAX-native: moju.piratio provides Groups, Models, Laws, and Operators, while moju.monitor provides ResidualEngine, build_loss, audit, and visualize. The moju.torch namespace exposes a PyTorch-facing engine and interop helpers.

Installation

pip install moju                 # residuals, audit, visualize, PDF
pip install "moju[io]"           # + state_ref file loaders
pip install "moju[studio]"       # Streamlit Moju Studio
pip install "moju[torch,io]"     # PyTorch + file loaders

API at a glance

Groups - Dimensionless numbers such as Re, Pr, Pe, Fo, Ma, Bi, Sc, Da, and related quantities.
Models - Constitutive relationships such as viscosity (including kinematic_viscosity_from_re for Burgers), density, diffusivity, heat transfer, friction, and turbulence closures.
Laws - Residuals for heat/diffusion, wave, mass, momentum, Darcy/Brinkman, Poisson, Burgers, Schrodinger, and related equations.
Operators - JAX autodiff helpers: gradient, jacobian, divergence, laplacian, curl, time derivatives, advection, and symmetric gradient.
ResidualEngine - Build residuals, physics losses, audit logs, and visual diagnostics.
moju.torch - PyTorch-facing residual engine, R_eff loss helpers, nondimensionalization helpers, and JAX-law wrappers.

High-level architecture

User config and state dictionaries flow through ResidualEngine. Groups materialize dimensionless quantities into merged state. Laws produce governing residuals. Constitutive audits compare catalog models against references or implied values. Optional data/ residuals run in eval mode when state_ref is supplied. build_loss uses law residuals for optimization; audit and visualize summarize the full log.

flowchart TB subgraph config [User config] Laws[Laws specs] Groups[Groups specs] Constitutive[Constitutive audits] Constants[Constants] end subgraph inputs [Inputs] StatePred[state_pred] StateRef[state_ref optional] end subgraph engine [ResidualEngine] Merge[Merge state constants groups] LawResiduals[Laws residuals] ConstitutiveResiduals[Constitutive residuals] DataResiduals[data residuals in eval] Log[Audit log] end StatePred --> Merge Constants --> Merge Groups --> Merge Merge --> LawResiduals Merge --> ConstitutiveResiduals StateRef --> DataResiduals Laws --> LawResiduals Constitutive --> ConstitutiveResiduals LawResiduals --> Log ConstitutiveResiduals --> Log DataResiduals --> Log LawResiduals --> BuildLoss[build_loss] Log --> Audit[audit] Log --> Visualize[visualize]

Training and monitoring (ResidualEngine)

Path A builds state_pred from (model, params, collocation) using a state_builder. Path B passes state_pred directly, which is convenient for CFD snapshots, saved neural-operator fields, or already-computed derivatives.

In the minimal 1D slab cooling example, Path B supplies T, T_t, T_laplacian, x, t, L, and material properties directly. No finite-difference inference is needed. The Fourier conduction law runs from those fields, and the default law-linked implied audit checks Models.thermal_diffusivity(k, rho, cp) against the diffusivity implied by the law fields. For Burgers, the default implied audit is Models.kinematic_viscosity_from_re(U, L, re) versus ν recovered by projecting inertial terms onto ∇²u (not dynamic viscosity μ).

compute_residuals - runs laws, groups, constitutive audits, and optional eval data comparisons.
build_loss - cascades over law residuals only, using the same robust R_eff reduction family as the monitor log.
audit - reports R_norm = R_eff / scale_k, per-key scores, category scores, and overall admissibility. Governing laws use RMS R_eff with law_scale_mode="auto" term-balance scale_k by default; constitutive implied_delta / ref_delta use worst-point r_max for admissibility (logged rms stays RMS for training). Path B: state_units="dimensional" for SI uploads.
visualize - Plotly dashboard for residuals, category scores, spatial fields, and constitutive divergence and consistency when closure debug data is available. The constitutive row includes a Divergence heatmap (normalised as (model − implied) / (|model| + ε)) and a Constitutive Consistency line plot with spatially varying ±0.1 % / ±0.5 % / ±1 % acceptability bands and tier boundary markers centred on the model prediction.

Path B can optionally fill missing law inputs from primitives on a structured grid. Default fill is finite-difference (auto_path_b_derivatives=True). For periodic Fourier-friendly fields (e.g. FNO grids), use PathBGridConfig(diff_method="spectral", periodic=True) or fill_path_b_spectral; temporal dt/dtt remain FD. Bare auto_path_b_derivatives=True never enables spectral. See path_b_derivatives.md.

Use monitor_training_vs_eval.md for run-mode details and law_implied_audits.md for the law-linked constitutive audit registry.

PyTorch support

Install moju[torch] to use the moju.torch namespace. It provides TorchResidualEngine, build_loss_torch, r_eff_scalar_torch, Torch-native nondimensionalization helpers, and wrap_law_torch for calling JAX laws from Torch tensors through jax2torch.

See scripts/torch_laws_jax2torch_example.py, tests/test_torch_engine.py, and tests/test_torch_interop.py for the current public surface.

Examples and guides

examples/slab_cooling_demo.py - full 1D slab cooling demo.
examples/cfd_snapshot_cookbook_heat_1d.py - xarray snapshot to eval audit.
examples/cookbook_path_b_fd_law_laplace.py - optional Path B finite-difference fill.
examples/cookbook_path_b_spectral_burgers.py - Path B periodic spectral fill for Burgers.
examples/cookbook_constitutive_divergence.py - constitutive divergence dashboard.
scripts/torch_laws_jax2torch_example.py - basic Torch interop.