Metadata-Version: 2.4
Name: moju
Version: 0.4.0
Summary: Supervised Physics scaling and modeling for SciML.
Author-email: Ifimo Lab <lab@ifimoanalytics.com>
Maintainer-email: Ifimo Analytics <info@ifimoanalytics.com>
License: MIT
Keywords: physics,jax,sciml,pinn,dimensionless,cfd,differentiable,residual
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.0
Requires-Dist: jaxlib>=0.4.0
Requires-Dist: numpy>=1.20.0
Provides-Extra: units
Requires-Dist: pint>=0.20; extra == "units"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Provides-Extra: torch
Requires-Dist: torch>=2.0.0; extra == "torch"
Requires-Dist: jax2torch>=0.0.7; extra == "torch"
Provides-Extra: report
Requires-Dist: reportlab>=4.0; extra == "report"
Requires-Dist: optax>=0.2.0; extra == "report"
Provides-Extra: ref
Requires-Dist: xarray>=2023.1.0; extra == "ref"
Provides-Extra: ref-vtk
Requires-Dist: meshio>=5.3.0; extra == "ref-vtk"
Provides-Extra: ref-foam
Requires-Dist: meshio>=5.3.0; extra == "ref-foam"
Provides-Extra: ref-hdf5
Requires-Dist: h5py>=3.8.0; extra == "ref-hdf5"
Provides-Extra: ref-all
Requires-Dist: xarray>=2023.1.0; extra == "ref-all"
Requires-Dist: meshio>=5.3.0; extra == "ref-all"
Requires-Dist: h5py>=3.8.0; extra == "ref-all"
Requires-Dist: scipy>=1.10.0; extra == "ref-all"
Dynamic: license-file

# moju: Physics-AI supervision for engineering-grade simulations

```bash
pip install moju
```

**moju** helps you use AI for flow, heat, and other physics while keeping the math you trust at the center: Reynolds number, viscosity, conservation of mass and momentum. It is JAX-native and fully differentiable so you can use it in training loops or as a standalone toolkit. Whether you're new to AI or an experienced simulation engineer, you can run the examples in minutes. One library gives you dimensionless scaling, physical models, and differentiable residuals that check whether your fields satisfy the governing equations.

*Physics you know, in the AI you train. Dimensionless scaling, constitutive models, and equation residuals in one JAX library.*

If you work with flow, heat transfer, or similar physics and want to try AI without leaving the math you trust, moju is for you. Beginners can run the examples; experts can plug it into their training loops.

## Quick Start

Get your first result in under two minutes.

1. **Install:** run `pip install moju` (or from source: `pip install -e .` in the repo root).
   - Optional (reference data adapters): `pip install moju[ref]` for `xarray`-based `state_ref` loaders.
   - Optional (VTK/VTU loaders): `pip install moju[ref_vtk]` for `meshio`-based VTU/VTK ingestion.
   - Optional (OpenFOAM loaders): `pip install moju[ref_foam]` for `meshio`-based OpenFOAM snapshot ingestion (often via VTK/VTU exports).
   - Optional (HDF5 loaders): `pip install moju[ref_hdf5]` for `h5py`-based HDF5 ingestion.
2. **Run it:** open Python and paste the block below. It computes a Reynolds number and air density so you can verify the install and see moju in action.

Optional (reproducible conda-forge environment via pixi):

```bash
# Install pixi: https://pixi.sh/
pixi run -e dev pytest
pixi run -e ref python -c "import xarray; import moju.monitor.state_ref; print('ref extras OK')"
```

```python
import moju
from moju.piratio import Groups, Models

print("moju", moju.__version__)

# Reynolds number for water in a pipe (velocity 1 m/s, diameter 0.1 m)
Re = Groups.re(u=1.0, L=0.1, rho=1000.0, mu=1e-3)
print("Reynolds number:", Re)

# Air density at 1 bar, 300 K (ideal gas)
rho = Models.ideal_gas_rho(P=101325.0, R=287.0, T=300.0)
print("Air density (kg/m³):", rho)
```

## What's included

The package provides two namespaces: **moju.piratio** (Groups, Models, Laws, Operators) and **moju.monitor** (ResidualEngine, build_loss, audit, visualize).

**moju.piratio** — four modules:

| Module    | Core Function         | Example Output        |
| --------- | --------------------- | --------------------- |
| Operators | Differential Calculus | ∇u, ∇²T, ∇×u          |
| Models    | Physical Properties   | μ(T), ρ(P,T), k(T)    |
| Groups    | Dimensionless Scaling | Re, Pr, Pe, Ma        |
| Laws      | Conservation Logic    | R_momentum, R_energy  |

**Groups.** Scale your problem with the numbers you already use: Reynolds, Prandtl, Nusselt, Mach, and more (Re, Pr, Nu, Ma, …). JIT-compiled and differentiable; single values or batched.

**Models.** Ready-made physical relationships: viscosity (Sutherland, power-law), density (ideal gas, Boussinesq), heat transfer (Stefan-Boltzmann, Fourier), friction (Darcy-Weisbach). All differentiable for use in loss functions and training.

