Overview

moju is a JAX-native library for physics-informed and physics-supervised SciML. The package provides moju.piratio (Groups, Models, Laws, Operators) and moju.monitor (ResidualEngine, build_loss, audit, visualize) so you can train or constrain models with the equations you trust.

Installation

pip install moju

Optional: pip install moju[units] for Pint-based units; pip install moju[dev] for pytest, black, ruff; pip install moju[ref] for xarray-based state_ref loaders (CFD snapshots → dict with optional interpolation); pip install moju[ref_vtk] for VTU/VTK loaders; pip install moju[ref_foam] for OpenFOAM loaders; pip install moju[ref_hdf5] for HDF5 loaders.

Quick start

import moju
from moju.piratio import Groups, Models

# Reynolds number and air density
Re = Groups.re(u=1.0, L=0.1, rho=1000.0, mu=1e-3)
rho = Models.ideal_gas_rho(P=101325.0, R=287.0, T=300.0)

API at a glance

Groups — Dimensionless numbers: Re, Pr, Nu, Ma, Gr, Pe, St, Bi, Fo, Sc, Le, Kn, Bo, Ca, Eu, Da, Ec.
Models — Viscosity (Sutherland, VFT, power-law), density (ideal gas, Boussinesq), heat (Stefan-Boltzmann, Fourier), friction (Darcy-Weisbach, Colebrook), and more.
Laws — Residuals for mass, momentum (Navier-Stokes, Stokes, Euler), heat/diffusion, Darcy/Brinkman, Poisson, wave, Burgers, Schrödinger, etc.
Operators — gradient, jacobian, divergence, laplacian, curl (2D/3D), time_derivative, time_derivative_second, advection, symmetric_gradient.
ResidualEngine (moju.monitor) — compute_residuals, build_loss, audit, visualize; single log for RMS and metrics.

All functions are JIT-compiled and support single-point or batched inputs (leading batch dimension). Use them inside loss functions and training loops with full JAX autodiff.

High-level architecture

User config (Laws, Groups, optional constitutive_audit / scaling_audit, Constants) and inputs (state_pred with predicted fields, optional state_ref for data residual and ref_delta) feed into ResidualEngine. By default, law-linked implied audits prepend extra audit rows derived from selected laws (see law_implied_audits.md); disable with MonitorConfig(law_implied_audits=False). Merged state = state_pred + constants + group outputs (output_key). Residuals: laws; constitutive / scaling closures (ref_delta, constitutive implied_delta, scaling π-constant on Path A); and data when state_ref is set. build_loss uses laws only; audit / visualize use the full log.

flowchart TB subgraph config [User config] Laws[Laws specs] Groups[Groups materialize Fo Bi etc] ConstAudit[constitutive_audit] ScaleAudit[scaling_audit] Constants[Constants] end subgraph inputs [Inputs] StatePred[state_pred fields] StateRef[state_ref optional] end subgraph core [ResidualEngine] Merge[Merge + run groups] LawsR[Laws residuals] ConstR[Constitutive closures] ScaleR[Scaling closures] DataR[data residual if state_ref] end StatePred --> Merge Constants --> Merge Groups --> Merge Merge --> LawsR Merge --> ConstR Merge --> ScaleR StateRef --> DataR config --> LawsR ConstAudit --> ConstR ScaleAudit --> ScaleR LawsR --> ResidualDict[Flat RMS keys laws constitutive scaling data] ConstR --> ResidualDict ScaleR --> ResidualDict DataR --> ResidualDict ResidualDict --> BuildLoss[build_loss laws only] BuildLoss --> PhysicsLoss[Physics loss scalar] PhysicsLoss -->|"user adds data loss"| TotalLoss[total_loss] ResidualDict --> Log[Log object] Log --> Audit[audit] Log --> Visualize[visualize]

Training and monitoring (ResidualEngine)

ResidualEngine: from moju.monitor import ResidualEngine, MonitorConfig, AuditSpec, audit_spec_to_engine_dict, build_loss, audit, visualize. Configure laws, groups, and optional constitutive_audit / scaling_audit audit specs tied to Models.* and Groups.* functions. Prefer the typed config (MonitorConfig, AuditSpec) for IDE autocompletion; serialize with to_dict()/from_dict() (AuditSpec.implied_fn is not serialized). ResidualEngine(config=MonitorConfig(...)) merges implied_fn from live AuditSpec objects; for hand-built dict specs use audit_spec_to_engine_dict(spec). Two entry paths: Path A builds state_pred via state_builder; Path B passes state_pred directly. Closure keys: constitutive/<name>/..., scaling/<name>/... (e.g. ref_delta, constitutive implied_delta, scaling pi_constant).

