Moju

What is Moju?

Moju turns predicted state fields into governing-law residuals, physics losses, constitutive consistency checks, and audit reports. It is JAX-native at the core, with a PyTorch-facing interface through moju.torch.

• Compose physics from Models, Groups, and Laws.
• Use build_loss for governing-law training losses.
• Use audit and visualize for training/eval diagnostics.
• Infer law-linked constitutive checks, such as Fourier conduction implying thermal diffusivity.
• Apply the same residual and audit ideas to JAX and PyTorch workflows.

Install

pip install moju
pip install "moju[io]"
pip install "moju[torch,io]"

5-Minute Example: 1D Slab Cooling

Minimal Path B flow: pass state_pred directly. The derivatives are already provided, so no finite-difference inference is needed.

import jax.numpy as jnp
from moju.monitor import audit, build_loss, build_minimal_residual_engine, visualize

L = 0.02
rho, cp, k = 2700.0, 900.0, 200.0
alpha = k / (rho * cp)

x = jnp.linspace(0.0, L, 64)
t = jnp.ones_like(x) * 10.0
T = 300.0 + 20.0 * (1.0 - x / L) ** 2
T_laplacian = jnp.ones_like(x) * (40.0 / (L**2))
T_t = alpha * T_laplacian

state_pred = {
    "x": x,
    "t": t,
    "T": T,
    "T_t": T_t,
    "T_laplacian": T_laplacian,
    "L": jnp.ones_like(x) * L,
    "k": jnp.ones_like(x) * k,
    "rho": jnp.ones_like(x) * rho,
    "cp": jnp.ones_like(x) * cp,
    "alpha": jnp.ones_like(x) * alpha,
}

engine = build_minimal_residual_engine(
    law_names=["fourier_conduction"],
    coord_dimension=1,
)
residuals = engine.compute_residuals(state_pred, run_mode="training")
loss = build_loss(residuals)
report = audit(engine.log)
fig = visualize(engine.log, engine=engine)

The Fourier law automatically adds the law-linked thermal_diffusivity implied audit. fig is the Plotly dashboard.

Core Concepts

• Models - constitutive relationships.
• Groups - dimensionless quantities materialized into state.
• Laws - governing-equation residuals.
• ResidualEngine - runs laws, groups, constitutive audits, and optional data comparisons.
• audit - admissibility scoring from the log (RMS for laws, worst-point for constitutive closure).
• visualize - Plotly dashboards for residuals, categories, spatial fields, and constitutive diagnostics.
• moju.torch - PyTorch-facing residual engine, loss helpers, and JAX-law wrappers.

Learn More

• API and architecture overview
• Training vs eval behavior
• Law-linked implied constitutive audits
• Repository README
• Examples

Philosophy

Moju does not define physics for you. It provides a structured way to apply, measure, and audit the physics you already trust.