API reference¶

The public API is re-exported from the top-level mimirax package. The modules below are documented from their source docstrings.

Contracts and types¶

The contracts mimirax is written against, and nothing else.

WHY PROTOCOLS. The package’s reason to exist is to outlive any single solver: its core must never import jaccpot, nornax or odisseo. So the core is written against structural interfaces – anything with the right methods satisfies them, whether it is jaccpot’s FMM, an ODISSEO rollout, an analytic potential, or the doubles in mimirax.testing. This is how the ecosystem already works: jaccpot/nornax_adapter.py satisfies nornax’s MutualForceModel without importing nornax, and nornax’s forces/base.py is the template these follow.

Every protocol here is runtime_checkable so that isinstance (and the optional beartype hook) can validate an argument structurally. Optional capabilities are deliberately not protocol members: adding one would make isinstance reject every implementation that predates it. An implementation advertises an optional capability by having the attribute; a consumer probes with hasattr or inspect.signature, never by widening the protocol.

Differentiability is part of every contract below: unless a docstring says otherwise, an implementation must be traceable by jax.grad in its array arguments, because the core numerical paths differentiate through it.

class mimirax.protocols.Constraint(*args, **kwargs)[source]¶

Bases: Protocol

An equality constraint, expressed as a defect that is zero when satisfied.

nornax’s multiple-shooting shooting_defect is one; a fixed total mass is another. Consumed by mimirax.parameters.quadratic_penalty().

defect(params)[source]¶

Return the constraint defect at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: Any shape; all zeros means the constraint holds.
Return type:: Array

class mimirax.protocols.FoldableObservable(*args, **kwargs)[source]¶

Bases: Protocol

An Observable that can also be accumulated one snapshot at a time.

WHY THIS EXISTS. A time-averaged observable is a reduction over a trajectory, and the cheap way to evaluate a reduction is to fold it into the integration rather than to stack the trajectory and average afterwards. Stacking costs O(t * n) memory in the forward pass alone – for a made-to-measure fit over many dynamical times that is the term that decides whether the problem fits on a machine, and it has nothing to do with autodiff. Folding costs O(m).

So this is a capability, not a replacement. It is a separate protocol because Observable is runtime_checkable and widening it would make isinstance reject every observable that predates this; and it is narrower than Observable because not every observable is foldable. A mean folds. Exponential smoothing folds, through the recursion state <- a * state + value. A median does not, and neither does anything needing the whole sample at once.

An implementation must satisfy the fold law: for any trajectory,

result(fold(… fold(initial(s_0), value(s_0)) …, value(s_{t-1})), t)

equals self(stacked_trajectory), the plain Observable call. The two paths are then interchangeable and a caller can choose on memory grounds alone. mimirax.observables.TimeAverage satisfies both protocols and a test pins the law to round-off.

value is separated from fold on purpose: a rollout can then evaluate value inside its own loop – where the snapshot already is – and stack only the (m,) results instead of the (n, 3) state.

value(snapshot)[source]¶

Evaluate the observable at one snapshot.

Parameters:: snapshot (PyTree) – A single time slice, with no leading time axis.
Returns:: (m,) observable values at that snapshot.
Return type:: Array

initial(snapshot)[source]¶

Return the zero accumulator, shaped from a template snapshot.

Must not depend on the snapshot’s values, only on its shapes and dtypes, so that it can be built with jax.eval_shape() and carries no gradient.

Parameters:: snapshot (PyTree) – A template time slice.
Returns:: The accumulator’s identity element.
Return type:: PyTree

fold(state, value)[source]¶

Accumulate one snapshot’s value into the running state.

Called in trajectory order, so an implementation may weight by recency. Must be a fixed-shape pytree map, since it runs as a lax.scan carry.

Parameters:

state (PyTree) – The running accumulator.
value (Array) – (m,) result of value() at the next snapshot.

Returns:

The updated accumulator.

Return type:

PyTree

result(state, num_steps)[source]¶

Finish the reduction, given how many snapshots went into it.

Parameters:

state (PyTree) – The final accumulator.
num_steps (int) – How many snapshots were folded. Static, so a normalization may depend on it.

Returns:

(m,) reduced observable, equal to the plain Observable call on the stacked trajectory.

Return type:

Array

class mimirax.protocols.ForceModel(*args, **kwargs)[source]¶

Bases: Protocol

A gravitational (or any pairwise) acceleration provider.

The smallest interface a solver has to offer for mimirax to build a forward model on it: accelerations at positions due to masses. jaccpot’s FastMultipoleMethod.compute_accelerations(positions, masses) satisfies it through mimirax.adapters.jaccpot; an analytic potential satisfies it directly; mimirax.testing.DirectSumGravity is the double.

Differentiable in positions and masses. A tree or FMM backend is differentiable at fixed topology only; where its topology is rebuilt is the backend’s business and must be documented on the adapter.

accelerations(positions, masses, *, args=None)[source]¶

Return the acceleration at every particle.

Parameters:

positions (Vec3) – (n, 3) particle positions.
masses (PerParticle) – (n,) particle masses.
args (object) – Caller-owned channel for backend options; the core never sets it.

Returns:

(n, 3) accelerations.

Return type:

Vec3

class mimirax.protocols.ForwardModel(*args, **kwargs)[source]¶

Bases: Protocol

Parameters in, a physical state (or prediction) out.

This is the seam between inference and physics: everything above it (observables, likelihoods, priors, samplers) is mimirax’s; everything below it (forces, integrators) is a solver’s. A forward model may be a static map (positions -> field at tracers, as in Paper I section 7), a rollout (initial conditions -> final state, as in ODISSEO), or the identity.

Must be differentiable in params.

class mimirax.protocols.Likelihood(*args, **kwargs)[source]¶

Bases: Protocol

Compares a prediction with an observation and returns a log-probability.

A misfit is the negative of this; the convention is fixed to log-densities so that likelihoods and priors add. Registered in mimirax.likelihoods.LIKELIHOODS. Must be differentiable in predicted.

log_prob(predicted, observed)[source]¶

Return the log-likelihood of observed given predicted.

Parameters:

predicted (Array) – What an Observable produced from the model.
observed (Array) – The data, same shape as predicted.

Returns:

The log-likelihood; additive constants may be dropped, and the implementation’s docstring must say if they are.

Return type:

Scalar

class mimirax.protocols.Observable(*args, **kwargs)[source]¶

Bases: Protocol

An observation operator: a state in, a data-space array out.

Kinematic maps, projections, binned densities and stream-track coordinates are all observables. Registered in mimirax.observables.OBSERVABLES. Must be differentiable in state.

class mimirax.protocols.Optimizer(*args, **kwargs)[source]¶

Bases: Protocol

Minimizes a scalar objective over a parameter pytree.

Registered in mimirax.inference.METHODS.

minimize(objective, params)[source]¶

Minimize objective starting from params.

Parameters:

objective (Callable[[PyTree], Scalar]) – A differentiable scalar function of the parameters.
params (PyTree) – The starting point; its structure is preserved in the result.

Returns:

Final parameters and the per-step objective trace.

Return type:

FitResult

class mimirax.protocols.Prior(*args, **kwargs)[source]¶

Bases: Protocol

A log-density over parameters. Regularizers are priors that say so.

Registered in mimirax.priors.PRIORS. Must be differentiable in params.

log_prob(params)[source]¶

Return the log prior density at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: The log density, up to an additive constant.
Return type:: Scalar

class mimirax.protocols.Reparameterization(*args, **kwargs)[source]¶

Bases: Protocol

A bijection between an unconstrained and a constrained parameter space.

Optimizers and samplers work in the unconstrained space; the forward model sees the constrained one. Implementations live in mimirax.parameters.

forward(unconstrained)[source]¶

Map unconstrained values to the constrained space.

Parameters:: unconstrained (Array) – Values in the unconstrained space.
Returns:: The constrained values, same shape.
Return type:: Array

inverse(constrained)[source]¶

Map constrained values back to the unconstrained space.

Parameters:: constrained (Array) – Values in the constrained space.
Returns:: The unconstrained values, same shape.
Return type:: Array

log_abs_det_jacobian(unconstrained)[source]¶

Return log abs det J of forward() at unconstrained.

Parameters:: unconstrained (Array) – Where to evaluate the Jacobian.
Returns:: The summed log absolute Jacobian determinant, the term a density picks up under the change of variables.
Return type:: Scalar

class mimirax.protocols.Sampler(*args, **kwargs)[source]¶

Bases: Protocol

Draws parameter samples from an unnormalized log-density.

Registered in mimirax.inference.METHODS.

sample(log_density, params, *, key, num_samples)[source]¶

Draw num_samples samples of log_density starting from params.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target density.
params (PyTree) – The initial state of the chain.
key (Array) – A jax.random key.
num_samples (int) – How many samples to return. Static.

Returns:

The samples and their log-densities.

Return type:

SampleResult

Result containers shared by the inference protocols and their implementations.

Kept in their own module (as yggdrax/types.py is) so that mimirax.protocols can reference them without importing any implementation, and so that an adapter satisfying a protocol structurally can return them without importing mimirax.inference.

class mimirax.types.FitResult(params, objective_trace)[source]¶

Bases: NamedTuple

What an optimizer returns.

Variables:

params (PyTree) – The final parameters, in the same structure as the initial ones.
objective_trace (Array) – (k,) objective value at every step, so convergence can be judged from the result alone.

Parameters:

params (Any)
objective_trace (Array)

params: Any¶: Alias for field number 0

objective_trace: Array¶: Alias for field number 1

class mimirax.types.SampleResult(samples, log_density)[source]¶

Bases: NamedTuple

What a sampler returns.

Variables:

samples (PyTree) – Parameter samples with a leading k axis on every leaf.
log_density (Array) – (k,) log target density at each sample.

Parameters:

samples (Any)
log_density (Array)

samples: Any¶: Alias for field number 0

log_density: Array¶: Alias for field number 1

Parameters¶

Parameter pytrees, reparameterizations, gauge fixing and constraints.

The free parameters of an inference problem are whatever pytree the caller chose – a bare (n, 3) array of positions, a dict of potential parameters, a nested structure. Nothing here assumes a structure; every function works through jax.tree_util or through a flat vector obtained with ravel().

Three concerns live here because they are about the parameter space, not about any model:

Reparameterizations – bijections between an unconstrained space, where optimizers and samplers move freely, and the constrained space a forward model expects (positive masses, bounded angles). Each is a Reparameterization and is registered in REPARAMETERIZATIONS; ODISSEO’s Plummer_sphere_reparam is the ecosystem’s existing instance of the pattern.
Gauge fixing – removing directions the data cannot see. A self-gravitating system’s centre of mass and net momentum are the standing examples: a fit over particle positions that leaves them free wanders along a flat direction of the objective.
Constraints – equalities expressed as a defect that is zero when the constraint holds (nornax’s multiple-shooting shooting_defect is one), and the quadratic penalty that turns a defect into an objective term.

Everything is differentiable in its array arguments; nothing here has a static-argument requirement under jit except the transforms mapping, whose keys select code paths.

class mimirax.parameters.AffineTransform(scale=1.0, shift=0.0)[source]¶

Bases: object

A fixed rescaling and shift: constrained = scale * u + shift.

Used to put parameters of very different magnitude on one footing before an optimizer sees them.

Variables:

scale (float) – Multiplicative factor; must be non-zero.
shift (float) – Additive offset.

Parameters:

scale (float)
shift (float)

scale: float = 1.0¶

shift: float = 0.0¶

forward(unconstrained)[source]¶

Apply the affine map.

Parameters:: unconstrained (Array) – Real values.
Returns:: scale * unconstrained + shift.
Return type:: Array

inverse(constrained)[source]¶

Undo the affine map.

Parameters:: constrained (Array) – Real values.
Returns:: (constrained - shift) / scale.
Return type:: Array

log_abs_det_jacobian(unconstrained)[source]¶

Return size * log|scale|.

Parameters:: unconstrained (Array) – Only its size is used.
Returns:: The summed log-Jacobian.
Return type:: Scalar

class mimirax.parameters.FunctionConstraint(defect_fn)[source]¶

Bases: object

A Constraint wrapping a plain defect function.

Variables:: defect_fn (Callable[[PyTree], Array]) – Maps parameters to a defect array that is zero when the constraint holds.
Parameters:: defect_fn (Callable[[Any], Array])

defect_fn: Callable[[Any], Array]¶

defect(params)[source]¶

Evaluate the wrapped defect function.

Parameters:: params (PyTree) – The free parameters.
Returns:: The defect.
Return type:: Array

class mimirax.parameters.IdentityTransform[source]¶

Bases: object

The trivial reparameterization: constrained equals unconstrained.

forward(unconstrained)[source]¶

Return unconstrained unchanged.

Parameters:: unconstrained (Array) – Values in the unconstrained space.
Returns:: The same values.
Return type:: Array

inverse(constrained)[source]¶

Return constrained unchanged.

Parameters:: constrained (Array) – Values in the constrained space.
Returns:: The same values.
Return type:: Array

log_abs_det_jacobian(unconstrained)[source]¶

Return zero: the identity has unit Jacobian.

Parameters:: unconstrained (Array) – Ignored except for its dtype.
Returns:: 0.
Return type:: Scalar

class mimirax.parameters.LogTransform[source]¶

Bases: object

Positive constrained values: constrained = exp(unconstrained).

The natural choice for masses, scale radii and noise levels.

forward(unconstrained)[source]¶

Exponentiate.

Parameters:: unconstrained (Array) – Real values.
Returns:: Strictly positive values, same shape.
Return type:: Array

inverse(constrained)[source]¶

Take the logarithm.

Parameters:: constrained (Array) – Strictly positive values.
Returns:: Real values, same shape.
Return type:: Array

log_abs_det_jacobian(unconstrained)[source]¶

Return sum(unconstrained), since d exp(u) / du = exp(u).

Parameters:: unconstrained (Array) – Where to evaluate.
Returns:: The summed log-Jacobian.
Return type:: Scalar

class mimirax.parameters.SoftplusTransform[source]¶

Bases: object

Positive constrained values with a linear tail: softplus(u).

Unlike LogTransform this does not blow up for large u, which makes it the safer choice under aggressive optimizer steps.

forward(unconstrained)[source]¶

Apply log(1 + exp(u)).

Parameters:: unconstrained (Array) – Real values.
Returns:: Strictly positive values, same shape.
Return type:: Array

inverse(constrained)[source]¶

Apply log(exp(c) - 1).

Parameters:: constrained (Array) – Strictly positive values.
Returns:: Real values, same shape.
Return type:: Array

log_abs_det_jacobian(unconstrained)[source]¶

Return sum(log sigmoid(u)), the log-derivative of softplus.

Parameters:: unconstrained (Array) – Where to evaluate.
Returns:: The summed log-Jacobian.
Return type:: Scalar

class mimirax.parameters.TotalMassConstraint(total)[source]¶

Bases: object

Fix the sum of a mass vector – the made-to-measure normalization.

Variables:: total (float) – The required total mass.
Parameters:: total (float)

total: float¶

defect(params)[source]¶

Return sum(params) - total for a mass-vector parameter.

Parameters:: params (PyTree) – The masses, as an array or a pytree whose leaves are summed.
Returns:: A scalar defect.
Return type:: Array

mimirax.parameters.centre_of_mass(positions, masses)[source]¶

Return the mass-weighted mean position.

Parameters:

positions (Vec3) – (n, 3) positions.
masses (PerParticle) – (n,) masses.

Returns:

(3,) centre of mass.

Return type:

Array

mimirax.parameters.constrain(unconstrained, transforms)[source]¶

Map a dict of unconstrained parameters to the constrained space.

