From f7fd4468f188ff6808f44ac2b4221b471ca5f0de Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 11:31:54 +0800 Subject: [PATCH 01/13] feat: add correlation-stability entropy example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Demonstrates entropy of a rolling correlation series as a correlation-regime stability gauge (idea from an r/dataisbeautiful thread). Ships alongside a rolling-std baseline so the example shows that an entropy signal must be checked against a trivial baseline before being trusted — on this synthetic data, std separates the regimes and entropy of the smooth correlation swing does not. - _synthetic.correlated_assets(): two return streams, stable shared-factor regime then a correlation breakdown - examples/correlation_stability.py: rolling corr -> permutation entropy vs rolling-std baseline, prints regime separation for both - tests: smoke test for the new example + a behavioral test asserting the generator's stable half is more correlated than its broken half Uses only the existing public API; no library changes. Co-Authored-By: Claude Opus 4.8 --- examples/_synthetic.py | 32 +++++++++++ examples/correlation_stability.py | 88 +++++++++++++++++++++++++++++++ tests/test_examples.py | 24 ++++++++- 3 files changed, 143 insertions(+), 1 deletion(-) create mode 100644 examples/correlation_stability.py diff --git a/examples/_synthetic.py b/examples/_synthetic.py index 22e1a69..1d02139 100644 --- a/examples/_synthetic.py +++ b/examples/_synthetic.py @@ -131,3 +131,35 @@ def qc_sensor(seed=13): out_of_control = 50 + np.linspace(0, 3, 400) + 1.5 * rng.randn(400) values = np.concatenate([in_control, out_of_control]) return pd.Series(values, name="dimension_mm") + + +def correlated_assets(seed=14): + """Two asset return streams whose correlation regime breaks down halfway. + + First half: both are driven by a shared market factor, so their rolling + correlation sits in a stable, high band. Second half: the shared driver + fades and a sign-flipping component takes over, so the correlation wanders + all over the place. The *correlation series itself* goes from steady to + erratic — that is what entropy of the rolling correlation is meant to catch. + + Returns a 2-column DataFrame of daily returns (``a``, ``b``). + """ + rng = _rng(seed) + n = 400 + + # Stable regime: strong shared factor + small idiosyncratic noise. + factor = rng.randn(n) + a_stable = factor + 0.4 * rng.randn(n) + b_stable = factor + 0.4 * rng.randn(n) + + # Breakdown regime: shared factor weakens and a slow sign-flipping driver + # drags the correlation up and down, so it never settles. + factor2 = rng.randn(n) + flip = np.sin(np.linspace(0, 6 * np.pi, n)) # swings the coupling sign + a_unstable = factor2 + 0.8 * rng.randn(n) + b_unstable = flip * factor2 + 0.8 * rng.randn(n) + + a = np.concatenate([a_stable, a_unstable]) + b = np.concatenate([b_stable, b_unstable]) + idx = pd.date_range("2025-01-01", periods=len(a), freq="B") + return pd.DataFrame({"a": a, "b": b}, index=idx) diff --git a/examples/correlation_stability.py b/examples/correlation_stability.py new file mode 100644 index 0000000..b872d24 --- /dev/null +++ b/examples/correlation_stability.py @@ -0,0 +1,88 @@ +"""Correlation-stability entropy example. + +Run: python examples/correlation_stability.py + +Idea (from an r/dataisbeautiful thread): instead of taking the entropy of a +price series directly, take a *rolling correlation* between two assets and feed +THAT series into entroscope. Low entropy of the correlation series means the +correlation sits in a stable band; high entropy means the relationship is +wandering — a correlation-regime breakdown. + +This example ships the honest comparison alongside it: the rolling standard +deviation of the same correlation series. Std-dev already measures "how much the +correlation moves". Entropy only earns its place if it separates the stable and +broken regimes at least as cleanly. The numbers below let you judge that — they +are not a claim that entropy wins. + +Synthetic data lives in ``_synthetic`` so this runs with no external files. Swap +``data.correlated_assets()`` for a 2-column frame of your own returns to try it +on real data. +""" + +import pandas as pd + +from entroscope import permutation + +try: + from . import _synthetic as data +except ImportError: + import _synthetic as data + + +def _split_means(series, half): + """Mean of a derived series over its stable vs broken halves. + + The derived series is shorter than the raw data (rolling windows drop + leading NaNs), so ``half`` is taken on the derived series' own length. + """ + clean = series.dropna() + cut = len(clean) // 2 + return clean.iloc[:cut].mean(), clean.iloc[cut:].mean() + + +def correlation_stability(): + """Entropy of rolling correlation vs a rolling-std baseline.""" + returns = data.correlated_assets() + + # 1. Rolling correlation between the two assets — the derived series. + corr = returns["a"].rolling(window=30).corr(returns["b"]) + + # 2. Entropy of that correlation series (ordinal, no binning needed). + corr_entropy = permutation.rolling(corr.dropna(), window=40, order=3) + + # 3. Honest baseline: rolling std of the same correlation series. + corr_std = corr.rolling(window=40).std() + + half = len(returns) // 2 + ent_stable, ent_broken = _split_means(corr_entropy, half) + std_stable, std_broken = _split_means(corr_std, half) + + print("Correlation stability — two assets, regime break at the midpoint") + print(f" {'measure':<28}{'stable':>10}{'broken':>10}{'separation':>12}") + print( + f" {'permutation entropy of corr':<28}" + f"{ent_stable:>10.3f}{ent_broken:>10.3f}{ent_broken - ent_stable:>12.3f}" + ) + print( + f" {'rolling std of corr (baseline)':<28}" + f"{std_stable:>10.3f}{std_broken:>10.3f}{std_broken - std_stable:>12.3f}" + ) + print( + " -> on THIS data the plain std baseline separates the regimes and the\n" + " entropy of the smooth correlation swing does not — a reminder to\n" + " always check an entropy signal against a trivial baseline before\n" + " trusting it. Entropy earns its place when the broken regime is\n" + " erratic/disordered, not smoothly oscillating.\n" + ) + return corr, corr_entropy, corr_std + + +def main(): + print("=" * 60) + print("CORRELATION-STABILITY ENTROPY") + print("=" * 60) + correlation_stability() + + +if __name__ == "__main__": + main() diff --git a/tests/test_examples.py b/tests/test_examples.py index 4d80071..ba59bd3 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -24,7 +24,7 @@ def _load(module_name): return module -@pytest.mark.parametrize("name", ["medical", "business"]) +@pytest.mark.parametrize("name", ["medical", "business", "correlation_stability"]) def test_example_runs(name, capsys): module = _load(name) module.main() # must not raise @@ -49,3 +49,25 @@ def test_synthetic_generators_return_series(): series = gen() assert isinstance(series, pd.Series) assert len(series) > 0 + + +def test_correlated_assets_has_regime_breakdown(): + """The generator must actually deliver the stable->broken regime it claims. + + The correlation-stability example is only meaningful if the first half is + genuinely more correlated than the second. Assert that behavioral contract, + not just the type — if the regime logic is broken later, this fails loudly. + """ + import pandas as pd + + data = _load("_synthetic") + df = data.correlated_assets() + + assert isinstance(df, pd.DataFrame) + assert list(df.columns) == ["a", "b"] + assert len(df) > 0 + + half = len(df) // 2 + stable_corr = df["a"].iloc[:half].corr(df["b"].iloc[:half]) + broken_corr = df["a"].iloc[half:].corr(df["b"].iloc[half:]) + assert stable_corr > broken_corr From ab90b96d14d9fbc504899a2fbe8ff42e6bd00043 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 11:53:49 +0800 Subject: [PATCH 02/13] feat: add shared delay-embedding for transfer entropy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements embed() in entroscope/_transfer_estimators.py — the single source of (y_future, y_past, x_past) aligned columns consumed by both the binned and KSG estimators. Verified in isolation against hand-computed values (Gate C). Co-Authored-By: Claude Sonnet 4.6 --- entroscope/_transfer_estimators.py | 31 ++++++++++++++++++++++++++++++ tests/test_transfer.py | 30 +++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) create mode 100644 entroscope/_transfer_estimators.py create mode 100644 tests/test_transfer.py diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py new file mode 100644 index 0000000..d97f87e --- /dev/null +++ b/entroscope/_transfer_estimators.py @@ -0,0 +1,31 @@ +"""Transfer-entropy estimators (history length 1) and shared delay-embedding. + +Private module: all the math lives here so the public ``transfer`` module stays +thin. Both estimators consume the SAME embedded vectors from ``embed`` — the +embedding is verified in isolation (see tests) so a bug here cannot silently fool +both estimators. + +NOTE: imports for later estimators (digamma from scipy, knn from .utils) are +added in the tasks that implement the binned and KSG estimators respectively. +""" + +import numpy as np + + +def embed(x, y, lag=1): + """Build aligned (y_future, y_past, x_past) sample columns, history length 1. + + For each valid time t (from ``lag`` to n-1): y_future=y[t], y_past=y[t-lag], + x_past=x[t-lag]. Returns three 1-D arrays of length ``n - lag``. + """ + x = np.asarray(x, dtype=float) + y = np.asarray(y, dtype=float) + if lag < 1: + raise ValueError("lag must be >= 1") + n = len(y) + if n - lag < 1: + raise ValueError(f"series too short ({n}) for lag {lag}") + y_future = y[lag:] + y_past = y[: n - lag] + x_past = x[: n - lag] + return y_future, y_past, x_past diff --git a/tests/test_transfer.py b/tests/test_transfer.py new file mode 100644 index 0000000..5704aa9 --- /dev/null +++ b/tests/test_transfer.py @@ -0,0 +1,30 @@ +import numpy as np +import pytest + +from entroscope import _transfer_estimators as est + + +def test_embed_hand_computed_lag1(): + # 6 points, lag=1 -> 5 aligned samples. + x = np.array([10.0, 11.0, 12.0, 13.0, 14.0, 15.0]) + y = np.array([20.0, 21.0, 22.0, 23.0, 24.0, 25.0]) + y_future, y_past, x_past = est.embed(x, y, lag=1) + # t = 1..5 + np.testing.assert_array_equal(y_future, [21.0, 22.0, 23.0, 24.0, 25.0]) + np.testing.assert_array_equal(y_past, [20.0, 21.0, 22.0, 23.0, 24.0]) + np.testing.assert_array_equal(x_past, [10.0, 11.0, 12.0, 13.0, 14.0]) + + +def test_embed_lag2(): + x = np.arange(6, dtype=float) + y = np.arange(10, 16, dtype=float) + y_future, y_past, x_past = est.embed(x, y, lag=2) + # t = 2..5 -> 4 samples + np.testing.assert_array_equal(y_future, [12.0, 13.0, 14.0, 15.0]) + np.testing.assert_array_equal(y_past, [10.0, 11.0, 12.0, 13.0]) + np.testing.assert_array_equal(x_past, [0.0, 1.0, 2.0, 3.0]) + + +def test_embed_too_short_raises(): + with pytest.raises(ValueError): + est.embed(np.array([1.0]), np.array([2.0]), lag=1) From f4925af2d3180cc9a62108c282f2f022d53b3851 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 11:57:07 +0800 Subject: [PATCH 03/13] feat: add Chebyshev-norm k-NN helpers for KSG Co-Authored-By: Claude Sonnet 4.6 --- entroscope/utils/knn.py | 40 ++++++++++++++++++++++++++++++++++++++++ tests/test_transfer.py | 19 +++++++++++++++++++ 2 files changed, 59 insertions(+) create mode 100644 entroscope/utils/knn.py diff --git a/entroscope/utils/knn.py b/entroscope/utils/knn.py new file mode 100644 index 0000000..00c13c1 --- /dev/null +++ b/entroscope/utils/knn.py @@ -0,0 +1,40 @@ +"""k-nearest-neighbor helpers for KSG-style estimators (Chebyshev / max norm).""" + +import numpy as np +from scipy.spatial import cKDTree + + +def kth_neighbor_distance(points, k): + """Chebyshev distance to the k-th nearest neighbor of each point. + + `points` is (n, d). Returns length-n array. Self (distance 0) is excluded by + querying k+1 neighbors and dropping the first. + """ + points = np.asarray(points, dtype=float) + tree = cKDTree(points) + # query returns the point itself first (distance 0); take the (k+1)-th column. + dists, _ = tree.query(points, k=k + 1, p=np.inf) + return np.asarray(dists)[:, k] + + +def count_within_radius(points, radii): + """For each point, count OTHER points with Chebyshev distance strictly < radius. + + `points` is (n, d); `radii` is length n. Excludes the point itself. + """ + points = np.asarray(points, dtype=float) + radii = np.asarray(radii, dtype=float) + tree = cKDTree(points) + counts = np.empty(len(points), dtype=int) + for i, (p, r) in enumerate(zip(points, radii)): + # ball_point with p=inf, count neighbors strictly inside r, minus self. + idx = tree.query_ball_point(p, r=r, p=np.inf) + # exclude self and any point exactly at radius r (strict inequality). + c = 0 + for j in idx: + if j == i: + continue + if np.max(np.abs(points[j] - p)) < r: + c += 1 + counts[i] = c + return counts diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 5704aa9..dae6dfa 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -28,3 +28,22 @@ def test_embed_lag2(): def test_embed_too_short_raises(): with pytest.raises(ValueError): est.embed(np.array([1.0]), np.array([2.0]), lag=1) + + +from entroscope.utils import knn + + +def test_kth_neighbor_distance_chebyshev(): + # Points on a line, spacing 1. For point at 0 with k=1, nearest is at 1 -> dist 1. + pts = np.array([[0.0], [1.0], [2.0], [5.0]]) + d = knn.kth_neighbor_distance(pts, k=1) + assert d[0] == pytest.approx(1.0) + assert d[3] == pytest.approx(3.0) # nearest to 5 is 2 + + +def test_count_within_radius_excludes_self_and_boundary(): + pts = np.array([[0.0], [1.0], [2.0]]) + # radius strictly greater than distance: count points with dist < r (excl self). + counts = knn.count_within_radius(pts, radii=np.array([1.5, 1.5, 1.5])) + # point 0: neighbor at 1 (dist1<1.5) -> 1 ; point 1: 0 and 2 -> 2 ; point2: 1 ->1 + np.testing.assert_array_equal(counts, [1, 2, 1]) From d20bc63c8f84847e26149e4173c63995023a2670 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 12:02:25 +0800 Subject: [PATCH 04/13] feat: add binned transfer-entropy estimator --- entroscope/_transfer_estimators.py | 31 ++++++++++++++++++++++++++++++ tests/test_transfer.py | 22 +++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py index d97f87e..48f65d9 100644 --- a/entroscope/_transfer_estimators.py +++ b/entroscope/_transfer_estimators.py @@ -29,3 +29,34 @@ def embed(x, y, lag=1): y_past = y[: n - lag] x_past = x[: n - lag] return y_future, y_past, x_past + + +def _hist_prob(*cols, bins): + """Joint probability table over the given equal-width-binned columns.""" + sample = np.column_stack(cols) + counts, _ = np.histogramdd(sample, bins=bins) + total = counts.sum() + return counts / total if total > 0 else counts + + +def te_binned(y_future, y_past, x_past, bins=6): + """Transfer entropy X->Y via equal-width histogram probabilities (base 2).""" + p_fpx = _hist_prob(y_future, y_past, x_past, bins=bins) # p(yf, yp, xp) + p_fp = p_fpx.sum(axis=2) # p(yf, yp) + p_px = p_fpx.sum(axis=0) # p(yp, xp) + p_p = p_fpx.sum(axis=(0, 2)) # p(yp) + + te = 0.0 + nf, npq, nx = p_fpx.shape + for i in range(nf): + for j in range(npq): + for k in range(nx): + pijk = p_fpx[i, j, k] + if pijk <= 0: + continue + denom = p_fp[i, j] * p_px[j, k] + numer = pijk * p_p[j] + if denom <= 0 or numer <= 0: + continue + te += pijk * np.log2(numer / denom) + return float(te) diff --git a/tests/test_transfer.py b/tests/test_transfer.py index dae6dfa..5726058 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -47,3 +47,25 @@ def test_count_within_radius_excludes_self_and_boundary(): counts = knn.count_within_radius(pts, radii=np.array([1.5, 1.5, 1.5])) # point 0: neighbor at 1 (dist1<1.5) -> 1 ; point 1: 0 and 2 -> 2 ; point2: 1 ->1 np.testing.assert_array_equal(counts, [1, 2, 1]) + + +def test_te_binned_independent_near_zero(): + rng = np.random.RandomState(0) + x = rng.randn(4000) + y = rng.randn(4000) + yf, yp, xp = est.embed(x, y, lag=1) + te = est.te_binned(yf, yp, xp, bins=6) + assert abs(te) < 0.05 + + +def test_te_binned_directional_coupling(): + rng = np.random.RandomState(1) + x = rng.randn(4000) + y = np.empty_like(x) + y[0] = rng.randn() + y[1:] = x[:-1] + 0.1 * rng.randn(3999) # y_t driven by x_{t-1} + yf, yp, xp = est.embed(x, y, lag=1) + te_xy = est.te_binned(yf, yp, xp, bins=6) + yf2, yp2, xp2 = est.embed(y, x, lag=1) # reverse direction + te_yx = est.te_binned(yf2, yp2, xp2, bins=6) + assert te_xy > te_yx + 0.1 From 292cbbabdbaac2b95eb5189fcec20505d1c36ab4 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 12:17:34 +0800 Subject: [PATCH 05/13] feat: add KSG (Kraskov) transfer-entropy estimator --- entroscope/_transfer_estimators.py | 32 ++++++++++++++++++++++++++++++ tests/test_transfer.py | 22 ++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py index 48f65d9..1ea32c1 100644 --- a/entroscope/_transfer_estimators.py +++ b/entroscope/_transfer_estimators.py @@ -10,6 +10,8 @@ """ import numpy as np +from scipy.special import digamma +from .utils import knn def embed(x, y, lag=1): @@ -60,3 +62,33 @@ def te_binned(y_future, y_past, x_past, bins=6): continue te += pijk * np.log2(numer / denom) return float(te) + + +_LN2 = np.log(2.0) + + +def te_ksg(y_future, y_past, x_past, k=4): + """Transfer entropy X->Y via the KSG (Kraskov) k-NN estimator, in bits. + + Joint space Z = (y_future, y_past, x_past). eps = Chebyshev distance to the + k-th neighbor in Z; marginal neighbor counts within eps give the digamma + terms. Returned in bits (nats / ln 2) to match the library's base-2 output. + """ + yf = np.asarray(y_future, dtype=float).reshape(-1, 1) + yp = np.asarray(y_past, dtype=float).reshape(-1, 1) + xp = np.asarray(x_past, dtype=float).reshape(-1, 1) + + z = np.column_stack([yf, yp, xp]) + eps = knn.kth_neighbor_distance(z, k=k) + + sp_yp = yp + sp_fp = np.column_stack([yf, yp]) + sp_px = np.column_stack([yp, xp]) + + n_yp = knn.count_within_radius(sp_yp, eps) + n_fp = knn.count_within_radius(sp_fp, eps) + n_px = knn.count_within_radius(sp_px, eps) + + terms = digamma(n_yp + 1) - digamma(n_fp + 1) - digamma(n_px + 1) + te_nats = digamma(k) + np.mean(terms) + return float(te_nats / _LN2) diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 5726058..c66a61d 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -69,3 +69,25 @@ def test_te_binned_directional_coupling(): yf2, yp2, xp2 = est.embed(y, x, lag=1) # reverse direction te_yx = est.te_binned(yf2, yp2, xp2, bins=6) assert te_xy > te_yx + 0.1 + + +def test_te_ksg_independent_near_zero(): + rng = np.random.RandomState(2) + x = rng.randn(2000) + y = rng.randn(2000) + yf, yp, xp = est.embed(x, y, lag=1) + te = est.te_ksg(yf, yp, xp, k=4) + assert abs(te) < 0.05 + + +def test_te_ksg_directional_coupling(): + rng = np.random.RandomState(3) + x = rng.randn(2000) + y = np.empty_like(x) + y[0] = rng.randn() + y[1:] = x[:-1] + 0.1 * rng.randn(1999) + yf, yp, xp = est.embed(x, y, lag=1) + te_xy = est.te_ksg(yf, yp, xp, k=4) + yf2, yp2, xp2 = est.embed(y, x, lag=1) + te_yx = est.te_ksg(yf2, yp2, xp2, k=4) + assert te_xy > te_yx + 0.1 From 4c211fa2ba64edfe7a823daea72ac6dd282ed1a8 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 12:38:05 +0800 Subject: [PATCH 06/13] test: add Gaussian closed-form correctness gate for transfer entropy Co-Authored-By: Claude Sonnet 4.6 --- tests/test_transfer.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_transfer.py b/tests/test_transfer.py index c66a61d..0caed14 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -91,3 +91,43 @@ def test_te_ksg_directional_coupling(): yf2, yp2, xp2 = est.embed(y, x, lag=1) te_yx = est.te_ksg(yf2, yp2, xp2, k=4) assert te_xy > te_yx + 0.1 + + +def _gaussian_system(n, a=0.5, b=0.7, sigma=0.5, seed=7): + rng = np.random.RandomState(seed) + x = rng.randn(n) + y = np.empty(n) + y[0] = rng.randn() + eps = sigma * rng.randn(n) + for t in range(1, n): + y[t] = a * y[t - 1] + b * x[t - 1] + eps[t] + return x, y + + +def test_gate_a_gaussian_closed_form_ksg(): + b, sigma = 0.7, 0.5 + true_te = 0.5 * np.log2(1 + (b ** 2) / (sigma ** 2)) + x, y = _gaussian_system(20000, a=0.5, b=b, sigma=sigma) + yf, yp, xp = est.embed(x, y, lag=1) + te = est.te_ksg(yf, yp, xp, k=4) + assert te == pytest.approx(true_te, abs=0.05) + + +def test_gate_a_gaussian_closed_form_binned(): + b, sigma = 0.7, 0.5 + true_te = 0.5 * np.log2(1 + (b ** 2) / (sigma ** 2)) + x, y = _gaussian_system(40000, a=0.5, b=b, sigma=sigma) + yf, yp, xp = est.embed(x, y, lag=1) + # bins=16 reduces discretisation bias enough to land within 0.1 bits of + # truth; bins=12 has ~0.11 bias (histogram underestimation of continuous MI) + te = est.te_binned(yf, yp, xp, bins=16) + assert te == pytest.approx(true_te, abs=0.1) + + +def test_binned_ksg_cross_check(): + # consistency, not correctness — backs up Gate A. + x, y = _gaussian_system(20000, a=0.4, b=0.6, sigma=0.6, seed=11) + yf, yp, xp = est.embed(x, y, lag=1) + te_b = est.te_binned(yf, yp, xp, bins=12) + te_k = est.te_ksg(yf, yp, xp, k=4) + assert te_b == pytest.approx(te_k, abs=0.1) From 37ffa81cc8a3e4d147dad762d4f903c2299702b6 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 12:45:16 +0800 Subject: [PATCH 07/13] test: add Kraskov-2004 MI correctness gate (Gate B) Co-Authored-By: Claude Sonnet 4.6 --- entroscope/_transfer_estimators.py | 17 +++++++++++++++++ tests/test_transfer.py | 14 ++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py index 1ea32c1..c5c68fc 100644 --- a/entroscope/_transfer_estimators.py +++ b/entroscope/_transfer_estimators.py @@ -92,3 +92,20 @@ def te_ksg(y_future, y_past, x_past, k=4): terms = digamma(n_yp + 1) - digamma(n_fp + 1) - digamma(n_px + 1) te_nats = digamma(k) + np.mean(terms) return float(te_nats / _LN2) + + +def ksg_mutual_information(a, b, k=4): + """KSG (Kraskov algorithm 1) mutual information between a and b, in NATS. + + a, b are (n, da) and (n, db). Returns I(a; b) in nats so it can be checked + against the analytic -0.5 ln(1 - r^2) for correlated Gaussians. + """ + a = np.asarray(a, dtype=float) + b = np.asarray(b, dtype=float) + n = len(a) + joint = np.column_stack([a, b]) + eps = knn.kth_neighbor_distance(joint, k=k) + n_a = knn.count_within_radius(a, eps) + n_b = knn.count_within_radius(b, eps) + mi = digamma(k) + digamma(n) - np.mean(digamma(n_a + 1) + digamma(n_b + 1)) + return float(mi) diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 0caed14..4c23c2e 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -131,3 +131,17 @@ def test_binned_ksg_cross_check(): te_b = est.te_binned(yf, yp, xp, bins=12) te_k = est.te_ksg(yf, yp, xp, k=4) assert te_b == pytest.approx(te_k, abs=0.1) + + +def test_gate_b_ksg_mutual_information_matches_analytic(): + # Correlated bivariate Gaussian; analytic MI = -0.5 ln(1 - r^2) nats. + r = 0.6 + rng = np.random.RandomState(13) + n = 20000 + a = rng.randn(n) + b = r * a + np.sqrt(1 - r ** 2) * rng.randn(n) + true_mi_nats = -0.5 * np.log(1 - r ** 2) + mi_nats = est.ksg_mutual_information( + a.reshape(-1, 1), b.reshape(-1, 1), k=4 + ) + assert mi_nats == pytest.approx(true_mi_nats, abs=0.03) From ac7373803f4476237c377cbeeb2fab6f6aa03e6c Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 13:52:19 +0800 Subject: [PATCH 08/13] feat: add public transfer.compute with sample-size guard --- entroscope/__init__.py | 3 ++- entroscope/transfer.py | 51 ++++++++++++++++++++++++++++++++++++++++++ tests/test_transfer.py | 35 +++++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 1 deletion(-) create mode 100644 entroscope/transfer.py diff --git a/entroscope/__init__.py b/entroscope/__init__.py index 4920d72..bc9fc27 100644 --- a/entroscope/__init__.py +++ b/entroscope/__init__.py @@ -1,6 +1,6 @@ """entroscope: the definitive entropy toolkit for time series data.""" -from . import shannon, permutation, spectral, sample, approximate, differential, multiscale +from . import shannon, permutation, spectral, sample, approximate, differential, multiscale, transfer from .utils import plot __version__ = "0.1.1" @@ -12,5 +12,6 @@ "approximate", "differential", "multiscale", + "transfer", "plot", ] diff --git a/entroscope/transfer.py b/entroscope/transfer.py new file mode 100644 index 0000000..8873032 --- /dev/null +++ b/entroscope/transfer.py @@ -0,0 +1,51 @@ +"""Transfer entropy — directional information flow X->Y (Schreiber 2000). + +The library's first bivariate measure: two series in, one directional scalar out. +Two estimators: ``method="ksg"`` (Kraskov k-NN, no binning, default) and +``method="binned"`` (histogram). Math lives in ``_transfer_estimators``. +""" + +import warnings + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +from . import _core +from . import _transfer_estimators as est + +_EMBED_DIM = 3 # history-length-1 joint space (y_future, y_past, x_past) + + +def _coerce_pair(x, y): + xa, _ = _core.as_array(x) + ya, yindex = _core.as_array(y) + if len(xa) != len(ya): + raise ValueError(f"x and y must be equal length ({len(xa)} != {len(ya)})") + return xa, ya, yindex + + +def _guard(n_samples, k, method): + if n_samples < 10 * _EMBED_DIM: + warnings.warn( + f"sample size ({n_samples}) is small for embedding dimension " + f"{_EMBED_DIM}; transfer-entropy estimate may be unreliable", + UserWarning, + stacklevel=3, + ) + + +def _estimate(xa, ya, k, lag, method, bins): + yf, yp, xp = est.embed(xa, ya, lag=lag) + _guard(len(yf), k, method) + if method == "ksg": + return est.te_ksg(yf, yp, xp, k=k) + if method == "binned": + return est.te_binned(yf, yp, xp, bins=bins) + raise ValueError(f"unknown method {method!r}; use 'ksg' or 'binned'") + + +def compute(x, y, *, k=4, lag=1, method="ksg", bins=6): + """Transfer entropy TE(X->Y) as a float (bits).""" + xa, ya, _ = _coerce_pair(x, y) + return _estimate(xa, ya, k, lag, method, bins) diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 4c23c2e..8f7ca9a 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -145,3 +145,38 @@ def test_gate_b_ksg_mutual_information_matches_analytic(): a.reshape(-1, 1), b.reshape(-1, 1), k=4 ) assert mi_nats == pytest.approx(true_mi_nats, abs=0.03) + + +import warnings +import pandas as pd +from entroscope import transfer + + +def test_compute_returns_float_ksg(): + rng = np.random.RandomState(5) + x = rng.randn(2000) + y = np.empty_like(x); y[0] = 0.0 + y[1:] = x[:-1] + 0.1 * rng.randn(1999) + val = transfer.compute(x, y, method="ksg", k=4) + assert isinstance(val, float) + assert val > 0.1 + + +def test_compute_method_binned(): + rng = np.random.RandomState(6) + x = rng.randn(3000) + y = rng.randn(3000) + val = transfer.compute(x, y, method="binned", bins=6) + assert isinstance(val, float) + + +def test_compute_length_mismatch_raises(): + with pytest.raises(ValueError): + transfer.compute(np.zeros(10), np.zeros(9)) + + +def test_compute_warns_on_thin_sample(): + x = np.arange(20, dtype=float) + y = np.arange(20, dtype=float) + 0.1 + with pytest.warns(UserWarning, match="sample size"): + transfer.compute(x, y, method="ksg", k=2) From 45310ccd09b85c29baf48c29ec58234c37c5d559 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 14:04:01 +0800 Subject: [PATCH 09/13] feat: add transfer.rolling/delta/plot bivariate driver --- entroscope/transfer.py | 43 ++++++++++++++++++++++++++++++++++++ tests/test_transfer.py | 49 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+) diff --git a/entroscope/transfer.py b/entroscope/transfer.py index 8873032..9db9dca 100644 --- a/entroscope/transfer.py +++ b/entroscope/transfer.py @@ -49,3 +49,46 @@ def compute(x, y, *, k=4, lag=1, method="ksg", bins=6): """Transfer entropy TE(X->Y) as a float (bits).""" xa, ya, _ = _coerce_pair(x, y) return _estimate(xa, ya, k, lag, method, bins) + + +def rolling(x, y, window=120, *, k=4, lag=1, method="ksg", bins=6): + """Rolling TE(X->Y) over sliding windows. Series-in -> Series-out.""" + xa, ya, yindex = _coerce_pair(x, y) + n = len(xa) + if window <= _EMBED_DIM + lag: + raise ValueError( + f"window ({window}) must exceed embedding+lag ({_EMBED_DIM + lag})" + ) + if window > n: + raise ValueError(f"window ({window}) larger than series length ({n})") + out = np.full(n, np.nan) + for end in range(window, n + 1): + wx = xa[end - window : end] + wy = ya[end - window : end] + out[end - 1] = _estimate(wx, wy, k, lag, method, bins) + return _core.wrap(out, yindex) + + +def delta(x, y, window=120, *, k=4, lag=1, method="ksg", bins=6): + """First difference of the rolling transfer entropy.""" + roll = rolling(x, y, window, k=k, lag=lag, method=method, bins=bins) + if isinstance(roll, pd.Series): + return roll.diff() + out = np.full_like(roll, np.nan) + out[1:] = np.diff(roll) + return out + + +def plot(x, y, window=120, *, k=4, lag=1, method="ksg", bins=6, title=None): + """Figure of rolling TE(X->Y). Never calls plt.show().""" + roll = rolling(x, y, window, k=k, lag=lag, method=method, bins=bins) + fig, ax = plt.subplots(figsize=(10, 4)) + if isinstance(roll, pd.Series): + ax.plot(roll.index, roll.to_numpy()) + else: + ax.plot(range(len(roll)), roll) + ax.set_title(title or f"Rolling transfer entropy (window={window})") + ax.set_xlabel("position") + ax.set_ylabel("transfer entropy (bits)") + fig.tight_layout() + return fig diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 8f7ca9a..1438c18 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -180,3 +180,52 @@ def test_compute_warns_on_thin_sample(): y = np.arange(20, dtype=float) + 0.1 with pytest.warns(UserWarning, match="sample size"): transfer.compute(x, y, method="ksg", k=2) + + +import matplotlib + + +def _coupled(n, seed=9): + rng = np.random.RandomState(seed) + x = rng.randn(n) + y = np.empty_like(x); y[0] = 0.0 + y[1:] = x[:-1] + 0.1 * rng.randn(n - 1) + return x, y + + +def test_rolling_array_in_array_out(): + x, y = _coupled(400) + out = transfer.rolling(x, y, window=120, method="ksg", k=4) + assert isinstance(out, np.ndarray) + assert len(out) == len(y) + assert np.isnan(out[:119]).all() + assert not np.isnan(out[-1]) + + +def test_rolling_series_in_series_out_preserves_index(): + x, y = _coupled(400) + idx = pd.date_range("2025-01-01", periods=400, freq="D") + sx, sy = pd.Series(x, index=idx), pd.Series(y, index=idx) + out = transfer.rolling(sx, sy, window=120, method="ksg", k=4) + assert isinstance(out, pd.Series) + assert out.index.equals(idx) + + +def test_rolling_window_too_small_raises(): + x, y = _coupled(50) + with pytest.raises(ValueError): + transfer.rolling(x, y, window=2) + + +def test_delta_is_first_difference(): + x, y = _coupled(400) + roll = transfer.rolling(x, y, window=120, method="ksg", k=4) + d = transfer.delta(x, y, window=120, method="ksg", k=4) + assert len(d) == len(roll) + np.testing.assert_allclose(d[121:], np.diff(roll)[120:], rtol=1e-9) + + +def test_plot_returns_figure(): + x, y = _coupled(400) + fig = transfer.plot(x, y, window=120, method="ksg", k=4) + assert isinstance(fig, matplotlib.figure.Figure) From 804fe7d9229c07f5a5fddda8d0a7323be78178f0 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 14:07:12 +0800 Subject: [PATCH 10/13] feat: add transfer-entropy information-flow example Co-Authored-By: Claude Sonnet 4.6 --- examples/_synthetic.py | 16 +++++++++++++++ examples/information_flow.py | 40 ++++++++++++++++++++++++++++++++++++ tests/test_examples.py | 14 ++++++++++++- 3 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 examples/information_flow.py diff --git a/examples/_synthetic.py b/examples/_synthetic.py index 1d02139..4585678 100644 --- a/examples/_synthetic.py +++ b/examples/_synthetic.py @@ -163,3 +163,19 @@ def correlated_assets(seed=14): b = np.concatenate([b_stable, b_unstable]) idx = pd.date_range("2025-01-01", periods=len(a), freq="B") return pd.DataFrame({"a": a, "b": b}, index=idx) + + +def lead_lag_assets(seed=15): + """Two assets where A leads B by one step, plus reverse-direction noise. + + A drives B (B_t depends on A_{t-1}), so transfer entropy should be clearly + larger in the A->B direction than B->A. Returns a 2-column DataFrame (a, b). + """ + rng = _rng(seed) + n = 600 + a = rng.randn(n) + b = np.empty(n) + b[0] = rng.randn() + b[1:] = a[:-1] + 0.3 * rng.randn(n - 1) + idx = pd.date_range("2025-01-01", periods=n, freq="B") + return pd.DataFrame({"a": a, "b": b}, index=idx) diff --git a/examples/information_flow.py b/examples/information_flow.py new file mode 100644 index 0000000..9bd6122 --- /dev/null +++ b/examples/information_flow.py @@ -0,0 +1,40 @@ +"""Transfer-entropy example: directional information flow between two assets. + +Run: python examples/information_flow.py + +Transfer entropy is directional: TE(A->B) measures how much A's past helps +predict B's future beyond B's own past. Here A leads B by construction, so +TE(A->B) should clearly exceed TE(B->A). Synthetic data lives in ``_synthetic``; +swap ``data.lead_lag_assets()`` for a 2-column frame of your own returns. +""" + +from entroscope import transfer + +try: + from . import _synthetic as data +except ImportError: + import _synthetic as data + + +def information_flow(): + """Directional transfer entropy between a leading and a lagging asset.""" + df = data.lead_lag_assets() + te_ab = transfer.compute(df["a"], df["b"], method="ksg", k=4) + te_ba = transfer.compute(df["b"], df["a"], method="ksg", k=4) + print("Information flow (KSG transfer entropy, bits)") + print(f" TE(A -> B): {te_ab:6.3f}") + print(f" TE(B -> A): {te_ba:6.3f}") + print(f" net A->B : {te_ab - te_ba:6.3f}") + print(" -> A leads B by construction, so TE(A->B) should dominate\n") + return df, te_ab, te_ba + + +def main(): + print("=" * 60) + print("TRANSFER ENTROPY — DIRECTIONAL INFORMATION FLOW") + print("=" * 60) + information_flow() + + +if __name__ == "__main__": + main() diff --git a/tests/test_examples.py b/tests/test_examples.py index ba59bd3..4ad7041 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -24,7 +24,9 @@ def _load(module_name): return module -@pytest.mark.parametrize("name", ["medical", "business", "correlation_stability"]) +@pytest.mark.parametrize( + "name", ["medical", "business", "correlation_stability", "information_flow"] +) def test_example_runs(name, capsys): module = _load(name) module.main() # must not raise @@ -71,3 +73,13 @@ def test_correlated_assets_has_regime_breakdown(): stable_corr = df["a"].iloc[:half].corr(df["b"].iloc[:half]) broken_corr = df["a"].iloc[half:].corr(df["b"].iloc[half:]) assert stable_corr > broken_corr + + +def test_lead_lag_assets_is_directional_frame(): + import pandas as pd + + data = _load("_synthetic") + df = data.lead_lag_assets() + assert isinstance(df, pd.DataFrame) + assert list(df.columns) == ["a", "b"] + assert len(df) > 0 From 75e906ba013ea59086d8d38dd2f93fd3182c83b0 Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 14:14:02 +0800 Subject: [PATCH 11/13] chore: lint cleanup for transfer entropy (ruff) Remove unused imports and split semicolon statements flagged by ruff in the transfer tests; drop a pre-existing unused pandas import in the correlation-stability example. Full suite: 132 passed. Co-Authored-By: Claude Opus 4.8 --- examples/correlation_stability.py | 2 -- tests/test_transfer.py | 7 ++++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/examples/correlation_stability.py b/examples/correlation_stability.py index b872d24..01880f4 100644 --- a/examples/correlation_stability.py +++ b/examples/correlation_stability.py @@ -19,8 +19,6 @@ on real data. """ -import pandas as pd - from entroscope import permutation try: diff --git a/tests/test_transfer.py b/tests/test_transfer.py index 1438c18..c4db263 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -147,7 +147,6 @@ def test_gate_b_ksg_mutual_information_matches_analytic(): assert mi_nats == pytest.approx(true_mi_nats, abs=0.03) -import warnings import pandas as pd from entroscope import transfer @@ -155,7 +154,8 @@ def test_gate_b_ksg_mutual_information_matches_analytic(): def test_compute_returns_float_ksg(): rng = np.random.RandomState(5) x = rng.randn(2000) - y = np.empty_like(x); y[0] = 0.0 + y = np.empty_like(x) + y[0] = 0.0 y[1:] = x[:-1] + 0.1 * rng.randn(1999) val = transfer.compute(x, y, method="ksg", k=4) assert isinstance(val, float) @@ -188,7 +188,8 @@ def test_compute_warns_on_thin_sample(): def _coupled(n, seed=9): rng = np.random.RandomState(seed) x = rng.randn(n) - y = np.empty_like(x); y[0] = 0.0 + y = np.empty_like(x) + y[0] = 0.0 y[1:] = x[:-1] + 0.1 * rng.randn(n - 1) return x, y From a2419619afd8a8af3e3dee8901c54bc03854001a Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 14:15:30 +0800 Subject: [PATCH 12/13] docs: clarify transfer-entropy correctness authorities and KSG rolling perf Co-Authored-By: Claude Opus 4.8 --- entroscope/_transfer_estimators.py | 6 ++++-- entroscope/transfer.py | 6 +++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py index c5c68fc..21383d4 100644 --- a/entroscope/_transfer_estimators.py +++ b/entroscope/_transfer_estimators.py @@ -5,8 +5,10 @@ embedding is verified in isolation (see tests) so a bug here cannot silently fool both estimators. -NOTE: imports for later estimators (digamma from scipy, knn from .utils) are -added in the tasks that implement the binned and KSG estimators respectively. +Correctness is established by three independent test authorities (see +tests/test_transfer.py): the bivariate-Gaussian closed form (magnitude), the +Kraskov-2004 analytic mutual information (KSG core), and a hand-checked isolated +embedding test. """ import numpy as np diff --git a/entroscope/transfer.py b/entroscope/transfer.py index 9db9dca..a2e3e4d 100644 --- a/entroscope/transfer.py +++ b/entroscope/transfer.py @@ -52,7 +52,11 @@ def compute(x, y, *, k=4, lag=1, method="ksg", bins=6): def rolling(x, y, window=120, *, k=4, lag=1, method="ksg", bins=6): - """Rolling TE(X->Y) over sliding windows. Series-in -> Series-out.""" + """Rolling TE(X->Y) over sliding windows. Series-in -> Series-out. + + Note: with method="ksg" this runs a k-NN search per window, so it is + noticeably slower than the single-series rolling measures on long inputs. + """ xa, ya, yindex = _coerce_pair(x, y) n = len(xa) if window <= _EMBED_DIM + lag: From 1190cc37c34f06b599882a3a830c4c5d718779ac Mon Sep 17 00:00:00 2001 From: jjscripts Date: Wed, 3 Jun 2026 14:20:05 +0800 Subject: [PATCH 13/13] style: apply ruff format to transfer-entropy files CI runs 'ruff format --check'; the new files weren't formatter-clean. Cosmetic only (import wrapping, comment spacing, line reflow), no logic change. Co-Authored-By: Claude Opus 4.8 --- entroscope/__init__.py | 11 ++++++++++- entroscope/_transfer_estimators.py | 8 ++++---- entroscope/transfer.py | 4 +--- tests/test_transfer.py | 20 +++++++++----------- 4 files changed, 24 insertions(+), 19 deletions(-) diff --git a/entroscope/__init__.py b/entroscope/__init__.py index bc9fc27..19c5370 100644 --- a/entroscope/__init__.py +++ b/entroscope/__init__.py @@ -1,6 +1,15 @@ """entroscope: the definitive entropy toolkit for time series data.""" -from . import shannon, permutation, spectral, sample, approximate, differential, multiscale, transfer +from . import ( + shannon, + permutation, + spectral, + sample, + approximate, + differential, + multiscale, + transfer, +) from .utils import plot __version__ = "0.1.1" diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py index 21383d4..0e5565a 100644 --- a/entroscope/_transfer_estimators.py +++ b/entroscope/_transfer_estimators.py @@ -45,10 +45,10 @@ def _hist_prob(*cols, bins): def te_binned(y_future, y_past, x_past, bins=6): """Transfer entropy X->Y via equal-width histogram probabilities (base 2).""" - p_fpx = _hist_prob(y_future, y_past, x_past, bins=bins) # p(yf, yp, xp) - p_fp = p_fpx.sum(axis=2) # p(yf, yp) - p_px = p_fpx.sum(axis=0) # p(yp, xp) - p_p = p_fpx.sum(axis=(0, 2)) # p(yp) + p_fpx = _hist_prob(y_future, y_past, x_past, bins=bins) # p(yf, yp, xp) + p_fp = p_fpx.sum(axis=2) # p(yf, yp) + p_px = p_fpx.sum(axis=0) # p(yp, xp) + p_p = p_fpx.sum(axis=(0, 2)) # p(yp) te = 0.0 nf, npq, nx = p_fpx.shape diff --git a/entroscope/transfer.py b/entroscope/transfer.py index a2e3e4d..09d50c3 100644 --- a/entroscope/transfer.py +++ b/entroscope/transfer.py @@ -60,9 +60,7 @@ def rolling(x, y, window=120, *, k=4, lag=1, method="ksg", bins=6): xa, ya, yindex = _coerce_pair(x, y) n = len(xa) if window <= _EMBED_DIM + lag: - raise ValueError( - f"window ({window}) must exceed embedding+lag ({_EMBED_DIM + lag})" - ) + raise ValueError(f"window ({window}) must exceed embedding+lag ({_EMBED_DIM + lag})") if window > n: raise ValueError(f"window ({window}) larger than series length ({n})") out = np.full(n, np.nan) diff --git a/tests/test_transfer.py b/tests/test_transfer.py index c4db263..6cf2b17 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -11,8 +11,8 @@ def test_embed_hand_computed_lag1(): y_future, y_past, x_past = est.embed(x, y, lag=1) # t = 1..5 np.testing.assert_array_equal(y_future, [21.0, 22.0, 23.0, 24.0, 25.0]) - np.testing.assert_array_equal(y_past, [20.0, 21.0, 22.0, 23.0, 24.0]) - np.testing.assert_array_equal(x_past, [10.0, 11.0, 12.0, 13.0, 14.0]) + np.testing.assert_array_equal(y_past, [20.0, 21.0, 22.0, 23.0, 24.0]) + np.testing.assert_array_equal(x_past, [10.0, 11.0, 12.0, 13.0, 14.0]) def test_embed_lag2(): @@ -21,8 +21,8 @@ def test_embed_lag2(): y_future, y_past, x_past = est.embed(x, y, lag=2) # t = 2..5 -> 4 samples np.testing.assert_array_equal(y_future, [12.0, 13.0, 14.0, 15.0]) - np.testing.assert_array_equal(y_past, [10.0, 11.0, 12.0, 13.0]) - np.testing.assert_array_equal(x_past, [0.0, 1.0, 2.0, 3.0]) + np.testing.assert_array_equal(y_past, [10.0, 11.0, 12.0, 13.0]) + np.testing.assert_array_equal(x_past, [0.0, 1.0, 2.0, 3.0]) def test_embed_too_short_raises(): @@ -106,7 +106,7 @@ def _gaussian_system(n, a=0.5, b=0.7, sigma=0.5, seed=7): def test_gate_a_gaussian_closed_form_ksg(): b, sigma = 0.7, 0.5 - true_te = 0.5 * np.log2(1 + (b ** 2) / (sigma ** 2)) + true_te = 0.5 * np.log2(1 + (b**2) / (sigma**2)) x, y = _gaussian_system(20000, a=0.5, b=b, sigma=sigma) yf, yp, xp = est.embed(x, y, lag=1) te = est.te_ksg(yf, yp, xp, k=4) @@ -115,7 +115,7 @@ def test_gate_a_gaussian_closed_form_ksg(): def test_gate_a_gaussian_closed_form_binned(): b, sigma = 0.7, 0.5 - true_te = 0.5 * np.log2(1 + (b ** 2) / (sigma ** 2)) + true_te = 0.5 * np.log2(1 + (b**2) / (sigma**2)) x, y = _gaussian_system(40000, a=0.5, b=b, sigma=sigma) yf, yp, xp = est.embed(x, y, lag=1) # bins=16 reduces discretisation bias enough to land within 0.1 bits of @@ -139,11 +139,9 @@ def test_gate_b_ksg_mutual_information_matches_analytic(): rng = np.random.RandomState(13) n = 20000 a = rng.randn(n) - b = r * a + np.sqrt(1 - r ** 2) * rng.randn(n) - true_mi_nats = -0.5 * np.log(1 - r ** 2) - mi_nats = est.ksg_mutual_information( - a.reshape(-1, 1), b.reshape(-1, 1), k=4 - ) + b = r * a + np.sqrt(1 - r**2) * rng.randn(n) + true_mi_nats = -0.5 * np.log(1 - r**2) + mi_nats = est.ksg_mutual_information(a.reshape(-1, 1), b.reshape(-1, 1), k=4) assert mi_nats == pytest.approx(true_mi_nats, abs=0.03)