mixle.inference package

The inference concern — fit a model and quantify its parameters.

One home for turning data into a fitted/posterior model. Every entry point is the same idea — infer parameters from data — and differs only in what it requires of its input:

  • closed-form conjugate Bayes (conjugate_posterior) — needs a ConjugateUpdatable family;

  • MLE / EM / MAP (fit / optimize / run_em) — needs an Estimator for the model;

  • sampling-based inference (nuts / advi) — needs a sampleable / differentiable target (a log-density callable, or a model that can be sampled). This is exactly why it belongs here and not in a separate package: it is inference under a capability precondition on the target, nothing more.

Everything physically lives in this package: the estimation / EM / fit / objectives / Fisher machinery (mixle.inference.{estimation,em,fit,objectives,fisher}), the MCMC samplers (mixle.inference.mcmc), the engine-agnostic NUTS/ADVI target facade (mixle.inference.target + .backends + .diagnostics). Conjugate Bayes is re-exported from its canonical home mixle.stats.bayes. mixle.infer remains as a deprecated shim onto this package.

These imports are eager and cycle-free: the machinery’s only mixle.stats dependency is the compute layer (mixle.stats.compute.{pdist,sequence}), never the mixle.stats package surface — the vectorized seq_* drivers were moved out of mixle.stats.__init__ into compute.sequence for exactly this reason.

class Explanation(total, parts=<factory>, responsibilities=None, component=None)[source]

Bases: object

Exact additive attribution of log p(x) (plus the latent posterior for mixtures).

Parameters:
total: float
parts: list[tuple[str, float]]
responsibilities: ndarray | None = None
component: int | None = None
most_anomalous(k=3)[source]
Parameters:

k (int)

Return type:

list[tuple[str, float]]

summary()[source]
Return type:

str

explain(model, x)[source]

Exact per-part attribution of model.log_density(x) (see module docstring).

Parameters:
Return type:

Explanation

class InterventionalNetwork(net, interventions)[source]

Bases: object

A Bayesian network under do(...): sample and summarize the post-intervention world.

Parameters:

  • net (Any)

  • interventions (dict[int, Any])

sample(size=1, *, seed=None)[source]

Ancestral sampling with the intervened fields clamped (their factors are never consulted).

Parameters:

  • size (int)

  • seed (int | None)

Return type:

list[tuple]

expectation(field, *, n=4000, seed=0)[source]

Monte-Carlo E[field | do(...)] for a numeric field.

Parameters:
Return type:

float

distribution(field, *, n=4000, seed=0)[source]

Interventional marginal of a discrete field as {value: probability}.

Parameters:
Return type:

dict[Any, float]

average_causal_effect(net, treatment, a, b, outcome, *, n=4000, seed=0)[source]

E[outcome | do(treatment=a)] - E[outcome | do(treatment=b)] (numeric outcome).

Parameters:
Return type:

float

counterfactual(net, observed, interventions)[source]

What THIS observed record would have been under the intervention (abduction-action-prediction).

Per Pearl’s three steps, walked in topological order:

  • abduction – a linear-Gaussian field’s exogenous noise is point-identified from the row: its residual eps = observed - coef @ parents_observed;

  • action – intervened fields take their do values;

  • prediction – the SAME residual replays through the counterfactual parents: cf = coef @ parents_cf + eps.

Honest boundaries: (1) a field that is not linear-Gaussian keeps its observed value only while its parents are unchanged under the intervention (that much IS identified); if its parents change, its exogenous noise cannot be recovered from one observation and this raises — use average_causal_effect() for the population answer instead of a guessed individual one. (2) The counterfactual is relative to the network’s DAG as given: purely observational structure learning cannot orient Markov-equivalent edges (x -> y and y -> x fit equally well), so if the causal direction matters, assert it from domain knowledge rather than trusting the learned arrow.

Parameters:
Return type:

tuple

do(net, interventions)[source]

Return the network under Pearl’s do operator (see module docstring).

Parameters:
Return type:

InterventionalNetwork

class Forecast(mean, lo, hi, level, state_probs, samples=None)[source]

Bases: object

Per-step predictive summaries plus the state-marginal trajectory.

Parameters:
mean: Any
lo: Any
hi: Any
level: float
state_probs: ndarray
samples: Any = None
forecast(model, history, horizon, *, level=0.9, n=4000, seed=0, keep_samples=False)[source]

Forecast horizon steps beyond history under a fitted HMM.

Parameters:

  • model (Any) – a fitted HiddenMarkovModelDistribution (any emission family with a sampler).

  • history (Any) – the observed sequence to condition on (one sequence).

  • horizon (int) – number of future steps to predict.

  • level (float) – central-interval mass (0.9 -> the 5%..95% band).

  • n (int) – Monte-Carlo draws per step for the emission quantiles (state marginals are exact).

  • seed (int) – reproducibility.

  • keep_samples (bool) – also return the raw (H, n) predictive draws (scalar emissions only).

Return type:

Forecast

class ParameterEstimator[source]

Bases: ABC, Generic[SS]

Estimates a distribution from accumulated sufficient statistics.

accumulator_factory() supplies accumulators that gather sufficient statistics of type SS, and estimate(nobs, suff_stat) maps those statistics (plus optional regularization configured on the estimator) to a new distribution.

to_dict()[source]

Return a safe JSON-compatible representation of this estimator.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Reconstruct an estimator from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

ParameterEstimator

to_json(**kwargs)[source]

Serialize this estimator as safe strict JSON.

Parameters:

kwargs (Any)

Return type:

str

classmethod from_json(text)[source]

Deserialize an estimator from to_json output.

Parameters:

text (str)

Return type:

ParameterEstimator

abstractmethod estimate(nobs, suff_stat)[source]
Parameters:

  • nobs (float | None)

  • suff_stat (SS)

Return type:

SequenceEncodableProbabilityDistribution

abstractmethod accumulator_factory()[source]
Return type:

StatisticAccumulatorFactory

resident_accumulation_supported()[source]

Return whether engine-resident (fixed-width) sufficient statistics suffice for estimate.

Most exponential-family M-steps consume only the resident sufficient statistics, so the default is True. Estimators whose M-step needs more than that (e.g. a full count histogram for the negative-binomial dispersion solve) override this to False so stacked/generated kernels fall back to the host accumulator, keeping every backend’s fixed point identical.

Return type:

bool

get_prior()[source]

Return the parameter prior configured on this estimator, if any.

The unified estimation contract treats the prior as the single regularization concept: None gives maximum likelihood, a conjugate prior gives the Bayesian posterior update inside estimate. The default reads the prior attribute (None when unset).

Return type:

ProbabilityDistribution | None

model_log_density(model)[source]

Return the prior log-density of model’s parameters (the ELBO global term).

Used by the variational/MAP objective in fit. The default is 0.0 (no prior); conjugate estimators override this to evaluate their prior at the model’s parameters, mapping to the prior’s parameterization first.

Parameters:

model (ProbabilityDistribution)

Return type:

float

estimate(data, estimator, prev_estimate=None)[source]

Perform E-step in EM algorithm by iterating over all observations in ‘data’.

Arg estimator must be consistent with prev_estimate. That is, prev_estimate must be an estimate that could be obtained from estimator.

Data must type consistent with estimator and prev_estimate.

Returns the next iteration of EM algorithm by iterating over each observation of data. See seq_estimate() for a more computationally efficient implementation.

Parameters:

  • data (Union[Sequence[T], pyspark.rdd.RDD]) – Sequence of iid observations of data type consistent with ‘estimator’ and/or ‘prev_estimate’.

  • estimator (ParameterEstimator) – Model to be estimated from ‘data’.

  • prev_estimate (Optional[SequenceEncodableProbabilityDistribution]) – Previous estimate of EM algorithm. Must be included for distributions that require initialization.

Returns:

SequenceEncodableProbabilityDistribution object.

Return type:

SequenceEncodableProbabilityDistribution

initialize(data, estimator, rng, p=0.1)[source]

Randomly initialize a model corresponding to ParameterEstimator for iid observations data.

Note: ParameterEstimator must be of data type T, matching the input data.

This function sequentially iterates over the entire data set ‘data’, repeatedly calling initialize() method of the SequenceEncodableStatisticAccumulator object created from ‘estimator’. Data points are weighted 0 or 1 with probability p.

Seq_initialize() is much more efficient, and should produce the same initialized model for the same data sets.

Parameters:

  • data (Union[Sequence[T], pyspark.rdd.RDD]) – Set of iid observations compatible with ‘estimator’.

  • estimator (ParameterEstimator) – ParameterEstimator object for desired model to be estimated from data.

  • rng (RandomState) – RandomState object for setting seed.

  • p (float) – Proportion of data to randomly sample for initializing model.

Returns:

SequenceEncodableProbabilityDistribution object consistent with ‘estimator’.

Return type:

SequenceEncodableProbabilityDistribution

seq_estimate(enc_data, estimator, prev_estimate)[source]

Perform vectorized E-step in EM algorithm for encoded sequence of observations in ‘enc_data’.

Arg estimator must be consistent with prev_estimate. That is, prev_estimate must be an estimate that could be obtained from estimator.

Arg enc_data must type consistent with estimator and prev_estimate (result of seq_encode() call).

Returns the next iteration of EM algorithm with vectorized calls to “seq_update()” of the corresponding SequenceEncodableStatsiticAccumulator objects.

Parameters:

  • enc_data (Union[List[Tuple[int, T]], 'pyspark.rdd.RDD']) – Sequence encoded data of format matching output of seq_encode() function.

  • estimator (ParameterEstimator) – Model to be estimated from ‘enc_data’.

  • prev_estimate (SequenceEncodableProbabilityDistribution) – Previous estimate of EM algorithm.

Returns:

SequenceEncodableProbabilityDistribution object.

Return type:

T_D

seq_initialize(enc_data, estimator, rng, p=0.1)[source]
Vectorized initialization of a model corresponding to ParameterEstimator for encoded sequences of iid data

observations.

Arg enc_data must type consistent with estimator (result of seq_encode() call). Arg estimator must be of data type consistent with encoded sequence data type in ‘enc_data’.

Vectorized initialization of SequenceEncodableProbabilityDistribution corresponding to ‘estimator’ from enc_data. Observations in the encoded sequence enc_data are kept with probability p.

This functions relies on calls to SequenceEncodableStatisticAccumulator.seq_initialize(), which is a vectorized initialization of the SequenceEncodableStatisticAccumulator object.

This method should produce the same initialized model as a call to initialize() if the data sets are the same.

Parameters:

  • enc_data (Union[List[Tuple[int, T]], 'pyspark.rdd.RDD']) – Sequence encoded data of format matching output of seq_encode() function.

  • estimator (ParameterEstimator) – Model to be estimated from ‘enc_data’.

  • rng (RandomState) – RandomState object for setting seed.

  • p (float) – Proportion of data to randomly sample for initializing model.

Returns:

SequenceEncodableProbabilityDistribution object consistent with ‘estimator’.

Return type:

SequenceEncodableProbabilityDistribution

optimize(data, estimator=None, max_its=10, delta=1.0e-9, init_estimator=None, init_p=0.1, rng=None, prev_estimate=None, vdata=None, enc_data=None, enc_vdata=None, out=sys.stdout, print_iter=1, num_chunks=1, engine=None, precision=None, fields=None, resources=None, placement=None, sub_chunks=1, chunk_size=None, backend='local', num_workers=None, client=None, comm=None, root=0, root_only=False, strategy=None, reuse_estep_ll=True, objective='auto', on_step=None, structure='auto')[source]
Estimation of ‘estimator’ via EM algorithm for max_its iterations or until

new_loglikelihood - old_loglikelihood < delta.

Parameters:

  • data (Optional[List[T]]) – List of data type T containing observed data. Must be compatible with data type of estimator.

  • estimator (ParameterEstimator | ProbabilityDistribution | None) – What to fit. A ParameterEstimator is used as-is; a distribution prototype (any ProbabilityDistribution) is coerced to its matching estimator via proto.estimator() so you build the model shape only once; None infers an estimator from raw data (mixle.utils.automatic.get_estimator).

  • max_its (int) – Maximum number of EM iterations to be performed. Default value is 10 iterations.

  • delta (Optional[float]) – Stopping criteria for EM algorithm used if max_its is not set: Iterate until |old_loglikelihood - new_loglikelihood| < delta or iterations == max_its.

  • init_estimator (Optional[ParameterEstimator]) – ParameterEstimator to used to initialize EM algorithm parameters. If None, estimator is used. Must be consistent with estimator.

  • init_p (float) – Value in (0.0,1.0] for randomizing the proportion of data points used in initialization.

  • rng (RandomState) – RandomState used to set seed for initializing EM algorithm.

  • vdata (Optional[Sequence[T]]) – Optional validation set.

  • prev_estimate (Optional[SeqeuenceEncodableProbabilityDistribution]) – Optional model estimate used from prior fitting. Must be consistent with estimator.

  • enc_data (Optional[List[Tuple[int, E]]]) – Optional encoded data of form List[Tuple[int, E]]. Formed from data if None.

  • enc_vdata (Optional[List[Tuple[int, E0]]]) – Optional sequence encoded validation set.

  • out (IO) – IO stream to write out iterations of EM algorithm. Pass out=None to silence all output.

  • print_iter (int) – Print the log-likelihood difference every print_iter iterations; the final converged iteration is always reported. Pass print_iter=0 to suppress the periodic lines (keeping only the converged line), or out=None to silence entirely.

  • num_chunks (int) – Number of chunks for encoded data.

  • engine (Optional[Any]) – Optional ComputeEngine for local kernel scoring/accumulation. Distributed engine placement is intentionally deferred to the orchestrator/planner layer.

  • precision (Optional[Any]) – Optional floating-point precision such as 'float32' or np.float64. Pass 'auto' to let mixle.engines.auto_precision choose from the data and engine: float32 only on a GPU torch engine with well-conditioned numeric data, else float64. Pass 'minimal' for the data-aware CPU allocator (mixle.inference.precision_plan): it inspects the data magnitude and the model’s leaf families/conditioning and runs the reduced float32 fused kernel where verified safe (accumulation stays float64), else float64 – the “preserve accuracy with minimal compute” default for local fits.

  • fields (Optional[Any]) – DataFrame column/field selection. A single field yields scalar observations; several fields yield tuple observations unless the estimator/model is record-shaped, in which case dict records are produced by source column name.

  • resources (Optional[Any]) – Optional planner resources. When supplied with raw data, optimize encodes through the shared encoded-data factory so placement, sub-chunks, and per-shard engines use the orchestrator contract.

  • placement (Optional[Any]) – Optional explicit placement produced by mixle.utils.parallel.planner.plan.

  • sub_chunks (int) – Number of sub-chunks per placement shard when resources or placement is supplied.

  • chunk_size (Optional[int]) – Approximate chunk size for ordinary local sequence encoding.

  • backend (str) – Encoded-data backend for raw data. 'local' keeps the historical local encoding unless resources/placement are supplied; 'mp' and 'mpi' use the shared encoded-data factory.

  • num_workers (Optional[int]) – Worker count for backend='mp' and optional partition count hint for backend='dask'.

  • client (Optional[Any]) – Existing dask.distributed client for backend='dask'. If omitted, the dask backend uses an active default client or starts a local threaded client.

  • comm (Optional[Any]) – MPI communicator for backend='mpi'.

  • root (int) – MPI root rank for backend='mpi'.

  • root_only (bool) – MPI root-only data mode for backend='mpi'.

  • strategy (Optional[Any]) – Optional EM strategy from mixle.inference.em (e.g. AnnealedEM, HardEM, MonteCarloEM) or any callable (enc, estimator, model) -> model to use in place of the standard exact E/M step. None uses the standard step.

  • reuse_estep_ll (bool) – Default True. Reuse the data log-likelihood computed during the E-step (the posterior normalizer / forward pass / variational ELBO) for convergence instead of running a separate scoring pass each iteration – typically ~1.5-2x faster per iteration for latent models (mixtures, HMMs and variants, topic models, associations, IBP, …) on the default local engine. Convergence then lags by one iteration (same fixed point) and the best-likelihood model is returned; fixed-iteration fits (delta=None) are identical to the standard loop. Automatically falls back to the standard loop for engines/strategies/ distributed backends or models that can’t report the LL (no slowdown there). Set False to force the exact historical per-iteration scoring behavior.

  • objective (str) – Convergence/selection objective. 'auto' (default) makes the prior the single switch – a model exposing a variational ELBO (seq_local_elbo) is fit by variational Bayes ('vb'), an estimator carrying a parameter prior by penalized log-likelihood ('map'), and everything else by plain maximum likelihood ('mle'). Pass 'mle' / 'map' / 'vb' to force a specific objective. fit accepts the same argument; both share this resolution so a Bayesian estimator is fit on the correct objective regardless of the verb used. (Only 'mle' is compatible with the fused E-step shortcut; reuse_estep_ll is ignored for 'map'/'vb'.)

  • on_step (Optional[Callable[[EMStep], None]]) – Optional per-iteration callback receiving an EMStep (iter, model, log_density, delta) for the accepted model. Use it to checkpoint a long run – e.g. on_step=registry.checkpointer('run', every=5) – and resume with prev_estimate=. Called on every iteration regardless of print_iter.

  • structure (str) – 'auto' (default) makes the tagline literal for flat tuple records fit with no estimator: the cross-field dependency graph is discovered (mixle.inference.learn_bayesian_network()) and returned WHEN it beats the independent composite by BIC — otherwise (no edges, non-record data, or any failure) the historical automatic-composite path proceeds untouched. 'off' restores the unconditional historical behavior. Only consulted when estimator is None and no prev_estimate/init_estimator/strategy/enc_data is supplied.

Returns:

SequenceEncodableProbabilityDistribution corresponding to estimator when stopping criteria of EM algorithm

is met.

Return type:

SequenceEncodableProbabilityDistribution

jit_seq_log_density(model, engine=None)[source]

Return a JittedScorer: the whole-tree log-density of model compiled to one XLA program.

jit_seq_log_density(model)(data) is bit-identical to model.seq_log_density(encode(data)) but runs as a single jax.jit XLA program over the entire composite tree – fast for repeated scoring of a fixed model. Requires the JAX optional extra and that every leaf in model declares JAX support (raises EngineNotSupportedError-style errors from the engine layer otherwise).

Parameters:
Return type:

JittedScorer

jit_em_mixture(model, data, *, max_its=100, engine=None)[source]

Fit a finite mixture of same-family scalar exponential-family leaves by EM, with the ENTIRE EM loop (every E-step + closed-form weighted M-step iteration) compiled to ONE jax.jit XLA program via lax.scan – the parameters are traced inputs threaded through the loop on-device, so there is no per-iteration recompile and no per-iteration host sync. This is roadmap A2 bullet 1 (“repeated EM iterations run as one XLA program”) realized literally.

model is the initial mixture (its components seed the EM); supported leaves: Gaussian, Poisson, Exponential. Runs a fixed max_its iterations (no host-side early stop – that is the point: the loop stays on-device). Returns a fitted MixtureDistribution, bit-close to the host EM from the same start (it is the same EM update). Raises NotImplementedError for unsupported structure.

SPEED – measured, honest: the payoff is GPU/TPU and large scale, where XLA parallelizes the E-step over millions of points and many components. On an Apple M4 GPU (via jax-metal) this kernel runs ~21x faster than mixle’s vectorized NumPy EM (K=10, N=1e6, 50 iters: 156 ms vs 3254 ms) with identical estimates, and ~12x faster than the same jitted loop on CPU. On CPU it is *not* a speedup – mixle’s host EM is already vectorized NumPy (+ a fused path) and a long sequential scan of small steps loses to it (~0.1-0.8x for K up to 40). So: GPU -> big win, CPU -> use the host EM. (The CPU win from A2 is the single-pass scoring jit, jit_seq_log_density(), ~8x.) Note: jax-metal is version-pinned – the GPU result above used Python 3.11 + jax/jaxlib 0.4.34 + jax-metal 0.1.1; newer jaxlib emits StableHLO that jax-metal 0.1.1 cannot compile.

Parameters:
class JittedScorer(model, engine=None)[source]

Bases: object

A jax.jit-compiled whole-tree log-density scorer for a fixed model.

Calling the scorer encodes data, runs the engine-neutral backend_seq_log_density over the whole model tree on the JAX engine under jax.jit, and returns host log-densities. The compiled program is cached and reused across calls with the same data shape.

Parameters:

  • model (Any)

  • engine (Any)

laplace_posterior(model, data, *, eps=1e-4, ridge=1e-6)[source]

Laplace posterior over model’s parameters from a finite-difference Hessian of its own seq_log_density – works for ANY model whose parameters _flatten() covers (the scalar exponential-family leaves, Composite and Mixture, recursively), conjugate or not, with no per-model inference code. model should be the fitted (MLE/MAP) model – its parameters are the Laplace mode. Returns a LaplacePosterior you can .sample() (a fitted model per draw).

Parameters:
Return type:

LaplacePosterior

class LaplacePosterior(mode_model, u_mode, cov, rebuild)[source]

Bases: object

Gaussian Laplace posterior over a model’s parameters (in the unconstrained space), with draws rebuilt back into fitted models. mean_model is the mode; cov the unconstrained covariance.

sample(n=1, rng=None)[source]
Parameters:

n (int)

summary()[source]
Return type:

dict

