Optimization

It is often useful to construct a distribution \(d^\prime\) which is consistent with some marginal aspects of \(d\), but otherwise optimizes some information measure. For example, perhaps we are interested in constructing a distribution which matches pairwise marginals with another, but otherwise has maximum entropy:

In [1]: from dit.algorithms.distribution_optimizers import MaxEntOptimizer

In [2]: xor = dit.example_dists.Xor()

In [3]: meo = MaxEntOptimizer(xor, [[0,1], [0,2], [1,2]])

In [4]: meo.optimize()
Out[4]: 
     message: Optimization terminated successfully
     success: True
      status: 0
         fun: -3.000001615505876
           x: [ 1.250e-01  1.250e-01  1.250e-01  1.250e-01  1.250e-01
                1.250e-01  1.250e-01  1.250e-01]
         nit: 87
         jac: [-3.000e+00 -3.000e+00 -3.000e+00 -3.000e+00 -3.000e+00
               -3.000e+00 -3.000e+00 -3.000e+00]
        nfev: 175
        njev: 87
 multipliers: [-1.771e+06]

In [5]: dp = meo.construct_dist()

In [6]: print(dp)
Class:    Distribution
Alphabet: (('0', '1'), ('0', '1'), ('0', '1'))
Base:     linear

x                 p(X0,X1,X2)
('0', '0', '0')   1/8
('0', '0', '1')   1784702/14277617
('0', '1', '0')   2075854/16606833
('0', '1', '1')   2309353/18474823
('1', '0', '0')   1897263/15178103
('1', '0', '1')   1/8
('1', '1', '0')   5109551/40876407
('1', '1', '1')   1/8

Helper Functions

There are three special functions to handle common optimization problems:

In [7]: from dit.algorithms import maxent_dist, marginal_maxent_dists

The first is maximum entropy distributions with specific fixed marginals. It encapsulates the steps run above:

In [8]: print(maxent_dist(xor, [[0,1], [0,2], [1,2]]))
Class:    Distribution
Alphabet: (('0', '1'), ('0', '1'), ('0', '1'))
Base:     linear

x                 p(X0,X1,X2)
('0', '0', '0')   1/8
('0', '0', '1')   1/8
('0', '1', '0')   1/8
('0', '1', '1')   1/8
('1', '0', '0')   1/8
('1', '0', '1')   1/8
('1', '1', '0')   1/8
('1', '1', '1')   1/8

The second constructs several maximum entropy distributions, each with all subsets of variables of a particular size fixed:

In [9]: k0, k1, k2, k3 = marginal_maxent_dists(xor)

where k0 is the maxent dist corresponding the same alphabets as xor; k1 fixes \(p(x_0)\), \(p(x_1)\), and \(p(x_2)\); k2 fixes \(p(x_0, x_1)\), \(p(x_0, x_2)\), and \(p(x_1, x_2)\) (as in the maxent_dist example above), and finally k3 fixes \(p(x_0, x_1, x_2)\) (e.g. is the distribution we started with).

M-Flat m-Projections

The dual of the MaxEnt / e-flat ladder is Amari’s m-flat hierarchy [Ama01]. Instead of matching \(k\)-way marginals, one truncates the additive (ANOVA) expansion of the pmf and takes the reverse-KL m-projection onto that mixture family. Sparse supports use the symmetric smooth \(P_\varepsilon=(1-\varepsilon)P+\varepsilon U\) (or the \(\varepsilon\downarrow 0\) schedule via m_projection_eps_limit):

In [10]: from dit.algorithms import m_projection, m_projection_eps_limit, mflat_mprojection_dists

In [11]: q2 = m_projection(xor, 2, eps=1e-6)

In [12]: q2_lim = m_projection_eps_limit(xor, order=2)

In [13]: ladder = mflat_mprojection_dists(xor, eps_schedule=(1e-4, 1e-6, 1e-8))

See Information Profiles (MFlatConnectedInformations and DualDependencyDecomposition) for reverse-KL gaps along the order chain and the full dependency lattice.

Maximum Entropy Solver (IPF)

By default, maxent_dist computes the maximum entropy distribution using Iterative Proportional Fitting (IPF), the classic algorithm from reconstructability analysis and log-linear modeling. Starting from the uniform distribution, IPF cyclically rescales the working distribution so each constrained marginal matches the data, iterating until convergence. It is typically far faster than the general scipy convex optimizer. IPF converges only linearly on cyclic structures with induced structural zeros, however, so when it fails to converge within its iteration budget maxent_dist automatically falls back to the scipy optimizer to preserve accuracy. The scipy backend can also be requested explicitly via method='scipy':

In [14]: print(maxent_dist(xor, [[0,1], [0,2], [1,2]], method='scipy'))
Class:    Distribution
Alphabet: (('0', '1'), ('0', '1'), ('0', '1'))
Base:     linear

