fix: break ties in the k-NN estimators and use Chebyshev when Z is None (follow-up to #14)

causationentropy/core/information/conditional_mutual_information.py

@@ -9,6 +9,7 @@
     poisson_joint_entropy,
 )
 from causationentropy.core.information.mutual_information import (
+    _add_tie_breaking_noise,
     gaussian_mutual_information,
     geometric_knn_mutual_information,
     kde_mutual_information,
@@ -140,23 +141,23 @@ def kde_conditional_mutual_information(
 
 
 def knn_conditional_mutual_information(X, Y, Z, metric=None, k=1):
-    """
+    r"""
     Estimate conditional mutual information using k-nearest neighbor method.
 
-    This function implements conditional mutual information estimation using
-    the relationship:
+    This function implements the k-nearest-neighbour estimator of Frenzel and
+    Pompe (2007) and Vejmelka and Paluš (2008), the conditional extension of
+    the KSG estimator:
 
     .. math::
 
-        I(X; Y | Z) = I(X, Y) - I(X, Y; Z)
-
-    where both mutual information terms are estimated using the KSG k-NN estimator.
-
-    The approach leverages the fact that:
-
-    .. math::
+        I(X; Y | Z) = \psi(k) - \langle \psi(n_{xz} + 1) + \psi(n_{yz} + 1)
+        - \psi(n_z + 1) \rangle
 
-        I(X; Y | Z) = I(X; Y) - I(X; Y | Z)
+    where :math:`\epsilon_i` is the distance from sample :math:`i` to its k-th
+    neighbour in the joint :math:`(X, Y, Z)` space, and :math:`n_{xz}`,
+    :math:`n_{yz}` and :math:`n_z` count the samples strictly within
+    :math:`\epsilon_i` in the :math:`(X, Z)`, :math:`(Y, Z)` and :math:`Z`
+    subspaces. If ``Z`` is None, the KSG mutual information estimator is used.
 
     Parameters
     ----------
@@ -167,8 +168,9 @@ def knn_conditional_mutual_information(X, Y, Z, metric=None, k=1):
     Z : array-like of shape (n_samples, n_features_z) or None
         Conditioning variable. If None, computes marginal mutual information.
     metric : str or None, default=None
-        Distance metric for k-NN calculations. Uses Minkowski for the
-        marginal-MI case (Z=None) and Chebyshev for conditional MI.
+        Distance metric for k-NN calculations. If None, uses Chebyshev
+        distance (the maximum norm assumed by the estimator) for both the
+        conditional case and the marginal-MI case (Z=None).
     k : int, default=1
         Number of nearest neighbors.
 
@@ -179,23 +181,30 @@ def knn_conditional_mutual_information(X, Y, Z, metric=None, k=1):
 
     Notes
     -----
-    This implementation uses the decomposition approach rather than direct
-    conditional MI estimation. The accuracy depends on:
+    The joint-space distance and all marginal neighbour counts use the same
+    metric. Tiny noise (relative amplitude 1e-10, fixed seed) is added to X, Y
+    and Z before the neighbour search so that repeated values, such as count
+    data or bootstrap resamples, do not produce infinite or NaN estimates.
+    The accuracy depends on:
 
-    - Quality of marginal MI estimates
     - Dimensionality of the joint space
     - Sample size relative to effective dimensionality
+    - The choice of k
 
     References
     ----------
     .. [1] Kraskov, A., Stögbauer, H., Grassberger, P. Estimating mutual information.
            Physical Review E 69, 066138 (2004).
+    .. [2] Frenzel, S., Pompe, B. Partial mutual information for coupling analysis
+           of multivariate time series. Physical Review Letters 99, 204101 (2007).
+    .. [3] Vejmelka, M., Paluš, M. Inferring the directionality of coupling with
+           conditional mutual information. Physical Review E 77, 026214 (2008).
     """
+    effective_metric = "chebyshev" if metric is None else metric
     if Z is None:
-        effective_metric = "minkowski" if metric is None else metric
         return knn_mutual_information(X, Y, metric=effective_metric, k=k)
     else:
-        effective_metric = "chebyshev" if metric is None else metric
+        X, Y, Z = _add_tie_breaking_noise([X, Y, Z])
         JS = np.column_stack((X, Y, Z))
         # Find the K-th smallest distance in the joint space using the
         # same metric that is used for all marginal neighbor counts.

causationentropy/core/information/mutual_information.py

@@ -106,7 +106,45 @@ def kde_mutual_information(X, Y, bandwidth="silverman", kernel="gaussian"):
     return mi
 
 
-def knn_mutual_information(X, Y, metric="euclidean", k=1):
+def _add_tie_breaking_noise(arrays, scale=1e-10, seed=0):
+    """
+    Add tiny independent noise to each array so that no two samples coincide.
+
+    The KSG-type estimators count neighbours strictly inside the distance to
+    the k-th neighbour. When a point has k or more exact copies, that distance
+    is 0, the counts become -1 and ``digamma(0)`` returns ``-inf``. Adding
+    low-amplitude noise breaks such ties, as recommended by Kraskov et al.
+    (2004) and done in Tigramite's CMIknn and scikit-learn's
+    ``mutual_info_regression``.
+
+    Parameters
+    ----------
+    arrays : sequence of array-like, each of shape (n_samples, n_features)
+        Variables to perturb. Every array receives its own noise.
+    scale : float, default=1e-10
+        Noise amplitude relative to each column's standard deviation (or to
+        ``max(|mean|, 1)`` for a constant column).
+    seed : int, default=0
+        Seed for the noise, so repeated calls on the same input give the same
+        estimate.
+
+    Returns
+    -------
+    list of np.ndarray
+        Perturbed float copies of ``arrays``, in the same order.
+    """
+    rng = np.random.default_rng(seed)
+    noisy = []
+    for A in arrays:
+        A = np.asarray(A, dtype=float)
+        spread = np.std(A, axis=0)
+        fallback = np.maximum(np.abs(np.mean(A, axis=0)), 1.0)
+        amplitude = scale * np.where(spread > 0, spread, fallback)
+        noisy.append(A + amplitude * rng.standard_normal(A.shape))
+    return noisy
+
+
+def knn_mutual_information(X, Y, metric="chebyshev", k=1):
     r"""
     Estimate mutual information using k-nearest neighbor (KSG) method.
 
@@ -127,8 +165,10 @@ def knn_mutual_information(X, Y, metric="euclidean", k=1):
         First variable.
     Y : array-like of shape (n_samples, n_features_y)
         Second variable.
-    metric : str, default='euclidean'
-        Distance metric for neighborhood calculations.
+    metric : str, default='chebyshev'
+        Distance metric for neighborhood calculations. The KSG estimator
+        assumes the maximum norm (Chebyshev distance); other metrics give
+        biased estimates.
     k : int, default=1
         Number of nearest neighbors to consider.
 
@@ -150,11 +190,17 @@ def knn_mutual_information(X, Y, metric="euclidean", k=1):
     - Small k: Lower bias, higher variance
     - Large k: Higher bias, lower variance
 
+    Tiny noise (relative amplitude 1e-10, fixed seed) is added to X and Y
+    before the neighbour search so that repeated values, such as count data,
+    do not produce infinite or NaN estimates.
+
     References
     ----------
     .. [1] Kraskov, A., Stögbauer, H., Grassberger, P. Estimating mutual information.
            Physical Review E 69, 066138 (2004).
     """
+    X, Y = _add_tie_breaking_noise([X, Y])
+
     # construct the joint space
     n = X.shape[0]
     JS = np.column_stack((X, Y))

causationentropy/core/stats.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Tuple, Union
 
 import numpy as np
@@ -416,7 +417,7 @@ def bootstrap_cmi_confidence_interval(
     block_length: int = None,
     mean_block_length: float = 10.0,
     use_stationary: bool = False,
-    metric: str = "euclidean",
+    metric: Union[str, None] = None,
     k: int = 5,
     bandwidth: str = "silverman",
     seed: Union[int, np.random.Generator, None] = 42,
@@ -456,8 +457,9 @@ def bootstrap_cmi_confidence_interval(
     use_stationary : bool, default=False
         If True, use the stationary bootstrap; otherwise the moving block
         bootstrap.
-    metric : str, default='euclidean'
-        Distance metric for k-NN based estimators.
+    metric : str or None, default=None
+        Distance metric for k-NN based estimators. None uses each estimator's
+        own default (Chebyshev for ``"knn"``, Euclidean for ``"geometric_knn"``).
     k : int, default=5
         Number of neighbors for k-NN based estimators.
     bandwidth : str, default='silverman'
@@ -482,6 +484,13 @@ def bootstrap_cmi_confidence_interval(
         If ``X``/``Y``/``Z`` disagree on ``T``, or bootstrap sizes are out
         of range.
 
+    Warns
+    -----
+    UserWarning
+        If ``method`` is ``"knn"`` or ``"geometric_knn"``. Resampling with
+        replacement repeats rows, which biases nearest-neighbour estimates,
+        so these intervals should not be relied on.
+
     Examples
     --------
     >>> import numpy as np
@@ -517,6 +526,15 @@ def bootstrap_cmi_confidence_interval(
         if Z.ndim != 2 or Z.shape[0] != X.shape[0]:
             raise ValueError("Z must be None or 2-D with the same rows as X.")
     n = X.shape[0]
+    if method in ("knn", "geometric_knn"):
+        warnings.warn(
+            "Bootstrap resamples repeat rows, and repeated rows shrink "
+            "nearest-neighbour distances, so intervals from k-NN based "
+            f"estimators (method={method!r}) are biased. Prefer "
+            "method='gaussian' for bootstrap intervals.",
+            UserWarning,
+            stacklevel=2,
+        )
 
     if use_stationary:
         indices = stationary_bootstrap_indices(

causationentropy/tests/core/information/test_knn_ties_and_defaults.py

@@ -0,0 +1,123 @@
+"""Tie handling and default metrics of the k-NN (KSG) estimators."""
+
+import warnings
+
+import numpy as np
+import pytest
+
+from causationentropy.core.information.conditional_mutual_information import (
+    conditional_mutual_information,
+    knn_conditional_mutual_information,
+)
+from causationentropy.core.information.mutual_information import (
+    knn_mutual_information,
+)
+from causationentropy.core.stats import bootstrap_cmi_confidence_interval
+
+
+class TestKNNRepeatedValues:
+    """Repeated samples used to make the k-th neighbour distance 0."""
+
+    def test_repeated_values_give_finite_estimates(self):
+        X = np.repeat(np.arange(10.0), 20)[:, None]
+        Y = X[::-1]
+
+        assert np.isfinite(knn_mutual_information(X, Y, k=3))
+        assert np.isfinite(knn_conditional_mutual_information(X, Y, X, k=3))
+        assert np.isfinite(
+            conditional_mutual_information(X, Y, None, method="knn", k=3)
+        )
+        assert np.isfinite(conditional_mutual_information(X, Y, X, method="knn", k=3))
+
+    def test_independent_count_data_is_close_to_zero(self):
+        rng = np.random.default_rng(0)
+        X, Y, Z = [rng.poisson(3, (300, d)).astype(float) for d in (1, 1, 2)]
+
+        mi = knn_mutual_information(X, Y, k=3)
+        cmi = knn_conditional_mutual_information(X, Y, Z, k=3)
+
+        assert np.isfinite(mi) and np.isfinite(cmi)
+        assert abs(mi) < 0.1
+        assert abs(cmi) < 0.1
+
+    def test_estimates_are_deterministic(self):
+        rng = np.random.default_rng(1)
+        X, Y, Z = [rng.poisson(3, (200, 1)).astype(float) for _ in range(3)]
+
+        first = knn_conditional_mutual_information(X, Y, Z, k=3)
+        second = knn_conditional_mutual_information(X, Y, Z, k=3)
+
+        assert first == second
+
+    def test_inputs_are_not_modified(self):
+        X = np.repeat(np.arange(5.0), 10)[:, None]
+        Y = X[::-1].copy()
+        X_before, Y_before = X.copy(), Y.copy()
+
+        knn_conditional_mutual_information(X, Y, X, k=3)
+
+        np.testing.assert_array_equal(X, X_before)
+        np.testing.assert_array_equal(Y, Y_before)
+
+
+class TestKNNDefaultMetric:
+    """The KSG estimators assume the maximum norm, also when Z is None."""
+
+    def test_marginal_case_defaults_to_chebyshev(self):
+        rng = np.random.default_rng(2)
+        X = rng.standard_normal((300, 1))
+        Y = 0.5 * X + rng.standard_normal((300, 1))
+
+        assert knn_conditional_mutual_information(
+            X, Y, None, k=5
+        ) == knn_conditional_mutual_information(X, Y, None, metric="chebyshev", k=5)
+        assert knn_mutual_information(X, Y, k=5) == knn_mutual_information(
+            X, Y, metric="chebyshev", k=5
+        )
+
+    def test_marginal_default_is_close_to_known_value(self):
+        rng = np.random.default_rng(1)
+        r = 0.5
+        truth = -0.5 * np.log(1 - r**2)  # about 0.144
+        dependent, independent = [], []
+        for _ in range(10):
+            X = rng.standard_normal((500, 1))
+            Y_dep = r * X + np.sqrt(1 - r**2) * rng.standard_normal((500, 1))
+            Y_ind = rng.standard_normal((500, 1))
+            dependent.append(knn_conditional_mutual_information(X, Y_dep, None, k=5))
+            independent.append(knn_conditional_mutual_information(X, Y_ind, None, k=5))
+
+        assert abs(np.mean(dependent) - truth) < 0.04
+        assert abs(np.mean(independent)) < 0.04
+
+
+class TestBootstrapKNNWarning:
+    """Bootstrap resamples repeat rows, which biases k-NN estimates."""
+
+    def _data(self):
+        rng = np.random.default_rng(0)
+        X = rng.standard_normal((100, 1))
+        Z = rng.standard_normal((100, 1))
+        Y = 0.6 * X + 0.8 * Z + rng.standard_normal((100, 1))
+        return X, Y, Z
+
+    def test_knn_bootstrap_warns_and_returns_finite_values(self):
+        X, Y, Z = self._data()
+
+        with pytest.warns(UserWarning, match="biased"):
+            estimate, lower, upper, replicates = bootstrap_cmi_confidence_interval(
+                X, Y, Z, method="knn", n_bootstraps=10, seed=0
+            )
+
+        assert np.isfinite(estimate)
+        assert np.all(np.isfinite(replicates))
+        assert lower <= upper
+
+    def test_gaussian_bootstrap_does_not_warn(self):
+        X, Y, Z = self._data()
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", UserWarning)
+            bootstrap_cmi_confidence_interval(
+                X, Y, Z, method="gaussian", n_bootstraps=10, seed=0
+            )