Moju

Make AI models physically admissible and auditable.

View on GitHub

What is Moju?

Moju is a lightweight framework for Physics AI and Scientific Machine Learning that enables you to:

Why Moju?

Most Physics-Informed models stop at adding physics loss.

Moju goes further.

The Big Idea

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.

5-Minute Example


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"]
    

Core Concepts

Use Cases

Installation

pip install moju

Philosophy

Moju does not define physics.
Moju provides a structured way to enforce and audit it.