x                 p(X0,X1,X2)
('0', '0', '0')   1/8
('0', '0', '1')   1784702/14277617
('0', '1', '0')   2075854/16606833
('0', '1', '1')   2309353/18474823
('1', '0', '0')   1897263/15178103
('1', '0', '1')   1/8
('1', '1', '0')   5109551/40876407
('1', '1', '1')   1/8

Reconstructability Analysis

The maximum entropy reconstruction underlies reconstructability analysis, which decomposes a distribution into a structure of marginals and assesses each structure by two quantities: its error (transmission, the information lost relative to the data) and its complexity (degrees of freedom, the number of free parameters). The dependency decomposition (see Information Profiles) evaluates these over the whole lattice of structures, yielding the “decomposition spectrum”:

In [15]: from dit.algorithms import degrees_of_freedom

In [16]: from dit.multivariate import transmission

In [17]: from dit.profiles import DependencyDecomposition

In [18]: from dit.multivariate import entropy

In [19]: print(DependencyDecomposition(xor, measures={'H': entropy, 'T': transmission, 'df': degrees_of_freedom}))
+-----------------------------------+
|      Dependency Decomposition     |
+------------+--------+--------+----+
| dependency |   H    |   T    | df |
+------------+--------+--------+----+
|    012     |  2.000 |  0.000 | 7  |
|  01:02:12  |  3.000 |  1.000 | 6  |
|   01:02    |  3.000 |  1.000 | 5  |
|   01:12    |  3.000 |  1.000 | 5  |
|   02:12    |  3.000 |  1.000 | 5  |
|    01:2    |  3.000 |  1.000 | 4  |
|    02:1    |  3.000 |  1.000 | 4  |
|    12:0    |  3.000 |  1.000 | 4  |
|   0:1:2    |  3.000 |  1.000 | 3  |
+------------+--------+--------+----+

Mixtures of product distributions

Shared-randomness profiles fit mixtures of fully factorized distributions (naive-Bayes / latent-class models) by EM. Maximum likelihood equals the forward-KL projection onto

\[\mathcal{F}_k = \Bigl\{ Q : Q(x) = \sum_{\alpha=1}^{k} \pi_\alpha \prod_i Q_i(x_i \mid \alpha) \Bigr\}.\]

Helpers fit_mixture_of_products() and mixture_of_products_dists() feed BindingMixtureProfile (see Information Profiles).

For additive building blocks (Amari ANOVA / fixed marginal lifts), see m_projection() (criteria jsd, forward_kl, reverse_kl) and fit_marginal_lift_mixture().

In [20]: from dit.algorithms import fit_mixture_of_products

In [21]: gb = dit.Distribution(['000', '111'], [0.5, 0.5])

In [22]: q2 = fit_mixture_of_products(gb, k=2, n_init=8, seed=0)['dist']

In [23]: print(round(dit.divergences.kullback_leibler_divergence(gb, q2), 8))
0.0

Marginal lifts

lift_marginal() extends a marginal pmf on a subset of coordinates to a full joint (uniformly, by default, over the remaining coordinates). fit_marginal_lift_mixture() fits the best convex combination of such lifts of all marginals up to a given order, and marginal_lift_dists() returns the whole ladder of fits for orders \(0, \ldots, k_{\max}\), which drives MarginalLiftProfile.

Convex Maximization over Polytopes

Some measures, such as the \(\Iprec{\bullet}\) partial information decomposition (PID_Prec), need the maximum of a convex function over a polytope. maximize_convex_function() exploits the fact that the maximum is attained at a vertex: it enumerates the vertices of \(\{x : A_{\text{ineq}} x \leq b_{\text{ineq}},\ A_{\text{eq}} x = b_{\text{eq}}\}\) with pypoman (from the optional extra) and evaluates the objective at each one:

import numpy as np
from dit.algorithms import maximize_convex_function

# maximize the squared norm over the probability simplex in R^3
A_ineq, b_ineq = -np.eye(3), np.zeros(3)
A_eq, b_eq = np.ones((1, 3)), np.ones(1)
x, value = maximize_convex_function(lambda x: x @ x, A_ineq, b_ineq, A_eq, b_eq)
# x is (up to rounding) a vertex of the simplex, and value is 1

Other Algorithms in dit

The rest of dit.algorithms is documented alongside the measures that use it; every function is listed in the API reference.

  • Channel capacity, capacity–cost, and Sibson capacity: Channels.

  • Lattice joins and meets, minimal sufficient statistics, and information trimming: Operations.

  • Statistics of numeric distributions (mean, variance, quantiles, …): Numeric Outcomes.

  • Support spectra (support_biadjacency, support_singular_values, spectral_entanglement_bound): The Gray-Wyner Network.

  • Sample-space pruning and expansion (pruned_samplespace, expanded_samplespace): the API reference.