**Laws.** Check if a flow or temperature field satisfies the physics. You pass velocities, pressures, gradients; moju returns a residual. Zero means the conservation law is satisfied. Differentiable residuals for physics-informed loss terms. Covers mass, momentum (Navier-Stokes, Stokes, Euler), heat diffusion, Darcy flow, and more.

**Operators.** Derivatives for fields defined by a neural network: gradient, divergence, Laplacian, curl, time derivatives. Pass your network and collocation points; moju returns the derivatives via JAX autodiff. Single points or batched.

**moju.monitor** — **ResidualEngine**, build_loss, audit, visualize.

- **Residuals**: `compute_residuals(...)` returns residuals under `laws/…`, optional `constitutive/…`, `scaling/…`, and `data/…` (when `state_ref` is provided).
- **Two entry paths**:
  - **Path A (recommended)**: provide `(model, params, collocation)` and a `state_builder` so moju can build `state_pred` (and derivative keys) consistently.
  - **Path B (advanced)**: provide `state_pred` directly (including derivative keys like `d_T_dx`, `d_mu_dt` when chain closures apply).
- **Constitutive and scaling/similarity audits** are tied to **Models.* and Groups.* functions** via audit specs, with three closure types:
  - `ref_delta` (requires `state_ref`)
  - `chain_dx` (spatial chain rule; requires spatially varying inputs + derivative keys)
  - `chain_dt` (temporal chain rule; requires time-varying inputs + derivative keys)
  Items with **no spatial and no temporal variation** are omitted from the report.
- **Physics loss**: **build_loss** uses **laws only**.
- **Audit**: per-key **R_norm(k) = RMS(r_k) / scale_k** and admissibility(k) = 1/(1 + R_norm(k)). Scale is **state-derived by default** (from merged state and specs for laws/constitutive/scaling/data keys); each log entry stores `entry["scale"]`. Passing **r_ref** to `audit(log, r_ref=...)` overrides scale for those keys (e.g. baseline or training-curve reference). Category scores for Governing laws / Constitutive / Scaling-similarity; overall score is a **geometric mean** across present categories (empty categories excluded).
- **Typed config**: use `MonitorConfig` and `AuditSpec` to define audits with IDE-friendly autocompletion, and serialize with `to_dict()` / `from_dict()`.
- **Introspection**: use `engine.required_state_keys()` and `engine.required_derivative_keys()` to see exactly which keys must be present in `state_pred`.

**Custom laws and groups.** In law or group specs, add optional `"fn": your_callable`; kwargs come from `state_map`. **Laws**: return the PDE residual (used in `build_loss`). **Groups**: return value is written to `output_key` in merged state (e.g. for Fo, Bi).

### Monitor in minutes

Minimal typed config:

```python
import jax.numpy as jnp
from moju.monitor import AuditSpec, MonitorConfig, ResidualEngine

cfg = MonitorConfig(
    laws=[{"name": "laplace_equation", "state_map": {"phi_laplacian": "phi_xx"}}],
    scaling_audit=[
        AuditSpec(
            name="pe",
            output_key="Pe",
            state_map={"re": "Re", "pr": "Pr"},
            predicted_spatial=["Re", "Pr"],
        )
    ],
)
engine = ResidualEngine(config=cfg)

state_pred = {"phi_xx": jnp.array(0.0), "Re": 10.0, "Pr": 2.0, "Pe": 20.0, "d_Re_dx": 1.0, "d_Pr_dx": 0.0, "d_Pe_dx": 2.0}
residuals = engine.compute_residuals(state_pred)
print(sorted(engine.required_state_keys()))
print(sorted(engine.required_derivative_keys()))
```

Run the minimal chain demos:

```bash
python examples/monitor_chain_spatial_demo.py
python examples/monitor_chain_temporal_demo.py
```

Run end-to-end examples (NN → state → engine → PDF):

```bash
python examples/monitor_heat_end_to_end.py
python examples/monitor_burgers_end_to_end.py
```

### Using high-fidelity CFD as `state_ref` (xarray)

If you have CFD or experimental data on a labeled grid, you can ingest it into a `state_ref` dict and optionally interpolate it onto your collocation grid.

```python
import numpy as np
import xarray as xr
from moju.monitor.state_ref import from_xarray

ds = xr.Dataset(
    data_vars={
        "T_cfd": (("t", "x"), np.random.randn(5, 8)),
    },
    coords={"t": np.linspace(0.0, 1.0, 5), "x": np.linspace(0.0, 1.0, 8)},
)

state_ref = from_xarray(
    ds,
    var_map={"T": "T_cfd"},
    target={"t": np.linspace(0.0, 1.0, 11), "x": np.linspace(0.0, 1.0, 21)},
    method="linear",
)
```

### Reference loaders (VTK/VTU, OpenFOAM, HDF5)

Additional thin adapters are available via optional extras:

