dit 2.1__tar.gz → 2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dit
-Version: 2.1
+Version: 2.2
 Summary: Python package for information theory on discrete random variables.
 Project-URL: Homepage, http://dit.io
 Project-URL: Repository, https://github.com/dit/dit
@@ -22,7 +22,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 +59,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 +71,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 +94,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 +469,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.2}/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.2}/dit/__init__.py

@@ -37,3 +37,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.2/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.2}/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.1 → dit-2.2}/dit/algorithms/optimization.py

@@ -18,16 +18,21 @@ from boltons.iterutils import pairwise
 from loguru import logger
 from scipy.optimize import Bounds, basinhopping, brute, differential_evolution, dual_annealing, minimize, shgo
 
+from ..algorithms.caekl_psp import (
+    caekl_mutual_information_psp_pmf,
+    caekl_mutual_information_psp_pmf_grad_data,
+)
 from ..algorithms.channelcapacity import channel_capacity
 from ..distconst import insert_rvf, modify_outcomes
 from ..distribution import Distribution
 from ..exceptions import OptimizationException, ditException
 from ..helpers import flatten, normalize_rvs, parse_rvs
 from ..math import prod, sample_simplex
-from ..utils import partitions, powerset
+from ..utils import powerset
 from ..utils.optimization import (
     BasinHoppingCallBack,
     BasinHoppingInnerCallBack,
+    BoundedDisplacement,
     Uniquifier,
     accept_test,
     basinhop_status,
@@ -610,6 +615,53 @@ class BaseOptimizer(metaclass=ABCMeta):
 
         return dual_total_correlation
 
+    def _residual_entropy(self, rvs, crvs=None):
+        """
+        Compute the residual entropy (variation of information).
+
+        Parameters
+        ----------
+        rvs : set
+            The random variables to compute the residual entropy of.
+        crvs : set
+            The random variables to condition on.
+
+        Returns
+        -------
+        re : func
+            The residual entropy.
+        """
+        if crvs is None:
+            crvs = set()
+        idx_joint = tuple(self._all_vars - (rvs | crvs))
+        idx_margs = [tuple(self._all_vars - ((rvs - {rv}) | crvs)) for rv in rvs]
+        idx_crvs = tuple(self._all_vars - crvs)
+
+        def residual_entropy(pmf):
+            """
+            Compute the specified residual entropy.
+
+            Parameters
+            ----------
+            pmf : np.ndarray
+                The joint probability distribution.
+
+            Returns
+            -------
+            re : float
+                The residual entropy.
+            """
+            pmf_joint = pmf.sum(axis=idx_joint, keepdims=True)
+            pmf_margs = [pmf_joint.sum(axis=marg, keepdims=True) for marg in idx_margs]
+            pmf_crvs = pmf_joint.sum(axis=idx_crvs, keepdims=True)
+
+            h_crvs = self._h(pmf_crvs)
+            h_margs = [self._h(marg) - h_crvs for marg in pmf_margs]
+
+            return sum(h_margs)
+
+        return residual_entropy
+
     ###########################################################################
     # Analytic gradients of the objective building blocks.
     #
@@ -755,6 +807,21 @@ class BaseOptimizer(metaclass=ABCMeta):
 
         return grad
 
+    def _residual_entropy_grad(self, rvs, crvs=None):
+        """Gradient builder for :meth:`_residual_entropy`."""
+        if crvs is None:
+            crvs = set()
+        idx_margs = [tuple(self._all_vars - ((rvs - {rv}) | crvs)) for rv in rvs]
+        idx_crvs = tuple(self._all_vars - crvs)
+        n = len(rvs)
+
+        def grad(pmf):
+            g = sum(self._marginal_entropy_grad(pmf, marg) for marg in idx_margs)
+            g = g - n * self._marginal_entropy_grad(pmf, idx_crvs)
+            return self._full_grad(g, pmf)
+
+        return grad
+
     def _caekl_mutual_information(self, rvs, crvs=None):
         """
         Compute the CAEKL mutual information.
@@ -773,15 +840,6 @@ class BaseOptimizer(metaclass=ABCMeta):
         """
         if crvs is None:
             crvs = set()
-        parts = [p for p in partitions(rvs) if len(p) > 1]
-        idx_parts = {}
-        for part in parts:
-            for p in part:
-                if p not in idx_parts:
-                    idx_parts[p] = tuple(self._all_vars - (p | crvs))
-        part_norms = [len(part) - 1 for part in parts]
-        idx_joint = tuple(self._all_vars - (rvs | crvs))
-        idx_crvs = tuple(self._all_vars - crvs)
 
         def caekl_mutual_information(pmf):
             """
@@ -797,19 +855,14 @@ class BaseOptimizer(metaclass=ABCMeta):
             caekl : float
                 The CAEKL mutual information.
             """
-            pmf_joint = pmf.sum(axis=idx_joint, keepdims=True)
-            pmf_parts = {p: pmf_joint.sum(axis=idx, keepdims=True) for p, idx in idx_parts.items()}
-            pmf_crvs = pmf_joint.sum(axis=idx_crvs, keepdims=True)
-
-            h_crvs = self._h(pmf_crvs)
-            h_joint = self._h(pmf_joint) - h_crvs
-
-            pairs = zip(parts, part_norms, strict=True)
-            candidates = [(sum(self._h(pmf_parts[p]) - h_crvs for p in part) - h_joint) / norm for part, norm in pairs]
-
-            caekl = min(candidates)
-
-            return caekl
+            return caekl_mutual_information_psp_pmf(
+                pmf,
+                all_vars=self._all_vars,
+                rvs=rvs,
+                crvs=crvs,
+                h=self._h,
+                sum_axes=lambda p, idx: p.sum(axis=idx, keepdims=True),
+            )
 
         return caekl_mutual_information
 
@@ -824,34 +877,20 @@ class BaseOptimizer(metaclass=ABCMeta):
         """
         if crvs is None:
             crvs = set()
-        parts = [p for p in partitions(rvs) if len(p) > 1]
-        idx_parts = {}
-        for part in parts:
-            for p in part:
-                if p not in idx_parts:
-                    idx_parts[p] = tuple(self._all_vars - (p | crvs))
-        part_norms = [len(part) - 1 for part in parts]
-        idx_joint = tuple(self._all_vars - (rvs | crvs))
-        idx_crvs = tuple(self._all_vars - crvs)
 
         def grad(pmf):
-            # Recompute the candidates to find the active (argmin) partition.
-            pmf_joint = pmf.sum(axis=idx_joint, keepdims=True)
-            pmf_parts = {p: pmf_joint.sum(axis=idx, keepdims=True) for p, idx in idx_parts.items()}
-            pmf_crvs = pmf_joint.sum(axis=idx_crvs, keepdims=True)
-
-            h_crvs = self._h(pmf_crvs)
-            h_joint = self._h(pmf_joint) - h_crvs
-
-            pairs = list(zip(parts, part_norms, strict=True))
-            candidates = [(sum(self._h(pmf_parts[p]) - h_crvs for p in part) - h_joint) / norm for part, norm in pairs]
-
-            idx_min = min(range(len(candidates)), key=lambda i: candidates[i])
-            part, norm = pairs[idx_min]
+            _, partition, idx_joint, idx_crvs, idx_parts = caekl_mutual_information_psp_pmf_grad_data(
+                pmf,
+                all_vars=self._all_vars,
+                rvs=rvs,
+                crvs=crvs,
+                h=self._h,
+                sum_axes=lambda p, idx: p.sum(axis=idx, keepdims=True),
+            )
+            norm = len(partition) - 1
 
-            # d/dpmf of (sum_p H(p) - (|part|-1) H(crvs) - H(joint)) / norm.
-            g = sum(self._marginal_entropy_grad(pmf, idx_parts[p]) for p in part)
-            g = g - (len(part) - 1) * self._marginal_entropy_grad(pmf, idx_crvs)
+            g = sum(self._marginal_entropy_grad(pmf, idx_parts[block]) for block in partition)
+            g = g - norm * self._marginal_entropy_grad(pmf, idx_crvs)
             g = g - self._marginal_entropy_grad(pmf, idx_joint)
             g = g / norm
             return self._full_grad(g, pmf)
@@ -1375,6 +1414,7 @@ class BaseNonConvexOptimizer(BaseOptimizer):
             minimizer_kwargs=minimizer_kwargs,
             niter=niter,
             accept_test=accept_test,
+            take_step=BoundedDisplacement(rng=self._rng),
             callback=self._callback,
             **bh_kwargs,
         )
@@ -1384,6 +1424,20 @@ class BaseNonConvexOptimizer(BaseOptimizer):
         if not success:  # pragma: no cover
             result = self._callback.minimum() or res_shotgun
 
+        if result is None:
+            # Basin hopping reported failure and neither the callback minima nor
+            # the (optional) seeding shotgun yielded a usable point. On
+            # near-degenerate distributions (e.g. probabilities ~1e-8 that are
+            # not subnormal, so they survive the test strategy's denormal
+            # filter) this happens intermittently depending on the random
+            # starts, producing flaky "No optima found" failures even though the
+            # problem is solvable. Fall back to a many-start shotgun, which
+            # tolerates SLSQP's spurious ``success=False`` at feasible optima via
+            # ``_best_feasible``. This never overrides a result that already
+            # succeeded.
+            logger.debug("Basin hopping produced no result; retrying with shotgun fallback")
+            result = self._optimize_shotgun(x0, minimizer_kwargs, niter=max(25, niter))
+
         return result
 
     def _optimization_diffevo(self, x0, minimizer_kwargs, niter):  # pragma: no cover