API

maxent_dist(dist, rvs, x0=None, maxiter=1000, sparse=True, method='ipf')[source]

Return the maximum entropy distribution consistent with the marginals from dist specified in rvs.

Parameters:

  • dist (Distribution) – The distributions whose marginals should be matched.

  • rvs (list of lists) – The marginals from dist to constrain.

  • x0 (np.ndarray) – Initial condition for the scipy optimizer. Ignored when method is ‘ipf’, which always starts from the uniform distribution.

  • maxiter (int) – The number of optimization iterations to perform. Used only by the scipy optimizer; IPF uses its own convergence tolerance.

  • sparse (bool) – Whether the returned distribution should be sparse or dense.

  • method (str) – The algorithm used to compute the maximum entropy distribution. ‘ipf’ (the default) uses Iterative Proportional Fitting, which is typically far faster than the convex optimizer. IPF converges only linearly on cyclic structures with induced structural zeros, so when it fails to converge within its iteration budget this falls back to the scipy optimizer to preserve accuracy. ‘scipy’ (also accepts ‘slsqp’) forces the scipy convex optimizer.

Returns:

me – The maximum entropy distribution.

Return type:

Distribution

marginal_maxent_dists(dist, k_max=None)[source]

Return the marginal-constrained maximum entropy distributions.

Parameters:

  • dist (distribution) – The distribution used to constrain the maxent distributions.

  • k_max (int) – The maximum order to calculate.

Returns:

dists – A list of distributions, where the i`th element is the maxent distribution with the `i-size marginals fixed.

Return type:

list

ipf_dist(dist, rvs, tol=1e-10, maxiter=10000, sparse=True, return_status=False)[source]

Return the maximum entropy distribution consistent with the marginals from dist specified in rvs, computed via Iterative Proportional Fitting.

Parameters:

  • dist (Distribution) – The distribution whose marginals should be matched.

  • rvs (list of lists) – The marginals from dist to constrain. Each inner list specifies a set of random variables (a marginal “projection”) to hold fixed, e.g. [[0, 1], [1, 2]] for the structure AB:BC.

  • tol (float) – The convergence tolerance. Iteration stops once the largest deviation of any constrained marginal from its target falls below tol.

  • maxiter (int) – The maximum number of IPF cycles to perform.

  • sparse (bool) – Whether the returned distribution should be sparse or dense.

  • return_status (bool) – If True, return a (me, converged) tuple, where converged is True iff the final marginal deviation fell below tol within maxiter cycles. IPF converges only linearly on cyclic structures with induced structural zeros, so callers that need an exact answer should check this flag and fall back to a convex optimizer when it is False.

Returns:

  • me (Distribution) – The maximum entropy distribution.

  • converged (bool) – Only returned when return_status is True (see above).

m_projection(dist, order, nrestarts=16, maxiter=3000, warm_start=None, criterion='reverse_kl', eps=None)[source]

Project dist onto the order-order m-flat mixture family.

Parameters:

  • dist (Distribution) – The target distribution \(P\).

  • order (int) – Maximum ANOVA order kept in \(\mathcal{M}_{\mathrm{order}}\).

  • nrestarts (int) – Optimizer controls.

  • maxiter (int) – Optimizer controls.

  • warm_start (Distribution or None) – Optional warm start.

  • criterion ({'reverse_kl', 'forward_kl', 'jsd'}) – Projection objective.

  • eps (float or None) – For reverse_kl, weight in \(P_\varepsilon=(1-\varepsilon)P+\varepsilon U\). Defaults to 1e-8 when None.

Returns:

q

Return type:

Distribution

m_projection_from_subsets(dist, subsets, nrestarts=16, maxiter=3000, warm_start=None, criterion='reverse_kl', eps=None)[source]

Project dist onto the m-flat family spanned by subsets.

mflat_mprojection_dists(dist, k_max=None, nrestarts=16, maxiter=3000, criterion='reverse_kl', eps=None, eps_schedule=None)[source]

Return the m-flat projection ladder \(Q^{(0)},\ldots,Q^{(k_{\max})}\).

fit_mixture_of_products(dist, k, *, n_init=12, max_iter=200, tol=1e-10, seed=0)[source]

Fit a \(k\)-mixture of product distributions to dist by EM.

Parameters:

  • dist (Distribution) – Target joint.

  • k (int) – Number of mixture components.

  • n_init (int) – Random restarts and EM iteration cap.

  • max_iter (int) – Random restarts and EM iteration cap.

  • tol (float) – Log-likelihood convergence tolerance.

  • seed (int) – RNG seed for restarts.