Parameters:

unconstrained (Mapping[str, Array]) – Parameters keyed by name.
transforms (Mapping[str, Reparameterization]) – Which reparameterization applies to which key. Keys absent from transforms pass through unchanged.

Returns:

The constrained parameters, same keys.

Return type:

dict[str, Array]

mimirax.parameters.constraint_violation(constraint, params)[source]¶

Return the largest absolute defect component, a convergence diagnostic.

Parameters:

constraint (Constraint) – The constraint.
params (PyTree) – Where to evaluate it.

Returns:

max |defect|.

Return type:

Scalar

mimirax.parameters.log_abs_det_jacobian(unconstrained, transforms)[source]¶

Sum the log-Jacobians of every transform applied by constrain().

Parameters:

unconstrained (Mapping[str, Array]) – Parameters keyed by name, in the unconstrained space.
transforms (Mapping[str, Reparameterization]) – Which reparameterization applies to which key.

Returns:

The total, which a density over the unconstrained space adds to the constrained-space log-density.

Return type:

Scalar

mimirax.parameters.net_momentum(velocities, masses)[source]¶

Return the total linear momentum.

Parameters:

velocities (Vec3) – (n, 3) velocities.
masses (PerParticle) – (n,) masses.

Returns:

(3,) momentum.

Return type:

Array

mimirax.parameters.num_parameters(params)[source]¶

Count the scalar degrees of freedom in a parameter pytree.

Parameters:: params (PyTree) – Any pytree of arrays.
Returns:: The total number of array elements across all leaves.
Return type:: int

mimirax.parameters.quadratic_penalty(constraint, params, weight=1.0)[source]¶

Turn a constraint defect into an objective term, weight * sum(defect^2).

Parameters:

constraint (Constraint) – The constraint.
params (PyTree) – Where to evaluate it.
weight (float) – The penalty strength.

Returns:

The penalty; add it to a negative log-posterior.

Return type:

Scalar

mimirax.parameters.ravel(params)[source]¶

Flatten a parameter pytree into one vector, with its inverse.

Parameters:: params (PyTree) – Any pytree of arrays.
Returns:: The concatenated leaves as a 1-d array, and a function that rebuilds the original structure from a vector of the same length.
Return type:: tuple[Array, Callable[[Array], PyTree]]

mimirax.parameters.remove_centre_of_mass(positions, masses)[source]¶

Translate positions so their centre of mass is at the origin.

Parameters:

positions (Vec3) – (n, 3) positions.
masses (PerParticle) – (n,) masses.

Returns:

The recentred positions.

Return type:

Vec3

mimirax.parameters.remove_net_momentum(velocities, masses)[source]¶

Boost velocities so the total linear momentum vanishes.

Parameters:

velocities (Vec3) – (n, 3) velocities.
masses (PerParticle) – (n,) masses.

Returns:

The velocities in the zero-momentum frame.

Return type:

Vec3

mimirax.parameters.unconstrain(constrained, transforms)[source]¶

Map a dict of constrained parameters back to the unconstrained space.

Parameters:

constrained (Mapping[str, Array]) – Parameters keyed by name.
transforms (Mapping[str, Reparameterization]) – Which reparameterization applies to which key. Keys absent from transforms pass through unchanged.

Returns:

The unconstrained parameters, same keys.

Return type:

dict[str, Array]

Observables, likelihoods and priors¶

Observation operators and their registry.

An observable maps a forward model’s state to data space; see mimirax.protocols.Observable. New observables are registered here with OBSERVABLES.register(name, cls) – no core module changes.

class mimirax.observables.GaussianRadialBins(centres, width, moments=('mass', 'v2'))[source]¶

Bases: object

A reference kernel: soft radial bins, and their mass-weighted v^2 moment.

The doubles’ counterpart on the observable side. Real made-to-measure kernels – line-of-sight velocity distributions, IFU pixels, stream tracks – belong to the science that uses them (D-027 keeps domain observables in the domain packages); this exists so the method has something to run on today, the way ProjectedPositions does.

The bins are Gaussian rather than top-hat because a top-hat’s derivative with respect to a position is zero almost everywhere and a delta on a measure-zero set: a fit that also moves the particles would see no gradient from it. The Gaussians overlap, which makes the recovered weights genuinely degenerate in a controlled, reportable way rather than accidentally well-conditioned – read the degeneracy off mimirax.diagnostics.fisher_information(), do not tune it away.

Each requested moment contributes len(centres) columns, in the order moments lists them, so the kernel is (n, len(centres) * len(moments)).

Variables:

centres (Array) – (b,) bin centres in radius.
width (float) – Gaussian bin width, in the same units as centres.
moments (tuple[str, ...]) – Which moments to stack: "mass" for the binned weight itself and "v2" for the binned weight times |v|^2 (a mass-weighted kinetic moment, linear in the weights as the method requires – a ratio such as a velocity dispersion is not).

Raises:

ValueError – If moments is empty or names a moment that is not supported.

Parameters:

centres (Array)
width (float)
moments (tuple[str, ...])

centres: Array¶

width: float¶

moments: tuple[str, ...] = ('mass', 'v2')¶

class mimirax.observables.IdentityObservable[source]¶

Bases: object

The state itself, as one array: the data are the model output.

class mimirax.observables.ProjectedPositions(axes=(0, 1))[source]¶

Bases: object

Positions projected onto a subset of axes and flattened.

The simplest mock observable: an (n, 3) state seen in projection, as an imaging survey sees a stellar system. Differentiable in state.

Variables:: axes (tuple[int, ...]) – Which coordinate axes survive the projection. Static.
Parameters:: axes (tuple[int, ...])

axes: tuple[int, ...] = (0, 1)¶

class mimirax.observables.TimeAverage(inner, decay_steps=None)[source]¶

Bases: object

Average an inner observable over a recorded trajectory’s leading axis.

state is a pytree whose every array leaf has a leading t axis – the stacked records a rollout returns. The inner observable is vmap-ed over that axis and the results are averaged, so the inner observable is written for a single snapshot and never has to know it is being time-averaged.

Two averages are available. The default is the plain mean, which is the right thing when the recorded window is what the fit is meant to average over. With decay_steps set, the weights are exponential in the step index and normalized to sum to one, w_t propto exp(-(t_last - t) / decay_steps): the discrete counterpart of Syer & Tremaine’s temporal smoothing d y~ / dt = alpha (y - y~), with alpha dt = 1 / decay_steps. It is in units of recorded steps, not of time, so the observable stays independent of the integrator’s step size; a driver that wants a physical timescale converts once, at the call site.

Neither choice makes a claim about whether the window is long enough for the average to represent a stationary state. That is measured – vary the number of recorded steps and watch the fitted weights – not assumed.

Both averages fold, so this also satisfies FoldableObservable: a caller can either stack the trajectory and average it (__call__()) or accumulate it one snapshot at a time inside the integration (value() / initial() / fold() / result()), and the two give the same answer to round-off. The second path costs O(m) memory where the first costs O(t * n), which is what mimirax.adapters.nornax.FoldedRollout exists to exploit. The mean folds by summing; the exponential average folds through the recursion state <- exp(-1 / decay_steps) * state + value, whose value after t snapshots is sum_s a^(t - 1 - s) value_s – exactly the unnormalized weights weights() returns – so result() divides by (1 - a^t) / (1 - a).

Variables:

inner (Callable[[PyTree], Array]) – The per-snapshot observable, an Observable.
decay_steps (float | None) – Exponential-smoothing timescale in recorded steps; None (the default) gives the plain mean.

Raises:

ValueError – If decay_steps is given and is not strictly positive.

Parameters:

inner (Callable[[Any], Array])
decay_steps (float | None)

inner: Callable[[Any], Array]¶

decay_steps: float | None = None¶

step_values(state)[source]¶

Return the inner observable at every recorded step, unaveraged.

Parameters:: state (PyTree) – A pytree whose array leaves all carry a leading t axis.
Returns:: (t, m) per-step predictions.
Return type:: Array

weights(num_steps)[source]¶

Return the (t,) averaging weights, which sum to one.

Parameters:: num_steps (int) – How many recorded steps to weight.
Returns:: Uniform 1 / t weights by default, exponential weights ending at the last step when decay_steps is set.
Return type:: Array

value(snapshot)[source]¶

Evaluate the inner observable at one snapshot.

Parameters:: snapshot (PyTree) – A single time slice, with no leading time axis.
Returns:: (m,) inner observable at that snapshot.
Return type:: Array

initial(snapshot)[source]¶

Return the zero accumulator, shaped by tracing the inner observable.

Uses jax.eval_shape(), so the inner observable is never evaluated and the accumulator carries no gradient from snapshot.

Parameters:: snapshot (PyTree) – A template time slice.
Returns:: (m,) zeros of the inner observable’s dtype.
Return type:: Array

fold(state, value)[source]¶

Accumulate one snapshot’s value.

Parameters:

state (Array) – The running accumulator.
value (Array) – (m,) value at the next snapshot.

Returns:

state + value for the plain mean; the decayed recursion exp(-1 / decay_steps) * state + value for exponential smoothing.

Return type:

Array

result(state, num_steps)[source]¶

Normalize the accumulator into the averaged observable.

Parameters:

state (Array) – The final accumulator.
num_steps (int) – How many snapshots were folded. Static.

Returns:

(m,) time-averaged observable, equal to __call__() on the stacked trajectory.

Return type:

Array

class mimirax.observables.WeightedKernelSum(kernel, weights_key='weights')[source]¶

Bases: object

The made-to-measure observable y_j = sum_i w_i K_j(z_i).

Differentiable in the weights and in whatever the kernel reads, so a fit over the weights alone and a fit that also moves the initial conditions use the same observable.

The state is a mapping: weights_key selects the weight leaf and the kernel is handed the whole state, so a kernel is free to use positions, velocities or anything else the forward model recorded.

Variables:

kernel (Callable[[PyTree], Array]) – state -> (n, m): particle i’s contribution to observable j. Must be differentiable in the state’s array leaves.
weights_key (str) – Which leaf of the state holds the (n,) weights.

Parameters:

kernel (Callable[[Any], Array])
weights_key (str)

kernel: Callable[[Any], Array]¶

weights_key: str = 'weights'¶

kernel_values(state)[source]¶

Return the raw (n, m) kernel, unweighted.

The classic force-of-change update needs this, and so does any hand-written check of the residual gradient sum_j Delta_j K_ji / sigma_j^2 against the one autodiff produces.

Parameters:: state (PyTree) – A state the kernel accepts.
Returns:: (n, m) kernel values.
Return type:: Array

mimirax.observables.make_observable(name, **kwargs)[source]¶

Instantiate a registered observable by name.

Parameters:

name (str) – A key of OBSERVABLES.
**kwargs (object) – Constructor arguments for the registered class.

Returns:

The instance.

Return type:

Observable

Reference observation operators: the trivial ones every pipeline test needs.

Real observables – kinematic maps, binned densities, stream-track coordinates – are the work of the science papers and are added to mimirax.observables.OBSERVABLES as they are written. These two exist so that every core numerical path has an observable to run through today.

class mimirax.observables.reference.IdentityObservable[source]¶

Bases: object

The state itself, as one array: the data are the model output.

class mimirax.observables.reference.ProjectedPositions(axes=(0, 1))[source]¶

Bases: object

Positions projected onto a subset of axes and flattened.

The simplest mock observable: an (n, 3) state seen in projection, as an imaging survey sees a stellar system. Differentiable in state.

Variables:: axes (tuple[int, ...]) – Which coordinate axes survive the projection. Static.
Parameters:: axes (tuple[int, ...])

axes: tuple[int, ...] = (0, 1)¶

The made-to-measure observable, a reference kernel for it, and time averaging.

Made-to-measure is a statement about one shape of observable. Every quantity the method fits is linear in the particle weights,

\[y_j = \sum_i w_i K_j(z_i),\]

with z_i particle i’s phase-space coordinates and K_j the kernel of observable j – a binned density, a mass-weighted velocity moment, a surface-brightness pixel. WeightedKernelSum is that form, with the kernel supplied by the caller, and it is the seam the classic algorithm needs: Syer & Tremaine’s force of change is driven by K_j(z_i) per particle, not by y_j alone, so the kernel has to be reachable and not buried inside a closure. WeightedKernelSum.kernel_values() is that access.

TimeAverage is the second half of the method. A made-to-measure fit compares the data not with an instantaneous model value but with one averaged along the orbits, which is what turns a snapshot of N particles into a statement about a stationary distribution function. It wraps any inner observable and averages it over the leading axis of a recorded trajectory – mimirax.adapters.nornax.NornaxRollout’s output, or any pytree whose leaves all carry a time axis.

What is not claimed here: that a given number of recorded steps is a long enough average, or that a given kernel makes the weights identifiable. Both are measurements, and mimirax.diagnostics is where they are made.

class mimirax.observables.m2m.GaussianRadialBins(centres, width, moments=('mass', 'v2'))[source]¶

Bases: object

A reference kernel: soft radial bins, and their mass-weighted v^2 moment.

Each requested moment contributes len(centres) columns, in the order moments lists them, so the kernel is (n, len(centres) * len(moments)).

Variables:

centres (Array) – (b,) bin centres in radius.
width (float) – Gaussian bin width, in the same units as centres.
moments (tuple[str, ...]) – Which moments to stack: "mass" for the binned weight itself and "v2" for the binned weight times |v|^2 (a mass-weighted kinetic moment, linear in the weights as the method requires – a ratio such as a velocity dispersion is not).

Raises:

ValueError – If moments is empty or names a moment that is not supported.

Parameters:

centres (Array)
width (float)
moments (tuple[str, ...])

centres: Array¶

width: float¶

moments: tuple[str, ...] = ('mass', 'v2')¶

class mimirax.observables.m2m.TimeAverage(inner, decay_steps=None)[source]¶

Bases: object

Average an inner observable over a recorded trajectory’s leading axis.

Variables:

inner (Callable[[PyTree], Array]) – The per-snapshot observable, an Observable.
decay_steps (float | None) – Exponential-smoothing timescale in recorded steps; None (the default) gives the plain mean.

Raises:

ValueError – If decay_steps is given and is not strictly positive.

Parameters:

inner (Callable[[Any], Array])
decay_steps (float | None)

inner: Callable[[Any], Array]¶

decay_steps: float | None = None¶

step_values(state)[source]¶

Return the inner observable at every recorded step, unaveraged.

Parameters:: state (PyTree) – A pytree whose array leaves all carry a leading t axis.
Returns:: (t, m) per-step predictions.
Return type:: Array

weights(num_steps)[source]¶

Return the (t,) averaging weights, which sum to one.

Parameters:: num_steps (int) – How many recorded steps to weight.
Returns:: Uniform 1 / t weights by default, exponential weights ending at the last step when decay_steps is set.
Return type:: Array

value(snapshot)[source]¶

Evaluate the inner observable at one snapshot.

Parameters:: snapshot (PyTree) – A single time slice, with no leading time axis.
Returns:: (m,) inner observable at that snapshot.
Return type:: Array

initial(snapshot)[source]¶

Return the zero accumulator, shaped by tracing the inner observable.

Uses jax.eval_shape(), so the inner observable is never evaluated and the accumulator carries no gradient from snapshot.

Parameters:: snapshot (PyTree) – A template time slice.
Returns:: (m,) zeros of the inner observable’s dtype.
Return type:: Array

fold(state, value)[source]¶

Accumulate one snapshot’s value.

Parameters:

state (Array) – The running accumulator.
value (Array) – (m,) value at the next snapshot.

Returns:

state + value for the plain mean; the decayed recursion exp(-1 / decay_steps) * state + value for exponential smoothing.

