mbi.extensions

Extensions for mbi providing alternative estimation approaches.

class mbi.extensions.MixtureOfProducts(_logits: dict[str | int, Array], domain: Domain, total: float)[source]

Bases: object

A discrete distribution as a mixture of product distributions.

domain: Domain

total: float

classmethod random(domain: Domain, total: float, num_components: int = 100, seed: int = 0, scale: float = 0.25) → MixtureOfProducts[source]: Initialize with random logits.

property num_components: int: Number of mixture components K.

property products: dict[str | int, Array]: Per-attribute categorical probabilities (softmax of logits).

project(attrs: str | int | Sequence[str | int]) → Factor[source]: Compute the marginal over attrs via einsum.

supports(attrs: str | int | Sequence[str | int]) → bool[source]: Any subset of domain attributes is supported.

synthetic_data(rows=None) → Dataset[source]: Generate synthetic data via randomized rounding.

class mbi.extensions.MixtureOfProductsEstimator(*, marginal_oracle: marginal_oracles.MarginalOracle | None = None, optimizer: optax.GradientTransformation, num_components: int = 100, seed: int = 0)[source]

Bases: Estimator

Estimates a distribution as a mixture of product distributions.

Inspired by RAP (https://arxiv.org/abs/2103.06641) and RAP^{softmax} (https://arxiv.org/abs/2106.07153). Instead of parameterizing via a graphical model with potentials, the distribution is represented as a mixture of K product distributions optimized via gradient descent.

num_components

Number of mixture components K.

Type:: int

optimizer

An optax optimizer (e.g. optax.adam(0.1)).

Type:: optax._src.base.GradientTransformation

seed

Random seed for parameter initialization.

Type:: int

optimizer: GradientTransformation

num_components: int = 100

seed: int = 0

mbi.extensions.precompute_marginals(dataset: Dataset | JaxDataset, cliques: Sequence[tuple[str | int, ...]], *, transpose: bool = True) → CliqueVector[source]

Compute all marginals from a dataset using JIT-compiled bincount.

Compiles one jnp.bincount program per unique padded shape (after power-of-2 bucketing and canonical reordering) and executes it for every clique that maps to that shape. All compilations are submitted to a background thread upfront so they overlap with execution.

Accepts both Dataset (numpy-backed) and JaxDataset (JAX-backed). Numpy arrays are converted to JAX arrays internally.

Parameters:

dataset – The dataset to compute marginals from.
cliques – Cliques (tuples of attribute names) to compute.
transpose – If True (default), transpose each marginal so its axes match the original clique attribute order. Set to False to skip the transpose and keep the canonical sorted order.

Returns:

A CliqueVector mapping each clique to its marginal Factor.

class mbi.extensions.ReweightedDatasetEstimator(seed_data: Dataset, optimizer: optax.GradientTransformation, *, marginal_oracle: marginal_oracles.MarginalOracle | None = None)[source]

Bases: Estimator

Estimates a distribution by reweighting seed records.

Inspired by PMW^Pub (https://arxiv.org/abs/2102.08598). Given a seed set of records, the distribution is represented by learning a weight for each record to minimize the marginal loss.

seed_data

A Dataset of seed records to reweight.

Type:: mbi.dataset.Dataset

optimizer

An optax optimizer (e.g. optax.adam(0.1)).

Type:: optax._src.base.GradientTransformation

seed_data: Dataset

optimizer: GradientTransformation

mbi.extensions.precompile(domain: Domain, cliques: Sequence[tuple[str | int, ...]], rows: int, constraints: Sequence[Constraint] = ()) → Future[source]

Warm the JIT cache for synthetic_data asynchronously.

Only requires domain and clique structure — both available after measurement selection. Fire and forget; synthetic_data benefits from whatever has compiled so far.

No concrete arrays are allocated (except a PRNG key); all inputs use abstract ShapeDtypeStruct values via Factor.abstract so the compiler sees shapes and dtypes without materializing memory.

Parameters:

domain – The Domain over which the model is defined.
cliques – The cliques of the model (known after measurement selection).
rows – Number of records that will be generated.
constraints – Structural constraints the model is fit under. Their cliques are added to the plan so the warmed JIT cache matches the constraint-aware synthetic_data call.

mbi.extensions.synthetic_data(model: MarkovRandomField, rows: int, seed: int = 0) → Dataset[source]

Generate synthetic data from a fitted model using Gumbel rounding.

Output marginals match the model within ±2 counts per cell. If precompile was called earlier with matching domain, cliques, and rows, the JIT cache will be warm and this call is fast.

Parameters:

model – A fitted MarkovRandomField.
rows – Number of records to generate.
seed – Random seed for the JAX PRNG.

Returns:

A Dataset whose marginals closely match the model.