mbi.marginal_loss

Defines loss functions based on linear measurements of marginals.

This module provides structures and functions for defining and calculating loss based on potentially noisy linear measurements of marginal distributions. Key components include the LinearMeasurement class to represent individual measurements and the MarginalLossFn class to define loss functions over CliqueVector objects, enabling the evaluation of model fit against observed or noisy data. Utilities for clique manipulation and feasibility checks are also included.

Functions

calculate_l2_lipschitz	Estimate the Lipschitz constant of L(x) = || f(x) - y ||_2^2 where f is a linear function.
calculate_l2_lipschitz_from_metadata	Lipschitz constant of the (unnormalized) l2 loss gradient from metadata.
from_linear_measurements	Construct a MarginalLossFn from a list of LinearMeasurements.
mle_loss_fn	MLE loss: -marginals.dot(mu.log()).
primal_feasibility	Calculates the average L1 distance between overlapping marginals in mu (consistency).

Classes

DatavectorQuery	Identity query: f.datavector().
LinearMeasurement	A class for representing a private linear measurement of a marginal.
MarginalLossFn	A loss function over the concatenated vector of marginals.
SlicedQuery	Datavector with leading elements removed: f.datavector()[start:].
WeightedQuery	Datavector scaled by fixed weights: f.datavector() * weights.

class mbi.marginal_loss.DatavectorQuery(use_for_total_estimation: bool = True)

Bases: object

Identity query: f.datavector().

use_for_total_estimation

If False, exclude measurements using this query from minimum_variance_unbiased_total. Defaults to True.

Type:

bool

use_for_total_estimation: bool = True

op_norm_sq() → float

class mbi.marginal_loss.WeightedQuery(weights: ndarray)

Bases: object

Datavector scaled by fixed weights: f.datavector() * weights.

weights: ndarray

op_norm_sq() → float

class mbi.marginal_loss.SlicedQuery(start: int = 1)

Bases: object

Datavector with leading elements removed: f.datavector()[start:].

start: int = 1

op_norm_sq() → float

class mbi.marginal_loss.LinearMeasurement(noisy_measurement: Array | ndarray | bool | number | bool | int | float | complex, clique: tuple[str | int, ...], stddev: float = 1.0, query: Callable[[Factor], Array] = DatavectorQuery(use_for_total_estimation=True))

Bases: object

A class for representing a private linear measurement of a marginal.

noisy_measurement

The noisy measurement of the marginal.

Type:

jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex

clique

The clique (a tuple of attribute names) defining the marginal.

Type:

tuple[str | int, …]

stddev

The standard deviation of the noise added to the measurement.

Type:

float

query

A linear function that, when applied to a Factor, extracts a

Type:

collections.abc.Callable[[mbi.factor.Factor], jax.Array]

a vector with the same shape and interpretation as `noisy_measurement`.

noisy_measurement: Array | ndarray | bool | number | bool | int | float | complex

clique: tuple[str | int, ...]

stddev: float = 1.0

query: Callable[[Factor], Array] = DatavectorQuery(use_for_total_estimation=True)

compress(mapping: Mapping[Any, ndarray], domain: Domain) → LinearMeasurement

Compress this measurement by merging domain values.

Parameters:

• mapping – Maps attribute names to 1D integer arrays. mapping[attr][i] gives the compressed value for original value i.

• domain – The full dataset domain (needed to reshape the measurement vector into the clique’s shape).

Returns:

A new LinearMeasurement over the compressed domain.

class mbi.marginal_loss.MarginalLossFn(cliques: Sequence[tuple[str | int, ...]], loss_fn: Callable[[CliqueVector, Any], Array | ndarray | bool | number | bool | int | float | complex], data: Any = (), lipschitz: float = 1.0)

Bases: object

A loss function over the concatenated vector of marginals.