Return type:

Array

result(state, num_steps)[source]¶

Normalize the accumulator into the averaged observable.

Parameters:

state (Array) – The final accumulator.
num_steps (int) – How many snapshots were folded. Static.

Returns:

(m,) time-averaged observable, equal to __call__() on the stacked trajectory.

Return type:

Array

class mimirax.observables.m2m.WeightedKernelSum(kernel, weights_key='weights')[source]¶

Bases: object

The made-to-measure observable y_j = sum_i w_i K_j(z_i).

Differentiable in the weights and in whatever the kernel reads, so a fit over the weights alone and a fit that also moves the initial conditions use the same observable.

The state is a mapping: weights_key selects the weight leaf and the kernel is handed the whole state, so a kernel is free to use positions, velocities or anything else the forward model recorded.

Variables:

kernel (Callable[[PyTree], Array]) – state -> (n, m): particle i’s contribution to observable j. Must be differentiable in the state’s array leaves.
weights_key (str) – Which leaf of the state holds the (n,) weights.

Parameters:

kernel (Callable[[Any], Array])
weights_key (str)

kernel: Callable[[Any], Array]¶

weights_key: str = 'weights'¶

kernel_values(state)[source]¶

Return the raw (n, m) kernel, unweighted.

The classic force-of-change update needs this, and so does any hand-written check of the residual gradient sum_j Delta_j K_ji / sigma_j^2 against the one autodiff produces.

Parameters:: state (PyTree) – A state the kernel accepts.
Returns:: (n, m) kernel values.
Return type:: Array

Likelihoods and misfits, and their registry.

A likelihood compares a prediction with an observation and returns a log-probability; see mimirax.protocols.Likelihood. Register new ones with LIKELIHOODS.register(name, cls).

class mimirax.likelihoods.GaussianLikelihood(sigma=1.0)[source]¶

Bases: object

Independent Gaussian noise on every data point, of one width or of many.

log_prob is -chi_squared / 2. The normalization constant -sum_j log(sigma_j sqrt(2 pi)) is dropped: it does not depend on the prediction, so neither optimization nor sampling sees it. It does depend on sigma, so this log-density may not be compared across two different sigma – which is a statement about model comparison and not about fitting.

HETEROSCEDASTIC IS THE USUAL CASE HERE. A scalar sigma is only meaningful when every entry of the data vector is in the same units and measured equally well. A made-to-measure vector usually is not: stack a binned mass moment and a binned |v|^2 moment and their numerical scales differ by whatever |v|^2 happens to be, so a scalar width silently weights the fit by unit choice. Pass an array of per-point widths instead, one per entry of observed.

Variables:: sigma (SigmaLike) – The noise level. A scalar applies one width to every point; an (m,) array – or anything broadcasting against the data – applies one width per point. Must be strictly positive; not checked, because sigma may be a traced array.
Parameters:: sigma (Array | float)

sigma: Array | float = 1.0¶

log_prob(predicted, observed)[source]¶

Return the log-likelihood up to its additive constant.

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.

Returns:

-0.5 * chi_squared(predicted, observed, sigma).

Return type:

Scalar

mimirax.likelihoods.chi_squared(predicted, observed, sigma=1.0)[source]¶

Return sum(((predicted - observed) / sigma)^2).

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.
sigma (SigmaLike) – Noise level: a scalar for homoscedastic noise, or an array broadcasting against the data for one width per data point.

Returns:

The chi-squared misfit.

Return type:

Scalar

The Gaussian likelihood and its chi-squared misfit.

The width may be a scalar or one width per data point. That is not a convenience: a made-to-measure data vector routinely stacks moments with incommensurate units – a binned mass and a binned mass-weighted |v|^2 in the same vector – and a scalar sigma then weights them by accident of units rather than by how well either is measured. The array form is the supported way to say what each observable’s uncertainty actually is.

class mimirax.likelihoods.gaussian.GaussianLikelihood(sigma=1.0)[source]¶

Bases: object

Independent Gaussian noise on every data point, of one width or of many.

Variables:: sigma (SigmaLike) – The noise level. A scalar applies one width to every point; an (m,) array – or anything broadcasting against the data – applies one width per point. Must be strictly positive; not checked, because sigma may be a traced array.
Parameters:: sigma (Array | float)

sigma: Array | float = 1.0¶

log_prob(predicted, observed)[source]¶

Return the log-likelihood up to its additive constant.

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.

Returns:

-0.5 * chi_squared(predicted, observed, sigma).

Return type:

Scalar

mimirax.likelihoods.gaussian.chi_squared(predicted, observed, sigma=1.0)[source]¶

Return sum(((predicted - observed) / sigma)^2).

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.
sigma (SigmaLike) – Noise level: a scalar for homoscedastic noise, or an array broadcasting against the data for one width per data point.

Returns:

The chi-squared misfit.

Return type:

Scalar

Priors and regularizers, and their registry.

A prior is a log-density over parameters; a regularizer is a prior that says so. See mimirax.protocols.Prior. Register new ones with PRIORS.register(name, cls).

class mimirax.priors.EntropyPrior(mu=1.0, reference=None, key='weights', include_linear_term=True)[source]¶

Bases: object

The made-to-measure entropy prior over a positive weight vector.

Syer & Tremaine’s (1996) regularizer keeps the particle weights w near a reference w0. Two spellings of it are in the literature and they do not have the same mode, so both are here and the default is stated:

include_linear_term=True (the default, Dehnen 2009’s form):
```
S = -sum_i w_i * (log(w_i / w0_i) - 1)
```
whose gradient is dS/dw_i = -log(w_i / w0_i). It vanishes at w = w0, so the prior’s mode is the reference unconditionally.
include_linear_term=False (Syer & Tremaine’s own S = -sum w log(w / w0)), whose gradient is -(log(w_i / w0_i) + 1) and therefore vanishes at w = w0 / e. Its mode is the reference only on the constant-total-weight surface that mimirax.parameters.TotalMassConstraint fixes, where the Lagrange multiplier absorbs the constant.

Both are strictly concave on w > 0 – the Hessian is the negative diagonal -mu / w_i for either – so the term can only ever pull a fit toward the reference, never create a second optimum. log_prob returns mu * S, the log density, so it adds to a likelihood the way every other Prior here does; the mu S term of the made-to-measure objective is exactly this, and mu lives here and nowhere else (MadeToMeasure deliberately has no second copy of it).

Domain. The entropy is defined on w > 0 and this returns nan outside it, deliberately: the package’s rule is that positivity is enforced by a reparameterization (mimirax.parameters.LogTransform), not by clipping inside a density, and a silent where-guard here would hide a parameterization mistake instead of failing on it.

Variables:

mu (float) – Regularization strength; scales the whole term linearly.
reference (Array | float | None) – w0. A scalar or an (n,) array; None means the uniform reference 1, so the mode is the all-ones weight vector. A system of total weight M on n particles wants reference=M / n.
key (str | None) – Which leaf of a mapping params holds the weights. With None the prior is applied to every leaf, as GaussianPrior is; a non-None key (the default "weights") is what keeps the prior off the position and velocity leaves of a made-to-measure parameter dict, where a logarithm would be a domain error rather than a regularizer.
include_linear_term (bool) – Whether to use the -1 shifted form whose mode is the reference; see above.

Parameters:

mu (float)
reference (Array | float | None)
key (str | None)
include_linear_term (bool)

mu: float = 1.0¶

reference: Array | float | None = None¶

key: str | None = 'weights'¶

include_linear_term: bool = True¶

log_prob(params)[source]¶

Return mu * S at params, the entropy prior’s log density.

Parameters:: params (PyTree) – The particle weights: a mapping carrying key, or any pytree whose leaves are all weights when key is None.
Returns:: -mu * sum(w * (log(w / w0) - 1)) with the default include_linear_term, and -mu * sum(w * log(w / w0)) without it. nan where any weight is non-positive.
Return type:: Scalar
Raises:: KeyError – If params is a mapping and key is not one of its keys. The message lists the keys that are there, because the usual cause is a parameter dict that spells the weights differently.

divergence(params)[source]¶

Return the prior’s norm: the KL divergence from the reference.

sum_i [w_i log(w_i / w0_i) - w_i + w0_i], which is >= 0 and zero exactly at w = w0. Note mu does not appear: this is the functional mu weights, not the weighted term.

WHY THIS EXISTS AND -log_prob WILL NOT DO. Choosing mu by an L-curve means plotting the misfit against the regularizer’s norm, and -log_prob / mu = -S = sum_i w_i (log(w_i/w0_i) - 1) is not a norm: at the prior’s own mode w = w0 it equals -sum_i w0_i, which is negative, and its logarithm – which is what an L-curve is drawn in – does not exist. The two differ by the constant sum_i w0_i, so they have the same minimiser and a different L-curve. An L-curve drawn against -S is a curve of the wrong quantity; this is the right one.

With include_linear_term=False the log density is a different functional, but this divergence is unchanged – it is the natural norm for either spelling, and it is what mimirax.diagnostics.l_curve_corner() expects.

Parameters:: params (PyTree) – The particle weights, read the same way log_prob() reads them.
Returns:: The KL divergence from the reference. nan where any weight is non-positive, for the reason log_prob() gives.
Return type:: Scalar
Raises:: KeyError – If params is a mapping and key is not one of its keys.

class mimirax.priors.GaussianPrior(mean=0.0, sigma=1.0)[source]¶

Bases: object

Independent Gaussian prior of one mean and width on every parameter.

The normalization constant is dropped, as in GaussianLikelihood.

Variables:

mean (float) – Prior mean.
sigma (float) – Prior width.

Parameters:

mean (float)
sigma (float)

mean: float = 0.0¶

sigma: float = 1.0¶

log_prob(params)[source]¶

Return -0.5 * sum(((params - mean) / sigma)^2) over all leaves.

Parameters:: params (PyTree) – The free parameters.
Returns:: The log prior density up to a constant.
Return type:: Scalar

class mimirax.priors.L2Regularizer(weight=1.0)[source]¶

Bases: object

Ridge penalty as a prior: log_prob = -weight * sum(params^2).

A Gaussian prior of zero mean and sigma = 1 / sqrt(2 weight), spelled the way an optimizer user thinks of it.

Variables:: weight (float) – Penalty strength.
Parameters:: weight (float)

weight: float = 1.0¶

log_prob(params)[source]¶

Return the negative weighted sum of squares over all leaves.

Parameters:: params (PyTree) – The free parameters.
Returns:: The log prior density up to a constant.
Return type:: Scalar

Priors and regularizers: a Gaussian, an L2 penalty, and the M2M entropy prior.

WHERE THE JACOBIAN SITS. Every log_prob here is a density over the space it is handed, and none of them adds the log-Jacobian of a reparameterization. That is the scaffold’s convention for GaussianPrior and it is kept for EntropyPrior: a caller who moves in an unconstrained space and wants a density over that space adds mimirax.parameters.log_abs_det_jacobian() itself. The alternative – folding the Jacobian into the prior – would make the same prior mean two different things depending on which optimizer called it. This is friction 1 of the EDDA programme’s D-027 and the choice is made here, once, explicitly: priors are densities over the constrained parameters, and mimirax.inference.MadeToMeasure.minimize() evaluates its objective on constrained parameters for exactly that reason.

class mimirax.priors.reference.EntropyPrior(mu=1.0, reference=None, key='weights', include_linear_term=True)[source]¶

Bases: object

The made-to-measure entropy prior over a positive weight vector.

include_linear_term=True (the default, Dehnen 2009’s form):
```
S = -sum_i w_i * (log(w_i / w0_i) - 1)
```
whose gradient is dS/dw_i = -log(w_i / w0_i). It vanishes at w = w0, so the prior’s mode is the reference unconditionally.
include_linear_term=False (Syer & Tremaine’s own S = -sum w log(w / w0)), whose gradient is -(log(w_i / w0_i) + 1) and therefore vanishes at w = w0 / e. Its mode is the reference only on the constant-total-weight surface that mimirax.parameters.TotalMassConstraint fixes, where the Lagrange multiplier absorbs the constant.

Variables:

mu (float) – Regularization strength; scales the whole term linearly.
reference (Array | float | None) – w0. A scalar or an (n,) array; None means the uniform reference 1, so the mode is the all-ones weight vector. A system of total weight M on n particles wants reference=M / n.
key (str | None) – Which leaf of a mapping params holds the weights. With None the prior is applied to every leaf, as GaussianPrior is; a non-None key (the default "weights") is what keeps the prior off the position and velocity leaves of a made-to-measure parameter dict, where a logarithm would be a domain error rather than a regularizer.
include_linear_term (bool) – Whether to use the -1 shifted form whose mode is the reference; see above.

Parameters:

mu (float)
reference (Array | float | None)
key (str | None)
include_linear_term (bool)

mu: float = 1.0¶

reference: Array | float | None = None¶

key: str | None = 'weights'¶

include_linear_term: bool = True¶

log_prob(params)[source]¶

Return mu * S at params, the entropy prior’s log density.

Parameters:: params (PyTree) – The particle weights: a mapping carrying key, or any pytree whose leaves are all weights when key is None.
Returns:: -mu * sum(w * (log(w / w0) - 1)) with the default include_linear_term, and -mu * sum(w * log(w / w0)) without it. nan where any weight is non-positive.
Return type:: Scalar
Raises:: KeyError – If params is a mapping and key is not one of its keys. The message lists the keys that are there, because the usual cause is a parameter dict that spells the weights differently.

divergence(params)[source]¶

Return the prior’s norm: the KL divergence from the reference.

sum_i [w_i log(w_i / w0_i) - w_i + w0_i], which is >= 0 and zero exactly at w = w0. Note mu does not appear: this is the functional mu weights, not the weighted term.

Parameters:: params (PyTree) – The particle weights, read the same way log_prob() reads them.
Returns:: The KL divergence from the reference. nan where any weight is non-positive, for the reason log_prob() gives.
Return type:: Scalar
Raises:: KeyError – If params is a mapping and key is not one of its keys.

class mimirax.priors.reference.GaussianPrior(mean=0.0, sigma=1.0)[source]¶

Bases: object

Independent Gaussian prior of one mean and width on every parameter.

The normalization constant is dropped, as in GaussianLikelihood.

Variables:

mean (float) – Prior mean.
sigma (float) – Prior width.

Parameters:

mean (float)
sigma (float)

mean: float = 0.0¶

sigma: float = 1.0¶

log_prob(params)[source]¶

Return -0.5 * sum(((params - mean) / sigma)^2) over all leaves.

Parameters:: params (PyTree) – The free parameters.
Returns:: The log prior density up to a constant.
Return type:: Scalar

class mimirax.priors.reference.L2Regularizer(weight=1.0)[source]¶

Bases: object

Ridge penalty as a prior: log_prob = -weight * sum(params^2).

A Gaussian prior of zero mean and sigma = 1 / sqrt(2 weight), spelled the way an optimizer user thinks of it.

Variables:: weight (float) – Penalty strength.
Parameters:: weight (float)

weight: float = 1.0¶

log_prob(params)[source]¶

Return the negative weighted sum of squares over all leaves.

Parameters:: params (PyTree) – The free parameters.
Returns:: The log prior density up to a constant.
Return type:: Scalar

Inference¶

Inference methods – optimizers and samplers – and their registry.

Every entry satisfies Optimizer or Sampler. OptaxOptimizer and MadeToMeasure are implemented, as is DataResampling for uncertainties on made-to-measure weights; HMC, NUTS, MeanFieldVI and NelderMead are stubs that fix names and signatures. Register new methods with METHODS.register(name, cls).

class mimirax.inference.HMC(step_size=0.1, num_leapfrog_steps=10)[source]¶