collapse_mixture(mixture)[source]

Moment-match a mixture onto a single distribution, in closed form (exact for Gaussians).

Returns the Gaussian minimizing KL(mixture || Gaussian) – its mean and covariance are the mixture’s overall mean and covariance (law of total variance). This is the exact, sample-free counterpart to mixle.ops.project(mixture, GaussianDistribution(...).estimator()).

Parameters:

mixture (Any)

Return type:

Any

reduce_mixture(mixture, n_components, *, method='runnalls')[source]

Reduce a Gaussian mixture to n_components by greedily merging the least-costly pair (Runnalls).

Repeatedly merges the two components whose merge costs the least KL (_runnalls_cost()) until n_components remain. Every merge is moment-preserving, so the reduced mixture has the same overall mean and covariance as the original – only higher moments are lost. Returns a mixture of the same kind as the input (a GaussianMixtureDistribution for multivariate input).

Parameters:

  • mixture (Any) – a Gaussian mixture (GaussianMixtureDistribution or a mixture of Gaussian components).

  • n_components (int) – target number of components (>= 1); no-op if already <= the current count.

  • method (str) – only "runnalls" for now.

Return type:

Any

moment_project(teacher, target=None, *, exact=True, **sampling_kw)[source]

Project teacher onto a smaller student – exactly when possible, else by sampling.

If teacher is a Gaussian mixture and target is None (or a single Gaussian family), the projection is the closed-form collapse_mixture() – no samples, machine-precision. Otherwise (or when exact=False) it delegates to mixle.ops.project(), the sampling M-projection onto target’s family. This gives one entry point that is exact where the structure allows and honest (sampling, clearly) where it does not.

Parameters:
Return type:

Any

gaussian_kl(p, q)[source]

KL(p || q) between two Gaussians (uni- or multivariate), in nats. Accepts distributions.

0.5[ tr(Σq⁻¹ Σp) + (μq-μp)ᵀ Σq⁻¹ (μq-μp) - d + ln(det Σq / det Σp) ] – the analytic Gaussian KL.

Parameters:
Return type:

float

fisher_merge(estimates, fishers=None)[source]

Fisher-weighted merge of parameter estimates – the closed-form Laplace-posterior combination.

Given estimates θ_i (each a flat parameter vector) and their Fisher information F_i, returns θ* = (Σ F_i)⁻¹ (Σ F_i θ_i) – the point that maximizes the sum of the local Laplace log-posteriors Σ_i -½(θ-θ_i)ᵀ F_i (θ-θ_i). This is Matena & Raffel Fisher merging (diagonal F) and, in general, the precision-weighted mean; for Gaussians it coincides with the product-of-experts mean. No gradient steps – a single linear solve.

Parameters:

  • estimates (Any) – sequence of parameter vectors θ_i (each shape (p,)), or a stack (m, p).

  • fishers (Any) – per-estimate Fisher information. Each may be a scalar/1-D vector (diagonal Fisher, per-coordinate precision) or a (p, p) matrix (full Fisher). None uses unit Fisher (a plain average). A single value is broadcast to every estimate.

Returns:

The merged parameter vector θ* of shape (p,).

Return type:

ndarray

fit(data, estimator=None, max_its=10, delta=1.0e-6, init_estimator=None, **kwargs)[source]

Fit a model in the Bayesian (variational / MAP) sense, returning the posterior-bearing model.

This is the posterior-returning counterpart of optimize(). fit iterates the EM/VB update that maximizes the objective selected by objective (default 'auto'):

  • 'auto' – the prior is the single switch: 'vb' when the model exposes seq_local_elbo, 'map' when the estimator carries a parameter prior, else 'mle';

  • 'mle' – plain data log-likelihood (ignores any prior in the objective);

  • 'map' / 'vb' – penalized log-likelihood / ELBO obj = data term + prior term, where the data term is the observed-data LL (MAP) or local-ELBO contributions (variational), and the prior term is estimator.model_log_density(model).

Convergence is checked on the chosen objective, so under 'map'/'vb' the prior is part of the stopping rule and conjugate updates never decrease it. The returned model carries its conjugate posterior forward as model.get_prior(). With no prior anywhere, every objective reduces to plain EM, so fit and optimize agree for frequentist estimators. optimize accepts the same objective argument; the two share this resolution so a Bayesian estimator is fit correctly regardless of which verb the caller reaches for.

Args otherwise mirror optimize() (local encoded path). Returns the model with the best validation log-likelihood seen during the run.

fit is a thin wrapper over optimize() – they share the one EM/objective loop. fit adds only the opt-in data-structure check, a Bayesian-leaning default delta (1e-6), and the exact per-iteration-scored loop (reuse_estep_ll=False). Every other optimize() keyword – engines, precision, distributed backend, on_step, the fused E-step – is accepted here too and forwarded verbatim, so reaching for a heavier knob never means switching verbs. estimator accepts the same three spellings as optimize() (estimator, distribution prototype, or None to infer from data).

Parameters:

  • data (Sequence[T] | None)

  • estimator (ParameterEstimator | ProbabilityDistribution | None)

  • max_its (int)

  • delta (float | None)

  • init_estimator (ParameterEstimator | ProbabilityDistribution | None)

  • kwargs (Any)

Return type:

SequenceEncodableProbabilityDistribution

class EMStep(iter, model, log_density, delta)[source]

Bases: NamedTuple

One accepted EM iteration, handed to an optimize(on_step=...) callback.

iter is 1-based; model is the current accepted model – snapshot it to checkpoint, and resume with optimize(prev_estimate=...); log_density is the training objective at this step; delta is its gain over the previous step (inf on the first iteration).

Parameters:
iter: int

Alias for field number 0

model: Any

Alias for field number 1

log_density: float

Alias for field number 2

delta: float

Alias for field number 3

best_of(data, vdata, est, trials, max_its, init_p, delta, rng, init_estimator=None, enc_data=None, enc_vdata=None, out=sys.stdout, print_iter=1, reuse_estep_ll=True, objective='auto')[source]
Performs EM algorithm for trials-number of randomized initial conditions. Returns the best model fit in terms of

maximum log-likelihood value from validation data.

Parameters:

  • data (Optional[List[T]]) – List of data of type T. If None is given, enc_data must be provided as List[Tuple[int, enc_data_type]].

  • vdata (Optional[Sequence[T]]) – Optional validation set.

  • est (ParameterEstimator) – ParameterEstimator for model to be estimated.

  • trials (int) – Integer number >= 1, of randomized initial conditions to perform EM algorithm for.

  • max_its (int) – Integer value >=1, sets the maximum number of iterations of EM to be performed as stopping criteria.

  • init_p (float) – Value in (0.0,1.0] for randomizing the proportion of data points used in initialization.

  • delta (float) – Stopping criteria for EM when |old-log-likelihood - new-log-likelihood| < delta.

  • rng (RandomState) – RandomState for setting seed.

  • init_estimator (Optional[ParameterEstimator]) – Optional ParameterEstimator used for fitting.

  • enc_data (Optional[List[Tuple[int, E]]]) – Optional encoded data, if provided data need not be provided. If None, enc_data is set from data.

  • enc_vdata (Optional[List[Tuple[int, E0]]]) – Optional sequence encoded validation set.

  • out (I0) – Text output stream.

  • print_iter (int) – Print iterations (i.e. log-likelihood difference) every print_iter-iterations.

  • reuse_estep_ll (bool) – Default True. Forwarded to each trial’s optimize call – reuse the E-step likelihood for convergence instead of a separate scoring pass (see optimize). Set False to force the exact historical per-iteration scoring behavior.

  • objective (str) – Convergence/selection objective forwarded to each trial’s optimize call; 'auto' (default) selects MLE / MAP / variational Bayes from the prior (see optimize).

Returns:

Tuple of log-likelihood of best fitting model and the best fitting model from number of trials.

Return type:

tuple[float, SequenceEncodableProbabilityDistribution]

run_em(enc_data, estimator, initial_model, strategy=None, max_its=10, delta=1.0e-9, engine=None, objective=None, max_iter=None)[source]

Run an EM-family strategy until convergence or max_its.

objective takes the same values as optimize(): None (MLE), a selection string ('auto' / 'mle' / 'map' / 'vb'), or a ready model -> float callable. max_its is the canonical iteration-cap spelling (matching optimize / fit / best_of); max_iter is accepted as a back-compat alias and overrides max_its when given.

Parameters:

  • enc_data (Any)

  • estimator (ParameterEstimator)

  • initial_model (SequenceEncodableProbabilityDistribution)

  • strategy (EMStrategy | None)

  • max_its (int)

  • delta (float | None)

  • engine (Any | None)

  • objective (str | Callable[[Any], float] | None)

  • max_iter (int | None)

Return type:

SequenceEncodableProbabilityDistribution

class EMStrategy(*args, **kwargs)[source]

Bases: Protocol

Structural contract for an EM-family strategy consumed by run_em().

Every strategy object in this module (StandardEM, PosteriorTransformEM, AnnealedEM, …) satisfies this Protocol structurally by exposing a step(...) -> EMStepResult method. run_em and _em_step_fn dispatch on it polymorphically; membership is decided by isinstance().

step(enc_data, estimator, model, engine=..., objective=...)[source]
Parameters:

  • enc_data (Any)

  • estimator (ParameterEstimator)

  • model (SequenceEncodableProbabilityDistribution)

  • engine (Any | None)

  • objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class StreamingEstimator(estimator, schedule=None, model=None, init_estimator=None, init_p=0.1, rng=None, encoder=None, num_chunks=1)[source]

Bases: _StreamingBase

Decay-mode online estimator built from accumulator scaling and M-steps.

Parameters:

  • estimator (ParameterEstimator)

  • model (SequenceEncodableProbabilityDistribution | None)

  • init_estimator (ParameterEstimator | None)

  • init_p (float)

  • rng (RandomState | None)

  • num_chunks (int)

update(data=None, *, enc_data=None)[source]

Consume one batch and return the updated model.

Parameters:
Return type:

SequenceEncodableProbabilityDistribution

class IncrementalEstimator(estimator, model=None, init_estimator=None, init_p=0.1, rng=None, encoder=None, num_chunks=1)[source]

Bases: _StreamingBase

Neal-Hinton style incremental EM over replaceable data chunks.

Each chunk contributes a sufficient-statistic payload computed under the current model. Revisiting a chunk subtracts that chunk’s previous payload, adds the new payload, and runs the ordinary estimator M-step on the pooled statistics. No distribution-specific estimation code lives here; the class only uses scale(-1), combine(), and estimate().

Parameters:

  • estimator (ParameterEstimator)

  • model (SequenceEncodableProbabilityDistribution | None)

  • init_estimator (ParameterEstimator | None)

  • init_p (float)

  • rng (RandomState | None)

  • num_chunks (int)

update(data=None, *, enc_data=None, chunk_id=None)[source]

Replace one chunk contribution and return the updated model.

chunk_id is keyword-only so this matches StreamingEstimator.update()’s (data, *, enc_data) shape across the streaming surface; it is required (a None chunk_id raises) because the Neal-Hinton update keys each batch’s contribution by it.

Parameters:
Return type:

SequenceEncodableProbabilityDistribution

chunk_value(chunk_id)[source]

Return a copy of one stored chunk contribution.

Parameters:

chunk_id (Any)

reset()[source]

Drop all chunk contributions and fitted model state.

Return type:

None

class BayesianStreamingEstimator(estimator, mode='posterior_carry', schedule=None, model=None, init_estimator=None, init_p=0.1, rng=None, num_chunks=1)[source]

Bases: object

Streaming / recursive-Bayes driver over the mixle.stats estimator protocol.

mode='posterior_carry' (the default) performs exact recursive conjugate updating: each fitted posterior is carried forward as the next batch’s prior by rebuilding the estimator from the fitted model (model.estimator() returns an estimator whose prior is the model’s posterior).

mode='forgetting' applies a power-prior step: the current batch’s accumulated sufficient statistics are scaled by rho = schedule(step) (via the accumulator’s scale, which preserves structural support metadata such as a categorical’s min_val) before the ordinary conjugate estimate call, and the batch’s contribution to nobs is scaled to match.

The public surface is BayesianStreamingEstimator(estimator, mode=..., schedule=...) plus .update(data=None, enc_data=None) and .reset().

Parameters:

  • estimator (ParameterEstimator)

  • mode (str | None)

  • schedule (Any | None)

  • model (SequenceEncodableProbabilityDistribution | None)

  • init_estimator (ParameterEstimator | None)

  • init_p (float)

  • rng (RandomState | None)

  • num_chunks (int)

update(data=None, enc_data=None)[source]

Consume one batch and return the updated posterior-bearing model.

Parameters:
Return type:

SequenceEncodableProbabilityDistribution

reset()[source]

Drop the current model and stream counters.

Return type:

None

bayes_action(posterior, loss, actions, *, n=2000, seed=0, cvar_alpha=0.1, quantiles=(0.05, 0.5, 0.95))[source]

Pick the Bayes action: argmin_a E_{draw ~ posterior}[ loss(a, draw) ].

Parameters:

  • posterior (Any) – any object exposing samples(n, rng) – e.g. mixle.inference.posterior(model, data, over=...) (parameter, latent, or predictive).

  • loss (Callable[[Any, Any], float]) – loss(action, draw) -> float (or a numpy-vectorized loss(action, draws) -> array).

  • actions (Sequence[Any]) – the finite candidate-action set to minimise over.

  • n (int) – number of posterior draws for the Monte-Carlo expectation.

  • seed (int) – RNG seed for the posterior draw (reproducible).

  • cvar_alpha (float) – tail mass for the CVaR / VaR of the chosen action (0.1 -> worst 10%).

  • quantiles (Sequence[float]) – loss quantiles to report per action.

Returns:

{action, action_index, expected_loss, risk_profile, alternatives} – the chosen action, its expected loss, its tail-risk profile, and the expected loss of every candidate.

Raises:

  • ValueError – if actions is empty.

  • CapabilityError – if posterior does not expose the samples(n, rng) contract.

Return type:

dict[str, Any]

class RiskProfile(expected_loss, cvar, cvar_alpha, var, quantiles, std)[source]

Bases: object

The tail-risk summary of a single action’s posterior loss distribution.

Parameters:
expected_loss: float
cvar: float
cvar_alpha: float
var: float
quantiles: dict[float, float]
std: float
as_dict()[source]
Return type:

dict[str, Any]

certify(model, *, escape_tested=False, penalized=False)[source]

Return the EstimationCertificate for a fitted model (or distribution prototype).

Walks the model’s block structure and classifies each block’s estimation method + guarantee from its capability signals – no fitting is done here, only inspection. Pass escape_tested=True when the fit ran saddle-escape restarts (mixle.Model.fit() sets this automatically), which upgrades EM blocks from STATIONARY to STATIONARY_ESCAPE_TESTED.

Pass penalized (a reason string, or True) when the fit optimized a PENALIZED objective – soft constraints, conservation/PINN residual factors, potentials (E2). The optimum is then of the penalized surrogate, NOT the likelihood, so no block may claim more than STATIONARY however clean its own solver is: every stronger block is downgraded with the penalty named in its reason.

Parameters:
Return type:

EstimationCertificate

plan_estimation(model, *, escape_tested=False)[source]

Alias for certify() – the pre-fit planning view over a distribution prototype.

Parameters:
Return type:

EstimationCertificate

schedule(model, *, escape_tested=False)[source]

Plan the block-coordinate estimation schedule for model (planner v2, A3).

Built from the same block classification as certify(): EM blocks make the schedule a loop (E-step + per-block M-steps, repeated until convergence); without a latent block every block is one independent pass. Gradient blocks appear as explicit gradient passes with their pool placement, so the schedule is also the offload plan for the hybrid case.

Parameters:
Return type:

EstimationSchedule

class EstimationSchedule(passes=<factory>, latent=False)[source]

Bases: object

The block-coordinate schedule planner v2 produces (A3): ordered passes + the loop structure.

A fully-factorized model schedules one independent pass per block (no loop). A latent model schedules the EM loop explicitly: an E-step over the latent, then one M-step pass PER BLOCK per round – each M-step named with its own method, so the schedule shows exactly where the closed forms live inside the iteration (and which pass, if any, is the gradient block a pool would take).

Parameters:

  • passes (list[SchedulePass])

  • latent (bool)

passes: list[SchedulePass]
latent: bool = False
property per_round: list[SchedulePass]
property gradient_passes: list[SchedulePass]
describe()[source]
Return type:

str

class SchedulePass(order, kind, block, method, placement, repeat)[source]

Bases: object

One pass of the estimation schedule: what runs, on which block, how, where, and how often.

Parameters:
order: int
kind: str
block: str
method: str
placement: str
repeat: str
class EstimationCertificate(guarantee, blocks=<factory>, escape_tested=False)[source]

Bases: object

The auditable proof of how a model was (or would be) estimated: per-block plans + the aggregate.

guarantee is the minimum over blocks. why_not_adam summarizes where gradient optimization was used and why those blocks could not use a stronger closed-form or convex route.

Parameters:

  • guarantee (Guarantee)

  • blocks (list[BlockPlan])

  • escape_tested (bool)

guarantee: Guarantee
blocks: list[BlockPlan]
escape_tested: bool = False
property gradient_blocks: list[BlockPlan]
property closed_form_blocks: list[BlockPlan]
as_dict()[source]
Return type:

dict[str, Any]

why_not_adam()[source]

The audit: which blocks needed gradient descent and why – everything else got a stronger method.

Return type:

str

table()[source]
Return type:

str

class BlockPlan(name, kind, method, guarantee, gradient, placement, reason)[source]

Bases: object

The estimation plan for one block of a model: which method ran and how strong its guarantee is.

Parameters:

  • name (str)

  • kind (str)

  • method (str)

  • guarantee (Guarantee)

  • gradient (bool)

  • placement (str)

  • reason (str)

name: str
kind: str
method: str
guarantee: Guarantee
gradient: bool
placement: str
reason: str
class Guarantee(*values)[source]

Bases: IntEnum

How strong the solution to an estimation block is, as an ordered ladder (higher = stronger).

HEURISTIC = 1
STATIONARY = 2
STATIONARY_ESCAPE_TESTED = 3
GLOBAL = 4
GLOBAL_UNIQUE = 5
property label: str
uq(thing, data=None, *, alpha=0.1, equivalent=None)[source]

Quantify the uncertainty of thing, choosing the method from what thing is.

Parameters:

  • thing (Any) – a fitted mixle model, a torch module / point-predictor callable (or a list of them for a deep ensemble), or an LLM-style callable that maps a prompt to a generation.

  • data (Any) – for a mixle model, the fitting data (builds the Laplace posterior); for a point predictor, (X_cal, y_cal) calibration data; for an LLM, optional example prompts used to calibrate an abstention threshold.

  • alpha (float) – target miscoverage / abstention level (1 - alpha coverage).

  • equivalent (Callable[[Any, Any], bool] | None) – for the LLM path, an optional meaning-equivalence predicate over generations (default: exact string match after stripping).

Returns:

A UQResult exposing the method-appropriate accessors and its own calibration numbers.

Return type:

UQResult

class UQResult(kind, method, payload)[source]

Bases: object

The uncertainty of a predictor, with the method that produced it and receipts to check it.

Parameters:
kind: str
method: str
payload: dict[str, Any]
sample_models(n=200, *, seed=None)[source]

n fitted models drawn from the parameter posterior (epistemic ensemble).

Parameters:
Return type:

list[Any]

credible_interval(readout, alpha=0.1, *, n=400, seed=0)[source]

A 1-alpha credible interval on readout(model) over the parameter posterior.

Parameters:
Return type:

tuple[float, float]

interval(x, alpha=None)[source]

Calibrated prediction interval(s) at x. alpha overrides the calibrated level.

Parameters:
Return type:

tuple[ndarray, ndarray]

epistemic_std(x)[source]

Ensemble disagreement (std across members) at x – 0.0 for a single predictor.

Parameters:

x (Any)

Return type:

ndarray

semantic_entropy(prompt, *, n=8)[source]

Entropy (nats) over the meaning classes of n sampled generations for prompt.

Parameters:
Return type:

float

confident(prompt, *, n=8, max_entropy=None)[source]

True when semantic entropy is below the threshold – else the model disagrees with itself.

Parameters:
Return type:

bool

report()[source]
Return type:

dict[str, Any]

calibration_report(model, data)[source]

The calibration of model on held-out data (see module docstring).

data should be data the model was NOT fitted on – calibration measured on the training set is optimistic. Runs the PIT test when the model has a scalar predictive CDF; always reports the held-out mean log-density.

Parameters:
Return type:

CalibrationReport

class CalibrationReport(n, mean_log_density, pit_error=None, pit_histogram=None, bins=10, method='', note='')[source]

Bases: object

Whether a fitted model’s uncertainty is calibrated on held-out data.

pit_error is the total-variation distance of the PIT histogram from uniform (0 = perfectly calibrated). It has a finite-sample floor ~``sqrt(bins/n)`` even for a perfect model, so is_calibrated() judges against that floor rather than a fixed constant.

Parameters:
n: int
mean_log_density: float
pit_error: float | None = None
pit_histogram: dict[str, Any] | None = None
bins: int = 10
method: str = ''
note: str = ''
noise_floor()[source]

The PIT-error a perfectly calibrated model would show at this sample size (sampling noise).

Return type:

float

is_calibrated(tol=None)[source]

