mimirax¶

mimirax is the inference and reconstruction layer of the MIDGARD stack: it fits, samples and diagnoses the parameters of differentiable N-body forward models. It owns observation operators, likelihoods, priors, parameter handling and inference methods, and it consumes forces and integrators through protocols rather than importing any particular solver.

The pipeline is:

parameters (any pytree)
    -> reparameterization / gauge fixing     (parameters)
    -> forward model                          (a ForceModel or rollout behind a protocol)
    -> observable                             (observables: state -> data space)
    -> likelihood + priors                    (likelihoods, priors)
    -> optimizer / sampler                    (inference)
    -> residuals, traces, degeneracies        (diagnostics)

Installation¶

pip install -e ".[dev]"      # library + quality tooling; no solver needed
pip install -e ".[docs]"     # to build this documentation

Solver glue is behind extras — mimirax[jaccpot], mimirax[nornax], mimirax[odisseo] — and each needs the sibling package installed from GitHub first; see the README.

Quick start¶

import jax
import jax.numpy as jnp
import optax

from mimirax import GaussianLikelihood, IdentityObservable, InferenceProblem, OptaxOptimizer
from mimirax.testing import SoftenedPointMassField

tracers = 3.0 * jax.random.normal(jax.random.PRNGKey(0), (12, 3))

def forward(source):
    field = SoftenedPointMassField(source, jnp.ones(1), softening=0.2)
    return field.accelerations(tracers, jnp.ones(12)).reshape(-1)

observed = forward(jnp.asarray([[0.4, -0.3, 0.2]]))
problem = InferenceProblem(forward, IdentityObservable(), GaussianLikelihood(0.01), observed)
fit = OptaxOptimizer(optax.adam(0.02), num_steps=1500).minimize(
    problem.negative_log_posterior, jnp.zeros((1, 3))
)

Documentation