Bases: object

Hamiltonian Monte Carlo with a fixed leapfrog trajectory. Stub.

Variables:

step_size (float) – Leapfrog step.
num_leapfrog_steps (int) – Steps per proposal.

Parameters:

step_size (float)
num_leapfrog_steps (int)

step_size: float = 0.1¶

num_leapfrog_steps: int = 10¶

sample(log_density, params, *, key, num_samples)[source]¶

Not implemented yet.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target.
params (PyTree) – The initial state.
key (Array) – A jax.random key.
num_samples (int) – Samples to draw.

Returns:

Never returns.

Return type:

SampleResult

Raises:

NotImplementedError – Always.

class mimirax.inference.NUTS(step_size=0.1, max_tree_depth=10)[source]¶

Bases: object

The No-U-Turn Sampler. Stub.

Variables:

step_size (float) – Initial leapfrog step.
max_tree_depth (int) – Doubling limit.

Parameters:

step_size (float)
max_tree_depth (int)

step_size: float = 0.1¶

max_tree_depth: int = 10¶

sample(log_density, params, *, key, num_samples)[source]¶

Not implemented yet.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target.
params (PyTree) – The initial state.
key (Array) – A jax.random key.
num_samples (int) – Samples to draw.

Returns:

Never returns.

Return type:

SampleResult

Raises:

NotImplementedError – Always.

class mimirax.inference.DataResampling(optimizer, sequential=True)[source]¶

Bases: object

Sample the weight posterior by refitting perturbed data. Algorithm 1.

Deliberately not a Sampler. That protocol takes a scalar log_density, and this method cannot work from one: it has to reach inside the posterior to find the data and perturb it. Claiming the protocol would mean accepting a log_density and ignoring it, so sample() takes the InferenceProblem instead and the type system says what the method needs.

Variables:

optimizer (Optimizer) – What to run per sample. A MadeToMeasure is the intended choice; its num_steps multiplies the sample count, so this is where the cost lives.
sequential (bool) – Whether to draw the samples one at a time (jax.lax.map(), the default) or all at once (jax.vmap()). Each sample is an independent optimization, so vmap is embarrassingly parallel and num_samples times the memory – which for a fit that differentiates through an N-body rollout is the wrong trade by default. Turn it on for cheap forward models.

Parameters:

optimizer (Optimizer)
sequential (bool)

optimizer: Optimizer¶

sequential: bool = True¶

sample(problem, params, *, key, num_samples)[source]¶

Draw num_samples weight vectors by refitting resampled data.

Each sample perturbs the observations by their own uncertainty, Y~ ~ N(Y, S), refits from params, and keeps the fitted parameters. The returned log_density is the original problem’s log posterior at each sample – not the perturbed one’s – so it can be compared across samples and used to weight them.

The starting point is params for every sample, not the previous sample: the draws are independent by construction and chaining them would introduce a correlation the derivation does not have.

Parameters:

problem (InferenceProblem) – The posterior to sample. Its observed is perturbed and its likelihood must carry a sigma, since that is the S the perturbation is drawn from.
params (PyTree) – The starting point for every refit – normally the best fit.
key (Array) – A jax.random key.
num_samples (int) – How many samples to draw. Static.

Returns:

samples with a leading num_samples axis on every leaf, and log_density the original posterior at each.

Return type:

SampleResult

Raises:

TypeError – If the likelihood has no sigma. The algorithm is defined for Gaussian uncertainties and there is nothing to perturb by otherwise; a caller with a general likelihood needs a general sampler.

covariance(samples, key='weights')[source]¶

Return the sample covariance of one leaf, the estimate of V.

Parameters:

samples (PyTree) – The samples field of a SampleResult.
key (str | None) – Which leaf to take; None treats samples as the array itself.

Returns:

(d, d) covariance over the num_samples axis, with the num_samples - 1 denominator. Compare it with (K^T S^-1 K)^-1 where the model is linear and no prior is active; expect it to be too narrow otherwise, and by how much the prior contributes.

Return type:

Array

class mimirax.inference.InferenceProblem(forward, observable, likelihood, observed, priors=())[source]¶

Bases: object

A posterior over parameters, assembled from protocol-satisfying parts.

log_posterior(params) = likelihood.log_prob(observable(forward(params)), observed) + sum(prior.log_prob(params)). Every method is differentiable in params provided the parts are, which the protocols require.

Variables:

forward (ForwardModel) – Parameters to state.
observable (Observable) – State to data space.
likelihood (Likelihood) – Data-space comparison.
observed (Array) – The data.
priors (tuple[Prior, ...]) – Zero or more priors and regularizers, summed.

Parameters:

forward (ForwardModel)
observable (Observable)
likelihood (Likelihood)
observed (Array)
priors (tuple[Prior, ...])

forward: ForwardModel¶

observable: Observable¶

likelihood: Likelihood¶

observed: Array¶

priors: tuple[Prior, ...] = ()¶

predict(params)[source]¶

Run the forward model and the observable.

Parameters:: params (PyTree) – The free parameters.
Returns:: The prediction in data space.
Return type:: Array

log_likelihood(params)[source]¶

Return the log-likelihood of the data at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: likelihood.log_prob(predict(params), observed).
Return type:: Scalar

log_prior(params)[source]¶

Return the summed log prior density at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: Zero when there are no priors.
Return type:: Scalar

log_posterior(params)[source]¶

Return the unnormalized log posterior at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: log_likelihood + log_prior.
Return type:: Scalar

negative_log_posterior(params)[source]¶

Return -log_posterior, the objective an optimizer minimizes.

Parameters:: params (PyTree) – The free parameters.
Returns:: The negative log posterior.
Return type:: Scalar

class mimirax.inference.MadeToMeasure(optimizer=None, num_steps=100, epsilon=0.1, transform=<factory>, keys=('weights', ), also_fit=(), refine=None, refine_steps=0)[source]¶

Bases: object

The made-to-measure weight-adjustment method, in both of its variants.

Satisfies Optimizer through minimize(); force_of_change() is the second, classic iteration on the same objective and is deliberately not the protocol method, so a caller can never get the oracle by accident.

Which leaves are weights is keys (for a mapping params); every other leaf is carried through untouched by force_of_change() and optimized in its own space by minimize(). A bare-array params is taken to be the weights themselves.

Variables:

optimizer (optax.GradientTransformation | None) – The update rule minimize() descends with, in the unconstrained space; made_to_measure() builds an Adam one. None is the default and means minimize() is unavailable – no step size is guessed on the caller’s behalf – which leaves force_of_change() usable on its own with epsilon.
num_steps (int) – Weight-update steps, for either iteration. Static.
epsilon (float) – The classic force of change’s step size, Syer & Tremaine’s epsilon. Unused by minimize().
transform (Reparameterization) – The positivity reparameterization applied to the weight leaves. LogTransform by default, because gradient descent in ln w is the differentiable counterpart of the classic multiplicative update – see the module docstring.
keys (tuple[str, ...]) – Which leaves of a mapping params hold weights.
also_fit (tuple[str, ...]) –
Further leaves of params to treat as free parameters. Empty by default, and that default is the method’s definition rather than a convenience: made-to-measure fits weights, at initial conditions the caller fixed. Every other leaf of params – the positions and velocities a rollout needs – is held constant and returned unchanged.

The default was not always this. An earlier version optimized every leaf, so a fit handed {"positions", "velocities", "weights"} for 64 particles silently had 448 free parameters against ten observables instead of 64, and moved the positions by 17 % and the velocities by 47 % while its docstring claimed to be recovering weights. Fitting the initial conditions too is a legitimate thing to want – Syer & Tremaine’s method extends to it – but it is a different inference problem, and it now has to be asked for by name.
refine (optax.GradientTransformation | None) – An optional second stage, run from where the first one stopped: optax.lbfgs() is the intended use, and the reason the stage exists is that a quasi-Newton rule needs a starting point in the right basin more than it needs many steps. None (the default) runs one stage. Whether two stages beat one is a property of the problem and is measured, not assumed – on both of this module’s test problems the best single-stage fit beats the two-stage one; see reports/M2M_module.md.
refine_steps (int) – Steps for the second stage. Static. Ignored when refine is None; a quasi-Newton stage wants tens, not thousands.

Parameters:

optimizer (GradientTransformation | None)
num_steps (int)
epsilon (float)
transform (Reparameterization)
keys (tuple[str, ...])
also_fit (tuple[str, ...])
refine (GradientTransformation | None)
refine_steps (int)

optimizer: GradientTransformation | None = None¶

num_steps: int = 100¶

epsilon: float = 0.1¶

transform: Reparameterization¶

keys: tuple[str, ...] = ('weights',)¶

also_fit: tuple[str, ...] = ()¶

refine: GradientTransformation | None = None¶

refine_steps: int = 0¶

weight_keys(params)[source]¶

Return which leaves of params are weights.

Parameters:: params (PyTree) – The parameters.
Returns:: The subset of keys that params has, or () when params is not a mapping and is therefore the weights itself.
Return type:: tuple[str, …]
Raises:: KeyError – If params is a mapping and holds none of keys. Silently optimizing nothing is the one outcome worse than failing here.

free_parameters(params)[source]¶

Split params into the free leaves and a function that puts them back.

The frozen leaves never enter the optimizer’s state, so they cannot be moved by it, cannot pick up an update from a stale moment estimate, and cost nothing to carry. The rebuild preserves the caller’s key order, so FitResult comes back in the structure it went in.

Parameters:: params (PyTree) – The parameters as the caller supplied them.
Returns:: The free subtree, and a function mapping a free subtree back into a full parameter pytree. For a bare array the subtree is the array and the rebuild is the identity.
Return type:: tuple[PyTree, Callable[[PyTree], PyTree]]

to_unconstrained(params)[source]¶

Map the weight leaves into the unconstrained space.

Parameters:: params (PyTree) – Parameters with positive weight leaves.
Returns:: The same structure with the weight leaves transformed; a bare array is transformed whole.
Return type:: PyTree

to_constrained(unconstrained)[source]¶

Map the weight leaves back into the constrained space.

Parameters:: unconstrained (PyTree) – Parameters in the unconstrained space.
Returns:: The same structure with the weight leaves transformed back; a bare array is transformed whole.
Return type:: PyTree

minimize(objective, params)[source]¶

Descend objective in the unconstrained space; the differentiable variant.

The optimizer moves in u, the objective is always evaluated on the constrained parameters, and the result comes back constrained and in the caller’s own structure. The reparameterization’s log-Jacobian is not added: priors here are densities over the constrained parameters (see mimirax.priors.reference), so a caller who wants a density over u – for a sampler, say – adds mimirax.parameters.log_abs_det_jacobian() to the objective itself. Optimizing without it finds the maximum of the constrained posterior, which is what a made-to-measure fit is asking for.

Each stage is one lax.scan, as OptaxOptimizer’s loop is, so a fit that differentiates through an N-body rollout traces once per stage. The trace records the objective before each update, so objective_trace[0] is the starting value; with a refine stage the two traces are concatenated and the trace is num_steps + refine_steps long.

Any optax rule works, including a learning-rate schedule (optax.adam(optax.cosine_decay_schedule(...))) and the quasi-Newton and line-search rules whose update needs the objective’s value and the objective itself – see _descend(). Which rule to use is a property of the problem and is measured rather than asserted; the module’s report carries the comparison, and its short version is that no rule wins on both of the method’s two constructions.

Only the weight leaves move, plus anything named in also_fit. Everything else in params is a fixed input to the forward model – the initial conditions a rollout integrates from – and comes back unchanged. That is what made-to-measure means, and getting it wrong turns “recover 64 weights from ten observables” into a fit with 448 free parameters; see also_fit.

Parameters:

objective (Callable[[PyTree], Scalar]) – The time-averaged negative log posterior, a function of the constrained parameters.
params (PyTree) – The starting point, with positive weight leaves. The frozen leaves are still required here: the objective needs them.

Returns:

Constrained final parameters in the input’s structure, with the frozen leaves untouched, and the objective trace.

Return type:

FitResult

Raises:

ValueError – If optimizer is None; the step size is the caller’s choice and is not guessed here.

force_of_change(objective, params)[source]¶

Run the classic Syer & Tremaine force-of-change iteration.

w <- w * exp(-epsilon * dF/dw), the weight leaves only; every other leaf is returned exactly as it came in. This is the oracle minimize() is compared against, not a method to prefer: it is first-order in one hand-tuned step size, it descends in a fixed metric, and – the substantive point – it is only the right gradient when the orbits do not depend on the weights. Use it on the tracer construction (mimirax.adapters.nornax.NornaxRollout.tracer()). On the self-consistent construction dF/dw here still includes the orbits’ dependence on the weights, because jax.grad() cannot be told to forget it, so the iteration is not the classic algorithm there and the classic algorithm is not defined there either.

Its fixed point is minimize()’s: both stop exactly where dF/dw = 0. The multiplicative step is what enforces positivity – the package’s rule is reparameterization, not clipping – and it is identical to a unit-step gradient descent in ln w scaled by epsilon, which is the algebraic reason the two variants share their stationary points.

Parameters:

objective (Callable[[PyTree], Scalar]) – The time-averaged negative log posterior, a function of the constrained parameters.
params (PyTree) – The starting point, with positive weight leaves.

Returns:

Final parameters in the input’s structure, and the (num_steps,) objective trace, recorded before each update.

Return type:

FitResult

class mimirax.inference.MeanFieldVI(num_steps=1000, learning_rate=0.01)[source]¶

Bases: object

Mean-field Gaussian variational inference by gradient ascent on the ELBO. Stub.

Variables:

num_steps (int) – Optimization steps for the variational parameters.
learning_rate (float) – Step size.

Parameters:

num_steps (int)
learning_rate (float)

num_steps: int = 1000¶

learning_rate: float = 0.01¶

sample(log_density, params, *, key, num_samples)[source]¶

Not implemented yet.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target.
params (PyTree) – Initial mean of the variational family.
key (Array) – A jax.random key.
num_samples (int) – Draws to return from the fitted family.

Returns:

Never returns.

Return type:

SampleResult

Raises:

NotImplementedError – Always.

class mimirax.inference.NelderMead(num_steps=500)[source]¶

Bases: object

The Nelder-Mead simplex method. Stub.

Variables:: num_steps (int) – Simplex iterations.
Parameters:: num_steps (int)

num_steps: int = 500¶

minimize(objective, params)[source]¶

Not implemented yet.

Parameters:

objective (Callable[[PyTree], Scalar]) – The function to minimize.
params (PyTree) – The starting point.

Returns:

Never returns.

Return type:

FitResult

Raises:

NotImplementedError – Always.

class mimirax.inference.OptaxOptimizer(optimizer, num_steps=100)[source]¶

Bases: object

Minimize an objective with any optax.GradientTransformation.

The loop is a lax.scan over a static number of steps, so the whole fit traces once and runs on device. The trace records the objective before each update, so objective_trace[0] is the value at the starting point.

Variables:

optimizer (optax.GradientTransformation) – The update rule.
num_steps (int) – How many updates to take. Static.

Parameters:

optimizer (GradientTransformation)
num_steps (int)

optimizer: GradientTransformation¶

num_steps: int = 100¶

minimize(objective, params)[source]¶

Run num_steps updates from params.

Parameters:

objective (Callable[[PyTree], Scalar]) – A differentiable scalar function of the parameters.
params (PyTree) – The starting point.

Returns:

The final parameters and the (num_steps,) objective trace.

Return type:

FitResult

mimirax.inference.adam(learning_rate=0.01, num_steps=200)[source]¶

Build an Adam optimizer with the given settings.

Parameters:

learning_rate (float) – Adam’s step size.
num_steps (int) – Number of updates.

Returns:

Ready to minimize.

Return type:

OptaxOptimizer

mimirax.inference.made_to_measure(learning_rate=0.01, num_steps=200, *, epsilon=0.1, keys=('weights',), also_fit=(), refine=None, refine_steps=0)[source]¶

Build a made-to-measure method with an Adam rule, as adam() does.

Adam is the default because it is the default everywhere, not because it is the best rule for this objective – on the module’s tracer problem plain optax.sgd(1e-6) reaches a gradient norm an order of magnitude smaller, and optax.lbfgs() reaches a lower objective in fifty steps than Adam does in ten thousand. Neither transfers to the self-consistent problem, where a fixed-step SGD diverges outright. Pass optimizer= to MadeToMeasure directly for anything other than Adam, and read the measured comparison in reports/M2M_module.md before choosing.

Parameters:

learning_rate (optax.ScalarOrSchedule) – Adam’s step size in the log-weight space, or an optax schedule – the annotation is optax’s own alias rather than a hand-written float | Callable, because a schedule takes chex.Numeric and not Array, and CI’s pyright is stricter than the local one about the difference (a difference the scaffold report already records). – optax.exponential_decay(0.05, 2000, 0.3) measured better on the tracer problem than any constant rate tried.
num_steps (int) – Number of weight updates.
epsilon (float) – The classic force of change’s step size.
keys (tuple[str, ...]) – Which leaves of a mapping params hold weights.
also_fit (tuple[str, ...]) – Further leaves to treat as free parameters; empty means weights only.
refine (optax.GradientTransformation | None) – An optional second-stage rule, e.g. optax.lbfgs().
refine_steps (int) – Steps for the second stage.

Returns:

Ready to minimize() or to force_of_change().

Return type:

MadeToMeasure

Composing a forward model, an observable, a likelihood and priors into one objective.

This is the one place the pieces meet. Everything else in the package is either a pluggable implementation of one protocol or a method that consumes the scalar functions defined here.

class mimirax.inference.objective.InferenceProblem(forward, observable, likelihood, observed, priors=())[source]¶

Bases: object

A posterior over parameters, assembled from protocol-satisfying parts.

Variables:

forward (ForwardModel) – Parameters to state.
observable (Observable) – State to data space.
likelihood (Likelihood) – Data-space comparison.
observed (Array) – The data.
priors (tuple[Prior, ...]) – Zero or more priors and regularizers, summed.

Parameters:

forward (ForwardModel)
observable (Observable)
likelihood (Likelihood)
observed (Array)
priors (tuple[Prior, ...])

forward: ForwardModel¶

observable: Observable¶

likelihood: Likelihood¶

observed: Array¶

priors: tuple[Prior, ...] = ()¶

predict(params)[source]¶

Run the forward model and the observable.

Parameters:: params (PyTree) – The free parameters.
Returns:: The prediction in data space.
Return type:: Array

log_likelihood(params)[source]¶

Return the log-likelihood of the data at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: likelihood.log_prob(predict(params), observed).
Return type:: Scalar

log_prior(params)[source]¶

Return the summed log prior density at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: Zero when there are no priors.
Return type:: Scalar

log_posterior(params)[source]¶

Return the unnormalized log posterior at params.

Parameters:: params (PyTree) – The free parameters.
Returns:: log_likelihood + log_prior.
Return type:: Scalar

negative_log_posterior(params)[source]¶

Return -log_posterior, the objective an optimizer minimizes.

Parameters:: params (PyTree) – The free parameters.
Returns:: The negative log posterior.
Return type:: Scalar

Gradient-based optimization through optax.

The one inference method the scaffold implements rather than stubs, because every other method’s tests need a working baseline to compare against and because it is the method Paper I section 7’s fit uses.

class mimirax.inference.optimizers.OptaxOptimizer(optimizer, num_steps=100)[source]¶

Bases: object

Minimize an objective with any optax.GradientTransformation.

Variables:

optimizer (optax.GradientTransformation) – The update rule.
num_steps (int) – How many updates to take. Static.

Parameters:

optimizer (GradientTransformation)
num_steps (int)

optimizer: GradientTransformation¶

num_steps: int = 100¶

minimize(objective, params)[source]¶

Run num_steps updates from params.

Parameters:

objective (Callable[[PyTree], Scalar]) – A differentiable scalar function of the parameters.
params (PyTree) – The starting point.

Returns:

The final parameters and the (num_steps,) objective trace.

Return type:

FitResult

mimirax.inference.optimizers.adam(learning_rate=0.01, num_steps=200)[source]¶

Build an Adam optimizer with the given settings.

Parameters:

learning_rate (float) – Adam’s step size.
num_steps (int) – Number of updates.

Returns:

Ready to minimize.

Return type:

OptaxOptimizer

Hamiltonian Monte Carlo and NUTS. Stubs.

Both are Sampler implementations by signature; the bodies raise until a method is written against them. Whether that is a hand-rolled leapfrog integrator or a wrapper over an external library is left open on purpose – the registry takes either.

class mimirax.inference.hmc.HMC(step_size=0.1, num_leapfrog_steps=10)[source]¶

Bases: object

Hamiltonian Monte Carlo with a fixed leapfrog trajectory. Stub.

Variables:

step_size (float) – Leapfrog step.
num_leapfrog_steps (int) – Steps per proposal.

Parameters:

step_size (float)
num_leapfrog_steps (int)

step_size: float = 0.1¶

num_leapfrog_steps: int = 10¶

sample(log_density, params, *, key, num_samples)[source]¶

Not implemented yet.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target.
params (PyTree) – The initial state.
key (Array) – A jax.random key.
num_samples (int) – Samples to draw.

Returns:

Never returns.

Return type:

SampleResult

Raises:

NotImplementedError – Always.

class mimirax.inference.hmc.NUTS(step_size=0.1, max_tree_depth=10)[source]¶

Bases: object

The No-U-Turn Sampler. Stub.

Variables:

step_size (float) – Initial leapfrog step.
max_tree_depth (int) – Doubling limit.

Parameters:

step_size (float)
max_tree_depth (int)

step_size: float = 0.1¶

max_tree_depth: int = 10¶

sample(log_density, params, *, key, num_samples)[source]¶

Not implemented yet.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target.
params (PyTree) – The initial state.
key (Array) – A jax.random key.
num_samples (int) – Samples to draw.

Returns:

Never returns.

Return type:

SampleResult

Raises:

NotImplementedError – Always.

Variational inference. Stub.

A fitted variational family is something one samples from, so the class satisfies Sampler: sample fits the family to log_density and returns draws from it.

class mimirax.inference.vi.MeanFieldVI(num_steps=1000, learning_rate=0.01)[source]¶

Bases: object

Mean-field Gaussian variational inference by gradient ascent on the ELBO. Stub.

Variables:

num_steps (int) – Optimization steps for the variational parameters.
learning_rate (float) – Step size.

Parameters:

num_steps (int)
learning_rate (float)

num_steps: int = 1000¶

learning_rate: float = 0.01¶

sample(log_density, params, *, key, num_samples)[source]¶

Not implemented yet.

Parameters:

log_density (Callable[[PyTree], Scalar]) – The unnormalized log target.
params (PyTree) – Initial mean of the variational family.
key (Array) – A jax.random key.
num_samples (int) – Draws to return from the fitted family.

Returns:

Never returns.

Return type:

SampleResult

Raises:

NotImplementedError – Always.

Derivative-free optimization. Stub.

Kept in the method registry from day one because the programme’s own plan (the gradient-reliability study ahead of Kessel Run) expects cases where the gradient through a chaotic rollout is not trustworthy and a derivative-free method is the honest comparison.

class mimirax.inference.derivative_free.NelderMead(num_steps=500)[source]¶

Bases: object

The Nelder-Mead simplex method. Stub.

Variables:: num_steps (int) – Simplex iterations.
Parameters:: num_steps (int)

num_steps: int = 500¶

minimize(objective, params)[source]¶

Not implemented yet.

Parameters:

objective (Callable[[PyTree], Scalar]) – The function to minimize.
params (PyTree) – The starting point.

Returns:

Never returns.

Return type:

FitResult

Raises:

NotImplementedError – Always.

Made-to-measure: particle weights adjusted against time-averaged observables.

WHERE THIS LIVES, AND WHY. The EDDA programme’s decision D-026 (2026-09-06) puts the made-to-measure module in mimirax rather than in the integrator package: its three parts are an orbit integration (nornax’s), a force (any ForceModel), and an observable-residual-driven weight update with an entropy prior – and the last is inference. The integrator enters through mimirax.adapters.nornax, behind the mimirax[nornax] extra; nothing here imports nornax.

THE OBJECTIVE. Both methods in this module minimize the same scalar,

\[F(w) = \tfrac{1}{2} \chi^2\big(\bar y(w), Y\big) - \mu S(w),\]

the chi-squared of the time-averaged observables against the data plus the entropy prior of EntropyPrior. An InferenceProblem built from mimirax.adapters.nornax.NornaxRollout, mimirax.observables.TimeAverage around mimirax.observables.WeightedKernelSum, GaussianLikelihood and an EntropyPrior is that objective, and problem.negative_log_posterior is what both methods below are handed. mu lives on the prior and has no second copy here, so the two methods cannot disagree about the regularization they share.

TWO ITERATIONS, ONE STATIONARY POINT. Syer & Tremaine’s (1996) force of change is

\[\dot w_i = -\varepsilon\, w_i \Big[ \sum_j \frac{\Delta_j K_{ji}}{\sigma_j} - \mu \frac{\partial S}{\partial w_i} \Big] = -\varepsilon\, w_i \frac{\partial F}{\partial w_i},\]

so the bracket the classic algorithm assembles by hand is exactly dF/dw – which is why this module computes it with jax.grad() and does not re-derive it (a test pins the two against each other on the tracer problem). The two methods then differ only in the metric they descend in:

MadeToMeasure.force_of_change() is the classic update, ln w <- ln w - eps dF/dw. That is a preconditioned gradient step in w with preconditioner diag(w), and the multiplicative spelling is what keeps the weights positive without a clip – the historical Euler form w <- w (1 - eps dF/dw) can step straight through zero and differs from this at O(eps^2).
MadeToMeasure.minimize() is the differentiable variant this package exists for: plain gradient descent (any optax rule) in u = ln w, i.e. a step in w with preconditioner diag(w^2), on an objective that is differentiated through the orbit integration.

Both are stationary exactly where dF/dw = 0, since w > 0. So they must agree on the answer and are free to disagree on the path – which is what makes the classic iteration a usable oracle, and what a measured comparison between them is testing. When the weights are the masses of a self-gravitating system the orbits themselves depend on the weights and only the second method’s gradient is the true one; the first is then not merely slower but wrong about dF/dw, which is why force_of_change() is documented for the tracer case.

WHAT IS NOT CLAIMED. Nothing here says a time average over a given window is long enough, that a chosen mu is right, or that the method converges at FMM scale. Those are experiments, and the numbers this module was verified at are in the EDDA programme’s reports/M2M_module.md.

class mimirax.inference.m2m.MadeToMeasure(optimizer=None, num_steps=100, epsilon=0.1, transform=<factory>, keys=('weights', ), also_fit=(), refine=None, refine_steps=0)[source]¶

Bases: object

The made-to-measure weight-adjustment method, in both of its variants.

Variables:

optimizer (optax.GradientTransformation | None) – The update rule minimize() descends with, in the unconstrained space; made_to_measure() builds an Adam one. None is the default and means minimize() is unavailable – no step size is guessed on the caller’s behalf – which leaves force_of_change() usable on its own with epsilon.
num_steps (int) – Weight-update steps, for either iteration. Static.
epsilon (float) – The classic force of change’s step size, Syer & Tremaine’s epsilon. Unused by minimize().
transform (Reparameterization) – The positivity reparameterization applied to the weight leaves. LogTransform by default, because gradient descent in ln w is the differentiable counterpart of the classic multiplicative update – see the module docstring.
keys (tuple[str, ...]) – Which leaves of a mapping params hold weights.
also_fit (tuple[str, ...]) –
Further leaves of params to treat as free parameters. Empty by default, and that default is the method’s definition rather than a convenience: made-to-measure fits weights, at initial conditions the caller fixed. Every other leaf of params – the positions and velocities a rollout needs – is held constant and returned unchanged.

The default was not always this. An earlier version optimized every leaf, so a fit handed {"positions", "velocities", "weights"} for 64 particles silently had 448 free parameters against ten observables instead of 64, and moved the positions by 17 % and the velocities by 47 % while its docstring claimed to be recovering weights. Fitting the initial conditions too is a legitimate thing to want – Syer & Tremaine’s method extends to it – but it is a different inference problem, and it now has to be asked for by name.
refine (optax.GradientTransformation | None) – An optional second stage, run from where the first one stopped: optax.lbfgs() is the intended use, and the reason the stage exists is that a quasi-Newton rule needs a starting point in the right basin more than it needs many steps. None (the default) runs one stage. Whether two stages beat one is a property of the problem and is measured, not assumed – on both of this module’s test problems the best single-stage fit beats the two-stage one; see reports/M2M_module.md.
refine_steps (int) – Steps for the second stage. Static. Ignored when refine is None; a quasi-Newton stage wants tens, not thousands.

Parameters:

optimizer (GradientTransformation | None)
num_steps (int)
epsilon (float)
transform (Reparameterization)
keys (tuple[str, ...])
also_fit (tuple[str, ...])
refine (GradientTransformation | None)
refine_steps (int)

optimizer: GradientTransformation | None = None¶

num_steps: int = 100¶

epsilon: float = 0.1¶

transform: Reparameterization¶

keys: tuple[str, ...] = ('weights',)¶

also_fit: tuple[str, ...] = ()¶

refine: GradientTransformation | None = None¶

refine_steps: int = 0¶

weight_keys(params)[source]¶

Return which leaves of params are weights.

Parameters:: params (PyTree) – The parameters.
Returns:: The subset of keys that params has, or () when params is not a mapping and is therefore the weights itself.
Return type:: tuple[str, …]
Raises:: KeyError – If params is a mapping and holds none of keys. Silently optimizing nothing is the one outcome worse than failing here.

free_parameters(params)[source]¶

Split params into the free leaves and a function that puts them back.

Parameters:: params (PyTree) – The parameters as the caller supplied them.
Returns:: The free subtree, and a function mapping a free subtree back into a full parameter pytree. For a bare array the subtree is the array and the rebuild is the identity.
Return type:: tuple[PyTree, Callable[[PyTree], PyTree]]

to_unconstrained(params)[source]¶

Map the weight leaves into the unconstrained space.

Parameters:: params (PyTree) – Parameters with positive weight leaves.
Returns:: The same structure with the weight leaves transformed; a bare array is transformed whole.
Return type:: PyTree

to_constrained(unconstrained)[source]¶

Map the weight leaves back into the constrained space.

Parameters:: unconstrained (PyTree) – Parameters in the unconstrained space.
Returns:: The same structure with the weight leaves transformed back; a bare array is transformed whole.
Return type:: PyTree

minimize(objective, params)[source]¶

Descend objective in the unconstrained space; the differentiable variant.

Parameters:

objective (Callable[[PyTree], Scalar]) – The time-averaged negative log posterior, a function of the constrained parameters.
params (PyTree) – The starting point, with positive weight leaves. The frozen leaves are still required here: the objective needs them.

Returns:

Constrained final parameters in the input’s structure, with the frozen leaves untouched, and the objective trace.

Return type:

FitResult

Raises:

ValueError – If optimizer is None; the step size is the caller’s choice and is not guessed here.

