dit 2.1__tar.gz → 2.3__tar.gz

{dit-2.1 → dit-2.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: dit
-Version: 2.1
+Version: 2.3
 Summary: Python package for information theory on discrete random variables.
 Project-URL: Homepage, http://dit.io
 Project-URL: Repository, https://github.com/dit/dit
@@ -10,8 +10,6 @@ License-Expression: BSD-3-Clause
 License-File: LICENSE.txt
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: BSD License
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
@@ -22,7 +20,7 @@ Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.11
 Requires-Dist: boltons
 Requires-Dist: debtcollector
-Requires-Dist: lattices>=0.5.0
+Requires-Dist: lattices>=0.5
 Requires-Dist: loguru
 Requires-Dist: networkx>=2.6
 Requires-Dist: numpy>=1.22
@@ -59,6 +57,7 @@ Requires-Dist: slsqp-jax; extra == 'dev'
 Requires-Dist: sphinx; extra == 'dev'
 Requires-Dist: sphinx-rtd-theme; extra == 'dev'
 Requires-Dist: sphinxcontrib-bibtex; extra == 'dev'
+Requires-Dist: sympy; extra == 'dev'
 Requires-Dist: torch; extra == 'dev'
 Requires-Dist: ty; extra == 'dev'
 Provides-Extra: docs
@@ -70,17 +69,19 @@ Requires-Dist: nbconvert; extra == 'docs'
 Requires-Dist: numdifftools; extra == 'docs'
 Requires-Dist: pint; extra == 'docs'
 Requires-Dist: pycddlib-standalone; extra == 'docs'
-Requires-Dist: pypoman; extra == 'docs'
 Requires-Dist: python-ternary; extra == 'docs'
 Requires-Dist: scikit-learn; extra == 'docs'
 Requires-Dist: sphinx; extra == 'docs'
 Requires-Dist: sphinx-rtd-theme; extra == 'docs'
 Requires-Dist: sphinxcontrib-bibtex; extra == 'docs'
+Requires-Dist: sympy; extra == 'docs'
 Provides-Extra: jax
 Requires-Dist: jax; extra == 'jax'
 Requires-Dist: jaxlib; extra == 'jax'
 Requires-Dist: optimistix; extra == 'jax'
 Requires-Dist: slsqp-jax; extra == 'jax'
+Provides-Extra: music
+Requires-Dist: music21; extra == 'music'
 Provides-Extra: optional
 Requires-Dist: colorama; extra == 'optional'
 Requires-Dist: matplotlib; extra == 'optional'
@@ -91,6 +92,8 @@ Requires-Dist: pypoman; extra == 'optional'
 Requires-Dist: scikit-learn; extra == 'optional'
 Provides-Extra: pytensor
 Requires-Dist: pytensor; extra == 'pytensor'
+Provides-Extra: symbolic
+Requires-Dist: sympy; extra == 'symbolic'
 Provides-Extra: test
 Requires-Dist: hypothesis>=6.0; extra == 'test'
 Requires-Dist: hypothesis[numpy]; extra == 'test'
@@ -464,7 +467,7 @@ If you'd like to get in contact about anything, you can reach us through our `sl
    :target: https://github.com/dit/dit/actions/workflows/build.yml
    :alt: Continuous Integration Status
 
-.. |codecov| image:: https://codecov.io/gh/dit/dit/branch/master/graph/badge.svg
+.. |codecov| image:: https://codecov.io/gh/dit/dit/branch/main/graph/badge.svg
    :target: https://codecov.io/gh/dit/dit
    :alt: Test Coverage Status
 

{dit-2.1 → dit-2.3}/README.rst

@@ -359,7 +359,7 @@ If you'd like to get in contact about anything, you can reach us through our `sl
    :target: https://github.com/dit/dit/actions/workflows/build.yml
    :alt: Continuous Integration Status
 
-.. |codecov| image:: https://codecov.io/gh/dit/dit/branch/master/graph/badge.svg
+.. |codecov| image:: https://codecov.io/gh/dit/dit/branch/main/graph/badge.svg
    :target: https://codecov.io/gh/dit/dit
    :alt: Test Coverage Status
 

{dit-2.1 → dit-2.3}/dit/__init__.py

@@ -10,7 +10,13 @@ Python package for sigma-algebras defined on finite sets.
 
 """
 
-__version__ = "1.5"
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+
+try:
+    __version__ = _pkg_version("dit")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.0.0"
 
 from loguru import logger as _logger
 
@@ -37,3 +43,5 @@ import dit.other  # noqa: E402
 import dit.pid  # noqa: E402
 import dit.profiles  # noqa: E402
 import dit.shannon  # noqa: E402
+import dit.symbolic  # noqa: E402
+import dit.visualization  # noqa: E402

{dit-2.1 → dit-2.3}/dit/algorithms/__init__.py

@@ -12,6 +12,17 @@ from .maxentropyfw import *
 from .minimal_sufficient_statistic import *
 from .optimization import *
 from .distribution_optimizers import *
+from .mixture_of_products import fit_mixture_of_products, mixture_of_products_dists
+from .marginal_lifts import fit_marginal_lift_mixture, lift_marginal, marginal_lift_dists
+from .mprojection import (
+    m_projection,
+    m_projection_eps_limit,
+    m_projection_from_subsets,
+    mflat_design_matrix,
+    mflat_mprojection_dists,
+    mflat_subsets_from_dependency,
+    symmetric_smooth,
+)
 from .prune_expand import expanded_samplespace, pruned_samplespace
 from .stats import (
     cdf,
@@ -34,6 +45,11 @@ from .stats import (
     standard_moment,
     variance,
 )
+from .support_spectrum import (
+    spectral_entanglement_bound,
+    support_biadjacency,
+    support_singular_values,
+)
 
 # Don't expose anything yet.
 # from . import pid_broja

dit-2.3/dit/algorithms/caekl_psp.py

@@ -0,0 +1,215 @@
+"""
+CAEKL via PSP on marginalized PMF arrays (optimization backends).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from functools import cache
+
+import numpy as np
+
+EntropyFn = Callable[..., float]
+
+
+def _group_rvs_from_sets(rvs: Iterable[int]) -> list[frozenset[int]]:
+    return [frozenset({rv}) for rv in sorted(rvs)]
+
+
+def _cluster_to_rv_part(
+    partition: tuple[frozenset[int], ...],
+    group_rvs: list[frozenset[int]],
+) -> tuple[frozenset[int], ...]:
+    blocks: list[frozenset[int]] = []
+    for cluster in partition:
+        block: set[int] = set()
+        for group_idx in cluster:
+            block |= set(group_rvs[group_idx])
+        blocks.append(frozenset(block))
+    return tuple(blocks)
+
+
+def caekl_mutual_information_psp_pmf(
+    pmf,
+    *,
+    all_vars: set[int],
+    rvs: set[int],
+    crvs: set[int],
+    h: EntropyFn,
+    sum_axes,
+) -> float:
+    """
+    CAEKL of ``rvs`` given ``crvs`` for a joint ``pmf`` using PSP.
+
+    Parameters
+    ----------
+    pmf : array
+        Full joint PMF including auxiliary variables.
+    all_vars : set[int]
+        All variable indices in ``pmf``.
+    rvs : set[int]
+        Target random-variable indices (one group per index).
+    crvs : set[int]
+        Conditioning variable indices.
+    h : callable
+        Shannon entropy ``h(marginal_pmf)``.
+    sum_axes : callable
+        ``sum_axes(pmf, axes)`` marginalizes ``pmf`` over ``axes``.
+    """
+    value, _ = caekl_mutual_information_psp_pmf_with_partition(
+        pmf,
+        all_vars=all_vars,
+        rvs=rvs,
+        crvs=crvs,
+        h=h,
+        sum_axes=sum_axes,
+    )
+    return value
+
+
+def caekl_mutual_information_psp_pmf_with_partition(
+    pmf,
+    *,
+    all_vars: set[int],
+    rvs: set[int],
+    crvs: set[int],
+    h: EntropyFn,
+    sum_axes,
+) -> tuple[float, tuple[frozenset[int], ...]]:
+    """
+    CAEKL via PSP on a PMF, returning an optimal partition over ``rvs``.
+
+    The partition is expressed as a tuple of frozensets of elements of ``rvs``.
+    """
+    group_rvs = _group_rvs_from_sets(rvs)
+    n_groups = len(group_rvs)
+    idx_crvs = tuple(all_vars - crvs)
+    pmf_crvs = sum_axes(pmf, idx_crvs)
+
+    def conditional_h(vars_set: set[int]) -> float:
+        idx = tuple(all_vars - (vars_set | crvs))
+        pmf_joint = sum_axes(pmf, idx)
+        return h(pmf_joint) - h(pmf_crvs)
+
+    @cache
+    def h_groups(group_indices: frozenset[int]) -> float:
+        if not group_indices:
+            return 0.0
+        vars_set: set[int] = set()
+        for group_idx in group_indices:
+            vars_set |= set(group_rvs[group_idx])
+        return conditional_h(vars_set)
+
+    from ..multivariate.mmi_psp import caekl_mutual_information_psp_with_partition
+
+    value, partition = caekl_mutual_information_psp_with_partition(h_groups, n_groups)
+    return value, _cluster_to_rv_part(partition, group_rvs)
+
+
+def caekl_mutual_information_psp_pmf_grad_data(
+    pmf,
+    *,
+    all_vars: set[int],
+    rvs: set[int],
+    crvs: set[int],
+    h: EntropyFn,
+    sum_axes,
+) -> tuple[float, tuple[frozenset[int], ...], tuple[int, ...], tuple[int, ...], dict[frozenset[int], tuple[int, ...]]]:
+    """
+    CAEKL value, optimal partition, and marginalization indices for gradients.
+
+    Returns
+    -------
+    value : float
+    partition : tuple of frozenset[int]
+        Optimal partition blocks (subsets of ``rvs``).
+    idx_joint : tuple[int, ...]
+    idx_crvs : tuple[int, ...]
+    idx_parts : dict
+        Maps each block in ``partition`` to axes to sum for its marginal PMF.
+    """
+    value, partition = caekl_mutual_information_psp_pmf_with_partition(
+        pmf,
+        all_vars=all_vars,
+        rvs=rvs,
+        crvs=crvs,
+        h=h,
+        sum_axes=sum_axes,
+    )
+    idx_joint = tuple(all_vars - (rvs | crvs))
+    idx_crvs = tuple(all_vars - crvs)
+    idx_parts = {block: tuple(all_vars - (set(block) | crvs)) for block in partition}
+    return value, partition, idx_joint, idx_crvs, idx_parts
+
+
+def labels_from_partition(
+    partition: tuple[frozenset[int], ...],
+    rvs_sorted: tuple[int, ...],
+) -> np.ndarray:
+    """Encode a CAEKL partition as per-RV block labels (for JAX host callbacks)."""
+    labels = np.zeros(len(rvs_sorted), dtype=np.int32)
+    rv_index = {rv: i for i, rv in enumerate(rvs_sorted)}
+    for block_id, block in enumerate(partition):
+        for rv in block:
+            labels[rv_index[rv]] = block_id
+    return labels
+
+
+def partition_from_labels(
+    labels: np.ndarray,
+    rvs_sorted: tuple[int, ...],
+) -> tuple[frozenset[int], ...]:
+    """Decode block labels into a CAEKL partition over ``rvs_sorted``."""
+    labels = np.asarray(labels, dtype=np.int32)
+    blocks: dict[int, set[int]] = {}
+    for i, lab in enumerate(labels):
+        blocks.setdefault(int(lab), set()).add(rvs_sorted[i])
+    return tuple(frozenset(blocks[k]) for k in sorted(blocks))
+
+
+def caekl_partition_indices(
+    pmf,
+    *,
+    all_vars: set[int],
+    rvs: set[int],
+    crvs: set[int],
+    h: EntropyFn,
+    sum_axes,
+) -> tuple[tuple[frozenset[int], ...], tuple[int, ...], tuple[int, ...], dict[frozenset[int], tuple[int, ...]]]:
+    """
+    PSP optimal partition and marginalization indices for autodiff backends.
+
+    Returns ``(partition, idx_joint, idx_crvs, idx_parts)``.
+    """
+    _, partition = caekl_mutual_information_psp_pmf_with_partition(
+        pmf,
+        all_vars=all_vars,
+        rvs=rvs,
+        crvs=crvs,
+        h=h,
+        sum_axes=sum_axes,
+    )
+    idx_joint = tuple(all_vars - (rvs | crvs))
+    idx_crvs = tuple(all_vars - crvs)
+    idx_parts = {block: tuple(all_vars - (set(block) | crvs)) for block in partition}
+    return partition, idx_joint, idx_crvs, idx_parts
+
+
+def caekl_from_partition_pmf(
+    pmf,
+    partition: tuple[frozenset[int], ...],
+    *,
+    idx_joint: tuple[int, ...],
+    idx_crvs: tuple[int, ...],
+    idx_parts: dict[frozenset[int], tuple[int, ...]],
+    h: EntropyFn,
+    sum_axes,
+) -> float:
+    """Partition information ``I_P`` for a fixed partition (used after PSP selects ``P``)."""
+    norm = len(partition) - 1
+    pmf_joint = sum_axes(pmf, idx_joint)
+    pmf_crvs = sum_axes(pmf, idx_crvs)
+    h_crvs = h(pmf_crvs)
+    h_joint = h(pmf_joint) - h_crvs
+    h_parts = sum(h(sum_axes(pmf, idx_parts[block])) - h_crvs for block in partition)
+    return (h_parts - h_joint) / norm

{dit-2.1 → dit-2.3}/dit/algorithms/distribution_optimizers.py

@@ -532,6 +532,87 @@ class MinDualTotalCorrelationOptimizer(BaseDistOptimizer, BaseNonConvexOptimizer
         return objective
 
 
+class MinResidualEntropyOptimizer(BaseDistOptimizer, BaseNonConvexOptimizer):
+    """
+    Compute minimum residual entropy distributions.
+    """
+
+    def _objective_gradient(self):
+        """Gradient of the ``+R`` objective w.r.t. the pmf."""
+        return self._residual_entropy_grad(self._rvs)
+
+    def _objective(self):
+        """
+        Compute the residual entropy.
+
+        Returns
+        -------
+        objective : func
+            The objective function.
+        """
+        residual_entropy = self._residual_entropy(self._rvs)
+
+        def objective(self, x):
+            """
+            Compute :math:`R[rvs]`
+
+            Parameters
+            ----------
+            x : np.ndarray
+                An optimization vector.
+
+            Returns
+            -------
+            obj : float
+                The value of the objective.
+            """
+            pmf = self.construct_joint(x)
+            return residual_entropy(pmf)
+
+        return objective
+
+
+class MaxCAEKLMutualInformationOptimizer(BaseDistOptimizer, BaseNonConvexOptimizer):
+    """
+    Compute maximum CAEKL mutual information distributions.
+    """
+
+    def _objective_gradient(self):
+        """Gradient of the ``-J`` objective w.r.t. the pmf."""
+        caekl_grad = self._caekl_mutual_information_grad(self._rvs)
+        return lambda pmf: -caekl_grad(pmf)
+
+    def _objective(self):
+        """
+        Compute the negative CAEKL mutual information.
+
+        Returns
+        -------
+        objective : func
+            The objective function.
+        """
+        caekl_mutual_information = self._caekl_mutual_information(self._rvs)
+
+        def objective(self, x):
+            """
+            Compute :math:`-J[rvs]`
+
+            Parameters
+            ----------
+            x : np.ndarray
+                An optimization vector.
+
+            Returns
+            -------
+            obj : float
+                The value of the objective.
+            """
+            pmf = self.construct_joint(x)
+            return -caekl_mutual_information(pmf)
+
+        return objective
+
+
 class BROJABivariateOptimizer(MaxCoInfoOptimizer):
     """
     An optimizer for constructing the maximum co-information distribution

dit-2.3/dit/algorithms/marginal_lifts.py

@@ -0,0 +1,182 @@
+"""
+Convex combinations of lifts of a joint's own marginals.
+
+At order :math:`k`, form building blocks
+
+.. math::
+
+    \\bigl\\{ U \\bigr\\} \\cup
+    \\bigl\\{ \\mathrm{Lift}(P_S) : 1 \\le |S| \\le k \\bigr\\}
+
+and fit nonnegative weights summing to one by least squares:
+
+.. math::
+
+    Q^{(k)} = \\arg\\min_{\\alpha \\ge 0,\\ \\sum \\alpha = 1}
+    \\bigl\\| P - \\textstyle\\sum_S \\alpha_S \\mathrm{Lift}(P_S) \\bigr\\|_2^2.
+
+Lifts:
+
+* ``uniform`` — :math:`P_S \\otimes U_{X \\setminus S}`
+* ``product`` — :math:`P_S \\otimes \\prod_{i \\notin S} P_i`
+"""
+
+from collections import defaultdict
+from copy import deepcopy
+from itertools import combinations, product
+
+import numpy as np
+from scipy.optimize import minimize
+
+from ..algorithms.optutil import prepare_dist
+from ..distribution import Distribution
+from ..exceptions import ditException
+
+__all__ = (
+    "lift_marginal",
+    "fit_marginal_lift_mixture",
+    "marginal_lift_dists",
+)
+
+
+def _cartesian(dist):
+    d = prepare_dist(deepcopy(dist))
+    n = d.outcome_length()
+    alph = [tuple(sorted({o[i] for o in d.outcomes})) for i in range(n)]
+    outs = list(product(*alph))
+    pmf_map = {tuple(o): float(p) for o, p in zip(d.outcomes, d.pmf, strict=True)}
+    pmf = np.array([pmf_map.get(o, 0.0) for o in outs], dtype=float)
+    pmf /= pmf.sum()
+    return outs, pmf, alph, d
+
+
+def lift_marginal(outs, pmf, alph, S, mode="uniform"):
+    """
+    Lift a marginal on coordinates ``S`` to a full joint pmf.
+
+    Parameters
+    ----------
+    outs, pmf, alph
+        Dense Cartesian table for the joint.
+    S : tuple of int
+        Variable indices of the marginal.
+    mode : {'uniform', 'product'}
+        How to extend off ``S``.
+    """
+    n = len(alph)
+    S = tuple(S)
+    marg = defaultdict(float)
+    for o, p in zip(outs, pmf, strict=True):
+        marg[tuple(o[i] for i in S)] += p
+    rest = [i for i in range(n) if i not in S]
+    if mode == "uniform":
+        rest_size = int(np.prod([len(alph[i]) for i in rest])) if rest else 1
+        return np.array([marg[tuple(o[i] for i in S)] / rest_size for o in outs], dtype=float)
+    if mode != "product":
+        msg = f"unknown lift mode {mode!r}"
+        raise ditException(msg)
+    ones = []
+    for i in range(n):
+        m1 = defaultdict(float)
+        for o, p in zip(outs, pmf, strict=True):
+            m1[o[i]] += p
+        ones.append(m1)
+    q = np.empty(len(outs), dtype=float)
+    for t, o in enumerate(outs):
+        val = marg[tuple(o[i] for i in S)]
+        for i in rest:
+            val *= ones[i][o[i]]
+        q[t] = val
+    return q / q.sum()
+
+
+def fit_marginal_lift_mixture(dist, order, mode="uniform", n_init=12, seed=0):
+    """
+    Fit a convex combination of lifts of marginals of order at most ``order``.
+
+    Returns
+    -------
+    result : dict
+        Keys ``dist``, ``labels``, ``alpha``, ``L2``.
+    """
+    if order < 0:
+        msg = "order must be nonnegative"
+        raise ditException(msg)
+    outs, pmf, alph, template = _cartesian(dist)
+    n = len(alph)
+    order = min(order, n)
+
+    blocks, labels = [], []
+    blocks.append(np.ones(len(outs)) / len(outs))
+    labels.append(())
+    for k in range(1, order + 1):
+        for S in combinations(range(n), k):
+            blocks.append(lift_marginal(outs, pmf, alph, S, mode=mode))
+            labels.append(S)
+    A = np.column_stack(blocks)
+    nb = A.shape[1]
+
+    def loss(x):
+        return float(np.sum((A @ x - pmf) ** 2))
+
+    cons = {"type": "eq", "fun": lambda x: float(x.sum() - 1.0)}
+    best = None
+    rng = np.random.default_rng(seed)
+    for _ in range(n_init):
+        x0 = rng.dirichlet(np.ones(nb))
+        res = minimize(
+            loss,
+            x0,
+            bounds=[(0.0, None)] * nb,
+            constraints=cons,
+            method="SLSQP",
+            options={"maxiter": 2000, "ftol": 1e-14, "disp": False},
+        )
+        if best is None or res.fun < best[0]:
+            best = (res.fun, res.x)
+
+    alpha = np.maximum(best[1], 0.0)
+    alpha = alpha / alpha.sum()
+    q = A @ alpha
+    q = np.maximum(q, 0.0)
+    q = q / q.sum()
+    qd = Distribution(outs, q, base="linear", validate=False)
+    qd.normalize()
+    if template.get_rv_names() is not None:
+        qd.set_rv_names(template.get_rv_names())
+    return {
+        "dist": qd,
+        "labels": labels,
+        "alpha": alpha,
+        "L2": float(np.sqrt(best[0])),
+    }
+
+
+def marginal_lift_dists(dist, k_max=None, mode="uniform", n_init=12, seed=0):
+    """
+    Ladder of marginal-lift mixtures for orders :math:`0,\\ldots,k_{\\max}`.
+
+    Order 0 is the uniform distribution. Order :math:`n` includes the full
+    joint as a block and recovers :math:`P` exactly.
+    """
+    n = dist.outcome_length()
+    if k_max is None:
+        k_max = n
+    k_max = min(int(k_max), n)
+
+    dists = []
+    metas = []
+    for k in range(0, k_max + 1):
+        if k == 0:
+            outs, _, _, template = _cartesian(dist)
+            q = Distribution(outs, np.ones(len(outs)) / len(outs), base="linear", validate=False)
+            q.normalize()
+            if template.get_rv_names() is not None:
+                q.set_rv_names(template.get_rv_names())
+            dists.append(q)
+            metas.append({"labels": [()], "alpha": np.array([1.0]), "L2": None})
+        else:
+            fit = fit_marginal_lift_mixture(dist, k, mode=mode, n_init=n_init, seed=seed + k)
+            dists.append(fit["dist"])
+            metas.append({key: fit[key] for key in ("labels", "alpha", "L2")})
+    return dists, metas