Make AI models physically admissible and auditable.
View on GitHubMoju is a lightweight framework for Physics AI and Scientific Machine Learning that enables you to:
Most Physics-Informed models stop at adding physics loss.
Moju goes further.
Moju treats physics as composable building blocks:
Predictions (state_pred)
↓
Constitutive models
↓
Dimensionless groups (π)
↓
Governing laws
↓
Residuals → build_loss(residuals) → audit(engine.log)
Instead of manually wiring physics losses, Moju builds residuals for you and produces admissibility diagnostics.
import jax.numpy as jnp
from moju.monitor import ResidualEngine, build_loss, audit, MonitorConfig, AuditSpec
from moju.piratio import Models, Groups
mu0 = jnp.array(1.8e-5)
T0 = jnp.array(273.0)
S = jnp.array(110.4)
T = jnp.array(300.0)
mu = Models.sutherland_mu(T=T, mu0=mu0, T0=T0, S=S)
Re = jnp.array(10.0)
Pr = jnp.array(2.0)
Pe = Groups.pe(re=Re, pr=Pr)
cfg = MonitorConfig(
laws=[{"name": "laplace_equation", "state_map": {"phi_laplacian": "phi_xx"}}],
constitutive_audit=[
AuditSpec(
name="sutherland_mu",
output_key="mu",
state_map={"T": "T", "mu0": "mu0", "T0": "T0", "S": "S"},
predicted_spatial=["T"],
)
],
scaling_audit=[
AuditSpec(
name="pe",
output_key="Pe",
state_map={"re": "Re", "pr": "Pr"},
predicted_spatial=["Re"],
)
],
)
engine = ResidualEngine(config=cfg)
state_pred = {
"phi_xx": jnp.array(0.0),
"T": T,
"mu0": mu0,
"T0": T0,
"S": S,
"mu": mu,
"d_T_dx": jnp.array(0.0),
"d_mu_dx": jnp.array(0.0),
"Re": Re,
"Pr": Pr,
"Pe": Pe,
"d_Re_dx": jnp.array(0.0),
"d_Pe_dx": jnp.array(0.0),
}
residuals = engine.compute_residuals(state_pred)
loss = build_loss(residuals)
report = audit(engine.log)
print("Overall admissibility:", report["overall_admissibility_score"], report["overall_admissibility_level"])
print("Categories:", report["per_category"])
Output (keys you’ll get):
Physics Admissibility Report
Overall score: report["overall_admissibility_score"]
Overall level: report["overall_admissibility_level"]
Category scores: report["per_category"]["laws"|"constitutive"|"scaling"]
Per-key metrics: report["per_key"]
ref_delta, constitutive implied_delta, scaling pi_constant on Path A; optional Path B FD for registered law inputs; see Overview)implied_delta rows from selected laws (e.g. Fourier → thermal_diffusivity vs α from fields); law_implied_audits.mdpip install moju
Moju does not define physics.
Moju provides a structured way to enforce and audit it.