diff --git a/entroscope/__init__.py b/entroscope/__init__.py index 4920d72..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 +from . import ( + shannon, + permutation, + spectral, + sample, + approximate, + differential, + multiscale, + transfer, +) from .utils import plot __version__ = "0.1.1" @@ -12,5 +21,6 @@ "approximate", "differential", "multiscale", + "transfer", "plot", ] diff --git a/entroscope/_transfer_estimators.py b/entroscope/_transfer_estimators.py new file mode 100644 index 0000000..0e5565a --- /dev/null +++ b/entroscope/_transfer_estimators.py @@ -0,0 +1,113 @@ +"""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. + +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 +from scipy.special import digamma +from .utils import knn + + +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 + + +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) + + +_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) + + +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/entroscope/transfer.py b/entroscope/transfer.py new file mode 100644 index 0000000..09d50c3 --- /dev/null +++ b/entroscope/transfer.py @@ -0,0 +1,96 @@ +"""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) + + +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. + + 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: + 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/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/examples/_synthetic.py b/examples/_synthetic.py index 22e1a69..4585678 100644 --- a/examples/_synthetic.py +++ b/examples/_synthetic.py @@ -131,3 +131,51 @@ 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) + + +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/correlation_stability.py b/examples/correlation_stability.py new file mode 100644 index 0000000..01880f4 --- /dev/null +++ b/examples/correlation_stability.py @@ -0,0 +1,86 @@ +"""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. +""" + +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/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 4d80071..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"]) +@pytest.mark.parametrize( + "name", ["medical", "business", "correlation_stability", "information_flow"] +) def test_example_runs(name, capsys): module = _load(name) module.main() # must not raise @@ -49,3 +51,35 @@ 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 + + +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 diff --git a/tests/test_transfer.py b/tests/test_transfer.py new file mode 100644 index 0000000..6cf2b17 --- /dev/null +++ b/tests/test_transfer.py @@ -0,0 +1,230 @@ +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 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]) + + +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 + + +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 + + +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) + + +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) + + +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) + + +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)