From 11f05f4da9aed56bc6b7c16e146e23982be383e3 Mon Sep 17 00:00:00 2001 From: Parva Mehta Date: Tue, 6 Oct 2026 16:06:32 +0530 Subject: [PATCH] fix: break ties in k-NN estimators, Chebyshev default for Z=None (follow-up to #14) --- .../conditional_mutual_information.py | 45 ++++--- .../core/information/mutual_information.py | 52 +++++++- causationentropy/core/stats.py | 24 +++- .../information/test_knn_ties_and_defaults.py | 123 ++++++++++++++++++ 4 files changed, 220 insertions(+), 24 deletions(-) create mode 100644 causationentropy/tests/core/information/test_knn_ties_and_defaults.py diff --git a/causationentropy/core/information/conditional_mutual_information.py b/causationentropy/core/information/conditional_mutual_information.py index b910271..3b3df55 100644 --- a/causationentropy/core/information/conditional_mutual_information.py +++ b/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. diff --git a/causationentropy/core/information/mutual_information.py b/causationentropy/core/information/mutual_information.py index 3131c85..94157cc 100644 --- a/causationentropy/core/information/mutual_information.py +++ b/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)) diff --git a/causationentropy/core/stats.py b/causationentropy/core/stats.py index 6e48c46..42d718d 100644 --- a/causationentropy/core/stats.py +++ b/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( diff --git a/causationentropy/tests/core/information/test_knn_ties_and_defaults.py b/causationentropy/tests/core/information/test_knn_ties_and_defaults.py new file mode 100644 index 0000000..ce1fe1c --- /dev/null +++ b/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 + )