Separates the computation (loss_fn, static) from the captured data (data, dynamic JAX pytree leaves). This allows MarginalLossFn to be passed through jax.jit as a regular traced argument rather than a static one, so that different measurement values with the same structure reuse the same compiled program.

cliques

Cliques defining the scope of the marginals.

Type:

collections.abc.Sequence[tuple[str | int, …]]

loss_fn

A pure function (marginals, data) -> loss. This is treated as static metadata for JIT compilation.

Type:

collections.abc.Callable[[mbi.clique_vector.CliqueVector, Any], jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex]

data

Arbitrary pytree of arrays captured by loss_fn. This is traced through JIT and can change without recompilation.

Type:

Any

lipschitz

Lipschitz constant of the gradient. Defaults to 1.0.

Type:

float

cliques: Sequence[tuple[str | int, ...]]

loss_fn: Callable[[CliqueVector, Any], Array | ndarray | bool | number | bool | int | float | complex]

data: Any = ()

lipschitz: float = 1.0

mbi.marginal_loss.calculate_l2_lipschitz(domain: Domain, cliques: Sequence[tuple[str | int, ...]], loss_fn: Callable[[CliqueVector], Array | ndarray | bool | number | bool | int | float | complex]) → float

Estimate the Lipschitz constant of L(x) = || f(x) - y ||_2^2 where f is a linear function.

The Lipschitz constant can usually be obtained via the largest eigenvalue of the Hessian, which for linear functions represented in matrix form is A^T A. This function computes the same value without materializing this n x n matrix by using power iteration and leveraging jax.jvp.

Parameters:

• domain – The domain over which the loss_fn is defined.

• loss_fn – The loss function, assumed to be of the form || f(x) - y ||_2^2 where f is linear.

Returns:

An estimate of the Lipschitz constant of the grad(L).

mbi.marginal_loss.calculate_l2_lipschitz_from_metadata(domain: Domain, measurements: Sequence[LinearMeasurement]) → float | None

Lipschitz constant of the (unnormalized) l2 loss gradient from metadata.

The loss is quadratic, so its gradient’s Lipschitz constant is lambda_max of a Hessian that is block-diagonal by maximal clique:

H = sum_M (1 / sigma_M^2) (Q_M P_{c_M})^T (Q_M P_{c_M})

where P_{c_M} marginalizes a maximal-clique table down to clique c_M, with ||P_c||^2 = |C| / |c| (the product of the summed-out domain sizes). Bounding each block by the sum of its terms’ norms gives:

L <= max_C sum_{M->C} ||Q_M||^2 / sigma_M^2 * |C| / |c_M|.

The uniform vector is the top eigenvector of every marginalization operator, so this is exact when all queries are identity marginals (the common case) and a safe overestimate otherwise (a larger L only shrinks the step size).

Returns None if any query’s operator norm is unknown (e.g. a nonlinear or user-supplied callable query), signalling the caller to fall back to the power-iteration estimate.

mbi.marginal_loss.from_linear_measurements(measurements: Sequence[LinearMeasurement], domain: Domain, norm: str = 'l2', normalize: bool = False) → MarginalLossFn

Construct a MarginalLossFn from a list of LinearMeasurements.

Parameters:

• measurements – A list of LinearMeasurements.

• norm – Either “l1” or “l2”.

• normalize – Flag determining if the loss function should be normalized by the length of linear measurements and estimated total.

• domain – The domain over which the measurements were made, necessary for calcualting the Lipschitz parameter.

Returns:

The MarginalLossFn L(mu) = sum_{c} || Q_c mu_c - y_c || (possibly squared or normalized).

mbi.marginal_loss.primal_feasibility(mu: CliqueVector) → Array | ndarray | bool | number | bool | int | float | complex

Calculates the average L1 distance between overlapping marginals in mu (consistency).

mbi.marginal_loss.mle_loss_fn(marginals: CliqueVector) → MarginalLossFn

MLE loss: -marginals.dot(mu.log()).