compute_residuals(...) — Path B: compute_residuals(state_pred, state_ref=None, ...). Path A: compute_residuals(state_pred=None, model=..., params=..., collocation=..., ...) with ResidualEngine(state_builder=...). Always runs laws. Optional Path B FD: auto_path_b_derivatives=True or a PathBGridConfig with fill_law_fd=True fills missing registered law inputs from primitives on the same grid (law_fd_recipes, list_law_fd_supported_laws()). For Models/Groups audits: ref_delta needs state_ref; constitutive implied_delta needs AuditSpec.implied_value_key or implied_fn (and a resolvable implied value); scaling audits do not use implied_delta. If nothing applies, the spec is skipped (optional omit log). Use required_state_keys() to see required keys. Logs per-key RMS. Use log_to_python=False inside jax.jit / jax.grad.
build_loss — Cascaded over laws only.
audit — R_norm(k) = RMS(r_k)/scale_k; scale is state-derived by default (stored in each log entry as entry["scale"]). Passing r_ref overrides scale for those keys (e.g. baseline or training curves). Admissibility: per-key scores in per_key; each per_category entry (laws, constitutive, scaling) is the geometric mean of scores in that category; overall_admissibility_score is the geometric mean of present categories. Optional PDF (pip install moju[report]). Report groups metrics under Governing laws, Constitutive relations, Scaling and similarity, and Data when applicable.
π-constant scaling closure (Path A) — On a scaling_audit spec, set invariance_pi_constant=True, non-empty invariance_compare_keys, and optional invariance_scale_c (default 10, must be > 1). Recipe parameters must live in ResidualEngine.constants. A built-in recipe exists for each registered scaling group (same ids as list_scaling_closure_ids()); discover them with list_pi_constant_group_names() and inspect moju.monitor.pi_constant_recipes.GROUP_PI_CONSTANT_RECIPES (rows use multiply_c / divide_c with optional integer power, e.g. We scales sigma by c²). Recipes are canonical one-parameter orbits, not unique. Grashof (gr): Groups.gr fixes g = 9.81 inside the function; the recipe only rescales user-supplied beta, dT, L, nu. Example scripts: examples/cookbook_pi_constant_reynolds.py, examples/cookbook_pi_constant_prandtl.py. Logs residual scaling/<name>/pi_constant (R_norm denominator uses mean absolute of the scaled-branch compare keys).
visualize — Multi-panel Plotly dashboard from the log alone (install plotly or moju[viz]): per-key RMS/admissibility, category and overall trajectories (including data/ keys), heatmap, top R_norm keys, closure-type summary, omitted/inferred counts, twin RMS/R_norm, category radar, and worst-key traces. Optional r_ref (same as audit).

Discover valid AuditSpec.name ids: AuditSpec.name must match the registry function id (e.g. sutherland_mu for Models.*, pe for Groups.*).

from moju.monitor import list_constitutive_models, list_scaling_closure_ids

print("Constitutive (Models.*) ids:", list_constitutive_models())
print("Scaling (Groups.*) ids:", list_scaling_closure_ids())

Derivative strategies (Path B / CFD): for governing laws, fill_path_b_derivatives / compute_residuals(..., auto_path_b_derivatives=...) with fill_law_fd=True can fill registered inputs such as phi_laplacian or u_grad from phi and u on structured grids when those keys are missing.

CFD snapshot cookbook: see examples/cfd_snapshot_cookbook_heat_1d.py for an end-to-end reference workflow (xarray snapshot → optional interpolation → finite-difference gradients + smoothing → weak-form audit → interpret RMS/admissibility).

Path B auto-FD cookbooks: examples/cookbook_path_b_fd_law_laplace.py (law input fill for laplace_equation).

Turbulence-related constitutive audits: examples/cookbook_turbulence_law_of_wall.py (Models.law_of_the_wall), examples/cookbook_turbulence_colebrook.py (Models.colebrook_friction), examples/cookbook_constitutive_smagorinsky.py (Models.smagorinsky_nu_t LES template), examples/cookbook_constitutive_k_epsilon.py (Models.k_epsilon_nu_t), and examples/cookbook_constitutive_k_omega.py (Models.k_omega_nu_t). The k–ε/k–ω examples audit algebraic νₜ only; transport PDEs for k, ε, ω are separate (Laws.* / custom).

Implied constitutive audit (implied_delta, no state_ref): examples/cookbook_constitutive_implied_ideal_gas_rho.py (implied_value_key), examples/cookbook_constitutive_implied_power_law_fn.py (implied_fn). Residual F(pred args) − implied; omitted if implied is missing.

Custom physics (optional `fn`)

In law or group specs, add "fn": your_callable; kwargs from state_map. Laws return the PDE residual. Groups return a value stored at output_key in merged state.

For extra closure residuals, use register-style audit lists on the engine.

import jax.numpy as jnp
from moju.monitor import ResidualEngine, build_loss

def my_residual(x):
    return x - 1.0

engine = ResidualEngine(
    constants={},
    laws=[{"name": "my_law", "state_map": {"x": "x"}, "fn": my_residual}],
    groups=[],
)
state = {"x": jnp.array(2.0)}
residuals = engine.compute_residuals(state)
loss = build_loss(residuals)

Before / After Moju

Corrected code snippets for the “Hardwiring reality into every iteration” comparison (pure JAX PINN vs physics-native moju): Before / After Moju (corrected).