force_of_change(objective, params)[source]¶

Run the classic Syer & Tremaine force-of-change iteration.

Parameters:

objective (Callable[[PyTree], Scalar]) – The time-averaged negative log posterior, a function of the constrained parameters.
params (PyTree) – The starting point, with positive weight leaves.

Returns:

Final parameters in the input’s structure, and the (num_steps,) objective trace, recorded before each update.

Return type:

FitResult

mimirax.inference.m2m.made_to_measure(learning_rate=0.01, num_steps=200, *, epsilon=0.1, keys=('weights',), also_fit=(), refine=None, refine_steps=0)[source]¶

Build a made-to-measure method with an Adam rule, as adam() does.

Parameters:

learning_rate (optax.ScalarOrSchedule) – Adam’s step size in the log-weight space, or an optax schedule – the annotation is optax’s own alias rather than a hand-written float | Callable, because a schedule takes chex.Numeric and not Array, and CI’s pyright is stricter than the local one about the difference (a difference the scaffold report already records). – optax.exponential_decay(0.05, 2000, 0.3) measured better on the tracer problem than any constant rate tried.
num_steps (int) – Number of weight updates.
epsilon (float) – The classic force of change’s step size.
keys (tuple[str, ...]) – Which leaves of a mapping params hold weights.
also_fit (tuple[str, ...]) – Further leaves to treat as free parameters; empty means weights only.
refine (optax.GradientTransformation | None) – An optional second-stage rule, e.g. optax.lbfgs().
refine_steps (int) – Steps for the second stage.

Returns:

Ready to minimize() or to force_of_change().

Return type:

MadeToMeasure

Uncertainties on made-to-measure weights, by resampling the data.

THE METHOD, AND WHOSE IT IS. This is Algorithm 1 of Bovy, Kawata & Hunt (2018, arXiv:1704.03884), the paper that added uncertainty quantification to made-to-measure. Their observation is that for a linear-Gaussian model Y = K W + delta with delta ~ N(0, S), the posterior over W under a uniform prior is Gaussian with mean V K^T S^-1 Y and covariance V = (K^T S^-1 K)^-1, so instead of computing that covariance one can sample it:

Y~  ~  N(Y, S)
W~  =  argmin  chi^2(model(W), Y~)

Each W~ is an exact draw from the posterior, because it is a linear transform of a Gaussian with the right mean and the right covariance. No Metropolis acceptance, no chain to converge, no step size – and one optimization per sample, all of them independent and therefore trivially parallel.

WHY IT SUITS THIS PROBLEM. Bovy et al. give the reasons and they are the same reasons a general-purpose sampler is the wrong tool here: the weight space is high-dimensional (one dimension per particle), the posterior at any given snapshot is noisy, and the weights’ uncertainties are strongly correlated. mimirax’s own HMC and NUTS would meet all three problems; this method meets none of them, because it never proposes a step in weight space at all.

WHERE IT IS EXACT, AND WHERE IT IS NOT. This matters more here than in the original paper, and the module measures it rather than repeating it:

Exact for a linear model with Gaussian noise and no prior. The made-to-measure observable is linear in the weights (mimirax.observables.WeightedKernelSum), and in the tracer construction the orbits do not depend on them, so the whole objective is quadratic and this is the posterior. A test compares the sample covariance with (K^T S^-1 K)^-1 directly.
Approximate with a prior. Bovy et al. say so: the method “does not properly deal with particle weights for which the penalty term … has a significant effect”. Resampling the data does not resample the prior, so the draws are too narrow along every direction the prior holds up – which, on this package’s own tracer problem, is 54 of 64 directions. Read mimirax.effective_parameters() before trusting a spread from this.
Approximate under the positivity constraint, for the same reason: the derivation assumes W is unconstrained, and a weight pinned near zero has a one-sided posterior a Gaussian draw cannot represent.
Not a posterior at all in the self-consistent construction, where the weights move the orbits and the model is no longer linear in them.

None of that makes it useless – it makes it a method with a stated domain, which is better than a chain whose convergence nobody checked. The two approximations are quantified in tests/unit/test_resampling.py and in the EDDA programme’s reports/M2M_module.md.

class mimirax.inference.resampling.DataResampling(optimizer, sequential=True)[source]¶

Bases: object

Sample the weight posterior by refitting perturbed data. Algorithm 1.

Variables:

optimizer (Optimizer) – What to run per sample. A MadeToMeasure is the intended choice; its num_steps multiplies the sample count, so this is where the cost lives.
sequential (bool) – Whether to draw the samples one at a time (jax.lax.map(), the default) or all at once (jax.vmap()). Each sample is an independent optimization, so vmap is embarrassingly parallel and num_samples times the memory – which for a fit that differentiates through an N-body rollout is the wrong trade by default. Turn it on for cheap forward models.

Parameters:

optimizer (Optimizer)
sequential (bool)

optimizer: Optimizer¶

sequential: bool = True¶

sample(problem, params, *, key, num_samples)[source]¶

Draw num_samples weight vectors by refitting resampled data.

The starting point is params for every sample, not the previous sample: the draws are independent by construction and chaining them would introduce a correlation the derivation does not have.

Parameters:

problem (InferenceProblem) – The posterior to sample. Its observed is perturbed and its likelihood must carry a sigma, since that is the S the perturbation is drawn from.
params (PyTree) – The starting point for every refit – normally the best fit.
key (Array) – A jax.random key.
num_samples (int) – How many samples to draw. Static.

Returns:

samples with a leading num_samples axis on every leaf, and log_density the original posterior at each.

Return type:

SampleResult

Raises:

covariance(samples, key='weights')[source]¶

Return the sample covariance of one leaf, the estimate of V.

Parameters:

samples (PyTree) – The samples field of a SampleResult.
key (str | None) – Which leaf to take; None treats samples as the array itself.

Returns:

Return type:

Array

Diagnostics¶

Residuals between prediction and data.

mimirax.diagnostics.residuals.normalized_residuals(predicted, observed, sigma=1.0)[source]¶

Return (predicted - observed) / sigma.

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.
sigma (float) – Noise level.

Returns:

Residuals in units of the noise.

Return type:

Array

mimirax.diagnostics.residuals.residuals(predicted, observed)[source]¶

Return predicted - observed.

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.

Returns:

The signed residuals.

Return type:

Array

mimirax.diagnostics.residuals.rms_residual(predicted, observed)[source]¶

Return the root-mean-square residual.

Parameters:

predicted (Array) – Model prediction.
observed (Array) – Data, same shape.

Returns:

sqrt(mean((predicted - observed)^2)).

Return type:

Scalar

Convergence judgements from an objective trace.

mimirax.diagnostics.convergence.has_converged(trace, *, rtol=1e-06, window=5)[source]¶

Whether the last window relative changes are all below rtol.

A host-side judgement (returns a Python bool), meant for after a fit, not inside a traced loop.

Parameters:

trace (Array) – (k,) objective values, k > window.
rtol (float) – Relative-change threshold.
window (int) – How many trailing steps must all be below it.

Returns:

True when converged by this criterion.

Return type:

bool

mimirax.diagnostics.convergence.relative_change(trace, eps=1e-300)[source]¶

Return |x[k] - x[k-1]| / (|x[k-1]| + eps) along a trace.

Parameters:

trace (Array) – (k,) objective values.
eps (float) – Guard against division by zero.

Returns:

(k - 1,) relative changes.

Return type:

Array

Degeneracy diagnostics: which parameter directions the data do not see.

Recovering discrete source positions from field samples has continuous degeneracies (Paper I section 7 says so in its first paragraph), and the diagnostics that name them are more useful than the ones that hide them. The Fisher information’s null-ish eigenvectors are the first such diagnostic.

mimirax.diagnostics.degeneracy.degenerate_directions(fisher, *, rtol=1e-08)[source]¶

Return the eigenvalues and the eigenvectors below rtol of the largest.

Parameters:

fisher (Array) – (d, d) symmetric matrix from fisher_information().
rtol (float) – Eigenvalues below rtol * max(eigenvalues) count as degenerate.

Returns:

All (d,) eigenvalues in ascending order, and the (d, r) matrix whose columns are the degenerate directions (r may be zero).

Return type:

tuple[Array, Array]

mimirax.diagnostics.degeneracy.effective_parameters(data_curvature, prior_curvature)[source]¶

Return how many parameters the data actually determined.

tr(H_data (H_data + H_prior)^-1), the standard effective number of degrees of freedom of a regularized fit. In the joint eigenbasis it is sum_i lambda_i / (lambda_i + pi_i): a direction the data constrain tightly (lambda >> pi) contributes 1, a direction only the prior holds up (lambda << pi) contributes 0, and the total interpolates.

WHY THIS AND NOT AN EIGENVALUE COUNT. degenerate_directions() needs an rtol and returns a count, so it reports a rank at a threshold and puts a marginal direction wholly on one side. This needs no threshold and is the quantity that actually answers “is my objective informative enough”: on the made-to-measure tracer problem of reports/M2M_module.md – 64 weights, 10 observables – it returns 10.000, exactly the number of observables, while degenerate_directions(rtol=1e-3) reports 55 of 64 and the fit reaches a chi-squared of 1.4e-08. Ten numbers went in and ten degrees of freedom came out; the other 54 are the prior’s answer, and no optimizer, step size or mu changes that. More information means more or better observables.

Read it against the parameter count, not on its own: 10 out of 64 says the fit is prior-dominated, and 10 out of 10 would say it is saturated.

WHERE THIS DENSE ROUTE FAILS. It solves against H_data + H_prior, and a H_data built from a Jacobian as J^T C^-1 J already has the squared condition number of J. Its absolute error is about eps * d_max, so a direction at the prior’s threshold is lost once d_max / threshold exceeds 1 / eps. Measured in the EDDA Fisher probe (reports/FISHER_fd_vs_ad_probe.md, section 5.7), on 32 weights:

kept in single precision, it returns values from -314 to +257 on Jacobians taken at the probe’s default steps, and 28.19 on a single-precision AD Jacobian that, analysed in double, gives 13.92;
in double precision it differs from the SVD route by 6.6e-5 to 9.7e-4 on the probe’s three radial sets – more than the effects that probe set out to measure.

The single-precision failure is not a slow loss of accuracy. On the unit test’s version of that set the prior’s 3.2e-6 lies below the float32 spacing (8e-3) of every diagonal entry of J^T C^-1 J, so H_data + H_prior is H_data, bit for bit, and what comes back is tr(F F^-1) of a numerically singular F: rounding, whatever mu.

It is also only as much a Fisher matrix as data_curvature is: a Hessian taken away from an exact fit of a nonlinear model is not one (see fisher_information()). Given the Jacobian, prefer effective_parameters_from_jacobian(), which has neither problem. This function is kept as it is so that existing results reproduce.

Parameters:

data_curvature (Array) – (d, d) Hessian of the negative log likelihood alone – the data’s contribution. Get it as fisher_information(problem.negative_log_ likelihood_only, params), or by subtracting the priors’ curvature from the full fisher_information().
prior_curvature (Array) – (d, d) Hessian of the negative log prior, same basis and ordering.

Returns:

The effective number of data-determined parameters, between 0 and d. Not an integer in general.

Return type:

Scalar

mimirax.diagnostics.degeneracy.effective_parameters_from_jacobian(jacobian, prior_curvature, *, sigma=None)[source]¶

Return effective_parameters() from the Jacobian, by an SVD.

The same quantity, tr(F (F + H_prior)^-1) with F = J^T C^-1 J the Gauss–Newton matrix – the Fisher information of a Gaussian likelihood – computed without ever forming F. Whiten the Jacobian by the noise and by the prior, B = C^(-1/2) J L^-T with H_prior = L L^T. Then F + H_prior = L (B^T B + I) L^T, and the trace is exactly

sum_k s_k^2 / (s_k^2 + 1)

over the singular values s_k of B. A diagonal prior p makes L = diag(sqrt(p)) and the whitening a column scaling, so no (d, d) matrix is formed at all. With p_i = mu / w_i – the curvature of EntropyPrior at w – the sum is sum_k d_k / (d_k + mu) over the squared singular values d_k of C^(-1/2) J diag(sqrt(w)), the prior-whitened form of reports/M2M_production_readiness.md (A2); at uniform w = w0 it is sum_k d_k / (d_k + mu / w0) over those of C^(-1/2) J. Either way mu is a threshold on a spectrum the observables alone determine.

WHY AN SVD. Forming F squares the condition number: its absolute error is about eps * d_max, and the dense route loses a direction at the threshold once d_max / threshold exceeds 1 / eps. The singular values of B carry an absolute error of about eps * s_max against a threshold at s = 1, so this holds until s_max reaches 1 / eps: d_max / threshold up to 1 / eps^2, the square of the dense route’s range. (Those are norm-wise bounds, i.e. reasoning.) Measured in the EDDA Fisher probe (reports/FISHER_fd_vs_ad_probe.md, section 5.7) on the sensitivity figure’s radial_wide set, 32 weights, d_max / threshold of 1.6e13, from a single-precision AD Jacobian: this route in single precision gives 13.937 against 13.920 from the same Jacobian analysed in double, where effective_parameters() gives 28.19. In double precision, forward- and reverse-mode Jacobians agree through this route to 6e-12, while the dense route differs from it by up to 9.7e-4. The margin does shrink toward 1 / eps^2: on the unit test’s circular-orbit version of that set, the single-precision error is 0.003 at d_max / threshold of 1.6e12 and 0.03 at 1.6e13. The arithmetic is in the precision the inputs promote to, so a float64 prior turns a float32 Jacobian’s analysis into a float64 one.

WHY A JACOBIAN AND NOT A HESSIAN. F here is J^T C^-1 J at whatever point J was taken – positive semi-definite and of rank at most min(m, d), so the count cannot exceed the number of observables. The Hessian of fisher_information() is that matrix only where the model is linear or the residuals vanish; elsewhere its residual curvature is counted as information (see there).

Get J by differentiating the prediction – observable(forward(x)) – over the flattened free parameters: jax.jacrev() costs m reverse passes and is the cheaper when there are more parameters than data, as in made-to-measure; jax.jacfwd() costs d tangents.

Parameters:

jacobian (Array) – (m, d) derivative of the prediction with respect to the flattened parameters, rows in data order and columns in the prior’s order. Already whitened, C^(-1/2) J, if sigma is None.
prior_curvature (Array) – Hessian of the negative log prior, which must be positive definite. Either (d,), the diagonal of a diagonal prior – exact, and the form to pass for EntropyPrior, whose curvature is mu / w – or (d, d), a dense one, whitened by its Cholesky factor. A direction with no prior curvature has no defined count, and a zero diagonal entry, or a dense prior that is not positive definite, gives nan or inf.
sigma (SigmaLike | None) – Noise width to whiten the rows by: one number, or (m,), as GaussianLikelihood takes it. None means jacobian is already whitened; for a correlated noise covariance C = K K^T, pass K^-1 J and None.

Returns:

The effective number of data-determined parameters, between 0 and min(m, d). Not an integer in general.

Return type:

Scalar

Raises:

ValueError – If jacobian is not two-dimensional, if prior_curvature is neither (d,) nor (d, d), or if sigma is neither one number nor one per row.

mimirax.diagnostics.degeneracy.fisher_information(negative_log_posterior, params)[source]¶

Return the Hessian of the negative log posterior over the flattened parameters.

At the maximum a posteriori point this is the observed Fisher information; elsewhere it is the local curvature and should be read as such.

WHEN THIS IS A FISHER MATRIX, AND WHEN IT IS NOT. For a Gaussian likelihood with prediction y(theta), residual r = y - observed and noise sigma, the Hessian of the negative log likelihood is

