Add transfer entropy (KSG + binned) — first bivariate information-flow measure

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",
 ]

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)

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

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

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)