True when the PIT error is within tolerance. Default tol = 2.5x the finite-sample noise floor (so genuine miscalibration, not sampling noise, is what fails). Unknown -> False, conservatively.

Parameters:

tol (float | None)

Return type:

bool

as_dict()[source]
Return type:

dict[str, Any]

plan_placement(certificate, pool=None, *, telemetry=None)[source]

Decide local vs pool for every block of a certified fit (see module docstring).

Closed-form / convex / EM blocks always stay local. A gradient (pool_eligible) block goes to the pool only when pool.available and its estimated work clears pool.flop_threshold_tflop; otherwise it stays local too. Each decision carries a priced reason, and (with a pool) a placement telemetry event is emitted per block.

Parameters:

  • certificate (EstimationCertificate)

  • pool (PoolSpec | None)

  • telemetry (Any)

Return type:

PlacementPlan

class PlacementPlan(placements=<factory>, pool_spec=None)[source]

Bases: object

The full local-vs-pool plan for a model’s blocks, derived from its certificate + a PoolSpec.

Parameters:

  • placements (list[BlockPlacement])

  • pool_spec (PoolSpec | None)

placements: list[BlockPlacement]
pool_spec: PoolSpec | None = None
property pool_blocks: list[BlockPlacement]
property local_blocks: list[BlockPlacement]
property est_pool_cost: float
as_dict()[source]
Return type:

dict[str, Any]

report()[source]
Return type:

str

class BlockPlacement(name, kind, placement, reason, est_tflop=0.0, est_cost=0.0)[source]

Bases: object

Where one block runs and why: ‘local’ or ‘pool’, with the priced justification.

Parameters:
name: str
kind: str
placement: str
reason: str
est_tflop: float = 0.0
est_cost: float = 0.0
class PoolSpec(available=False, cost_per_hour=0.5, flop_threshold_tflop=1.0, pool_speedup=10.0)[source]

Bases: object

What the shared GPU pool offers, and when offloading is worth it.

available: is a pool configured at all (else everything is local). cost_per_hour: the pool’s price. flop_threshold: the estimated work (in TFLOP) below which the round-trip is not worth it – small gradient blocks stay local. local_tflops / pool speedup drive the estimate.

Parameters:
available: bool = False
cost_per_hour: float = 0.5
flop_threshold_tflop: float = 1.0
pool_speedup: float = 10.0
learn_placement_policy(rows, static_policy, *, cost_key='cost', k=8, min_neighbors=4)[source]

Learn a placement policy from telemetry (features, choice, outcome) rows (see module docstring).

static_policy maps a feature dict to a choice and is the fall-back when history is too thin. cost_key names the outcome field to minimize (default "cost"). Feature standardization and the neighbor index are built from the rows; LearnedPolicy.decide() and evaluate follow.

Parameters:
Return type:

LearnedPolicy

class LearnedPolicy(keys, vecs, choices, costs, static, mean=<factory>, scale=<factory>, k=8, min_neighbors=4, margin=0.02)[source]

Bases: object

A history-based placement policy that defers to a static teacher where it lacks evidence.

Parameters:
keys: list[str]
vecs: ndarray
choices: list[str]
costs: ndarray
static: Callable[[dict[str, Any]], str]
mean: ndarray
scale: ndarray
k: int = 8
min_neighbors: int = 4
margin: float = 0.02
decide(features)[source]

Return (choice, learned) – the learned pick when confident, else the static fallback.

Parameters:

features (dict[str, Any])

Return type:

tuple[str, bool]

evaluate(rows, *, cost_key='cost')[source]

Realized-cost comparison on held-out rows: learned policy vs always-static, vs each fixed choice.

Parameters:
Return type:

dict[str, Any]

learn_action_policy(rows, static_scorer=None, *, value_key='value', k=8, min_neighbors=4)[source]

Learn a reasoner acquisition policy from route telemetry (features, kind, outcome) rows.

static_scorer is the fall-back when history is thin (default mixle.substrate.act.score_action()). value_key names the outcome field to MAXIMIZE (default "value" – did the action yield evidence). Returns a LearnedAcquisition usable directly as investigate(..., scorer=policy).

Parameters:
Return type:

LearnedAcquisition

class LearnedAcquisition(keys, vecs, values, static, mean=<factory>, scale=<factory>, k=8, min_neighbors=4)[source]

Bases: object

A history-based action scorer for the reasoner: learns which actions pay off, else defers.

Drop-in for mixle.substrate.act.score_action() (call it as scorer=policy in investigate). From route telemetry – each row a fired action’s (features={kind,cost,overlap}, value) – it estimates the expected yield of an action in a query’s feature region and scores it yield / cost. Where nearby history is too thin, it FALLS BACK to the static lexical scorer: the same never-worse discipline as LearnedPolicy, now on the reasoner’s acquisition decisions (J3).

Parameters:
keys: list[str]
vecs: ndarray
values: ndarray
static: Callable[[Any, str], float]
mean: ndarray
scale: ndarray
k: int = 8
min_neighbors: int = 4
expected_yield(features)[source]

Estimated yield of an action with these features, or None when history is too thin to say.

Parameters:

features (dict[str, Any])

Return type:

float | None

learn_schedule_policy(rows, static_policy, *, latency_key='latency', k=8, min_neighbors=4)[source]

Learned pool scheduling (J4): when work is pool-eligible, learn WHERE/WHEN it actually runs fastest.

The same never-worse shape as placement, keyed on realized LATENCY instead of dollar cost: rows are (features, choice, outcome) where features describe the moment (queue depth, job size, local load), choice is the scheduling decision (“run_local” / “queue_pool” / “defer”), and the outcome’s latency is what the decision actually cost in wall-clock. Where nearby history is thin, the returned policy defers to the static scheduler.

Parameters:
Return type:

LearnedPolicy

meta_improve(rows, static_policy, *, cost_key='cost', holdout_frac=0.3, seed=0, k=8, min_neighbors=4)[source]

The meta-improve loop (J5): learn from telemetry, PROMOTE only on a never-worse holdout receipt.

Splits the telemetry into train/holdout, learns a policy on the train slice, and evaluates it against the static policy on the HELD-OUT decisions (realized cost, the same currency the platform pays). The learned policy is promoted iff its held-out mean cost is <= the static policy’s – the receipt is returned either way, so a non-promotion is auditable, not silent. Returns:

{promoted, policy, receipt: {learned_mean_cost, static_mean_cost, deferred_fraction, n}}

policy is the learned policy when promoted, else the static one wrapped for the same call shape – callers can always use the result’s policy and get never-worse behavior by construction.

Parameters:
Return type:

dict[str, Any]

simulate(model)[source]

Package a fitted model as a Simulator (see module docstring).

Parameters:

model (Any)

Return type:

Simulator

class Simulator(model)[source]

Bases: object

A fitted model packaged as a data generator, runnable under a baseline or named scenarios.

Parameters:

model (Any)

scenario(name, interventions)[source]

Register a named intervention scenario (requires a learned Bayesian network to apply).

Parameters:
Return type:

Simulator

run(n=100, *, scenario=None, interventions=None, seed=0)[source]

Generate n synthetic records under the baseline, a registered scenario, or ad-hoc interventions.

Parameters:
Return type:

list[Any]

outcome_mean(field_index, *, scenario=None, n=2000, seed=0)[source]

The mean of a numeric field under a scenario – the quantity to compare across conditions.

Parameters:

  • field_index (int)

  • scenario (str | None)

  • n (int)

  • seed (int)

Return type:

float

compare(field_index, scenario_a, scenario_b, *, n=4000, seed=0)[source]

mean(field | scenario_a) - mean(field | scenario_b) – the simulated effect of A vs B.

Parameters:

  • field_index (int)

  • scenario_a (str | None)

  • scenario_b (str | None)

  • n (int)

  • seed (int)

Return type:

float

class Scenario(name, interventions=<factory>)[source]

Bases: object

A named simulation condition: which fields are clamped to which values (an intervention).

Parameters:
name: str
interventions: dict[int, Any]
synthesize(source, *, label=None, verify=None, n=100, max_tries=None, seed=0)[source]

Build a verified dataset of n accepted rows from a generative source (see module docstring).

source is a fitted model (sampled), a list of real inputs (a generator is inferred over them), or a callable draw function. label (optional) is the teacher applied to each input. verify (optional) accepts verify(x) or verify(x, label) and gates each row – rejected rows are resampled up to max_tries total draws. The verifier is attached to the returned Dataset so consumers can recheck() independently.

Parameters:
Return type:

Dataset

class Dataset(inputs, labels=None, verify=None, acceptance_rate=1.0, n_rejected=0, provenance=<factory>)[source]

Bases: object

A verified synthetic dataset: inputs, optional labels, and the verifier that vouched for them.

Parameters:
inputs: list[Any]
labels: list[Any] | None = None
verify: Callable[[...], bool] | None = None
acceptance_rate: float = 1.0
n_rejected: int = 0
provenance: dict[str, Any]
pairs()[source]

(input, label) pairs – raises if the dataset is unlabeled.

Return type:

list[tuple[Any, Any]]

recheck()[source]

Re-run the attached verifier over every row.

Returns True when every row still passes, or when no verifier is attached.

Return type:

bool

create(data, *, budget=None, device=None, calibrate=None, quantify_uq=False, max_its=25, seed=0)[source]

Create a certified model from data (see module docstring).

data is a list of records/scalars. calibrate (a fraction in (0,1)) reserves a holdout for a PIT calibration check. quantify_uq=True attaches an auto-selected UQ handle. budget / device (any object; recorded verbatim) constrain the fit toward a smaller model. Returns a CreatedModel bundling the fit, its certificate, and the requested post-conditions.

Parameters:

  • data (Any)

  • budget (Any | None)

  • device (Any | None)

  • calibrate (float | None)

  • quantify_uq (bool)

  • max_its (int)

  • seed (int)

Return type:

CreatedModel

class CreatedModel(model, certificate, strategy, calibration=None, uq=None, provenance=<factory>)[source]

Bases: object

A certified model artifact: the fitted model plus its guarantees and provenance.

Parameters:
model: Any
certificate: Any
strategy: str
calibration: Any | None = None
uq: Any | None = None
provenance: dict[str, Any]
property guarantee: Any

The aggregate estimation guarantee (MIN over blocks) – the artifact’s headline claim.

why()[source]

One line summarizing how the artifact was estimated.

Return type:

str

is_calibrated()[source]

Whether the calibration holdout judged the model calibrated (None if not checked).

Return type:

bool | None

skill(name, obj, *, description='', tags=(), call=None, provenance=None, registry=None)[source]

Package obj (a fitted model, a CreatedModel, or a function) as a reusable Skill and register it (see module docstring).

The estimation certificate is inherited from obj.certificate when present (so a certified model yields a certified skill). call overrides how the skill is invoked; otherwise a model verb (predict / __call__ / sampler) is used. The skill is added to registry (or the process default) and returned.

Parameters:
Return type:

Skill

class Skill(name, call, description='', tags=(), certificate=None, provenance=<factory>)[source]

Bases: object

A named, described, provenanced callable derived from a fitted artifact (or a plain function).

Parameters:
name: str
call: Callable[[...], Any]
description: str = ''
tags: tuple[str, ...] = ()
certificate: Any | None = None
provenance: dict[str, Any]
property guarantee: Any | None

The inherited estimation guarantee, if this skill wraps a certified model.

score(query)[source]

Lexical overlap of query with this skill’s name/description/tags, in [0, 1].

Parameters:

query (str)

Return type:

float

class SkillRegistry[source]

Bases: object

A findable collection of skills – register once, retrieve by name or by query.

add(sk)[source]
Parameters:

sk (Skill)

Return type:

Skill

get(name)[source]
Parameters:

name (str)

Return type:

Skill

all()[source]
Return type:

list[Skill]

find(query, k=5)[source]

The best k skills for query by lexical overlap (highest score first, ties by name).

Parameters:
Return type:

list[Skill]

best(query)[source]

The single best-matching skill for query (or None if nothing overlaps).

Parameters:

query (str)

Return type:

Skill | None

index(substrate)[source]

Mirror every skill into a Substrate for embedding-grade recall.

The registry’s own find is a lexical matcher; when a corpus grows past that, index the skills as artifact items and retrieve through the substrate instead. Returns the item ids.

Parameters:

substrate (Any)

Return type:

list[str]

default_registry()[source]

The process-wide default registry skill() writes to when no registry is given.

Return type:

SkillRegistry

posterior(model, data=None, *, over='predictive', prior=None, method='auto', **kwargs)[source]

Build the Posterior of model over the requested variables.

Parameters:

  • model (Any) – a fitted mixle distribution (or a latent-variable model for over='latent').

  • data (Any) – observations – required for over='params' and for over='latent' (the x the latent posterior conditions on); ignored for plug-in over='predictive'.

  • over (str) – 'latent' -> q(z | x) (needs the latent_posterior capability); 'params' -> q(theta | data); 'predictive' -> draws of new data.

  • prior (Any) – prior over parameters for over='params' (see conjugate_posterior / sample_parameter_posterior).

  • method (str) – for over='params' – 'auto' (conjugate when the family supports it, else MCMC), 'conjugate', or 'mcmc'.

  • **kwargs (Any) – forwarded to sample_parameter_posterior for the MCMC path (sampler, steps…).

Returns:

A Posterior – a LatentPosterior, ParameterPosterior, or PredictivePosterior.

Return type:

Posterior

class ParameterPosterior(draw_one, draw_many, *, mean_fn=None, chain=None, kind='')[source]

Bases: Posterior

q(theta | data) over a model family’s parameters – exact (conjugate) or MCMC.

Built by posterior(model, data, over="params"); the exact-vs-MCMC distinction is hidden behind the shared Posterior interface. A single sample() returns one parameter set, and samples() returns n of them; mean() and interval() summarize the posterior.

Parameters:

  • draw_one (Callable[[Any], Any])

  • draw_many (Callable[[int, Any], Any])

  • mean_fn (Callable[[], Any] | None)

  • chain (np.ndarray | None)

  • kind (str)

classmethod from_conjugate(cp)[source]

Wrap a closed-form ConjugatePosterior (each draw is a parameter dict).

Parameters:

cp (ConjugatePosterior)

Return type:

ParameterPosterior

classmethod from_mcmc(result)[source]

Wrap an MCMC MCMCResult; draws resample the retained chain, summaries use it directly.

Parameters:

result (Any)

Return type:

ParameterPosterior

sample(rng=None)[source]

One parameter draw from the posterior.

Parameters:

rng (Any)

Return type:

Any

samples(n, rng=None)[source]

n parameter draws (a dict of length-n arrays for conjugate; a list for MCMC).

Parameters:
Return type:

Any

mean()[source]

The posterior mean of the parameters.

Return type:

Any

interval(level=0.9)[source]

Central credible interval at level – [lo, hi] over the chain (MCMC) or 2000 draws.

Parameters:

level (float)

Return type:

Any

class PredictivePosterior(draw_one, draw_many)[source]

Bases: Posterior

The posterior-predictive: draws of new data from a fitted model.

plug_in() wraps a single fitted model’s sampler. from_parameter_posterior() integrates parameter uncertainty – each predictive draw first samples theta ~ q(theta | data), rebuilds the model, then draws data – so the spread reflects both sampling and parameter uncertainty.

Parameters:

  • draw_one (Callable[[Any], Any])

  • draw_many (Callable[[int, Any], Any])

classmethod plug_in(model)[source]

Plug-in predictive: draw new data from model at its fitted parameters.

Parameters:

model (Any)

Return type:

PredictivePosterior

classmethod from_parameter_posterior(param_post, build)[source]

Posterior-predictive integrating parameter uncertainty.

build maps one parameter draw (the object ParameterPosterior.sample() returns) to a distribution; each predictive draw rebuilds the model from a fresh theta and samples it.

Parameters:
Return type:

PredictivePosterior

sample(rng=None)[source]

One predictive draw of new data.

Parameters:

rng (Any)

Return type:

Any

samples(n, rng=None)[source]

n predictive draws of new data.

Parameters:
Return type:

Any

record_fit(model, data, *, seed, estimator=None)[source]

Record a ReproReceipt for a model fitted on data with seed (see module docstring).

Parameters:
Return type:

ReproReceipt

verify_reproducible(estimator, data, receipt, *, seed=None, max_its=25)[source]

Refit estimator on data and check the fit reproduces receipt (data + parameters).

Returns {reproducible, data_matches, params_match, refit_fingerprint}. reproducible is True iff BOTH the data fingerprint and the refit’s parameter fingerprint match the receipt – i.e. the exact fit can be recovered from the recorded recipe. seed defaults to the receipt’s seed.

Parameters:

  • estimator (Any)

  • data (Any)

  • receipt (ReproReceipt)

  • seed (int | None)

  • max_its (int)

Return type:

dict[str, Any]

class ReproReceipt(data_fingerprint, n, seed, estimator, param_fingerprint)[source]

Bases: object

The recipe to re-derive a fit: data + seed + estimator, plus the parameter fingerprint to check.

Parameters:

  • data_fingerprint (str)

  • n (int)

  • seed (int)

  • estimator (str)

  • param_fingerprint (str)

data_fingerprint: str
n: int
seed: int
estimator: str
param_fingerprint: str
as_dict()[source]
Return type:

dict[str, Any]

matches_data(data)[source]

Whether data is the exact dataset this fit was recorded on.

Parameters:

data (Any)

Return type:

bool

matches_model(model)[source]

Whether model has the exact parameters this receipt fingerprinted.

Parameters:

model (Any)

Return type:

bool

data_fingerprint(data, *, ndigits=_NDIGITS)[source]

A stable hash of a training dataset (order-sensitive; floats rounded) – identifies the exact input.

Parameters:
Return type:

str

param_fingerprint(model, *, ndigits=_NDIGITS)[source]

A stable hash of a fitted model’s parameters, via its to_json state (floats rounded).

Parameters:
Return type:

str

class BeliefState[source]

Bases: ABC

A distribution over a latent, exposing a uniform query + update interface.

Realizations answer where they are defined: mean(), cov(), var(), sd(), entropy(), interval(), sample(), marginal(), and – the point of a belief state – update(), which returns a new belief conditioned on fresh evidence.

abstractmethod mean()[source]

The posterior mean of the latent.

Return type:

ndarray

abstractmethod entropy()[source]

The differential/Shannon entropy H[q] (nats) – watch it shrink as evidence arrives.

Return type:

float

abstractmethod sample(n=1, rng=None)[source]

Draw n latent samples from the belief.

Parameters:
Return type:

ndarray

abstractmethod update(*args, **kwargs)[source]

Return a new belief conditioned on fresh evidence (the assimilation step).

Parameters:
Return type:

BeliefState

cov()[source]

The posterior covariance (not defined for every realization).

Return type:

ndarray

var()[source]

Per-coordinate posterior variance.

Return type:

ndarray

sd()[source]

Per-coordinate posterior standard deviation.

Return type:

ndarray

interval(level=0.9)[source]

Per-coordinate central credible interval at level – an (d, 2) array of [lo, hi].

Parameters:

level (float)

Return type:

ndarray

marginal(indices)[source]

The belief restricted to a subset of latent coordinates.

Parameters:

indices (Any)

Return type:

BeliefState

class GaussianBelief(mean, cov)[source]

Bases: BeliefState

A multivariate-Gaussian belief N(mean, cov) over a continuous latent.

Evidence is a linear-Gaussian observation y = H z + noise, noise ~ N(0, R); update() applies the exact Kalman measurement update (Joseph form, so the covariance stays symmetric positive-definite). fuse() combines two beliefs about the same latent as a product of Gaussian experts. condition() does noiseless Gaussian conditioning on a coordinate subset.

Parameters:

  • mean (Any)

  • cov (Any)

property dim: int
mean()[source]

The posterior mean of the latent.

Return type:

ndarray

cov()[source]

The posterior covariance (not defined for every realization).

Return type:

ndarray

entropy()[source]

The differential/Shannon entropy H[q] (nats) – watch it shrink as evidence arrives.

Return type:

float

interval(level=0.9)[source]

Per-coordinate central credible interval at level – an (d, 2) array of [lo, hi].

Parameters:

level (float)

Return type:

ndarray

sample(n=1, rng=None)[source]

Draw n latent samples from the belief.

Parameters:
Return type:

ndarray

update(H, y, R)[source]

Kalman measurement update: condition on y = H z + noise, noise ~ N(0, R).

Parameters:

  • H (Any) – (k, d) observation matrix (or (d,) / scalar for a single linear readout).

  • y (Any) – (k,) observed value (or scalar).

  • R (Any) – (k, k) observation-noise covariance (or (k,) diagonal / scalar).

Return type:

GaussianBelief

fuse(other)[source]

Product-of-experts fusion of two beliefs about the same latent (cross-modal fusion).

Equivalent to conditioning self on other treated as a direct Gaussian observation (H = I, R = other.cov), so it reuses the exact Kalman update.

Parameters:

other (GaussianBelief)

Return type:

GaussianBelief

condition(indices, values)[source]

Noiseless Gaussian conditioning: fix latent coordinates indices to values.

Returns the belief over the remaining coordinates. This is the exact R -> 0 limit of an observation that reads off those coordinates.

Parameters:
Return type:

GaussianBelief

marginal(indices)[source]

The belief restricted to a subset of latent coordinates.

Parameters:

indices (Any)