- **VTK/VTU** (`meshio`): `pip install moju[ref_vtk]` then use `from_vtu(...)` / `from_vtk(...)` (unstructured; no automatic interpolation).
- **OpenFOAM** (`meshio`): `pip install moju[ref_foam]` then use `from_openfoam(...)` (often easiest after exporting to VTK/VTU).
- **HDF5** (`h5py`): `pip install moju[ref_hdf5]` then use `from_hdf5(...)`.

Templates:

- `examples/monitor_state_ref_from_vtu_demo.py`
- `examples/monitor_state_ref_from_openfoam_demo.py`
- `examples/monitor_state_ref_from_hdf5_demo.py`

### Derivative strategies for Path B / CFD

For `chain_dx` / `chain_dt` closures, Path B workflows must provide derivative keys like `d_T_dx`, `d_mu_dt`, `d_Re_dx`.

- **Solver gradients (best when available)**: export gradients directly from the CFD solver / adjoint / postprocessing.
- **Finite differences**: central differences in the interior, one-sided at boundaries; use coordinate-aware spacing on nonuniform grids.
- **Smoothing before differencing**: denoise fields then differentiate (reduces noise, but can introduce bias/edge artifacts).

Planned extension point: **weak-form / integrated chain closures** can reduce noise sensitivity by integrating closure residuals over space/time windows using quadrature. A future API hook is `closure_mode="pointwise"|"weak"` plus `quadrature_weights`.

### CFD snapshot cookbook (end-to-end)

An end-to-end, copy/paste workflow for CFD snapshots (structured 1D) is provided in:

- `examples/cfd_snapshot_cookbook_heat_1d.py`

It demonstrates:
- **Load** a snapshot into `xarray.Dataset` (or `xarray.open_dataset(...)` for NetCDF)
- **Regrid/interpolate** to collocation points via `from_xarray(..., target=...)`
- **Compute gradients** via coordinate-aware finite differences (with optional smoothing)
- **Audit** using **weak-form** (`closure_mode="weak"`) with quadrature weights (`w_x`)
- **Interpret** the resulting RMS + admissibility score

## Examples

### First example

The Quick Start block above is enough to verify the install. Below are further examples.

### More scaling and physical models

```python
from moju.piratio import Groups, Models

# Dimensionless numbers (single values or arrays)
Re = Groups.re(u=1.0, L=0.1, rho=1000.0, mu=1e-3)   # Reynolds
Pr = Groups.pr(mu=1e-3, cp=4186.0, k=0.6)            # Prandtl (water)
Nu = Groups.nu(h=100.0, L=0.1, k=0.6)                # Nusselt
Ma = Groups.ma(u=100.0, a=343.0)                     # Mach number

# Physical models
mu_air = Models.sutherland_mu(T=300.0, mu0=1.8e-5, T0=273.0, S=110.4)  # Air viscosity
q_rad = Models.stefan_boltzmann_flux(epsilon=0.9, T=400.0)             # Radiative heat flux
nu = Models.kinematic_viscosity(mu=1e-3, rho=1000.0)                   # Kinematic viscosity
```

### Checking physics (Laws)

Use **Laws** to check whether a velocity field satisfies incompressible mass conservation (div u = 0). You pass the velocity gradient; moju returns a residual. Zero when the law is satisfied. In a full setup you obtain gradients from **Operators** and feed them into Laws to build physics-informed loss terms.

```python
import jax.numpy as jnp
from moju.piratio import Laws

# Velocity gradient for a flow that preserves volume (trace = 0)
# Example: constant velocity field -> gradient is zero
u_grad = jnp.array([[0.0, 0.0], [0.0, 0.0]])
residual = Laws.mass_incompressible(u_grad)
print("Mass residual (should be 0):", residual)
```

### Derivatives (Operators)

**Operators** compute derivatives of a function, e.g. a scalar or vector field from a neural network. Here we use a trivial scalar; in practice you pass your network and collocation points.

```python
import jax.numpy as jnp
from moju.piratio import Operators

# A simple scalar function of x (in practice this would be your neural network)
def scalar_field(params, x):
    return jnp.sum(x**2)

params = {}
x = jnp.array([1.0, 2.0])

grad = Operators.gradient(scalar_field, params, x)
print("Gradient of sum(x²) at [1, 2]:", grad)

lap = Operators.laplacian(scalar_field, params, x)
print("Laplacian at [1, 2]:", lap)
```

## Going further

moju is JAX-native, JIT-compiled, and fully differentiable. It supports a broad range of physics AI workflows: surrogate modeling, inverse problems, physics-informed training, digital twins, hybrid solvers, and anywhere else physics and machine learning meet. Residuals and operators integrate with JAX autodiff so you can train or constrain models to satisfy the equations. We build on the principle that **physics is the ground truth** and provide the "glass box" transparency needed to deploy AI in high-stakes settings (thermal management, flow simulation, and beyond). Versioning follows [VERSIONING.md](VERSIONING.md).

## License

MIT License. Open for the community. Developed by Ifimo Lab, a division of Ifimo Analytics.