J^T C^-1 J + sum_j (r_j / sigma_j^2) d^2 y_j / d theta^2

with J = dy / dtheta. The first term is the Fisher information – the Gauss–Newton matrix: positive semi-definite, of rank at most min(m, d), and what a Cramer–Rao forecast or effective_parameters() means by the data’s curvature. The second is residual curvature, and it is neither: it has no sign and is in general of full rank. So this Hessian is the Fisher matrix

at any point, when the prediction is linear in the parameters – the tracer made-to-measure weights, whose orbits do not depend on them;
wherever the residuals vanish – at the truth of noise-free mock data, or at an exact fit.

At the maximum a posteriori point of noisy data it is the observed information, which differs from the Fisher matrix by the residual term evaluated there. Anywhere else on a nonlinear model it is local curvature and not a Fisher matrix. Measured in the EDDA Fisher probe (reports/FISHER_fd_vs_ad_probe.md, section 1.8): in the self-consistent made-to-measure construction, where the weights are the masses and move the orbits, at uniform weights, with ten observables of 32 weights, the residual term fills the null space and effective_parameters() on this Hessian returns 32.000009 – more than the 32 parameters there are, which only an indefinite curvature can give – where J^T C^-1 J gives 10.00.

For the Fisher matrix away from an exact fit, differentiate the prediction rather than the objective and use effective_parameters_from_jacobian(). And passed the full negative log posterior, this includes the priors’ curvature too; pass the likelihood alone before reading it as the data’s information.

Parameters:

negative_log_posterior (Callable[[PyTree], Scalar]) – The objective, in the parameters’ own pytree structure.
params (PyTree) – Where to evaluate.

Returns:

(d, d) symmetric matrix over the flattened parameter vector, in the order mimirax.parameters.ravel() uses.

Return type:

Array

mimirax.diagnostics.degeneracy.l_curve_corner(regularization, misfit, norm)[source]¶

Pick a regularization weight at the L-curve’s corner.

The standard Tikhonov criterion, and for made-to-measure it answers “what should mu be” – a question the module previously answered with “1e-3, because it is small”. Given a fit repeated over a grid of mu, the corner of the log-log curve of misfit against regularizer norm is the point of maximum curvature: below it the misfit barely improves while the norm runs away, above it the norm barely improves while the misfit degrades.

Curvature is taken with respect to log(regularization), so an unevenly spaced grid is handled correctly, via

kappa = (rho' eta'' - rho'' eta') / (rho'^2 + eta'^2)^(3/2)

with eta = log(norm) and rho = log(misfit). The three points at each end are excluded. Two reasons, and the count is the second one: jnp.gradient falls back to one-sided differences at the ends, which routinely produce the largest spurious curvature in the array; and taking it twice makes the curvature at index i depend on the data at i +/- 2, so one bad endpoint reaches two points inward and excluding two would not contain it.

MEASURED, against the mu that actually minimizes the weight error – an oracle a real fit does not have – on three tracer problems over eight noise realizations each (reports/M2M_production_readiness.md, A4):

criterion	eff ~ n	eff << n	precise data
this function	1.04x	1.30x	1.02x
generalized cross-validation	1.01x	2.15x	1.01x
discrepancy, chi2 = m - eff	1.72x	2.39x	1.60x
`mu = 1e-3`	2.46x	5.30x	3.18x
NNLS, the honest mu -> 0 limit	4.70x	13.97x	6.73x

Read as a multiple of the oracle’s weight error, so 1.00x is the best achievable. GCV – which needs nothing this package does not already have, being misfit / (m - effective_parameters)^2 – beats this function on two of the three and is 2.15x on the degenerate problem, where this one is 1.30x. Neither dominates; both beat mu = 1e-3 on all three, and both beat driving mu to zero, which pins 21 %, 88 % and 51 % of the weights at exactly zero on the three problems.

No criterion here is universally reliable, and this one is not either. On synthetic Tikhonov problems chosen to be hostile, all three fail, in different places:

synthetic problem	this one	GCV	discrepancy
hard spectral gap, delta=1e-2	171x	16x	16x
hard spectral gap, delta=1e-4	5.9x	58x	3312x
smooth spectrum, s_i ~ 1/i	99x	1.0x	1.0x
smooth spectrum, s_i ~ 1/i^2	14x	3.5x	3.5x

So: run it per problem, run GCV beside it, and treat a disagreement between the two as a warning rather than picking one on faith. The curvature array is returned for exactly this reason – a curve with no corner has a flat curvature profile and no peak worth trusting, and a caller who only sees the returned mu cannot tell that apart from a sharp corner.

WHAT TO PASS AS norm. The regularizer’s norm, not minus its log density. For EntropyPrior that is divergence(): -S is negative at the prior’s own mode, so log(-S) does not exist and a curve drawn against it is a curve of the wrong quantity.

mu IS NOT DIMENSIONLESS and the number this returns does not transfer between problems: it carries the ratio of the misfit’s scale to the norm’s, so rescaling the data, sigma or the total weight rescales it. The criterion transfers; the value does not. Re-run it per problem.

Parameters:

regularization (Array) – (k,) regularization weights, strictly positive and strictly increasing – the mu grid the fit was repeated over. At least nine of them.
misfit (Array) – (k,) misfit at each mu: the achieved chi-squared, strictly positive.
norm (Array) – (k,) regularizer norm at each mu, strictly positive.

Returns:

The mu at the corner, and the (k,) curvature array, whose three end points on each side are nan because curvature was not evaluated there. Returning the curvature makes a bad corner visible: a curve with no corner has no clear peak, and a caller that only sees the mu cannot tell.

Return type:

tuple[Scalar, Array]

Raises:

ValueError – If the three arrays are not the same one-dimensional shape, if there are fewer than nine points, or if regularization is not strictly increasing and positive.

mimirax.diagnostics.degeneracy.profile_likelihood(negative_log_posterior, params, direction)[source]¶

Profile the objective along a direction, re-optimizing the rest. Stub.

Parameters:

negative_log_posterior (Callable[[PyTree], Scalar]) – The objective.
params (PyTree) – The starting point.
direction (Array) – A flattened-parameter direction.

Returns:

Never returns.

Return type:

Array

Raises:

NotImplementedError – Always; needs an optimizer in the loop and is left for a method session.

Adapters (optional extras)¶

jaccpot’s FMM as a ForceModel. mimirax[jaccpot].

FastMultipoleMethod.compute_accelerations(positions, masses) already has the protocol’s shape; this adapter only fixes the keyword contract and the lazy import. Gradients: compute_accelerations builds its tree on the host on every call and is not traceable by jax.grad. jaccpot’s differentiable seam is its differentiable_accelerations at frozen topology (jaccpot/docs/differentiable_fmm.md); wiring that in, with the prepare-once / evaluate-many split it needs, is the adapter’s next step and deliberately not guessed at here.

class mimirax.adapters.jaccpot.JaccpotForceModel(solver)[source]¶

Bases: object

Thin wrapper exposing a jaccpot solver through the mimirax protocol.

Variables:: solver (Any) – A jaccpot.FastMultipoleMethod (or anything with its compute_accelerations method).
Parameters:: solver (Any)

solver: Any¶

accelerations(positions, masses, *, args=None)[source]¶

Forward to solver.compute_accelerations.

Parameters:

positions (Vec3) – (n, 3) positions.
masses (PerParticle) – (n,) masses.
args (object) – Ignored; jaccpot’s options live on the solver.

Returns:

(n, 3) accelerations.

Return type:

Vec3

classmethod from_preset(preset='balanced', **solver_kwargs)[source]¶

Build a FastMultipoleMethod and wrap it.

Parameters:

preset (str) – A jaccpot preset name.
**solver_kwargs (object) – Forwarded to FastMultipoleMethod.

Returns:

A JaccpotForceModel.

Return type:

Any

Raises:

ImportError – If jaccpot is not installed; the message names the extra.

A nornax rollout as a ForwardModel. mimirax[nornax].

The integrator half of the made-to-measure module (EDDA decision D-026): initial conditions and particle weights in, a recorded trajectory out, through nornax.solvers.leapfrog_kdk.block_kdk_rollout. nornax is imported inside the call, so this module – and the protocol bridge below it – import and are testable with no solver installed, and a missing nornax fails at the rollout with a message naming the extra.

THE PROTOCOL MISMATCH, AND WHAT IT COSTS. nornax’s integrator wants a MutualForceModel: level_accelerations(positions, masses, *, rung, level), the acceleration contributed by pairs whose finer endpoint sits on rung level, applied antisymmetrically. mimirax’s ForceModel gives a total accelerations(positions, masses) and knows nothing about rungs. The two meet in exactly one place: with every particle on rung 0 and k_max = 0, level 0 is the total force, so a mimirax force model can drive a single-rung rollout exactly. SingleRungMutualForce is that bridge and it refuses any other level rather than returning something plausible. A multi-rung rollout needs a real MutualForceModel – nornax’s MutualDirectSumGravity, or jaccpot’s BlockStepFMM – passed straight through.

WHAT IS DIFFERENTIABLE. params["positions"], params["velocities"] and params["weights"] are all differentiable leaves of the returned trajectory. The weights reach the dynamics only in the self-consistent construction (NornaxRollout.self_consistent()), where they are the masses; in the tracer construction (NornaxRollout.tracer()) they reach the observable alone and d(trajectory)/d(weights) is zero on the position and velocity leaves by construction. What is not differentiable, in either: rung and the topology, both of which nornax severs from the gradient with stop_gradient. For a tree or FMM backend that is the substantive statement of the EDDA programme’s B4 – the gradient is the exact fixed-topology gradient of the numeric path on each segment between rebuilds, and a rebuild cadence of rebuild_every base steps is where the segments end.

class mimirax.adapters.nornax.FoldedRollout(rollout, observable, segments=1)[source]¶

Bases: object

A rollout with the observable folded in, in checkpointed segments.

params in, the reduced observable out – so this is a ForwardModel whose output is already in data space, and the InferenceProblem above it takes IdentityObservable.

WHAT IT BUYS, MEASURED. Two things cost O(t * n) in a made-to-measure fit, and only one of them is autodiff’s fault:

NornaxRollout stacks the trajectory – positions and velocities at every base step – because a plain Observable needs the whole thing at once. That is 6 t n doubles in the forward pass. Folding removes it: the per-step record becomes the (m,) observable value instead of the (n, 3) state.
lax.scan retains one state per step for the backward pass. Running the rollout as segments chained sub-rollouts, each wrapped in jax.checkpoint(), trades that for O(t / segments + segments) states and one extra forward pass per segment.

Peak scratch for one gradient, from XLA’s own memory_analysis, at n = 64:

t      segments=1 (NornaxRollout)   segments=sqrt(t)   reduction
   1.4 MB                       0.9 MB             1.6x
  5.6 MB                       1.0 MB             5.6x
 22.4 MB                      1.2 MB             18x

The segments = 1 column is linear in t; the sqrt(t) column is flat. segments = sqrt(num_steps) is the optimum of t / segments + segments and is what balanced() returns.

This is a memory schedule, not an approximation. The prediction matches the stacked path to 1.2e-16 and its gradient to 5.9e-15, which is what the fold law of FoldableObservable guarantees and what a test measures.

WHAT IT COSTS. The forward model now knows about the observable, which is exactly the seam mimirax.protocols otherwise keeps apart. That is why this is a separate, explicitly-named class rather than a flag on NornaxRollout: the trade is visible at the call site, and a caller who does not need the memory keeps the seam. The licence for chaining segments at all is nornax’s shooting_node, which recomputes acc and rung at each seam so that segments sub-rollouts are one trajectory of the same map – a property this package tests directly.

Variables:

rollout (NornaxRollout) – The integration to run. Its record attribute is ignored: this class installs its own record function, which returns the observable’s per-step value.
observable (FoldableObservable) – The reduction to fold in.
segments (int) – How many checkpointed segments. 1 folds the observable but keeps one scan, which already removes the stacked trajectory.

Raises:

ValueError – If segments is not positive or does not divide rollout.num_steps. An uneven split would weight the segments unequally in an exponential average, so it is refused rather than silently rounded.

Parameters:

rollout (NornaxRollout)
observable (FoldableObservable)
segments (int)

rollout: NornaxRollout¶

observable: FoldableObservable¶

segments: int = 1¶

classmethod balanced(rollout, observable)[source]¶

Build one with the memory-optimal segment count near sqrt(num_steps).

Peak storage goes as num_steps / segments + segments, minimized at sqrt(num_steps). The largest divisor of num_steps at or below that is used, since the split has to be even.

Parameters:

rollout (NornaxRollout) – The integration to run.
observable (FoldableObservable) – The reduction to fold in.

Returns:

With segments the best available divisor.

Return type:

FoldedRollout

class mimirax.adapters.nornax.NornaxRollout(force, dt, num_steps, k_max=0, weights_are_masses=True, masses=None, eta=0.1, eps=1.0, reassign_rungs=True, checkpoint=True, checkpoint_substeps=False, record=None, rebuild_fn=None, rebuild_every=1, base_index=0)[source]¶

Bases: object

Integrate a state forward with nornax and return the recorded trajectory.

params is a mapping with "positions" (n, 3), "velocities" (n, 3) and "weights" (n,). The returned trajectory is a mapping whose leaves all carry a leading num_steps axis, ready for mimirax.observables.TimeAverage:

"positions", "velocities" – (num_steps, n, 3), one row per base step, from nornax’s record_fn;
"weights" – (num_steps, n), the weights broadcast along the time axis. They are constant along a rollout (nothing in the integrator moves a mass), so this materializes nothing the scan had to carry; it is done so that the whole trajectory is one pytree an averaging observable can vmap over, which keeps Observable from needing a second argument.

Use self_consistent() or tracer() rather than the constructor: which of the two a fit is running is the single most consequential thing about a made-to-measure problem, and it should be named at the call site.

The initial state is built with nornax’s shooting_node, not by hand. That constructor exists precisely to recompute the two leaves that must not be free variables – acc from the positions and masses, and rung from that acc – so a fit that moves the weights cannot move the integration schedule by moving a derived quantity.

Variables:

force (Any) – A nornax MutualForceModel, or a mimirax ForceModel (bridged by SingleRungMutualForce, k_max = 0 only).
dt (float) – Base time step, nornax’s dt_max.
num_steps (int) – Base steps to take, nornax’s n_base. Static.
k_max (int) – Highest block-step rung. 0 is the single-rung reduced case. Static.
weights_are_masses (bool) – Whether params["weights"] is the dynamical mass vector (self-consistent M2M) or only enters the observable (tracer M2M).
masses (Array | None) – The dynamical masses when weights_are_masses is False; None means unit masses, which is what an external field wants since it ignores them.
eta (float) – nornax’s rung-assignment accuracy parameter.
eps (float) – nornax’s rung-assignment softening parameter.
reassign_rungs (bool) –
Whether rungs are recomputed at every base-step boundary (production) or frozen for the whole rollout. Freeze them to make the map globally smooth in the continuous state, which is what a finite-difference gradient check needs.