Return type:

GaussianBelief

as_belief(obj, node=None)[source]

Adapt any object exposing mean/cov (a FieldPosterior node, a fitted Gaussian, a ParameterPosterior) into a GaussianBelief.

node is forwarded when the source is node-addressable (e.g. FieldPosterior.mean(node) / .cov(node)); otherwise mean/cov are called with no argument.

Parameters:
Return type:

GaussianBelief

class UncertaintyDecomposition(total, aleatoric, epistemic, kind)[source]

Bases: object

A predictive uncertainty split into aleatoric + epistemic (summing to total).

kind is "entropy" (values in nats) or "variance" (values in the outcome’s squared units). Each field is a scalar for a single query point, or an array over query points.

Parameters:
total: ndarray
aleatoric: ndarray
epistemic: ndarray
kind: str
property fraction_epistemic: ndarray

Share of the total uncertainty that is epistemic (reducible by more data), in [0, 1].

item()[source]

Collapse size-1 arrays to Python floats (convenience for single-point decompositions).

Return type:

UncertaintyDecomposition

decompose_uncertainty(*, probs=None, means=None, variances=None)[source]

Front door: decompose a predictive into aleatoric + epistemic uncertainty.

Pass exactly one representation of the per-member predictive:

  • probs=(M, ..., K) – categorical predictives -> BALD entropy split (decompose_entropy());

  • means=(M, ...) (with optional variances=(M, ...)) – continuous predictives -> law-of-total-variance split (decompose_variance()).

To decompose a fitted model’s parameter uncertainty, build the members first with posterior_ensemble() (then predictive_distribution() for the discrete case).

Parameters:
Return type:

UncertaintyDecomposition

decompose_entropy(member_probs)[source]

BALD entropy split of a discrete predictive.

Parameters:

member_probs (Any) – array (M, ..., K) – M posterior draws / ensemble members, each a categorical predictive over K outcomes (optionally batched over query points in the middle axes). Rows need not be normalized; each is renormalized over the last axis.

Returns:

An UncertaintyDecomposition with kind="entropy" (nats). epistemic is the mutual information H(mean) - mean H and is clamped to >= 0 (it is non-negative in exact arithmetic; the clamp only removes tiny floating-point negatives).

Return type:

UncertaintyDecomposition

decompose_variance(member_means, member_vars=None)[source]

Law-of-total-variance split of a continuous predictive.

Parameters:

  • member_means (Any) – array (M, ...) – each member’s predictive mean E[y | theta_m].

  • member_vars (Any) – array (M, ...) – each member’s predictive variance Var[y | theta_m]. If None, aleatoric noise is taken as zero (members are point predictors) and the decomposition reports only the epistemic spread of the means.

Returns:

aleatoric = mean_m Var_m, epistemic = Var_m(mean_m), total their sum.

Return type:

An UncertaintyDecomposition with kind="variance"

predictive_distribution(members, support)[source]

Evaluate an iterable of fitted distributions over a discrete support -> (M, K) probs.

Each member’s log_density is evaluated at every point of support and softmax-normalized over the support, giving one categorical row per member. Feed the result to decompose_entropy().

Parameters:
Return type:

ndarray

posterior_ensemble(param_post, build, n=200, rng=None)[source]

Materialize n models from a parameter posterior – an ensemble representing q(theta|data).

build maps one parameter draw (whatever ParameterPosterior.sample() returns) to a fitted distribution, mirroring from_parameter_posterior(). The returned list is the “members” the decomposition integrates over – so epistemic uncertainty here is genuine parameter uncertainty, not just ensemble disagreement.

Parameters:
Return type:

list[Any]

class Clustering(representatives, probs, labels)[source]

Bases: object

Samples grouped into equivalence classes: representatives, class probs, per-sample labels.

Parameters:
representatives: list[Any]
probs: ndarray
labels: ndarray
cluster_samples(samples, equivalent=None)[source]

Group samples into equivalence classes under equivalent (default exact ==).

For discrete draws whose surface form varies but meaning does not – e.g. LLM generations (“Paris”, “It’s Paris.”, “The capital is Paris”) – pass a semantic equivalent (embedding similarity, an entailment check, normalized match). Greedy single-linkage against each cluster’s first member; returns the class distribution and per-sample assignments.

Parameters:
Return type:

Clustering

marginalize_meaning(items, equivalent=None, *, log_probs=None, weights=None)[source]

The distribution over meanings = the string distribution marginalized over each meaning class.

A generative model puts probability on strings; a meaning is an equivalence class of strings (equivalent decides sameness). The probability of a meaning c is the pushforward under the quotient – you sum the string probabilities over the class: P(c) = sum_{s in c} P(s). This returns that marginal (as a Clustering, probs = P(c)).

How the per-string probability enters:

  • log_probs – the model’s sequence log-probabilities log P(s) for each item; classes are combined by logsumexp (exact marginalization, numerically stable). Use this when the items are distinct strings whose probabilities you know – it corrects the counting form’s hidden “every string in a class is equiprobable” assumption.

  • weights – explicit non-negative masses per item (summed within class).

  • neither – uniform mass, i.e. counting: P(c) = count_c / N. For i.i.d. samples from the model this is the unbiased Monte-Carlo estimate of the same marginal.

Parameters:
Return type:

Clustering

semantic_entropy(samples, equivalent=None, *, log_probs=None, weights=None)[source]

Entropy (nats) over the meaning classes of samples – the model’s predictive uncertainty.

Sample a stochastic generator (an LLM at temperature) n times, marginalize the string distribution over meaning classes (marginalize_meaning()), and take the entropy of that marginal. High semantic entropy means the model disagrees with itself about what the answer is (a hallucination signal), as opposed to merely phrasing one answer many ways (which collapses to low entropy). Pass log_probs (the sequence log-likelihoods) to marginalize with the actual string probabilities rather than by sample counting. Feed the clusters’ per-member distributions to decompose_entropy() for an epistemic/aleatoric split.

Parameters:
Return type:

float

reliability_curve(prob, outcome, *, bins=10, strategy='uniform', ci=False, n_boot=1000, ci_level=0.95, seed=0)[source]

Reliability diagram: observed frequency vs mean forecast probability, per bin.

Bins the forecasts, then within each bin compares the mean predicted probability to the observed event frequency. A perfectly calibrated forecaster lies on the diagonal observed == predicted.

Parameters:

  • prob (ndarray) – (n,) predicted probabilities of the positive class (or top-label confidences).

  • outcome (ndarray) – (n,) 0/1 outcomes (or correctness indicators).

  • bins (int) – number of bins.

  • strategy (str) – "uniform" (equal-width) or "quantile" (equal-count) bins.

  • ci (bool) – if True attach a percentile bootstrap band on the observed frequency in each bin.

  • n_boot (int) – bootstrap resamples when ci is True.

  • ci_level (float) – central probability of the bootstrap band (e.g. 0.95).

  • seed (int | RandomState | None) – RNG seed for the bootstrap.

Returns:

{'mean_pred', 'obs_freq', 'count', 'bin_edges'} (one entry per non-empty bin), plus 'obs_lo' / 'obs_hi' when ci is True.

Return type:

dict[str, ndarray]

expected_calibration_error(prob, outcome, *, bins=10, strategy='uniform', norm='l1', ci=False, n_boot=1000, ci_level=0.95, seed=0)[source]

Expected Calibration Error: count-weighted average gap between confidence and accuracy.

ECE = sum_b (n_b / n) |obs_b - pred_b| over bins (norm='l2' uses the squared gap, square- rooted). Zero is perfect calibration. For multiclass classifiers reduce with top_label_confidence() first.

Parameters:

  • prob (ndarray) – (n,) predicted probabilities / confidences.

  • outcome (ndarray) – (n,) 0/1 outcomes / correctness indicators.

  • bins (int) – number of bins.

  • strategy (str) – "uniform" or "quantile" binning.

  • norm (str) – "l1" (mean absolute gap) or "l2" (root mean squared gap).

  • ci (bool) – if True also return a percentile bootstrap interval.

  • n_boot (int) – bootstrap controls.

  • ci_level (float) – bootstrap controls.

  • seed (int | RandomState | None) – bootstrap controls.

Returns:

The ECE (float), or (ece, lo, hi) when ci is True.

Return type:

float | tuple[float, float, float]

maximum_calibration_error(prob, outcome, *, bins=10, strategy='uniform')[source]

Maximum Calibration Error: the worst per-bin gap max_b |obs_b - pred_b|.

Unlike expected_calibration_error() this is not count-weighted, so it surfaces a small but badly-miscalibrated region that the average would hide.

Parameters:
Return type:

float

top_label_confidence(prob, labels)[source]

Reduce a multiclass classifier to the (confidence, correct) top-label calibration problem.

Parameters:

  • prob (ndarray) – (n, K) class-probability matrix.

  • labels (ndarray) – (n,) integer true labels.

Returns:

the max predicted probability per row and a 0/1 indicator of whether that argmax class was the true label. Feed these to reliability_curve() / expected_calibration_error().

Return type:

(confidence, correct)

class ProbabilityCalibrator(method='isotonic')[source]

Bases: object

Map raw scores to calibrated probabilities – fit against binary outcomes.

A raw score (a model’s confidence, a self-consistency fraction, a token likelihood) need not be a probability of anything: it can be monotone-but-miscalibrated, or have no relationship to the outcome at all. This learns the transform score -> P(outcome = 1 | score) from labeled data, so the output is a probability of the event you calibrated against.

  • method="isotonic" – monotone, non-parametric (pool-adjacent-violators). Assumes higher score => not-lower probability; flexible, needs enough calibration points.

  • method="platt" – logistic sigmoid(a * score + b). Two parameters, robust on little data, but assumes a sigmoidal relationship.

A near-flat fitted curve is itself the finding: it means the raw score carried little information about the outcome (its “likelihood” was unrelated to the event).

Parameters:

method (str)

fit(scores, outcomes)[source]

Fit the score->probability map on scores with binary outcomes (0/1).

Parameters:
Return type:

ProbabilityCalibrator

predict(scores)[source]

Calibrated probabilities for scores (clamped to [0, 1]).

Parameters:

scores (Any)

Return type:

ndarray

calibrate_probabilities(scores, outcomes, *, method='isotonic')[source]

Fit a ProbabilityCalibrator mapping scores to P(outcome=1 | score).

Parameters:
Return type:

ProbabilityCalibrator

pit_values(y, cdf)[source]

Probability Integral Transform values u_i = F_i(y_i).

Under a calibrated continuous predictive distribution the PIT values are Uniform(0, 1). Pass either the precomputed CDF values F_i(y_i) or a callable cdf(y) -> F(y).

Parameters:
Returns:

(n,) PIT values, clipped to [0, 1].

Return type:

ndarray

pit_ensemble(y, forecasts, *, randomize=True, seed=0)[source]

Rank-based PIT from a finite predictive ensemble.

u_i is the fraction of ensemble members <= y_i. With randomize=True ties are broken by a uniform jitter within the rank gap, which makes the PIT exactly Uniform(0, 1) under calibration even for discrete ensembles (the randomized PIT of Czado et al. 2009).

Parameters:

  • y (ndarray) – (n,) realised values.

  • forecasts (ndarray) – (n, m) ensemble (m draws per observation).

  • randomize (bool) – jitter ties for an exactly-uniform PIT.

  • seed (int | RandomState | None) – RNG seed when randomize is True.

Returns:

(n,) PIT values in [0, 1].

Return type:

ndarray

pit_histogram(pit, *, bins=10)[source]

Histogram of PIT values with the uniform reference level.

Parameters:

  • pit (ndarray) – (n,) PIT values.

  • bins (int) – number of equal-width bins on [0, 1].

Returns:

{'counts', 'density', 'edges', 'uniform'} where density integrates to 1 and uniform is the flat reference density (1.0) a calibrated forecast would match.

Return type:

dict[str, ndarray]

pit_calibration_error(pit, *, bins=10)[source]

Calibration error of a PIT histogram: mean absolute deviation from uniform mass.

sum_b |count_b/n - 1/bins| – 0 when the PIT histogram is perfectly flat (calibrated), larger when it is U-shaped (under-dispersed) or humped (over-dispersed).

Parameters:
Return type:

float

interval_coverage(lower, upper, y)[source]

Empirical coverage and mean width of a set of prediction intervals.

Parameters:

  • lower (ndarray) – (n,) lower endpoints.

  • upper (ndarray) – (n,) upper endpoints.

  • y (ndarray) – (n,) realised values.

Returns:

{'coverage', 'mean_width'} – the fraction of y inside [lower, upper] and the mean interval width. Compare coverage to the nominal level the interval was built for.

Return type:

dict[str, float]

coverage_curve(forecasts, y, *, levels=None)[source]

Empirical-vs-nominal coverage of central intervals across a grid of nominal levels.

For each nominal central level c the per-observation central interval [quantile((1-c)/2), quantile((1+c)/2)] is read off the predictive ensemble and its empirical coverage is measured. A calibrated forecast traces the diagonal empirical == nominal; bowing below the diagonal means the intervals are too narrow (over-confident).

Parameters:

  • forecasts (ndarray) – (n, m) predictive ensemble.

  • y (ndarray) – (n,) realised values.

  • levels (ndarray | None) – nominal central coverage levels in (0, 1); defaults to 0.05 .. 0.95 by 0.05.

Returns:

{'nominal', 'empirical'} arrays of equal length.

Return type:

dict[str, ndarray]

bonferroni(pvals, *, alpha=0.05)[source]

Bonferroni single-step FWER control: adjusted p_i = min(1, m p_i).

The simplest and most conservative correction; always valid regardless of dependence.

Parameters:

  • pvals (ndarray) – (m,) raw p-values.

  • alpha (float) – target family-wise error rate.

Returns:

{'reject', 'pvals_adjusted', 'n_reject', 'alpha'}.

Return type:

dict

holm(pvals, *, alpha=0.05)[source]

Holm step-down FWER control – uniformly more powerful than Bonferroni, equally valid.

Sorts p-values ascending and applies the running factor m - k with a cumulative maximum so the adjusted values stay monotone.

Parameters:
Return type:

dict

hochberg(pvals, *, alpha=0.05)[source]

Hochberg step-up FWER control (valid under independence / positive dependence).

Same per-step factor as Holm but applied step-up (from the largest p-value down), giving more rejections; requires the independence/PRDS assumption that Holm does not.

Parameters:
Return type:

dict

benjamini_hochberg(pvals, *, alpha=0.05)[source]

Benjamini–Hochberg FDR control; adjusted values are q-values.

Controls the expected false-discovery proportion at alpha under independence or positive regression dependence (PRDS). The standard choice for screening many hypotheses.

Parameters:
Return type:

dict

benjamini_yekutieli(pvals, *, alpha=0.05)[source]

Benjamini–Yekutieli FDR control, valid under arbitrary dependence.

Like benjamini_hochberg() but inflated by the harmonic factor c(m) = sum_{i=1}^m 1/i, so it holds for any dependence structure at the cost of power.

Parameters:
Return type:

dict

adjust_pvalues(pvals, *, method='bh', alpha=0.05)[source]

Unified dispatcher over the correction methods.

Parameters:

  • pvals (ndarray) – (m,) raw p-values.

  • method (str) – one of "bonferroni", "holm", "hochberg", "bh" (Benjamini–Hochberg), "by" (Benjamini–Yekutieli).

  • alpha (float) – target error rate (FWER for the first three, FDR for the last two).

Returns:

{'reject', 'pvals_adjusted', 'n_reject', 'alpha'}.

Return type:

dict

fisher_combine(pvals)[source]

Fisher’s method: combine independent p-values via -2 sum log p ~ chi^2_{2k}.

Sensitive to a few very small p-values. For combining evidence for the same hypothesis across independent tests.

Returns:

{'statistic', 'pvalue', 'df'}.

Parameters:

pvals (ndarray)

Return type:

dict[str, float]

stouffer_combine(pvals, *, weights=None)[source]

Stouffer’s Z method: combine p-values on the z-scale, optionally weighted.

Z = sum w_i Phi^{-1}(1 - p_i) / sqrt(sum w_i^2). Weights let more-precise studies (e.g. larger samples) count more; equal weights recover the unweighted combination.

Returns:

{'z', 'pvalue'} (one-sided combined p-value).

Parameters:
Return type:

dict[str, float]

tippett_combine(pvals)[source]

Tippett’s method (Sidak min-p): combined p = 1 - (1 - min p)^k.

Most powerful when a single strong signal exists among the tests.

Returns:

{'min_p', 'pvalue'}.

Parameters:

pvals (ndarray)

Return type:

dict[str, float]

bootstrap(data, statistic, *, n_boot=2000, method='bca', ci_level=0.95, seed=0, groups=None, clusters=None, block_length=None, m=None)[source]

Bootstrap confidence interval for statistic(data).

Parameters:

  • data (Any) – a single array (resampled along axis 0) or a tuple of arrays sharing their first axis (the statistic is then called as statistic(*parts)).

  • statistic (Callable[[...], Any]) – maps the data to a scalar or fixed-length vector.

  • n_boot (int) – number of bootstrap resamples.

  • method (str) – "percentile", "basic" (pivotal), or "bca" (bias-corrected & accelerated). "bca" is only second-order accurate for plain i.i.d. resampling; with groups / clusters / block_length / m set it falls back to "percentile".

  • ci_level (float) – central probability of the interval.

  • seed (int | RandomState | None) – RNG seed.

  • groups (ndarray | None) – (n,) labels for stratified resampling (resample within each group).

  • clusters (ndarray | None) – (n,) labels for cluster resampling (resample whole clusters with replacement).

  • block_length (int | None) – moving-block length for serially dependent (time-series) data.

  • m (int | None) – subsample size for m-out-of-n subsampling (without replacement).

Returns:

A BootstrapResult.

Return type:

BootstrapResult

class BootstrapResult(estimate, ci_low, ci_high, distribution, method, ci_level, standard_error)[source]

Bases: object

Result of a bootstrap() call.

Parameters:
estimate

the statistic on the original data (scalar or vector).

Type:

numpy.ndarray

ci_low / ci_high

confidence-interval endpoints (same shape as estimate).

distribution

(n_boot, ...) array of bootstrap replicates.

Type:

numpy.ndarray

method

the interval method used.

Type:

str

ci_level

the central probability of the interval.

Type:

float

standard_error

bootstrap standard error (std of the replicates).

Type:

numpy.ndarray

estimate: ndarray
ci_low: ndarray
ci_high: ndarray
distribution: ndarray
method: str
ci_level: float
standard_error: ndarray
block_bootstrap(data, statistic, block_length, *, n_boot=2000, ci_level=0.95, seed=0)[source]

Moving-block bootstrap for serially dependent (time-series) data.

Convenience wrapper over bootstrap() with block_length set: resamples contiguous blocks so within-block autocorrelation is preserved. Choose block_length on the order of the series’ correlation length.

Parameters:
Return type:

BootstrapResult

wild_bootstrap(fitted, residuals, statistic, *, n_boot=2000, kind='rademacher', ci_level=0.95, seed=0)[source]

Wild (residual-multiplier) bootstrap, robust to heteroscedasticity.

Builds synthetic responses y* = fitted + residual * v where v are mean-zero, unit-variance two-point multipliers drawn independently per observation, then recomputes the statistic on each y*. Because each residual keeps its own magnitude, the procedure preserves heteroscedasticity that an i.i.d. residual resample would destroy.

Parameters:

  • fitted (ndarray) – (n,) fitted values from the model.

  • residuals (ndarray) – (n,) residuals y - fitted.

  • statistic (Callable[[ndarray], Any]) – maps a synthetic response vector y* to a scalar or vector (e.g. refit and return coefficients).

  • n_boot (int) – number of resamples.

  • kind (str) – "rademacher" (v in {-1, +1}) or "mammen" (Mammen’s two-point distribution).

  • ci_level (float) – central probability of the percentile interval.

  • seed (int | RandomState | None) – RNG seed.

Returns:

A BootstrapResult (percentile interval).

Return type:

BootstrapResult

permutation_test(x, y, *, statistic=None, n_perm=10000, alternative='two-sided', paired=False, stratify=None, seed=0, exact_max=10000)[source]

Two-sample permutation test for an arbitrary statistic under a sharp null.

Under the null that the two samples are exchangeable, the labels can be shuffled freely; the statistic’s permutation distribution is the reference. For two-sided the statistic is centered at zero by construction (difference statistics) and compared on absolute value.

Parameters:

  • x (ndarray) – the two samples (1-D). For paired=True they must have equal length and pairing is preserved by sign-flipping the within-pair differences.

  • y (ndarray) – the two samples (1-D). For paired=True they must have equal length and pairing is preserved by sign-flipping the within-pair differences.

  • statistic (Callable[[ndarray, ndarray], float] | None) – f(x, y) -> float; defaults to the difference in means. For paired it is applied to (differences, zeros) so the default reduces to the mean paired difference.

  • n_perm (int) – number of random rearrangements (ignored if the exact set is enumerated).

  • alternative (str) – "two-sided", "greater", or "less".

  • paired (bool) – paired (sign-flip) permutation instead of label shuffling.

  • stratify (ndarray | None) – (n,) group labels (concatenated x-then-y) for restricted permutation – labels are shuffled only within each group, preserving group structure.

  • seed (int | RandomState | None) – RNG seed.

  • exact_max (int) – if the number of distinct rearrangements is <= exact_max they are enumerated for an exact p-value.