Returns:

result – Keys dist (fitted Distribution), pi, conds, I_xv (\(I(X;V)\) under data-weighted responsibilities), H_v, loglik.

Return type:

dict

mixture_of_products_dists(dist, k_max=None, *, n_init=12, max_iter=200, tol=1e-10, seed=0, early_stop=True, kl_tol=1e-08)[source]

Fit the mixture-of-products ladder \(Q^{(1)},\ldots,Q^{(k_{\max})}\).

Parameters:
Returns:

  • dists (list of Distribution) – dists[k-1] is the MLE in \(\mathcal{F}_k\).

  • meta (list of dict) – Per-k diagnostics (I_xv, H_v, loglik, pi, conds).

lift_marginal(outs, pmf, alph, S, mode='uniform')[source]

Lift a marginal on coordinates S to a full joint pmf.

Parameters:

  • outs – Dense Cartesian table for the joint.

  • pmf – Dense Cartesian table for the joint.

  • alph – Dense Cartesian table for the joint.

  • S (tuple of int) – Variable indices of the marginal.

  • mode ({'uniform', 'product'}) – How to extend off S.

fit_marginal_lift_mixture(dist, order, mode='uniform', n_init=12, seed=0)[source]

Fit a convex combination of lifts of marginals of order at most order.

Returns:

result – Keys dist, labels, alpha, L2.

Return type:

dict

marginal_lift_dists(dist, k_max=None, mode='uniform', n_init=12, seed=0)[source]

Ladder of marginal-lift mixtures for orders \(0,\ldots,k_{\max}\).

Order 0 is the uniform distribution. Order \(n\) includes the full joint as a block and recovers \(P\) exactly.

maximize_convex_function(f, A_ineq, b_ineq, A_eq=None, b_eq=None)[source]

Maximize a convex function over a polytope. This function uses the fact that the maximum of a convex function over a polytope will be achieved at one of the extreme points of the polytope.

The maximization is done by taking a system of linear inequalities, using the pypoman library to create a list of extreme points, and then evaluating the objective function on each point.

Parameters:

  • f (function) – Objective function to maximize

  • A_ineq (matrix) – Specifies inequalities matrix, should be num_inequalities x num_variables

  • b_ineq (array) – Specifies inequalities vector, should be num_inequalities long

  • A_eq (matrix) – Specifies equalities matrix, should be num_equalities x num_variables

  • b_eq (array) – Specifies equalities vector, should be num_equalities long

Returns:

  • optimal_extreme_point (np.ndarray) – The extreme point at which the maximum is attained.

  • maximum_function_value (float) – The maximum value of f.

Optimization Backends

By default, dit uses NumPy and SciPy for numerical optimization. Three additional backends are available that leverage automatic differentiation for computing exact gradients, which can improve convergence for large or complex problems:

JAX Backend

Install with pip install "dit[jax]". The JAX backend (dit.algorithms.optimization_jax) provides:

  • Automatic differentiation via jax.grad for exact gradient computation

  • JIT compilation for improved performance

  • GPU/TPU acceleration when available

PyTorch Backend

Install with pip install "dit[torch]". The PyTorch backend (dit.algorithms.optimization_torch) provides:

  • Automatic differentiation via torch.autograd for exact gradient computation

  • GPU acceleration via CUDA or MPS when available

  • torch.compile support for PyTorch 2.0+

PyTensor Backend

Install with pip install "dit[pytensor]". The PyTensor backend (dit.algorithms.optimization_pytensor) uses PyTensor (the maintained successor to the archived Aesara) and provides:

  • Symbolic graph compilation of the objective, its exact gradient via pytensor.grad, and each constraint Jacobian, compiled once and reused across the optimization

  • A native augmented-Lagrangian solver (compiled value/gradient with L-BFGS-B inner solves) for moderate problem sizes, with a SciPy SLSQP fallback

  • Optional Numba compilation of the compiled functions (set DIT_PYTENSOR_MODE=NUMBA)

Set DIT_PYTENSOR_COMPILEDIR to persist PyTensor’s compilation cache across process launches.

For measures that use the Markov variable optimizer (e.g. common informations), the backend can be selected via the backend parameter:

from dit.multivariate import wyner_common_information

# Default NumPy backend
wyner_common_information(d)

# JAX backend (requires jax)
wyner_common_information(d, backend='jax')

# PyTorch backend (requires torch)
wyner_common_information(d, backend='torch')

# PyTensor backend (requires pytensor)
wyner_common_information(d, backend='pytensor')

Logging

The optimization modules emit structured log messages via loguru. Logging is disabled by default. To enable it:

from loguru import logger
logger.enable("dit")

This will show optimization progress including problem dimensions, convergence status, and objective values.