This is not a free choice, and the difference is measured. With it on, the rollout integrates a schedule that depends on the weights through the acceleration, and nornax severs that dependence with stop_gradient. The gradient is therefore exact for the schedule the forward pass realised – reverse and forward mode agree on it to 2.6e-16 – but the objective is only piecewise smooth in the weights: a weight change large enough to move a particle across a rung boundary lands on a different map, with a kink between. On a 16-body, three-rung test system the production gradient differs from the frozen-schedule one by 65 % in norm (cosine 0.953), because the two realise different schedules. Neither is wrong; they are gradients of different maps. Freeze the schedule when a fit must descend one smooth objective, and expect the kinks when it must not.
checkpoint (bool) – Whether to wrap each base step in jax.checkpoint, bounding the retained backward-pass state to the base-step boundaries. Leave it on: it does not change the gradient (measured bit-identical, single- and multi-rung) and it is what keeps a long rollout’s memory linear in num_steps.
checkpoint_substeps (bool) – Whether to additionally remat each sub-step boundary’s kick. checkpoint bounds memory across base steps but still materializes all 2**k_max sub-step pair tensors while differentiating one base step; this bounds that to one boundary’s worth. nornax’s rollout documents it as needed for deep k_max gradients, which otherwise run out of memory – so it is the knob a made-to-measure fit at FMM scale and large k_max will need, and it is exposed here rather than left reachable only by bypassing the adapter. Off by default, as in nornax; it changes the memory schedule, not the result.
record (Callable[[Any], PyTree] | None) – nornax’s record_fn, which must return a mapping so that __call__() can add the "weights" leaf beside it; None records positions and velocities.
rebuild_fn (Callable[..., Any] | None) – A tree backend’s traceable (positions, masses) -> topology rebuild, e.g. jaccpot’s BlockStepFMM.rebuild_state.
rebuild_every (int) – Base steps per topology rebuild. Requires rebuild_fn.
base_index (int) – The global base-step index this rollout starts at, which anchors the rebuild cadence when a trajectory is run as chained segments.

Raises:

ValueError – If a mimirax ForceModel is paired with k_max > 0, or if num_steps is not positive.

Parameters:

force (Any)
dt (float)
num_steps (int)
k_max (int)
weights_are_masses (bool)
masses (Array | None)
eta (float)
eps (float)
reassign_rungs (bool)
checkpoint (bool)
checkpoint_substeps (bool)
record (Callable[[Any], Any] | None)
rebuild_fn (Callable[[...], Any] | None)
rebuild_every (int)
base_index (int)

force: Any¶

dt: float¶

num_steps: int¶

k_max: int = 0¶

weights_are_masses: bool = True¶

masses: Array | None = None¶

eta: float = 0.1¶

eps: float = 1.0¶

reassign_rungs: bool = True¶

checkpoint: bool = True¶

checkpoint_substeps: bool = False¶

record: Callable[[Any], Any] | None = None¶

rebuild_fn: Callable[[...], Any] | None = None¶

rebuild_every: int = 1¶

base_index: int = 0¶

classmethod self_consistent(force, dt, num_steps, *, k_max=0, **kwargs)[source]¶

Build the self-consistent construction: the weights are the masses.

The case the differentiable N-body stack was built for. Because the weights are the masses, the orbits themselves depend on the weights, so d(objective)/d(weights) flows through the force evaluations and through every step of the integration.

That orbit-response term is not unprecedented – Bovy, Kawata & Hunt (2018, arXiv:1704.03884) compute it for external-potential parameters by finite differencing the orbit integration, one extra integration per parameter – but they explicitly exclude self-gravity, and finite differencing costs n + 1 rollouts when the parameters are the n particle masses. Reverse mode costs one, whatever n is, which is what makes this construction tractable rather than merely expressible. force must be a real nornax MutualForceModel; a mimirax ForceModel is rejected above k_max = 0.

Parameters:

force (Any) – A nornax MutualForceModel.
dt (float) – Base time step.
num_steps (int) – Base steps to take.
k_max (int) – Highest block-step rung.
**kwargs (Any) – Any other attribute of NornaxRollout.

Returns:

Configured with weights_are_masses=True.

Return type:

NornaxRollout

classmethod tracer(force, dt, num_steps, *, masses=None, **kwargs)[source]¶

Build the classic tracer construction: the weights stay out of the dynamics.

The orbits are those of test particles in a potential the weights do not set, so the trajectory can be computed once and the weight fit runs on the recorded kernel – which is what makes the classic force-of-change iteration well defined (see mimirax.inference.MadeToMeasure.force_of_change()). Single-rung by construction: an external field is not antisymmetric, so the block schedule’s momentum-conserving level split does not apply to it.

Parameters:

force (Any) – A mimirax ForceModel – an analytic potential; mimirax.testing.SoftenedPointMassField is the double – or a nornax MutualForceModel.
dt (float) – Base time step.
num_steps (int) – Base steps to take.
masses (Array | None) – The dynamical masses. None means unit masses; an external field ignores them, so that is the right default there and the wrong one for a self-gravitating system whose masses are not the weights.
**kwargs (Any) – Any other attribute of NornaxRollout.

Returns:

Configured with weights_are_masses=False and k_max=0.

Return type:

NornaxRollout

mutual_force()[source]¶

Return the force in the form nornax’s integrator consumes.

Returns:: force itself when it already has level_accelerations, and SingleRungMutualForce around it otherwise.
Return type:: Any

dynamical_masses(params)[source]¶

Return the mass vector the integrator advances the state with.

Parameters:: params (PyTree) – The parameter mapping.
Returns:: params["weights"] in the self-consistent construction, and masses (or ones) in the tracer construction.
Return type:: PerParticle

rollout(params)[source]¶

Run the rollout and return nornax’s own (final_state, records).

The escape hatch for anything that needs the integrator’s state rather than an observable’s view of it – a shooting defect between chained segments, an energy diagnostic, a topology’s pair counts.

Parameters:: params (PyTree) – A mapping with "positions", "velocities" and "weights".
Returns:: The final BlockStepState and the stacked per-base-step records.
Return type:: tuple[Any, PyTree]

final_state(params)[source]¶

Return only the final BlockStepState.

Parameters:: params (PyTree) – A mapping with "positions", "velocities" and "weights".
Returns:: nornax’s BlockStepState at the end of the rollout.
Return type:: Any

folding(observable, *, segments=1)[source]¶

Return a ForwardModel that folds observable into the rollout.

The memory escape hatch. See FoldedRollout.

Parameters:

observable (FoldableObservable) – The reduction to accumulate; mimirax.observables.TimeAverage is one.
segments (int) – How many checkpointed segments to split the rollout into.

Returns:

Parameters in, the reduced observable out.

Return type:

FoldedRollout

class mimirax.adapters.nornax.SingleRungMutualForce(force)[source]¶

Bases: object

A mimirax ForceModel behind nornax’s per-level interface.

Valid at k_max = 0 only, and it says so by raising. With every particle on rung 0 every pair sits on level 0, so level 0’s antisymmetric contribution is the whole force and the wrapped model’s total is exactly right. At any other level the honest answer does not exist: a mimirax ForceModel has no notion of which pairs belong to which level, and both plausible guesses are wrong – returning the total at level 0 and zero above would kick the whole force with level 0’s weight at every sub-step boundary (the wrong equations, silently), and returning zero at level 0 would lose the force entirely.

What conformance can and cannot say about this bridge. nornax’s check_mutual_force_model at k_max = 0 checks shape, finiteness, the level partition (trivial with one level) and the momentum residual – and the last one is a property of the wrapped model, not of the bridge. A mutual model (mimirax.testing.DirectSumGravity) passes it to round-off; an external field (mimirax.testing.SoftenedPointMassField, the tracer construction’s double) does not and must not, because an external potential genuinely does not conserve the particles’ momentum. That is why the tracer construction is single-rung: the antisymmetry the block schedule relies on is exactly what an external field lacks. An external field could be split by level = rung_i instead of max(rung_i, rung_j) and driven multi-rung, but that is a different bridge for a different force, and guessing which one a caller meant from a model that cannot be asked is how the wrong equations get integrated.

Variables:: force (ForceModel) – The mimirax acceleration provider being bridged.
Parameters:: force (ForceModel)

force: ForceModel¶

level_accelerations(positions, masses, *, rung, level, args=None, topology=None)[source]¶

Return the wrapped model’s total acceleration, as level 0.

Parameters:

positions (Vec3) – (n, 3) positions.
masses (PerParticle) – (n,) masses.
rung (Array) – Ignored: at k_max = 0 every particle is on rung 0, so the rung partition carries no information the total does not already have.
level (int) – Must be 0.
args (object) – Passed through to the wrapped model’s caller-owned channel.
topology (object) – Accepted and ignored: a mimirax ForceModel has no frozen interaction structure to evaluate against.

Returns:

(n, 3) accelerations.

Return type:

Vec3

Raises:

ValueError – If level is not 0; see the class docstring for why there is no defensible answer above it.

An ODISSEO integration as a ForwardModel. mimirax[odisseo]. Stub.

ODISSEO’s odisseo.integrate(...) and its differentiable FMM lane (odisseo/differentiable.py) are the obvious body. Left as a stub because whether ODISSEO stays on the routing path between jaccpot and nornax is an open decision (O-4 in the EDDA programme), and an adapter written before that is settled would be written twice.

class mimirax.adapters.odisseo.OdisseoForwardModel(config, params)[source]¶

Bases: object

Run an ODISSEO simulation from parameters. Stub.

Variables:

config (Any) – An odisseo.option_classes.SimulationConfig.
params (Any) – An odisseo.option_classes.SimulationParams template whose fields the free parameters overwrite.

Parameters:

config (Any)
params (Any)

config: Any¶

params: Any¶

Test doubles¶

Analytic test doubles with closed-form gradients.

WHY THESE EXIST. Every core numerical path in mimirax has to be exercised on CPU with no solver installed, or the suite will not run often enough to be worth having. These doubles satisfy the protocols in mimirax.protocols and, unlike a solver, come with their derivatives written out by hand – so a test can check jax.grad against an exact expression rather than against finite differences whose tolerance has to be argued for.

LinearForwardModel – y = A x + b; the gradient of a sum-of-squares misfit is 2 A^T (A x + b - y).
SoftenedPointMassField – the softened Newtonian field of fixed sources at test-particle positions, with its tidal tensor d a / d x = -G sum_j m_j [ I / r^3 - 3 d d^T / r^5 ] in closed form.
DirectSumGravity – self-gravity of n bodies, the honest O(n^2) ForceModel any FMM is checked against.

All are float64-exact to round-off; tests compare at atol = 1e-12.

class mimirax.testing.doubles.DirectSumGravity(softening=0.0, gravitational_constant=1.0)[source]¶

Bases: object

Softened self-gravity of n bodies by direct summation.

The O(n^2) reference every hierarchical solver is measured against, and the double for a ForceModel that actually uses masses. Differentiable in positions and masses; the self-term is masked with where on both the distance and the weight so the gradient carries no 0 * inf.

Variables:

softening (float) – Plummer softening eps.
gravitational_constant (float) – G.

Parameters:

softening (float)
gravitational_constant (float)

softening: float = 0.0¶

gravitational_constant: float = 1.0¶

accelerations(positions, masses, *, args=None)[source]¶

Return the mutual accelerations.

Parameters:

positions (Vec3) – (n, 3) positions.
masses (PerParticle) – (n,) masses.
args (object) – Ignored.

Returns:

(n, 3) accelerations.

Return type:

Vec3

class mimirax.testing.doubles.LinearForwardModel(matrix, offset)[source]¶

Bases: object

y = A x + b: a ForwardModel whose Jacobian is A.

Variables:

matrix (Array) – (m, n) design matrix A.
offset (Array) – (m,) offset b.

Parameters:

matrix (Array)
offset (Array)

matrix: Array¶

offset: Array¶

exact_gradient_sum_of_squares(params, observed)[source]¶

Return d/dx sum((A x + b - y)^2) = 2 A^T (A x + b - y).

Parameters:

params (Array) – (n,) parameter vector.
observed (Array) – (m,) data y.

Returns:

(n,) exact gradient.

Return type:

Array

class mimirax.testing.doubles.SoftenedPointMassField(sources, source_masses, softening=0.0, gravitational_constant=1.0)[source]¶

Bases: object

Softened Newtonian field of fixed sources, evaluated at test particles.

a_i = -G sum_j m_j d_ij / (|d_ij|^2 + eps^2)^{3/2} with d_ij = x_i - s_j. A ForceModel whose masses argument is ignored (the targets are test particles), and the analytic stand-in for “accelerations at tracers” in Paper I section 7.

Variables:

sources (Array) – (s, 3) source positions.
source_masses (Array) – (s,) source masses.
softening (float) – Plummer softening eps.
gravitational_constant (float) – G.

Parameters:

sources (Array)
source_masses (Array)
softening (float)
gravitational_constant (float)

sources: Array¶

source_masses: Array¶

softening: float = 0.0¶

gravitational_constant: float = 1.0¶

accelerations(positions, masses, *, args=None)[source]¶

Return the field at every target.

Parameters:

positions (Vec3) – (n, 3) target positions.
masses (PerParticle) – Ignored: targets are test particles.
args (object) – Ignored.

Returns:

(n, 3) accelerations.

Return type:

Vec3

tidal_tensor(positions)[source]¶

Return d a_i / d x_i for every target, in closed form.

Parameters:: positions (Vec3) – (n, 3) target positions.
Returns:: (n, 3, 3) with [i, a, b] = d a_i[a] / d x_i[b]; equal to the diagonal blocks of jax.jacfwd(accelerations).
Return type:: Array

mimirax.testing.doubles.make_linear_problem(key, *, n_params=4, n_data=12, noise=0.0, sigma=1.0)[source]¶

Build a well-posed linear-Gaussian problem with a known truth.

Parameters:

key (Array) – A jax.random key.
n_params (int) – Number of free parameters. Static.
n_data (int) – Number of data points; must exceed n_params for a unique optimum. Static.
noise (float) – Standard deviation of Gaussian noise added to the synthetic data.
sigma (float) – The likelihood’s assumed noise level.

Returns:

The problem (linear forward model, identity observable, Gaussian likelihood, no priors) and the (n_params,) true parameters.

Return type:

tuple[InferenceProblem, Array]

Registry¶

The registry idiom every pluggable family in mimirax uses.

Observables, likelihoods, priors, reparameterizations and inference methods are all registries: a module-level mapping from a name to an implementation, a register call that refuses a silent overwrite, and an available query. This is yggdrax’s register_tree_builder / available_tree_types pattern (yggdrax/tree.py) made generic, so a reader of one package reads the other.

WHY A REGISTRY AND NOT A BRANCH. Extensibility is the package’s primary design axis: adding an observable or a sampler must not require editing a core module. A registry gives a new implementation one call to make instead of one elif to add, and it gives the error message for a typo the full list of what exists.

class mimirax._registry.Registry(kind)[source]¶

Bases: Generic[T]

A name -> implementation mapping with explicit registration.

Parameters:: kind (str) – What the registry holds, used in error messages – "observable", "likelihood", …

property kind: str¶

Return the kind of implementation this registry holds.

Returns:: The kind given at construction.
Return type:: str

register(name, item, *, overwrite=False)[source]¶

Parameters:

name (str) – Lookup key. Leading and trailing whitespace is stripped.
item (T) – The implementation to register.
overwrite (bool) – Whether an existing entry under name may be replaced. Off by default so two plugins cannot silently shadow each other.

Returns:

item, unchanged, so the call can wrap a definition.

Return type:

Raises:

ValueError – If name is empty, or already registered and overwrite is False.

decorate(name, *, overwrite=False)[source]¶

Return a decorator that registers the decorated object under name.

Parameters:

name (str) – Lookup key.
overwrite (bool) – Passed through to register().

Returns:

A decorator returning its argument unchanged.

Return type:

Callable[[T], T]

get(name)[source]¶

Look up the implementation registered under name.

Parameters:: name (str) – Lookup key.
Returns:: The registered implementation.
Return type:: T
Raises:: KeyError – If nothing is registered under name. The message lists every registered name.

available()[source]¶

Return the registered names, sorted.

Returns:: Registered names in sorted order.
Return type:: tuple[str, …]