Returns:

A PermutationResult.

Return type:

PermutationResult

class PermutationResult(statistic, pvalue, null_distribution, n_perm, exact, alternative)[source]

Bases: object

Result of a permutation_test().

Parameters:
statistic

the observed test statistic.

Type:

float

pvalue

the (one- or two-sided) p-value.

Type:

float

null_distribution

the statistic under each sampled/enumerated rearrangement.

Type:

numpy.ndarray

n_perm

number of rearrangements used.

Type:

int

exact

True if the full permutation set was enumerated.

Type:

bool

alternative

the alternative hypothesis.

Type:

str

statistic: float
pvalue: float
null_distribution: ndarray
n_perm: int
exact: bool
alternative: str
sandwich_covariance(scores, bread, *, clusters=None, n_params=None, small_sample=True)[source]

Generic sandwich covariance B M B' from per-observation scores and the bread.

Parameters:

  • scores (ndarray) – (n, p) per-observation contributions to the estimating equation (for OLS the i-th row is x_i * e_i; for a GLM the score of observation i).

  • bread (ndarray) – (p, p) inverse sensitivity B (e.g. (X'X)^{-1} for OLS, the inverse negative Hessian / inverse Fisher for an M-estimator).

  • clusters (ndarray | None) – optional (n,) cluster labels; scores are summed within cluster before forming the meat (the robust-to-within-cluster-correlation meat). If None, observations are independent.

  • n_params (int | None) – number of estimated parameters for the small-sample correction (defaults to p).

  • small_sample (bool) – apply the usual finite-sample correction (n/(n-p) for independent data, G/(G-1) * (n-1)/(n-p) for clustered).

Returns:

The (p, p) robust covariance matrix.

Return type:

ndarray

ols_robust_covariance(x, residuals, *, hc='hc1')[source]

Heteroscedasticity-consistent (White) covariance for OLS coefficients.

Cov = (X'X)^{-1} ( sum_i w_i x_i x_i' ) (X'X)^{-1} where the per-observation weight w_i is the squared residual, optionally adjusted for leverage h_ii:

  • hc0: e_i^2 (White’s original; biased down in small samples).

  • hc1: e_i^2 * n/(n-p) (degrees-of-freedom correction; the common default).

  • hc2: e_i^2 / (1 - h_ii).

  • hc3: e_i^2 / (1 - h_ii)^2 (best small-sample behaviour; ~jackknife).

Parameters:

  • x (ndarray) – (n, p) design matrix (include an intercept column if wanted).

  • residuals (ndarray) – (n,) OLS residuals y - X beta_hat.

  • hc (str) – one of "hc0", "hc1", "hc2", "hc3".

Returns:

The (p, p) robust covariance matrix.

Return type:

ndarray

cluster_robust_covariance(x, residuals, clusters, *, small_sample=True)[source]

Cluster-robust (one-way or multi-way) covariance for OLS coefficients.

Allows arbitrary correlation within clusters while assuming independence across them. Pass one label array for one-way clustering, or a list/tuple of label arrays for multi-way clustering, which uses the Cameron–Gelbach–Miller inclusion–exclusion V_A + V_B - V_{A∩B} (two-way) and its higher-way generalisation.

Parameters:

  • x (ndarray) – (n, p) design matrix.

  • residuals (ndarray) – (n,) OLS residuals.

  • clusters (ndarray | list) – (n,) labels, or a list of such arrays for multi-way clustering.

  • small_sample (bool) – apply the G/(G-1) * (n-1)/(n-p) finite-sample correction per term.

Returns:

The (p, p) cluster-robust covariance matrix.

Return type:

ndarray

newey_west_covariance(x, residuals, *, lags=None, small_sample=True)[source]

Newey–West HAC covariance for OLS coefficients (serially correlated errors).

Heteroscedasticity- and autocorrelation-consistent: the meat is the long-run covariance of the scores s_t = x_t e_t, estimated with Bartlett (triangular) weights so it stays positive semi-definite:

S = Gamma_0 + sum_{l=1}^{L} (1 - l/(L+1)) (Gamma_l + Gamma_l’), Gamma_l = sum_t s_t s_{t-l}’.

Rows of x (and residuals) are assumed to be in time order.

Parameters:

  • x (ndarray) – (n, p) design matrix, rows in time order.

  • residuals (ndarray) – (n,) residuals in time order.

  • lags (int | None) – truncation lag L; defaults to the Newey–West rule floor(4 (n/100)^{2/9}).

  • small_sample (bool) – apply the n/(n-p) degrees-of-freedom correction.

Returns:

The (p, p) HAC covariance matrix.

Return type:

ndarray

robust_standard_errors(cov)[source]

Standard errors sqrt(diag(cov)) from a covariance matrix.

Parameters:

cov (ndarray)

Return type:

ndarray

glm(x, y, *, family='gaussian', link=None, offset=None, weights=None, max_iter=100, tol=1e-8, robust=False)[source]

Fit a generalized linear model by iteratively reweighted least squares.

Parameters:

  • x (ndarray) – (n, p) design matrix (include an intercept column explicitly if wanted).

  • y (ndarray) – (n,) response (counts, 0/1 or proportions, positive reals, … per the family).

  • family (str | Family) – "gaussian", "binomial", "poisson", "gamma", "inverse_gaussian", "negativebinomial", or a Family.

  • link (str | None) – link name; defaults to the family’s canonical link.

  • offset (ndarray | None) – (n,) known additive term on the linear-predictor scale (e.g. log exposure).

  • weights (ndarray | None) – (n,) prior weights.

  • max_iter (int) – IRLS controls (convergence on the relative deviance change).

  • tol (float) – IRLS controls (convergence on the relative deviance change).

  • robust (bool) – if True report Huber–White sandwich standard errors instead of model-based ones.

Returns:

A GLMResult.

Return type:

GLMResult

class GLMResult(coef, se, fitted, deviance, dispersion, log_likelihood, n_iter, family, link, cov, _link=None)[source]

Bases: object

Fitted GLM.

Parameters:
coef

(p,) coefficient estimates.

Type:

numpy.ndarray

se

(p,) standard errors (model-based, or robust if requested).

Type:

numpy.ndarray

fitted

(n,) fitted means mu.

Type:

numpy.ndarray

deviance

residual deviance.

Type:

float

dispersion

estimated/assumed dispersion phi.

Type:

float

log_likelihood

maximised log-likelihood.

Type:

float

n_iter

IRLS iterations to convergence.

Type:

int

family / link

names.

cov

(p, p) coefficient covariance.

Type:

numpy.ndarray

coef: ndarray
se: ndarray
fitted: ndarray
deviance: float
dispersion: float
log_likelihood: float
n_iter: int
family: str
link: str
cov: ndarray
predict(x, *, offset=None)[source]

Predict the mean response mu at new design rows x.

Parameters:
Return type:

ndarray

property aic: float
property bic: float
z_values()[source]
Return type:

ndarray

p_values()[source]
Return type:

ndarray

class Family(name, variance, canonical, unit_deviance, estimate_dispersion, extra=1.0)[source]

Bases: object

An exponential-family error model: variance function, canonical link, deviance, dispersion.

Parameters:
name: str
variance: Callable[[ndarray], ndarray]
canonical: str
unit_deviance: Callable[[ndarray, ndarray], ndarray]
estimate_dispersion: bool
extra: float = 1.0
ridge_regression(x, y, alpha=1.0, *, fit_intercept=True)[source]

Ridge (L2-penalised) linear regression in closed form.

Minimises ||y - X b||^2 + alpha ||b||^2; the intercept (if fitted) is not penalised.

Parameters:
Return type:

PenalizedResult

elastic_net(x, y, alpha=1.0, l1_ratio=0.5, *, fit_intercept=True, max_iter=1000, tol=1e-7)[source]

Elastic-net linear regression by cyclic coordinate descent.

Minimises (1/2n) ||y - X b||^2 + alpha ( l1_ratio ||b||_1 + (1 - l1_ratio)/2 ||b||^2 ). l1_ratio = 1 is the lasso (sparse), l1_ratio = 0 is ridge.

Parameters:
Return type:

PenalizedResult

lasso(x, y, alpha=1.0, **kw)[source]

Lasso (L1) linear regression – elastic_net() with l1_ratio = 1.

Parameters:
Return type:

PenalizedResult

class PenalizedResult(coef, intercept, alpha, l1_ratio, n_iter)[source]

Bases: object

Fitted penalized linear regression.

Parameters:
coef

(p,) coefficients (excluding the intercept).

Type:

numpy.ndarray

intercept

fitted intercept.

Type:

float

alpha

overall penalty strength.

Type:

float

l1_ratio

elastic-net mixing (1 = lasso, 0 = ridge).

Type:

float

n_iter

coordinate-descent iterations (0 for the closed-form ridge).

Type:

int

coef: ndarray
intercept: float
alpha: float
l1_ratio: float
n_iter: int
predict(x)[source]
Parameters:

x (ndarray)

Return type:

ndarray

robust_regression(x, y, *, method='huber', c=None, max_iter=100, tol=1e-8)[source]

Robust (M-estimator) linear regression by IRLS with a robust scale.

Down-weights observations with large residuals so a few outliers cannot dominate the fit. huber uses the Huber weight (tuning c = 1.345 for 95% Gaussian efficiency); tukey uses the redescending Tukey biweight (c = 4.685), which rejects gross outliers entirely.

Parameters:
Return type:

RegressionFit

quantile_regression(x, y, tau=0.5, *, max_iter=200, tol=1e-7, eps=1e-6)[source]

Linear quantile regression: the conditional tau-quantile by IRLS on the check loss.

Minimises the pinball loss sum rho_tau(y - X b) via iteratively reweighted least squares with weights tau / |r| for positive residuals and (1 - tau) / |r| for negative ones (a smoothed Newton scheme; eps floors |r| for stability).

Parameters:
Return type:

RegressionFit

class RegressionFit(coef, fitted, scale, n_iter)[source]

Bases: object

Coefficients + fitted values from robust_regression() / quantile_regression().

Parameters:
coef: ndarray
fitted: ndarray
scale: float
n_iter: int
predict(x)[source]
Parameters:

x (ndarray)

Return type:

ndarray

kaplan_meier(time, event=None, *, ci_level=0.95)[source]

Kaplan–Meier product-limit estimate of the survival function S(t).

Parameters:

  • time (ndarray) – (n,) observed times (event or censoring).

  • event (ndarray | None) – (n,) 1 = event, 0 = right-censored (defaults to all events).

  • ci_level (float) – confidence level for the log–log survival band.

Returns:

{'time', 'survival', 'se', 'ci_low', 'ci_high', 'at_risk', 'n_events', 'median'}.

Return type:

dict[str, ndarray]

nelson_aalen(time, event=None)[source]

Nelson–Aalen estimate of the cumulative hazard H(t) = sum d_i / Y_i.

Returns:

{'time', 'cumhaz', 'se'} with the Poisson-type standard error of the cumulative hazard.

Parameters:
Return type:

dict[str, ndarray]

cox_ph(x, time, event, *, start=None, strata=None, ties='efron', max_iter=50, tol=1e-9)[source]

Cox proportional-hazards regression by Newton–Raphson on the partial likelihood.

The hazard is h(t | x) = h0(t) exp(x' beta); only the ordering of event times enters, so the baseline h0 is left unspecified (semi-parametric). Time-varying covariates are supported through the counting-process form: pass start so each row is an at-risk interval (start, stop] (a subject contributes several rows), and the risk set at an event time is every interval covering it.

Parameters:

  • x (ndarray) – (n, p) covariates (no intercept – it is absorbed into the baseline).

  • time (ndarray) – (n,) event/censoring times (the interval stop times).

  • event (ndarray) – (n,) 1 = event, 0 = censored.

  • start (ndarray | None) – optional (n,) interval start times for time-varying covariates / left truncation.

  • strata (ndarray | None) – optional (n,) labels; each stratum gets its own baseline hazard (coefficients shared).

  • ties (str) – "efron" (default, more accurate) or "breslow" tie handling.

  • max_iter (int) – Newton controls.

  • tol (float) – Newton controls.

Returns:

A CoxResult.

Return type:

CoxResult

class CoxResult(coef, se, cov, loglik, baseline_time, baseline_cumhaz, concordance, n_iter)[source]

Bases: object

Fitted Cox proportional-hazards model.

Parameters:
coef

(p,) log-hazard-ratio coefficients.

Type:

numpy.ndarray

se

(p,) standard errors (inverse observed information).

Type:

numpy.ndarray

cov

(p, p) covariance.

Type:

numpy.ndarray

loglik

maximised partial log-likelihood.

Type:

float

baseline_time / baseline_cumhaz

Breslow baseline cumulative hazard.

concordance

Harrell’s C-index.

Type:

float

n_iter

Newton iterations.

Type:

int

coef: ndarray
se: ndarray
cov: ndarray
loglik: float
baseline_time: ndarray
baseline_cumhaz: ndarray
concordance: float
n_iter: int
hazard_ratios()[source]
Return type:

ndarray

z_values()[source]
Return type:

ndarray

p_values()[source]
Return type:

ndarray

to_person_period(time, event, covariates=None)[source]

Expand right-censored durations into a person-period (long) array for discrete-time models.

Each subject contributes one row per discrete period they were at risk; the binary outcome is 1 in the period the event occurred and 0 otherwise. Integer time is the number of periods observed.

Returns:

{'period', 'outcome', 'subject', 'covariates'} (covariates repeated per period if given).

Parameters:
Return type:

dict[str, ndarray]

discrete_time_hazard(x, outcome, *, link='cloglog', offset=None)[source]

Discrete-time hazard model: a binary GLM on the person-period array.

Fit on the long-format data from to_person_period() (the design x typically holds period indicators / a time trend plus covariates). cloglog gives the grouped-proportional-hazards (interval-censored Cox) interpretation; logit gives the proportional-odds hazard.

Returns:

a mixle.inference.glm.GLMResult (binomial family with the chosen link).

Parameters:
aalen_johansen(time, event, *, causes=None)[source]

Aalen–Johansen cumulative incidence functions for competing risks.

With several mutually exclusive event types, the cause-specific CIF F_k(t) is the probability of failing from cause k by time t accounting for the competing causes (it is not 1 - KM on the cause, which overstates incidence).

Parameters:

  • time (ndarray) – (n,) event/censoring times.

  • event (ndarray) – (n,) integer cause label, 0 = censored, 1..K = causes.

  • causes (ndarray | None) – optional explicit list of cause labels; inferred from event if None.

Returns:

{'time', 'cif': {cause: array}, 'overall_survival'}.

Return type:

dict

aalen_additive(x, time, event, *, intercept=True)[source]

Aalen’s additive-hazards regression: cumulative regression functions B(t).

Models h(t | x) = b0(t) + sum_j x_j b_j(t) with time-varying additive effects. At each event time the increment dB is the least-squares solution over the risk set; the cumulative B(t) (returned) has interpretable slopes – a rising B_j means covariate j adds hazard.

Returns:

{'time', 'cum_coef'} where cum_coef is (n_event_times, p[+1]) cumulative coefficients (the first column is the baseline when intercept is True).

Parameters:
Return type:

dict

frailty_cox(x, time, event, groups, *, max_iter=50, tol=1e-5, ties='breslow')[source]

Shared gamma-frailty Cox model for clustered survival, by EM.

Subjects in the same group share an unobserved frailty w_g ~ Gamma(1/theta, 1/theta) (mean 1, variance theta) that multiplies the hazard, capturing within-group correlation. The E-step takes the posterior-mean frailties; the M-step refits Cox with log w_g as an offset and updates theta. theta -> 0 indicates no detectable clustering.

Returns:

A FrailtyCoxResult.

Parameters:
Return type:

FrailtyCoxResult

class FrailtyCoxResult(coef, se, theta, frailties, groups, n_iter=0)[source]

Bases: object

Shared gamma-frailty Cox result.

Parameters:
coef / se

fixed-effect log-hazard-ratios and standard errors.

theta

estimated frailty variance (0 means no clustering signal).

Type:

float

frailties

posterior mean random effect per group.

Type:

numpy.ndarray

groups

group labels aligned to frailties.

Type:

numpy.ndarray

n_iter

EM iterations.

Type:

int

coef: ndarray
se: ndarray
theta: float
frailties: ndarray
groups: ndarray
n_iter: int = 0
ordinal_regression(x, y, *, link='logit', max_iter=200)[source]

Fit a cumulative-link ordinal regression (ordered logit / probit) by maximum likelihood.

Parameters:

  • x (ndarray) – (n, p) covariates (no intercept – the thresholds play that role).

  • y (ndarray) – (n,) integer category labels 0..K-1 (ordered).

  • link (str) – "logit" (proportional odds) or "probit".

  • max_iter (int) – optimiser iterations.

Returns:

An OrdinalResult.

Return type:

OrdinalResult

class OrdinalResult(coef, thresholds, se, log_likelihood, link, n_categories)[source]

Bases: object

Fitted ordinal (cumulative-link) regression.

Parameters:
coef

(p,) slope coefficients (positive beta_j raises the latent score, shifting mass toward higher categories).

Type:

numpy.ndarray

thresholds

(K-1,) ordered cut points alpha.

Type:

numpy.ndarray

se

(p,) standard errors for coef.

Type:

numpy.ndarray

log_likelihood

maximised log-likelihood.

Type:

float

link

"logit" or "probit".

Type:

str

n_categories

number of ordered categories K.

Type:

int

coef: ndarray
thresholds: ndarray
se: ndarray
log_likelihood: float
link: str
n_categories: int
predict_proba(x)[source]

Per-category probabilities (n, K) at design rows x.

Parameters:

x (ndarray)

Return type:

ndarray

predict(x)[source]

Most-probable ordered category per row.

Parameters:

x (ndarray)

Return type:

ndarray

concordance_summary(x, y)[source]

All pairwise concordance measures between two ordinal variables.

Returns:

{'kendall_tau_b', 'gamma', 'somers_d_yx', 'somers_d_xy', 'concordant', 'discordant', 'tx', 'ty', 'txy'} – tx = pairs tied on x only, ty on y only, txy on both.

Parameters:
Return type:

dict[str, float]

kendall_tau(x, y)[source]

Kendall’s tau-b rank correlation (tie-corrected) between two ordinal variables.

Parameters:
Return type:

float

goodman_kruskal_gamma(x, y)[source]

Goodman–Kruskal gamma: (C - D) / (C + D) ignoring ties.

Parameters:
Return type:

float

somers_d(x, y, *, dependent='y')[source]

Somers’ D, the asymmetric rank association treating dependent as the response.

Parameters:
Return type:

float

mann_whitney_u(x, y, *, alternative='two-sided', use_continuity=True)[source]

Mann-Whitney U / Wilcoxon rank-sum test for two independent samples.

Tests whether x is stochastically greater/less than y. Uses mid-ranks for ties, the tie-corrected normal approximation, and (default) a continuity correction. alternative is 'two-sided', 'greater' (x > y), or 'less'. The rank-biserial correlation 2*U1/(n1 n2) - 1 is reported as the effect size.

Parameters:
Return type:

MannWhitneyResult

class MannWhitneyResult(statistic: 'float', statistic2: 'float', zscore: 'float', pvalue: 'float', rank_biserial: 'float', alternative: 'str')[source]

Bases: object

Parameters:
statistic: float
statistic2: float
zscore: float
pvalue: float
rank_biserial: float
alternative: str
wilcoxon_signed_rank(x, y=None, *, alternative='two-sided', zero_method='wilcox', correction=False)[source]

Wilcoxon signed-rank test for paired samples (or one sample vs 0).

Ranks |d| for d = x - y (mid-ranks for ties), splits into positive / negative rank sums, and uses the tie-corrected normal approximation. zero_method='wilcox' drops zero differences (and their ranks); 'pratt' keeps them in the ranking but drops them from the sums. The matched-pairs rank-biserial correlation is reported as the effect size.

Parameters:
Return type:

WilcoxonResult

class WilcoxonResult(statistic: 'float', zscore: 'float', pvalue: 'float', rank_biserial: 'float', alternative: 'str')[source]

Bases: object

Parameters:
statistic: float
zscore: float
pvalue: float
rank_biserial: float
alternative: str
sign_test(x, y=None, *, alternative='two-sided')[source]

Sign test for paired samples (or one sample vs 0): exact binomial test on the signs of x - y.

Only the directions of the differences are used (ties dropped), so it is maximally robust but less powerful than the signed-rank test. extra carries n_positive and n (non-zero pairs).

Parameters:
Return type:

TestResult

kruskal_wallis(*samples)[source]

Kruskal-Wallis H test: the rank-based k-sample generalization of Mann-Whitney (one-way ANOVA).

Tie-corrected H with a chi-square(k-1) reference. extra carries df and the epsilon_squared effect size (H - k + 1)/(N - k).

Parameters:

samples (Any)

Return type:

TestResult

friedman_test(*measurements)[source]

Friedman test for k related samples (repeated measures): the rank-based repeated-measures ANOVA.

Pass each treatment as a separate equal-length array (one value per block). Ranks within each block, tie-corrects, and uses a chi-square(k-1) reference. extra carries df and Kendall’s W concordance effect size.

Parameters:

measurements (Any)

Return type:

TestResult

brunner_munzel(x, y, *, alternative='two-sided', distribution='t')[source]

Brunner-Munzel test: the generalized Wilcoxon test that does NOT assume equal variances/shapes.

Tests the stochastic-equality null P(x < y) + 0.5 P(x = y) = 1/2. distribution='t' uses a Satterthwaite t reference (recommended for small samples); 'normal' the normal approximation. Reports the estimated relative effect p_hat = P(x < y) + 0.5 P(x = y) in extra.

Parameters:
Return type:

TestResult

mood_median_test(*samples, ties='below')[source]

Mood’s median test: chi-square test that k samples share a common median.

Cross-tabulates each observation as above / (at-or-below) the pooled grand median and runs a chi-square test of independence on the resulting 2xk table. extra carries the grand_median.

Parameters:
Return type:

TestResult

dunn_test(*samples, p_adjust='holm')[source]

Dunn’s post-hoc test: all pairwise rank-mean comparisons after a Kruskal-Wallis rejection.

Uses the pooled-rank z statistic with the shared tie-corrected variance, and adjusts the pairwise p-values by 'holm', 'bonferroni', or 'none'.

Parameters:

  • samples (Any)

  • p_adjust (str)

Return type:

DunnResult

class DunnResult(comparisons, zscores, pvalues, p_adjust)[source]

Bases: object

Post-hoc Dunn pairwise comparisons after Kruskal-Wallis.

Parameters:
comparisons: list[tuple[int, int]]
zscores: ndarray
pvalues: ndarray
p_adjust: str
jonckheere_terpstra(*samples, alternative='increasing')[source]

Jonckheere-Terpstra test for an ORDERED alternative across independent samples.

More powerful than Kruskal-Wallis when the groups are expected to shift monotonically in the given order. alternative='increasing' / 'decreasing' / 'two-sided'. Uses the tie-corrected normal approximation of the J statistic (sum of pairwise Mann-Whitney counts over ordered pairs).

Parameters:

  • samples (Any)

  • alternative (str)

Return type:

TestResult

page_trend_test(*measurements, decreasing=False)[source]

Page’s trend test for an ORDERED alternative in repeated measures.

Like Friedman but for a pre-specified ordering of the k treatments (the columns, in order). Tests L = sum_j j * R_j against the normal approximation. Set decreasing=True to predict the reverse ordering. extra carries the z-score.

Parameters:

  • measurements (Any)

  • decreasing (bool)

Return type:

TestResult

ks_1samp(x, cdf, *, alternative='two-sided')[source]

One-sample Kolmogorov-Smirnov goodness-of-fit test against a fully-specified cdf callable.

Parameters:
Return type:

TestResult

ks_2samp(x, y, *, alternative='two-sided')[source]

Two-sample Kolmogorov-Smirnov test: max gap between the two empirical CDFs (asymptotic p).

Parameters:
Return type:

TestResult

runs_test(x, *, cutoff='median')[source]

Wald-Wolfowitz runs test for randomness of a binary/dichotomized sequence.

Dichotomizes x about its median (or a supplied numeric cutoff) and tests whether the run count departs from what independence predicts (too few runs => clustering/trend; too many => over-alternation). Normal approximation, two-sided. extra carries the run count and z-score.

Parameters:
Return type:

TestResult

cliffs_delta(x, y)[source]

Cliff’s delta effect size in [-1, 1]: P(x > y) - P(x < y) (rank-based, ties count as 0).

Parameters:
Return type:

float

class TestResult(statistic, pvalue, extra=<factory>)[source]

Bases: object

Generic statistic + p-value result; extra carries test-specific fields (effect size, df, …).

Parameters:
statistic: float
pvalue: float
extra: dict[str, Any]
deming_regression(x, y, variance_ratio=1.0)[source]

Errors-in-variables (Deming) regression of y on x when both are noisy.

Parameters:

  • x – paired measurements; both may carry error.

  • y – paired measurements; both may carry error.

  • variance_ratio (float) – var(e_y) / var(e_x) – the ratio of output to input noise variance. 1.0 is total least squares (orthogonal regression); a large value -> ordinary least squares (no input error); a small value -> inverse regression (predictor dominated by error).

Returns:

A DemingFit with the unbiased slope / intercept and the recovered latent x*.

Return type:

DemingFit

class DemingFit(slope, intercept, variance_ratio, x, y)[source]

Bases: object

Result of deming_regression(): slope/intercept plus the recovered latent predictor values.

conditional_mean(x_star)[source]

The conditional mean E[y | x*] = a + b x* at true predictor values x*.

Parameters:

x_star (ndarray)

Return type:

ndarray

simex(fit_fn, x, y, sigma_u, *, lambdas=None, n_sims=100, extrapolation='quadratic', seed=0)[source]

SIMEX: simulation–extrapolation correction for a predictor measured with known error.

When a predictor x is observed as x = x* + u with u ~ N(0, sigma_u^2), naive estimates are biased (attenuation). SIMEX adds further noise of variance lambda sigma_u^2 for a grid of lambda >= 0, refits at each level (averaging over n_sims noise draws), then extrapolates the estimate back to lambda = -1 (zero measurement error). Works for any estimator returning a parameter vector.

Parameters:

  • fit_fn (Callable[[ndarray, ndarray], ndarray]) – f(x, y) -> theta returning the parameter vector for (possibly multi-column) x.

  • x (ndarray) – (n,) or (n, p) error-prone predictor(s).

  • y (ndarray) – (n,) response.

  • sigma_u (float | ndarray) – measurement-error standard deviation (scalar, or per-column for matrix x).

  • lambdas (ndarray | None) – extra-noise levels (defaults to 0, 0.5, 1.0, 1.5, 2.0).

  • n_sims (int) – noise replications per level.

  • extrapolation (str) – "quadratic" or "linear" extrapolant in lambda.

  • seed (int | RandomState | None) – RNG seed.

Returns:

{'estimate', 'naive', 'lambdas', 'curve'} – the SIMEX-corrected parameter vector, the naive lambda=0 estimate, and the per-level averaged estimates.

Return type:

dict

propagate_uncertainty(func, samples, *, quantiles=(0.025, 0.5, 0.975))[source]

Monte-Carlo propagation of an uncertainty set through an arbitrary functional.

Pushes input draws (a posterior sample, a bootstrap set, any uncertainty representation) through func and summarises the output distribution – the general “what is the uncertainty of g(theta)?” operation. func may be vectorised (accept the whole (n, ...) array) or act on a single draw; both are detected.

Parameters:

  • func (Callable[[ndarray], ndarray]) – the functional to propagate. Returns a scalar or fixed-length vector per input draw.

  • samples (ndarray) – (n, ...) input draws (rows are draws).

  • quantiles (tuple[float, ...]) – output quantiles to report.

Returns:

{'mean', 'std', 'quantiles', 'levels', 'samples'} over the propagated outputs.

Return type:

dict

paired_score_difference(scores_a, scores_b, *, lower_is_better=True, ci_level=0.95)[source]

Mean paired held-out score difference with a CI and a paired t-test.

Parameters:

  • scores_a (ndarray) – (n,) per-observation held-out scores for the two models (same observations, same order).

  • scores_b (ndarray) – (n,) per-observation held-out scores for the two models (same observations, same order).

  • lower_is_better (bool) – True for losses/scores (CRPS, log loss, pinball); False for higher-is-better metrics.

  • ci_level (float) – confidence level for the interval on the mean difference.

Returns:

{'mean_diff', 'se', 'ci_low', 'ci_high', 't', 'p_value', 'favored'} where mean_diff is mean(a - b) and favored is 'A' / 'B' / 'tie' at the given level.

Return type:

dict

vuong_test(loglik_a, loglik_b, *, k_a=0, k_b=0, correction='none')[source]

Vuong’s test for non-nested model selection.

Compares two models by their pointwise log-likelihoods. Under the null that both are equally close to the truth, the statistic sqrt(n) * mean(m) / sd(m) (with m_i = ll_a_i - ll_b_i, minus an optional complexity correction) is asymptotically standard normal. A large positive value favors A.

Parameters:

  • loglik_a (ndarray) – (n,) pointwise log-likelihoods of the two (non-nested) models.

  • loglik_b (ndarray) – (n,) pointwise log-likelihoods of the two (non-nested) models.

  • k_a (int) – parameter counts, used only if correction is set.

  • k_b (int) – parameter counts, used only if correction is set.

  • correction (str) – "none", "aic" (subtract k_a - k_b), or "bic" (subtract (k_a - k_b) log n / 2) from the log-likelihood ratio.

Returns:

{'statistic', 'p_value', 'favored'}.

Return type:

dict

clarke_test(loglik_a, loglik_b, *, k_a=0, k_b=0, correction='none')[source]

Clarke’s distribution-free paired sign test for non-nested models.

Counts how often model A’s pointwise log-likelihood beats B’s; under the null this count is Binomial(n, 0.5). More robust than vuong_test() when the per-observation log-ratio is heavy-tailed or skewed (where the normal approximation behind Vuong fails).

Returns:

{'statistic', 'p_value', 'favored', 'n'} – statistic is the number of points favoring A.

Parameters:
Return type:

dict

compare_elpd(pointwise_a, pointwise_b)[source]

Compare two models’ expected log pointwise predictive density (LOO/WAIC).

Takes the per-observation elpd contributions (the pointwise arrays returned by mixle.ppl.diagnostics.psis_loo() / waic) and returns the elpd difference with the standard error of the pointwise difference – the honest SE for model comparison (a difference within ~2 SE of zero is not decisive).

Parameters:

  • pointwise_a (ndarray) – (n,) per-observation elpd contributions (higher is better).

  • pointwise_b (ndarray) – (n,) per-observation elpd contributions (higher is better).

Returns:

{'elpd_diff', 'se', 'z', 'favored'} – elpd_diff = sum(a - b).

Return type:

dict

split_conformal(cal_pred, cal_y, test_pred, *, alpha=0.1, side='two-sided')[source]

Split (inductive) conformal interval from a calibration set.

Parameters:

  • cal_pred (ndarray) – (n,) model predictions on the calibration set.

  • cal_y (ndarray) – (n,) calibration responses.

  • test_pred (ndarray) – (m,) predictions at the test points.

  • alpha (float) – miscoverage level (1 - alpha coverage).

  • side (str) – "two-sided" (|y - yhat| score), "upper" (one-sided upper bound), or "lower" (one-sided lower bound).

Returns:

(lower, upper) arrays of length m (an unbounded side is -inf / +inf).

Return type:

tuple[ndarray, ndarray]

jackknife_plus(x, y, fit_predict, x_test, *, alpha=0.1)[source]

Jackknife+ intervals (leave-one-out), using all data for both fitting and calibration.

For each training point i the model is refit without i; R_i = |y_i - mu_{-i}(x_i)| is the LOO residual and mu_{-i}(x) the LOO prediction at a test point. The interval aggregates mu_{-i}(x) -/+ R_i across i (Barber et al. 2021), giving ~``1 - 2 alpha`` worst-case and ~``1 - alpha`` typical coverage without a data split. Costs n refits.

Returns:

(lower, upper) arrays of length len(x_test).

Parameters:
Return type:

tuple[ndarray, ndarray]

cv_plus(x, y, fit_predict, x_test, *, alpha=0.1, n_folds=10, seed=0)[source]

CV+ intervals: the K-fold analogue of jackknife_plus() (only n_folds refits).

Each point’s residual uses the model trained on the other folds, and the test prediction uses the same out-of-fold model. Much cheaper than Jackknife+ with nearly the same guarantee.

Returns:

(lower, upper) arrays of length len(x_test).

Parameters:
Return type:

tuple[ndarray, ndarray]

mondrian_conformal(cal_pred, cal_y, cal_groups, test_pred, test_groups, *, alpha=0.1)[source]

Mondrian (group-conditional) split conformal: a separate quantile per group.

Calibrates the conformal quantile within each group (taxonomy), so coverage holds conditional on the group rather than only marginally – the fix when error scale varies across known subpopulations.

Parameters:

  • cal_pred (ndarray) – calibration predictions, responses, and group labels.

  • cal_y (ndarray) – calibration predictions, responses, and group labels.

  • cal_groups (ndarray) – calibration predictions, responses, and group labels.

  • test_pred (ndarray) – test predictions and their group labels.

  • test_groups (ndarray) – test predictions and their group labels.

  • alpha (float) – miscoverage level.

Returns:

(lower, upper) arrays of length len(test_pred).

Return type:

tuple[ndarray, ndarray]

weighted_conformal(cal_pred, cal_y, test_pred, weights, *, alpha=0.1, test_weight=1.0)[source]

Covariate-shift-weighted split conformal (Tibshirani et al. 2019).

Under covariate shift the calibration and test inputs follow different distributions; reweighting the calibration scores by the likelihood ratio w(x) = p_test(x)/p_train(x) restores coverage. Uses the weighted empirical quantile of the calibration scores (each test point shares the same test_weight for its own potential score).

Parameters:

  • cal_pred (ndarray) – calibration predictions and responses.

  • cal_y (ndarray) – calibration predictions and responses.

  • test_pred (ndarray) – (m,) test predictions.

  • weights (ndarray) – (n,) likelihood-ratio weights for the calibration points (need not be normalised).

  • alpha (float) – miscoverage level.

  • test_weight (float) – the weight assigned to a test point (usually the mean test/train ratio; 1.0 when weights are self-normalised around the test density).

Returns:

(lower, upper) arrays of length m (a symmetric interval per test point).

Return type:

tuple[ndarray, ndarray]

conformal_label_threshold(cal_prob_true, *, alpha=0.1)[source]

Calibrate the LAC (least-ambiguous set-valued classifier) score threshold for 1 - alpha coverage.

The nonconformity score of a calibration point is 1 - p_model[true_class] – which needs the model’s class scores to rank well, not to be a true probability (the whole point: a softmax over a ReLU net is not a describable random process, but conformal still gives a finite-sample coverage guarantee from how those scores behave on held-out, exchangeable data). Returns the conformal quantile qhat of the calibration scores; a class is admitted at test time iff 1 - p[c] <= qhat (see conformal_label_sets()).

Parameters:

  • cal_prob_true (ndarray) – (n,) model score assigned to the true class of each calibration point.

  • alpha (float) – miscoverage level (1 - alpha marginal coverage of the returned sets).

Returns:

qhat – the score threshold (+inf when n is too small for the requested alpha).

Return type:

float

conformal_label_sets(cal_prob_true, test_prob, *, alpha=0.1, qhat=None)[source]

Split-conformal prediction sets for a classifier: distribution-free 1 - alpha label coverage.

Calibrates a LAC threshold (conformal_label_threshold()) on the held-out true-class scores, then admits every class whose score clears it. The returned boolean mask has guaranteed marginal coverage: the true label is in the set with probability >= 1 - alpha. A singleton set is a confident prediction; an empty or multi-label set is an honest “I’m not sure” – the signal a cost-aware cascade escalates on.

Parameters:

  • cal_prob_true (ndarray) – (n,) score assigned to the true class of each calibration point.

  • test_prob (ndarray) – (m, K) model class scores at the test points (rows need not sum to 1).

  • alpha (float) – miscoverage level.

  • qhat (float | None) – a precomputed threshold (e.g. from an earlier calibration); recomputed if None.

Returns:

(sets, qhat) – sets is an (m, K) boolean mask, qhat the threshold used.

Return type:

tuple[ndarray, float]

learn_structure(data, *, field_estimators=None, min_gain=0.0, max_levels=64, n_bins=4, max_its=30)[source]

Discover the dependency forest for heterogeneous data and return the fitted joint model.

Any field can be a parent: a discrete one conditions directly, a continuous one is quantile-binned into n_bins conditioning levels (so a real can drive a count, a category, or another real). Scores every (parent -> child) pair by dependency_gain(), greedily builds a maximum-gain acyclic forest (each field at most one parent), and fits each factor. Falls back to independent marginals where no dependence clears min_gain – never worse than a composite, much better when structure exists. This is “automatic inference for composable models of heterogeneous data” made real.

Parameters:
Return type:

DependencyTreeDistribution

learn_mixture_structure(data, n_components, *, restarts=3, max_iter=15, seed=0, min_gain=0.0, n_bins=4, max_its=30)[source]

Fit a MixtureOfDependencyTrees by hard EM – discover clusters AND each cluster’s dependency graph.

Each iteration re-learns a dependency forest per cluster on its currently-assigned points (M-step), then reassigns every record to its most-probable cluster (E-step), until assignments stabilize. Runs restarts random initializations and returns the highest-likelihood fit. Empty/tiny clusters are re-seeded so a component never collapses.

Parameters:
Return type:

MixtureOfDependencyTrees

learn_bayesian_network(data, *, max_parents=2, min_gain=0.0, max_its=30, weights=None)[source]

Discover a heterogeneous DAG for data and return the fitted network.

Each field greedily gains up to max_parents parents by BIC-penalized conditional likelihood, keeping the graph acyclic. Continuous children become conditional-linear-Gaussian factors (regression on continuous + one-hot discrete parents); discrete/count children condition on the joint config of their discrete parents, or become GLM nodes when a driver is continuous. With weights (soft-EM responsibilities), every factor fit and the BIC search itself are responsibility-weighted, with effective sample size sum(weights).

Parameters:
Return type:

HeterogeneousBayesianNetwork

learn_mixture_bayesian_network(data, n_components, *, restarts=3, max_iter=12, seed=0, max_parents=2, min_gain=0.0, max_its=30, em='hard')[source]

Fit a MixtureOfBayesianNetworks by EM – discover clusters AND each cluster’s DAG.

em="hard" (default): each iteration re-learns a network per cluster on its assigned points and reassigns every record to its most-probable cluster, until assignments stabilize. em="soft": proper EM – every record contributes to EVERY cluster with its responsibility, each component’s structure search and factor fits are responsibility-weighted (learn_bayesian_network(weights=)), and convergence is on the observed-data log-likelihood; boundary points shape both clusters instead of whipsawing between them. Both run from the same restarts (k-means-seeded + random); best final log-likelihood wins. Starved clusters are re-seeded (hard) / responsibility-floored (soft).

Parameters:
Return type:

MixtureOfBayesianNetworks

select_mixture_components(data, k_values=(1, 2, 3, 4), *, em='hard', seed=0, **kwargs)[source]

Model selection over the number of clusters by BIC.

Fits learn_bayesian_network() for k=1 and learn_mixture_bayesian_network() for each larger k, scores each by bayesian_network_bic(), and returns (best_model, report) where report = {"k": chosen, "bic": {k: score, ...}}. BIC’s n_params log n penalty is what stops a mixture of per-cluster DAGs from always preferring more clusters.

Parameters:
Return type:

tuple[Any, dict[str, Any]]

bayesian_network_bic(model, data)[source]

BIC (lower is better) for a fitted network or mixture: -2 LL + n_params log n.

Parameters:
Return type:

float

learn_structure_embedded(data, *, text_fields='auto', n_clusters=8, embed_dim=16, seed=0, max_parents=2, **embed_kw)[source]

Discover cross-field structure where some fields are free text (see module docstring).

Parameters:

  • data (Sequence[tuple]) – flat tuple records; text fields may hold arbitrary strings.

  • text_fields (Sequence[int] | str) – which field indices to embed, or "auto" – every string-valued field with too many distinct values to be a categorical.

  • n_clusters (int) – number of representative clusters surfaced by describe() (interpretability only; the model uses the full embedding vector, not a cluster code).

  • embed_dim (int) – embedding dimension for mixle.represent.fit_embedder() – the vector node’s dim.

  • **embed_kw (Any) – forwarded to fit_embedder (epochs, hidden, feature_dim, …).

  • seed (int)

  • max_parents (int)

  • **embed_kw

Return type:

EmbeddedStructureModel

class EmbeddedStructureModel(net, codecs)[source]

Bases: object

A discovered dependency graph over records whose text fields ride in as embedding VECTORS.

Parameters:

  • net (Any)

  • codecs (dict[int, EmbeddedFieldCodec])

encode_record(x)[source]

The record with each text field replaced by its embedding vector.

Parameters:

x (tuple)

Return type:

tuple

encode_records(rows)[source]
Parameters:

rows (Sequence[tuple])

Return type:

list[tuple]

edges()[source]
Return type:

list[tuple[int, int]]

log_density(x)[source]
Parameters:

x (tuple)

Return type:

float

seq_log_density(rows)[source]
Parameters:

rows (Sequence[tuple])

Return type:

ndarray

describe()[source]

The discovered structure with per-cluster representative examples for each text field.

Return type:

dict[str, Any]

class HeterogeneousBayesianNetwork(factors)[source]

Bases: object

A DAG joint over a heterogeneous record: log p(x) = sum_i log P(x_i | parents(i)) over fitted factors.

Parameters:

factors (Sequence[Any])

edges()[source]
Return type:

list[tuple[int, int]]

log_density(x)[source]
Parameters:

x (tuple)

Return type:

float

seq_log_density(encoded)[source]
Parameters:

encoded (Any)

Return type:

ndarray

dist_to_encoder()[source]
Return type:

Any

sampler(seed=None)[source]
Parameters:

seed (int | None)

Return type:

Any

class MixtureOfBayesianNetworks(components, weights)[source]

Bases: object

A latent mixture whose components each carry their own heterogeneous DAG (regression edges and all).

The fullest form of the moat: it discovers the clustering and each cluster’s cross-field graphical model, so the slope of one field on another (or the whole dependency graph) can differ between clusters – which a single network cannot represent. log p(x) = logsumexp_k(log w_k + log p_k(x)) over HeterogeneousBayesianNetwork components. Fit by learn_mixture_bayesian_network().

Parameters:

  • components (Sequence[HeterogeneousBayesianNetwork])

  • weights (Sequence[float])

log_density(x)[source]
Parameters:

x (tuple)

Return type:

float

seq_log_density(encoded)[source]
Parameters:

encoded (Any)

Return type:

ndarray

dist_to_encoder()[source]
Return type:

Any

responsibilities(data)[source]
Parameters:

data (Sequence[tuple])

Return type:

ndarray

sampler(seed=None)[source]
Parameters:

seed (int | None)

Return type:

Any

property n_components: int
dependency_gain(parent, child, child_estimator, *, max_its=30, penalty='bic')[source]

Description-length gain (nats) of modeling child conditioned on a discrete parent vs. independently.

Fits the marginal P(child) and the conditional P(child | parent) (a child model per parent value) on the same data and returns LL_cond - LL_marginal minus a complexity penalty for the extra parameters (BIC: 0.5 * (levels - 1) * k * ln n). Positive means the dependence is worth modeling. This is a model-based dependency test – it works across any pair of families, unlike a same-type MI estimate.

Parameters:
Return type:

float

class DependencyTreeDistribution(parents, factors, binners=None)[source]

Bases: object

A directed-forest joint over a heterogeneous record: each field is a marginal or a conditional on its parent.

log_density(record) = sum_root log P(f_root) + sum_child log P(f_child | f_parent). The dependence a CompositeDistribution assumes away is modeled here as per-parent-value conditionals – while it still scores, samples, and composes like any mixle distribution.

Parameters:

  • parents (Sequence[int | None])

  • factors (Sequence[Any])

  • binners (Sequence[Any] | None)

log_density(x)[source]
Parameters:

x (tuple)

Return type:

float

seq_log_density(encoded)[source]
Parameters:

encoded (Any)

Return type:

ndarray

dist_to_encoder()[source]
Return type:

Any

sampler(seed=None)[source]
Parameters:

seed (int | None)

Return type:

Any

edges()[source]

The learned dependency edges (parent_field, child_field).

Return type:

list[tuple[int, int]]

class MixtureOfDependencyTrees(components, weights)[source]

Bases: object

A latent mixture whose components each carry their own discovered dependency structure.

log p(x) = logsumexp_k ( log w_k + log p_k(x) ) where each p_k is a DependencyTreeDistribution. This is the deep form of the tagline: it discovers both the clustering and the within-cluster cross-field dependence – so the same category can map to different reals in different clusters, which neither a single dependency tree (one relationship) nor a mixture of independent composites (no within-cluster dependence) can represent. Fit by learn_mixture_structure().

Parameters:

  • components (Sequence[DependencyTreeDistribution])

  • weights (Sequence[float])

log_density(x)[source]
Parameters:

x (tuple)

Return type:

float

seq_log_density(encoded)[source]
Parameters:

encoded (Any)

Return type:

ndarray

dist_to_encoder()[source]
Return type:

Any

responsibilities(data)[source]

Posterior p(component | record) for each record – the E-step and a soft cluster assignment.

Parameters:

data (Sequence[tuple])

Return type:

ndarray

sampler(seed=None)[source]
Parameters:

seed (int | None)

Return type:

Any

property n_components: int
kfold(n, n_splits=5, *, shuffle=False, seed=0)[source]

Standard k-fold split of n rows.

Parameters:

  • n (int) – number of observations.

  • n_splits (int) – number of folds k (each row is in exactly one test fold).

  • shuffle (bool) – shuffle the row order before splitting (use for i.i.d. data; leave False to keep contiguous folds, i.e. blocked_kfold()).

  • seed (int | RandomState | None) – RNG seed when shuffle is True.

Returns:

k (train_index, test_index) pairs.

Return type:

list[tuple[ndarray, ndarray]]

blocked_kfold(n, n_splits=5)[source]

Contiguous-block k-fold (no shuffle) for serially dependent data.

Identical to kfold() with shuffle=False; named separately because keeping blocks contiguous is the point for time series, not an incidental default.

Parameters:
Return type:

list[tuple[ndarray, ndarray]]

leave_one_out(n)[source]

Leave-one-out CV: n folds, each holding out a single observation.

Parameters:

n (int)

Return type:

list[tuple[ndarray, ndarray]]

stratified_kfold(y, n_splits=5, *, shuffle=True, seed=0)[source]

Stratified k-fold preserving class proportions in every fold.

Distributes each class’s indices across the folds so that class frequencies in each test fold mirror the overall frequencies – important for imbalanced classification.

Parameters:

  • y (ndarray) – (n,) class labels.

  • n_splits (int) – number of folds.

  • shuffle (bool) – shuffle within each class before round-robin assignment.

  • seed (int | RandomState | None) – RNG seed when shuffle is True.

Returns:

n_splits (train_index, test_index) pairs.

Return type:

list[tuple[ndarray, ndarray]]

leave_one_group_out(groups)[source]

Leave-one-group-out CV: one fold per distinct group held out entirely.

Parameters:

groups (ndarray) – (n,) group labels (e.g. subject / site / genus id).

Returns:

One (train_index, test_index) pair per unique group.

Return type:

list[tuple[ndarray, ndarray]]

group_kfold(groups, n_splits=5)[source]

Group k-fold: partition groups into k folds so no group spans train and test.

Greedily assigns whole groups (largest first) to the currently-smallest fold, balancing fold sizes while keeping each group intact – the right scheme for repeated measures when there are more groups than folds.

Parameters:
Return type:

list[tuple[ndarray, ndarray]]

time_series_split(n, n_splits=5, *, gap=0, max_train_size=None)[source]

Forward-chaining time-series CV: train always precedes test (expanding window).

The series is cut into n_splits + 1 contiguous blocks; fold i tests on block i+1 and trains on everything strictly before it. A gap (buffer) drops the gap points just before each test block from the training set so leakage through short-range autocorrelation is avoided.

Parameters:

  • n (int) – number of time-ordered observations.

  • n_splits (int) – number of train/test splits.

  • gap (int) – number of observations to drop between the train and test segments.

  • max_train_size (int | None) – cap on the training-window length (sliding window); None keeps it expanding.

Returns:

n_splits (train_index, test_index) pairs.

Return type:

list[tuple[ndarray, ndarray]]

purged_kfold(n, n_splits=5, *, embargo=0)[source]

Purged/embargoed blocked k-fold (a.k.a. buffered CV).

Contiguous test blocks, but the embargo observations on each side of a test block are removed from the training set, so no training point sits within the autocorrelation reach of a test point. The standard guard for serially dependent data when you still want every block to serve as a test fold (cf. de Prado 2018).

Parameters:

  • n (int) – number of time-ordered observations.

  • n_splits (int) – number of contiguous folds.

  • embargo (int) – buffer width removed from train on both sides of each test block.

Returns:

n_splits (train_index, test_index) pairs.

Return type:

list[tuple[ndarray, ndarray]]

spatial_block_kfold(coords, n_splits=5, *, block_size=None, n_side=None, seed=0)[source]

Spatial-block k-fold: hold out contiguous spatial blocks, not scattered points.

Partitions space into a regular grid of blocks and assigns whole blocks to folds at random, so a test point’s spatial neighbours are held out with it rather than leaking into training. This is the right scheme for geostatistical / spatially autocorrelated data.

Parameters:

  • coords (ndarray) – (n, d) spatial coordinates (typically d = 2).

  • n_splits (int) – number of folds.

  • block_size (float | None) – grid-cell width in coordinate units; if None it is derived from n_side.

  • n_side (int | None) – number of grid cells per axis; defaults so there are comfortably more blocks than folds.

  • seed (int | RandomState | None) – RNG seed for assigning blocks to folds.

Returns:

n_splits (train_index, test_index) pairs.

Return type:

list[tuple[ndarray, ndarray]]

nested_kfold(n, *, outer_splits=5, inner_splits=4, shuffle=False, seed=0)[source]

Nested k-fold: an inner CV (for tuning) inside each outer fold (for evaluation).

The outer loop estimates generalisation; the inner loop selects hyper-parameters using only the outer-training data, so the reported score is not optimistically biased by tuning on the test set.

Parameters:

  • n (int) – number of observations.

  • outer_splits (int) – number of outer folds.

  • inner_splits (int) – number of inner folds within each outer training set.

  • shuffle (bool) – shuffle before splitting (i.i.d. data only).

  • seed (int | RandomState | None) – RNG seed.

Returns:

A list of NestedFold; inner folds index the original array.

Return type:

list[NestedFold]

class NestedFold(train, test, inner)[source]

Bases: object

One outer fold of nested_kfold().

Parameters:
train

outer training indices (used to build the inner CV).

Type:

numpy.ndarray

test

outer test indices (held out for the final, untuned-on evaluation).

Type:

numpy.ndarray

inner

inner (train_index, test_index) folds, indexing into the original array (not into train), so they can be applied directly.

Type:

list[tuple[numpy.ndarray, numpy.ndarray]]

train: ndarray
test: ndarray
inner: list[tuple[ndarray, ndarray]]
log_score(prob, *, mean=True)[source]

Logarithmic score (log loss) -log p(y) of the probability assigned to the outcome.

The log score is local (it depends only on the density at the realised value) and strictly proper. Pass the predictive probability/density evaluated at the observed outcome.

Parameters:

  • prob (ndarray) – array of predictive probabilities (or densities) at the realised outcomes. Values are clipped away from zero before the log so a single zero-probability event yields a large but finite penalty rather than inf.

  • mean (bool) – if True (default) return the mean log loss; otherwise the per-observation vector.

Returns:

Mean log loss (float) or the per-observation array.

Return type:

ndarray | float

brier_score(prob, outcome, *, mean=True)[source]

Brier score (mean squared error of probabilistic classification).

For binary forecasts pass 1-D prob (predicted probability of the positive class) and 0/1 outcome. For K-class forecasts pass prob shaped (n, K) and outcome either as integer class labels (length n) or as a one-hot (n, K) matrix; the multiclass score sums the squared error across classes, so it ranges in [0, 2].

Parameters:

  • prob (ndarray) – (n,) positive-class probabilities, or (n, K) class probabilities.

  • outcome (ndarray) – (n,) 0/1 or integer labels, or (n, K) one-hot.

  • mean (bool) – if True return the mean; otherwise the per-observation vector.

Returns:

Mean Brier score (float) or the per-observation array.

Return type:

ndarray | float

brier_decomposition(prob, outcome, *, bins=10)[source]

Murphy’s reliability–resolution–uncertainty decomposition of the (binary) Brier score.

Bins the forecast probabilities into bins equal-width bins on [0, 1] and computes

Brier = reliability - resolution + uncertainty,

where reliability (lower is better) is the mean squared gap between forecast probability and observed frequency within a bin, resolution (higher is better) is how far bin frequencies move from the base rate, and uncertainty is the base-rate variance p_bar (1 - p_bar) (a property of the data, not the forecaster).

Parameters:

  • prob (ndarray) – (n,) predicted probabilities of the positive class.

  • outcome (ndarray) – (n,) 0/1 outcomes.

  • bins (int) – number of equal-width probability bins.

Returns:

{'reliability', 'resolution', 'uncertainty', 'brier'}. The identity reliability - resolution + uncertainty == brier holds up to binning of the score’s in-bin variance term.

Return type:

dict[str, float]

crps_ensemble(forecasts, y, *, fair=False, mean=True)[source]

Continuous Ranked Probability Score from a finite predictive ensemble (sample).

Uses the energy form CRPS = E|X - y| - 1/2 E|X - X'| estimated from the ensemble draws. The CRPS generalises absolute error to distributional forecasts (a point forecast recovers |x-y|) and is reported in the same units as y.

Parameters:

  • forecasts (ndarray) – predictive draws. Either (n, m) (m draws for each of n observations) or (m,) for a single observation. Ragged ensembles are not supported – pad or call per observation.

  • y (ndarray) – (n,) realised values (or a scalar for the (m,) case).

  • fair (bool) – if True use the unbiased 1/(m(m-1)) spread estimator (the “fair”/almost-unbiased CRPS); if False (default) the standard 1/m^2 estimator.

  • mean (bool) – if True return the mean CRPS; otherwise the per-observation vector.

Returns:

Mean CRPS (float) or the per-observation array.

Return type:

ndarray | float

crps_gaussian(mu, sigma, y, *, mean=True)[source]

Closed-form CRPS for a Gaussian predictive distribution N(mu, sigma^2).

CRPS = sigma * [ z (2 Phi(z) - 1) + 2 phi(z) - 1/sqrt(pi) ] with z = (y - mu) / sigma (Gneiting & Raftery 2007). Exact, so it is the right reference when the forecast is Gaussian.

Parameters:

  • mu (ndarray) – predictive means, broadcastable with y.

  • sigma (ndarray) – predictive standard deviations (> 0), broadcastable with y.

  • y (ndarray) – realised values.

  • mean (bool) – if True return the mean; otherwise the per-observation vector.

Returns:

Mean CRPS (float) or the per-observation array.

Return type:

ndarray | float

interval_score(lower, upper, y, alpha, *, mean=True)[source]

Winkler interval score for a central (1 - alpha) prediction interval.

IS = (u - l) + (2/alpha)(l - y) 1[y < l] + (2/alpha)(y - u) 1[y > u] (Gneiting & Raftery 2007). It rewards narrow intervals but penalises misses by 2/alpha times the shortfall, so at matched nominal coverage the tighter, better-placed interval scores lower. This is the proper way to rank interval methods (conformal, quantile-regression, GP credible bands) head to head.

Parameters:

  • lower (ndarray) – (n,) interval lower endpoints.

  • upper (ndarray) – (n,) interval upper endpoints.

  • y (ndarray) – (n,) realised values.

  • alpha (float) – miscoverage level (e.g. 0.1 for a 90% interval), in (0, 1).

  • mean (bool) – if True return the mean; otherwise the per-observation vector.

Returns:

Mean interval score (float) or the per-observation array.

Return type:

ndarray | float

winkler_score(lower, upper, y, alpha, *, mean=True)[source]

Alias for interval_score() (the score is due to Winkler 1972).

Parameters:
Return type:

ndarray | float

pinball_loss(pred, y, tau, *, mean=True)[source]

Pinball (quantile / check) loss for a quantile forecast.

rho_tau(u) = max(tau * u, (tau - 1) * u) with u = y - pred; its expectation is minimised by the true tau-quantile, so it is the proper score for quantile forecasts and the loss that quantile regression and conformalised quantiles optimise.

Parameters:

  • pred (ndarray) – predicted tau-quantiles, shape (n,) (single level) or (n, q) (several levels, one per column of tau).

  • y (ndarray) – (n,) realised values.

  • tau (float | ndarray) – quantile level(s) in (0, 1) – a scalar for (n,) pred, or a length-q array matching the columns of a (n, q) pred.

  • mean (bool) – if True return the mean over all entries; otherwise the per-observation (or per-observation-per-level) array.

Returns:

Mean pinball loss (float) or the per-observation array.

Return type:

ndarray | float

energy_score(forecasts, y, *, fair=False, mean=True)[source]

Energy score: the multivariate generalisation of CRPS for vector-valued forecasts.

ES = E||X - y|| - 1/2 E||X - X'|| with Euclidean norms; for scalar y it equals crps_ensemble().

Parameters:

  • forecasts (ndarray) – predictive draws shaped (n, m, d) (m d-vectors per observation) or (m, d) for a single observation.

  • y (ndarray) – (n, d) realised vectors (or (d,) for the single-observation case).

  • fair (bool) – use the unbiased 1/(m(m-1)) spread estimator if True.

  • mean (bool) – if True return the mean; otherwise the per-observation vector.

Returns:

Mean energy score (float) or the per-observation array.

Return type:

ndarray | float

skill_score(score, reference, *, perfect=0.0)[source]

Skill score: fractional improvement of score over a reference score.

skill = (reference - score) / (reference - perfect) – 1 means the forecast hit the perfect score, 0 means it only matched the reference (e.g. climatology), and negative means it did worse than the reference. Works for any negatively-oriented score (lower is better) such as those above.

Parameters:

  • score (float) – the forecast’s score (lower is better).

  • reference (float) – the baseline/reference forecast’s score.

  • perfect (float) – the score of a perfect forecast (0.0 for the rules in this module).

Returns:

The skill score (float); nan if the reference already equals the perfect score.

Return type:

float

select_best(candidates, *, score, lower_is_better=False, conformal_alpha=None)[source]

Score each candidate and return the best, the verifier-based test-time-compute selector.

Parameters:

  • candidates (Any) – an iterable of candidate objects (anything score accepts).

  • score (Callable[[Any], float]) – a verifier score(candidate) -> float; the winner maximizes it (or minimizes it when lower_is_better).

  • lower_is_better (bool) – if True, the winner is the candidate with the smallest score.

  • conformal_alpha (float | None) – optional miscoverage level in (0, 1). When given, the result’s confident flag reports whether the winner’s lead over the runner-up exceeds a conformal/bootstrap band at confidence 1 - conformal_alpha (a z-scaled estimate of the score spread). None (default) leaves confident unset.

Returns:

A SelectionResult. It is also subscriptable (result["best"]), so callers may treat it as a small dict with keys best, best_index, scores, confident.

Raises:

ValueError – if candidates is empty, or conformal_alpha is outside (0, 1).

Return type:

SelectionResult

class SelectionResult(best, best_index, scores, confident=None, margin=None, band=None, _extras=<factory>)[source]

Bases: object

Result of a select_best() call.

Parameters:
best

the winning candidate (the actual object, not its index).

Type:

Any

best_index

the position of the winner in the input candidates.

Type:

int

scores

per-candidate scores in input order (a numpy float array).

Type:

numpy.ndarray

confident

whether the winner’s lead clears the conformal band — None when no conformal_alpha was supplied.

Type:

bool | None

margin

the winner’s lead over the runner-up (in score units, always >= 0); None when there is a single candidate.

Type:

float | None

band

the conformal/bootstrap band the margin was compared against — None when no conformal_alpha was supplied or there is a single candidate.

Type:

float | None

best: Any
best_index: int
scores: ndarray
confident: bool | None = None
margin: float | None = None
band: float | None = None
class ConjugateUpdatable[source]

Bases: PredicateCapability

Has a closed-form conjugate Bayesian update (the top tier of the inference ladder).

supports(x, ConjugateUpdatable) is True iff conjugate_posterior(x, data) returns an exact closed-form posterior; otherwise Bayesian inference falls back to the numerical fitters (MAP / Laplace / MCMC / VI). Detection is the single conjugate_posterior registry.

classmethod check(obj)[source]
Parameters:

obj (Any)

Return type:

bool

conjugate_posterior(dist, data, prior=None, weights=None)[source]

Closed-form conjugate posterior over the parameters of dist given data.

Every supported family returns a closed-form, full-Bayesian posterior: exact parameter samples, marginal likelihood, posterior mean / point estimate, and a posterior predictive. For families with a multi-parameter likelihood whose conjugate is conditional (Gamma, InverseGamma, InverseGaussian, Pareto, NegativeBinomial, vonMises), the non-target parameter (shape / location / scale / number-of-trials / concentration) is taken as known from dist – exactly as a Binomial’s number of trials is. Families with no closed-form conjugate (full Beta, full Gamma, LogSeries, …) raise a clear error rather than returning a partial answer.

Parameters:

  • dist – A mixle likelihood distribution instance whose type selects the conjugate family.

  • data – A sequence of observations of the kind dist scores.

  • prior (dict | None) – Optional dict of conjugate-prior hyperparameters (family specific). None uses a weak proper prior.

  • weights (ndarray | None) – Optional per-observation weights (e.g. EM responsibilities).

Returns:

A ConjugatePosterior exposing mean, sample, point_estimate, log_marginal_likelihood, posterior_predictive and summary.

Return type:

ConjugatePosterior

mixture_conjugate_posterior(dist, data, priors, prior_weights=None, weights=None)[source]

Posterior under a prior that is a mixture of conjugate priors (Diaconis-Ylvisaker).

Parameters:

  • dist – The likelihood distribution (must map to a closed-form conjugate realiser, since the reweighting needs each component’s marginal likelihood).

  • data – The observations.

  • priors (list[dict]) – One hyperparameter dict per mixture component (same keys conjugate_posterior() accepts for this family).

  • prior_weights (ndarray | None) – Prior mixing weights w_m (default uniform); normalised internally.

  • weights (ndarray | None) – Optional per-observation weights.

Returns:

the exact posterior, again a mixture of conjugate posteriors with weights w'_m proportional to w_m * Z_m.

Return type:

A MixtureConjugatePosterior

is_conjugate_family(dist)[source]

Return whether dist (instance or type) has a closed-form conjugate posterior.

Single source of truth: membership in the conjugate_posterior builder registry. This is the family-level capability — “can this distribution be updated in closed form?” — distinct from the instance-level has_conj_prior flag (whether a conjugate prior is currently attached for the MAP path). Backs mixle.capability.ConjugateUpdatable and ProbabilityDistribution.has_conjugate_prior().

Parameters:

dist (Any)

Return type:

bool

class ConjugatePosterior[source]

Bases: object

Base class for a closed-form conjugate posterior over a likelihood’s parameters.

Subclasses provide the family-specific hyperparameters and the four capabilities: mean (the posterior mean of the parameters), sample (exact draws of the parameters), point_estimate (a fitted distribution at the posterior mean), log_marginal_likelihood (the evidence of the observed data under the prior), and posterior_predictive (the distribution of a new draw).

family: str = 'conjugate'
log_base: float = 0.0
mean()[source]
Return type:

dict[str, Any]

sample(n=1, rng=None)[source]
Parameters:
Return type:

dict[str, ndarray]

sampler(seed=None)[source]

Return a sampler exposing the standard obj.sampler(seed).sample(size) API.

Mirrors the distribution sampling convention so conjugate posteriors read the same way as every other mixle object; here each draw is a parameter set from the posterior. size=None returns one parameter set (scalars), size=n a dict of length-n arrays. The explicit-rng form sample(n, rng) remains available.

Parameters:

seed (int | None)

Return type:

ConjugatePosteriorSampler

point_estimate()[source]
log_marginal_likelihood()[source]
Return type:

float

posterior_predictive()[source]
summary()[source]
Return type:

dict[str, Any]

hyper()[source]
Return type:

dict[str, Any]

class MixtureConjugatePosterior(components, post_weights, prior_weights, comp_log_evidence)[source]

Bases: ConjugatePosterior

Posterior under a prior that is itself a mixture of conjugate priors.

Diaconis-Ylvisaker (1979): if the prior is sum_m w_m * pi_m(theta) with each pi_m conjugate to the likelihood, the posterior is again a mixture of the component conjugate posteriors, sum_m w'_m * pi_m^post(theta), with the mixing weights reweighted by each component’s marginal likelihood Z_m:

w’_m proportional to w_m * Z_m.

Everything stays closed-form, so a prior can be multimodal (e.g. “the rate is near 1 OR near 10”) or robust (a heavy-tailed prior as a mixture of conjugates) without losing exact inference.

Parameters:

components (list[ConjugatePosterior])

family: str = 'MixtureOfConjugates'
components: list[ConjugatePosterior]
mean()[source]
Return type:

dict[str, Any]

sample(n=1, rng=None)[source]
Parameters:
Return type:

dict[str, ndarray]

point_estimate()[source]

The maximum-a-posteriori component’s point estimate (the dominant conjugate mode).

posterior_predictive()[source]

A mixle MixtureDistribution of the component predictives, weighted by the posterior weights.

log_marginal_likelihood()[source]
Return type:

float

hyper()[source]
Return type:

dict[str, Any]

class FisherView(dist, estimator=None)[source]

Bases: object

Accumulator-backed Fisher-geometry view of a distribution.

Parameters:

  • dist (Any) – mixle.stats or mixle.bstats distribution.

  • estimator (Any | None) – Optional estimator. When omitted, dist.estimator() is used.

Notes

The generic encoded-data path obtains per-row statistics by replaying seq_update with one-hot weights. That keeps the interface compatible with existing encoders, but it is a correctness fallback rather than a high-performance implementation. Important families should override to_fisher() with direct seq_expected_statistics kernels.

structured_statistics(x, estimate=None, weight=1.0)[source]

Structured sufficient stats for one observation.

For latent-variable distributions, this is the posterior-expected complete-data sufficient statistic under estimate (or this view’s distribution when estimate is omitted).

Parameters:
Return type:

Any

expected_structured_statistics(x, estimate=None, weight=1.0)[source]
Parameters:
Return type:

Any

sufficient_statistics(x, estimate=None, vectorizer=None)[source]
Parameters:

  • x (Any)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

Return type:

ndarray

expected_sufficient_statistics(x, estimate=None, vectorizer=None)[source]
Parameters:

  • x (Any)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

Return type:

ndarray

seq_structured_statistics(enc_data, estimate=None)[source]

Structured per-row stats from encoded data.

This generic implementation is intentionally conservative and may be slow for large data. It exists so every encoder-compatible model has a correct baseline.

Parameters:

  • enc_data (Any)

  • estimate (Any | None)

Return type:

list[Any]

statistics_matrix(data=None, enc_data=None, estimate=None, vectorizer=None, fit=True)[source]

Return an n x d matrix of per-observation sufficient statistics.

Parameters:

  • data (Sequence[Any] | None)

  • enc_data (Any | None)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

  • fit (bool)

Return type:

ndarray

expected_statistics_matrix(data=None, enc_data=None, estimate=None, vectorizer=None, fit=True)[source]
Parameters:

  • data (Sequence[Any] | None)

  • enc_data (Any | None)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

  • fit (bool)

Return type:

ndarray

seq_expected_statistics(enc_data, estimate=None, vectorizer=None, fit=True)[source]
Parameters:

  • enc_data (Any)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

  • fit (bool)

Return type:

ndarray

mean_statistics(stats=None, **kwargs)[source]
Parameters:
Return type:

ndarray

score_center(stats=None, **kwargs)[source]
Parameters:
Return type:

ndarray

fisher_information(stats=None, diagonal=False, ridge=1.0e-8, **kwargs)[source]

Empirical Fisher approximation from per-observation statistic vectors.

Parameters:
Return type:

ndarray

fisher_vectors(stats=None, metric='diagonal', center=None, fisher=None, ridge=1.0e-8, **kwargs)[source]

Return centered/whitened sufficient-statistic vectors.

metric=’identity’ returns centered statistics, ‘diagonal’ divides by per-coordinate Fisher standard deviations, and ‘full’ applies an empirical full-matrix whitening transform.

Parameters:
Return type:

ndarray

observed_fisher_information(stats=None, diagonal=False, center=None, ridge=1.0e-8, **kwargs)[source]

Observed Fisher estimate from score vectors for observed data.

For latent models the statistic rows are posterior-expected complete-data sufficient statistics, so centering them by their model expectation gives observed score vectors. Their covariance is the observed Fisher metric used by Fisher-vector embeddings. If a model center is unavailable, this falls back to the empirical mean.

Parameters:
Return type:

ndarray

observed_fisher_vectors(stats=None, metric='diagonal', center=None, fisher=None, ridge=1.0e-8, **kwargs)[source]
Parameters:
Return type:

ndarray

fisher_vector(x, estimate=None, metric='diagonal', center=None, fisher=None, vectorizer=None, ridge=1.0e-8)[source]
Parameters:

  • x (Any)

  • estimate (Any | None)

  • metric (str)

  • center (ndarray | None)

  • fisher (ndarray | None)

  • vectorizer (SufficientStatisticVectorizer | None)

  • ridge (float)

Return type:

ndarray

natural_parameters()[source]

Return natural parameters when a specialized view provides them.

Return type:

Any

class FixedFisherView(dist, labels)[source]

Bases: FisherView

Distribution-specific Fisher view with fixed vector coordinates.

Parameters:

  • dist (Any)

  • labels (Sequence[Path])

structured_statistics(x, estimate=None, weight=1.0)[source]

Structured sufficient stats for one observation.

For latent-variable distributions, this is the posterior-expected complete-data sufficient statistic under estimate (or this view’s distribution when estimate is omitted).

Parameters:
Return type:

Any

sufficient_statistics(x, estimate=None, vectorizer=None)[source]
Parameters:

  • x (Any)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

Return type:

ndarray

seq_structured_statistics(enc_data, estimate=None)[source]

Structured per-row stats from encoded data.

This generic implementation is intentionally conservative and may be slow for large data. It exists so every encoder-compatible model has a correct baseline.

Parameters:

  • enc_data (Any)

  • estimate (Any | None)

Return type:

list[Any]

statistics_matrix(data=None, enc_data=None, estimate=None, vectorizer=None, fit=True)[source]

Return an n x d matrix of per-observation sufficient statistics.

Parameters:

  • data (Sequence[Any] | None)

  • enc_data (Any | None)

  • estimate (Any | None)

  • vectorizer (SufficientStatisticVectorizer | None)

  • fit (bool)

Return type:

ndarray

mean_statistics(stats=None, model=True, **kwargs)[source]
Parameters:
Return type:

ndarray

fisher_information(stats=None, diagonal=False, ridge=1.0e-8, **kwargs)[source]

Empirical Fisher approximation from per-observation statistic vectors.

Parameters:
Return type:

ndarray

fisher_vectors(stats=None, metric='diagonal', center=None, fisher=None, ridge=1.0e-8, **kwargs)[source]

Return centered/whitened sufficient-statistic vectors.

metric=’identity’ returns centered statistics, ‘diagonal’ divides by per-coordinate Fisher standard deviations, and ‘full’ applies an empirical full-matrix whitening transform.

Parameters:
Return type:

ndarray

to_fisher(dist, **kwargs)[source]

Return a FisherView for dist via its own to_fisher hook.

Fisher views are co-located with each distribution: a distribution owns its view by overriding ProbabilityDistribution.to_fisher. This module keeps only the shared base machinery (FisherView/FixedFisherView/SufficientStatisticVectorizer and the reusable CountFisherView / EmpiricalMetricFixedFisherView helpers) plus _legacy_to_fisher(), the type-name dispatch for families not yet migrated to a per-file hook (and the generic fallback).

Parameters:
Return type:

FisherView

nuts(target, *, backend='auto', dim=None, init=None, num_samples=1000, warmup=1000, chains=1, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, rng=None, parallel=None, **backend_kwargs)[source]

No-U-Turn Sampler over an arbitrary log-target, dispatched to a registered engine.

The target contract depends on the backend (the kinds cannot be auto-converted across autodiff systems):

  • numpy / numba: a fused value_and_grad(theta) -> (logp, grad) (njit-jitted for numba). The caller supplies the (analytic) gradient.

  • torch / jax: a scalar logp(theta); the backend builds value_and_grad by autodiff (torch.func / NumPyro).

Parameters:

  • target (Callable[[...], Any]) – the log-target, per the contract above.

  • backend (str) – "auto" (default), or one of available_backends(). "auto" honors an explicit choice elsewhere, otherwise prefers the always-present numpy path for a plain numpy value_and_grad; pass backend="numba" to run an @njit target, "jax" / "torch" for an autodiff scalar logp.

  • dim (int | None) – parameter dimension. Required unless init is given.

  • init (Any) – initial state, shape (dim,) or (chains, dim) for per-chain starts. Defaults to zeros. A 1-D init is reused (jittered after the first) across chains.

  • num_samples (int) – retained post-warmup draws per chain.

  • warmup (int) – step-size adaptation / burn-in iterations per chain.

  • chains (int) – number of independent chains (>= 2 to get a meaningful R-hat).

  • mass (Any) – forwarded to the sampler.

  • target_accept (float) – forwarded to the sampler.

  • max_tree_depth (int) – forwarded to the sampler.

  • thin (int) – forwarded to the sampler.

  • rng (RandomState | int | None) – seed / RandomState for reproducibility.

  • parallel (bool | str | None) – run the chains independent chains concurrently. None/False (default) runs them in the backend’s usual single call; "thread" uses a thread pool (real speedup for the numba/torch backends, whose inner loops release the GIL); True or "process" uses a process pool (real speedup for the pure-numpy backend; requires a picklable target). Ignored for the jax backend, which already vectorizes chains internally. Each chain still gets an independent seed, so results are valid multi-chain draws (R-hat / ESS across chains).

  • **backend_kwargs (Any) – forwarded to the chosen backend (e.g. compile=, device= for torch; chain_method= for jax).

Returns:

NutsResult with pooled samples (chains*draws, d), per-chain chains (chains, draws, d), per-dimension rhat and ess, the total target-evaluation count, the adapted step_size, and extra={"backend": name}.

Return type:

NutsResult

nuts_torch(logp, *, dim=None, init=None, num_samples=1000, warmup=1000, chains=1, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, rng=None, compile=True, dtype=None, device=None)[source]

Torch-native NUTS over a torch scalar logp(theta) -> Tensor[()] (GPU / large autodiff targets).

The whole leapfrog/tree runs in torch tensors with a torch.compile``d ``value_and_grad (no numpy round-trip / graph re-trace per gradient). Same multi-chain interface and NutsResult as nuts(). Also registered as the "torch" backend.

Performance note (be deliberate about when to use this): on CPU this is typically slower than the numpy nuts() (per-op torch dispatch + host syncs in the tree dominate when the target is cheap). Its value is GPU (device=) and large autodiff targets. For CPU work prefer the numpy backend, numba (analytic gradient), or the jax/NumPyro backend (XLA, vectorized multi-chain). Chains run in a Python loop (per-chain latency).

Parameters:
Return type:

NutsResult

advi(target_batch, u0, s0, *, samples=1000, mc=16, steps=2000, lr=0.05, batch_size=None, family='meanfield', alpha=1.0, rng=None)[source]

Automatic-differentiation VI over a user batched torch target.

Parameters:

  • target_batch (Callable[[Any], Any]) – target_batch(U: Tensor(B, d)) -> Tensor(B,) returning the (unnormalized) joint log-target for each of B parameter draws. The caller owns any data minibatching/rescaling inside this callable; batch_size here is unused unless the caller wires it in (kept for signature parity with the internal ADVI).

  • u0 – initial variational mean and (marginal) scale in the unconstrained space.

  • s0 – initial variational mean and (marginal) scale in the unconstrained space.

  • samples (int) – number of posterior draws to return.

  • mc (int) – Monte-Carlo samples per step, optimizer steps, Adam learning rate.

  • steps (int) – Monte-Carlo samples per step, optimizer steps, Adam learning rate.

  • lr (float) – Monte-Carlo samples per step, optimizer steps, Adam learning rate.

  • family (str) – 'meanfield' (diagonal) or 'fullrank' (Cholesky covariance).

  • alpha (float) – Renyi/tilted objective exponent (1.0 = standard KL-ELBO).

  • rng (RandomState | int | None) – seed / RandomState.

  • batch_size (int | None)

Returns:

AdviResult with samples (samples, d) drawn from the fitted Gaussian q (in the same space target_batch consumes), plus the fitted mean and scale.

Return type:

AdviResult

class NutsResult(samples, chains, rhat, ess, num_target_evals, step_size, extra=<factory>)[source]

Bases: object

Draws and diagnostics from a multi-chain nuts() run.

samples is (chains * draws, d) pooled draws; chains is (n_chains, draws, d).

Parameters:
samples: ndarray
chains: ndarray
rhat: ndarray
ess: ndarray
num_target_evals: int
step_size: float
extra: dict
class AdviResult(samples, mean, scale, objective=None)[source]

Bases: object

Result of advi(): posterior draws plus the fitted variational parameters.

Parameters:
samples: ndarray
mean: ndarray
scale: ndarray
objective: float | None = None
rhat(chains)[source]

Gelman-Rubin potential scale reduction factor (R-hat) per parameter dimension.

R-hat compares the variance between chains to the variance within chains. Values near 1.0 indicate the chains have mixed and are sampling a common target; values above ~1.01-1.1 flag non-convergence. Standard multi-chain check (Gelman & Rubin 1992).

Parameters:

chains (Any) – array shaped (n_chains, n_draws) or (n_chains, n_draws, d). Needs at least two chains and two draws.

Returns:

A length-d array of R-hat values (one per parameter).

Return type:

ndarray

ess(samples, max_lag=None)[source]

Effective sample size per parameter from positive autocorrelation lags.

Accepts a single chain (n_draws,) / (n_draws, d) or a stack of chains (n_chains, n_draws, d) (the chains are pooled into one autocorrelation estimate after centering each chain by its own mean). Uses the initial-positive-sequence truncation (Geyer): sum autocorrelations until the first non-positive lag.

Returns:

A length-d array of ESS values.

Parameters:

  • samples (Any)

  • max_lag (int | None)

Return type:

ndarray

split_rhat(chains)[source]

Rank-normalized split-R-hat (Vehtari et al. 2021) – the robust, recommended R-hat.

Splits each chain in half (doubling the chain count, so within-chain non-stationarity is caught), rank-normalizes the pooled draws (robust to heavy tails / non-normality), then applies the Gelman-Rubin R-hat. Convergence is typically declared at < 1.01.

Parameters:

chains (Any)

Return type:

ndarray

folded_split_rhat(chains)[source]

Folded rank-normalized split-R-hat (Vehtari et al. 2021): split-R-hat on |x - median(x)|.

The plain split_rhat() compares chain locations; folding about the median makes it compare chain scales/tails, catching the case where chains share a mean but differ in spread. Use together with split_rhat() (or rhat_max()).

Parameters:

chains (Any)

Return type:

ndarray

rhat_max(chains)[source]

The recommended convergence R-hat: max(split_rhat, folded_split_rhat) (Vehtari et al. 2021).

A single number that flags non-convergence in either the location (bulk) or the scale (folded) of the chains; declare convergence at < 1.01.

Parameters:

chains (Any)

Return type:

ndarray

mcse_mean(chains)[source]

Monte Carlo standard error of the posterior mean: sd(x) / sqrt(ESS) per parameter.

The sampling error in the estimated mean from autocorrelated draws – the posterior standard deviation deflated by the (autocorrelation-based) effective sample size of the raw chains.

Parameters:

chains (Any)

Return type:

ndarray

geweke_z(chain, first=0.1, last=0.5)[source]

Geweke convergence z-score comparing the start and end of a single chain (per parameter).

Tests stationarity within one chain (Geweke 1992): if the chain has converged, the mean of its first first fraction and its last last fraction agree, so the z-statistic (mean_a - mean_b) / sqrt(se_a^2 + se_b^2) – with autocorrelation-adjusted standard errors (the segment variance divided by its effective sample size) – is approximately standard normal. |z| < 2 indicates convergence; a large |z| flags a chain still drifting. Returns one z per parameter. Complements the multi-chain split_rhat().

Parameters:
Return type:

ndarray

mcmc_summary(chains)[source]

Per-parameter posterior + convergence summary (the ArviZ-style summary table).

Returns one dict per parameter with the posterior mean/sd and 5/50/95% quantiles, plus the recommended convergence diagnostics: r_hat (= rhat_max(), the max of the bulk and folded rank-normalized split-R-hats), ess_bulk, ess_tail, and mcse_mean (Vehtari et al. 2021). Declare convergence at r_hat < 1.01 with ess_bulk/ess_tail comfortably in the hundreds.

Parameters:

chains (Any)

Return type:

list[dict[str, float]]

ess_bulk(chains)[source]

Bulk effective sample size: ESS of the rank-normalized split chains (mixing in the distribution body).

Parameters:

chains (Any)

Return type:

ndarray

ess_tail(chains, prob=0.05)[source]

Tail effective sample size: the smaller ESS of the lower-/upper-prob tail indicators.

Tail quantiles mix more slowly than the bulk; tail-ESS is min(ESS[1(x <= q_prob)], ESS[1(x >= q_{1-prob})]) over the split chains (Vehtari et al. 2021), surfacing poor tail exploration that bulk-ESS misses.

Parameters:
Return type:

ndarray

available_backends()[source]

Return the names of registered backends whose engine is importable, in registration order.

Return type:

list[str]

class InferenceBackend(name, available, target_kind, nuts)[source]

Bases: object

A registered inference engine: a name, an availability probe, a target contract, a sampler.

Parameters:

  • name (str)

  • available (Callable[[], bool])

  • target_kind (str)

  • nuts (Callable[..., NutsResult])

name: str
available: Callable[[], bool]
target_kind: str
nuts: Callable[..., NutsResult]
register_inference_backend(backend)[source]

Register (or replace) an inference backend under backend.name.

Parameters:

backend (InferenceBackend)

Return type:

None

class EventStudyResult(effect, se, z, p_value, ci, treated_mean, treated_se, control_mean, control_se, tau2_treated, n_treated, n_control, shrunk_treated)[source]

Bases: object

Pooled influence estimate. effect is the DiD ATT (treated minus control) when a control exists.

Parameters:
effect: float
se: float
z: float
p_value: float
ci: tuple[float, float]
treated_mean: float
treated_se: float
control_mean: float | None
control_se: float | None
tau2_treated: float
n_treated: int
n_control: int
shrunk_treated: ndarray
gaussian_effect(pre, post)[source]

Per-subject mean shift and its (Welch) sampling variance from pre/post activity samples.

Parameters:
Return type:

tuple[float, float]

poisson_lograte_effect(k_pre, t_pre, k_post, t_post)[source]

Per-subject log activity-rate shift log(rate_post) - log(rate_pre) for event counts over windows.

k_* are event counts, t_* the window durations (or exposures). Uses a Haldane 0.5 correction so zero-count windows are finite; variance is the delta-method log-rate variance 1/k_post + 1/k_pre.

Parameters:
Return type:

tuple[float, float]

hierarchical_event_study(treated_effects, treated_vars, control_effects=None, control_vars=None, *, alpha=0.05)[source]

Pool per-subject effects into a population influence estimate (DiD if a control group is given).

*_effects / *_vars are the per-subject shifts and their variances from stage 1. With a control group the reported effect is treated_mean - control_mean – the difference-in-differences ATT.

Parameters:
Return type:

EventStudyResult

tipping_drift(result)[source]

Sensitivity bound: the unmeasured differential drift that would explain the effect away.

Within-subject DiD is unbiased only if, absent treatment, treated and control would have drifted equally. This returns the differential drift delta (in effect units) that nullifies the estimate (= effect) and the value that pushes the 95% CI through zero – so a reader can judge whether a confound that large is plausible. Larger = more robust.

Parameters:

result (EventStudyResult)

Return type:

dict

Subpackages

Submodules