Add divergence module (KL + Jensen-Shannon) for distribution comparison

entroscope/__init__.py

@@ -9,6 +9,7 @@
     differential,
     multiscale,
     transfer,
+    divergence,
 )
 from .utils import plot
 
@@ -22,5 +23,6 @@
     "differential",
     "multiscale",
     "transfer",
+    "divergence",
     "plot",
 ]

entroscope/divergence.py

@@ -0,0 +1,91 @@
+"""Divergence between two distributions — KL and Jensen-Shannon.
+
+Distribution-vs-distribution measures: how far apart are two samples? Inputs are
+two raw sample arrays/Series; both are histogrammed over a SHARED range so the
+resulting probability vectors are aligned and comparable. Base-2 (bits) to match
+the rest of the library.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.special import rel_entr
+
+from . import _core
+
+_EPS = 1e-12
+
+
+def _binned_probs(p, q, bins):
+    """Histogram p and q over SHARED edges; return two aligned probability vectors.
+
+    Edges span the combined [min, max] of both samples so the bins line up. A
+    degenerate (zero-width) combined range falls back to a single bin.
+    """
+    if bins <= 0:
+        raise ValueError("bins must be a positive integer")
+    p = np.asarray(p, dtype=float)
+    q = np.asarray(q, dtype=float)
+    lo = min(p.min(), q.min())
+    hi = max(p.max(), q.max())
+    if hi <= lo:  # all values identical across both samples
+        return np.array([1.0]), np.array([1.0])
+    edges = np.linspace(lo, hi, bins + 1)
+    cp, _ = np.histogram(p, bins=edges)
+    cq, _ = np.histogram(q, bins=edges)
+    pp = cp / cp.sum()
+    qq = cq / cq.sum()
+    return pp, qq
+
+
+_LN2 = np.log(2.0)
+
+
+def _kl_bits(pp, qq):
+    """KL(pp || qq) in bits, with epsilon smoothing so the result stays finite."""
+    pp = pp + _EPS
+    qq = qq + _EPS
+    pp = pp / pp.sum()
+    qq = qq / qq.sum()
+    return float(np.sum(rel_entr(pp, qq)) / _LN2)
+
+
+def kl(p, q, bins=10):
+    """Kullback-Leibler divergence KL(p || q) in bits (directional).
+
+    p, q are raw samples; both are binned over a shared range. Epsilon smoothing
+    keeps the result finite even when q has an empty bin where p has mass.
+    """
+    pa, _ = _core.as_array(p)
+    qa, _ = _core.as_array(q)
+    pp, qq = _binned_probs(pa, qa, bins)
+    return _kl_bits(pp, qq)
+
+
+def js(p, q, bins=10):
+    """Jensen-Shannon divergence in bits, symmetric and bounded to [0, 1].
+
+    js(p, q) = 0.5*KL(p || m) + 0.5*KL(q || m) with m = (p + q) / 2.
+    """
+    pa, _ = _core.as_array(p)
+    qa, _ = _core.as_array(q)
+    pp, qq = _binned_probs(pa, qa, bins)
+    m = 0.5 * (pp + qq)
+    return 0.5 * _kl_bits(pp, m) + 0.5 * _kl_bits(qq, m)
+
+
+def plot(p, q, bins=10, title=None):
+    """Overlay the two binned distributions. Returns a Figure; never shows it."""
+    pa, _ = _core.as_array(p)
+    qa, _ = _core.as_array(q)
+    pp, qq = _binned_probs(pa, qa, bins)
+    x = np.arange(len(pp))
+    fig, ax = plt.subplots(figsize=(10, 4))
+    width = 0.4
+    ax.bar(x - width / 2, pp, width=width, label="p", alpha=0.8)
+    ax.bar(x + width / 2, qq, width=width, label="q", alpha=0.8)
+    ax.set_title(title or f"Binned distributions (bins={bins})")
+    ax.set_xlabel("bin")
+    ax.set_ylabel("probability")
+    ax.legend()
+    fig.tight_layout()
+    return fig

examples/_synthetic.py

@@ -179,3 +179,16 @@ def lead_lag_assets(seed=15):
     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)
+
+
+def distribution_drift(seed=16):
+    """A reference batch and a drifted batch of the same measurement.
+
+    The reference is centered and tight; the drifted batch shifts and widens, as
+    production data does when the world changes under a trained model. Returns a
+    (reference, drifted) tuple of pd.Series.
+    """
+    rng = _rng(seed)
+    reference = pd.Series(100 + 5 * rng.randn(2000), name="reference")
+    drifted = pd.Series(108 + 9 * rng.randn(2000), name="drifted")
+    return reference, drifted

examples/drift.py

@@ -0,0 +1,43 @@
+"""Distribution-drift example: KL and Jensen-Shannon between two batches.
+
+Run:  python examples/drift.py
+
+Compares a reference batch (e.g. training data) against a drifted batch (e.g.
+later production data). Higher divergence means the distributions have moved
+apart. JS is symmetric and bounded to [0, 1] bits; KL is directional. Synthetic
+data lives in ``_synthetic``; swap ``data.distribution_drift()`` for two batches
+of your own values.
+"""
+
+from entroscope import divergence
+
+try:
+    from . import _synthetic as data
+except ImportError:
+    import _synthetic as data
+
+
+def drift_detection():
+    """Divergence between a reference batch and a drifted batch."""
+    reference, drifted = data.distribution_drift()
+    kl = divergence.kl(reference, drifted, bins=30)
+    js = divergence.js(reference, drifted, bins=30)
+    js_self = divergence.js(reference, reference.copy(), bins=30)
+    print("Distribution drift (divergence, bits)")
+    print(f"  KL(reference || drifted): {kl:6.3f}")
+    print(f"  JS(reference,  drifted) : {js:6.3f}")
+    print(f"  JS(reference,  itself)  : {js_self:6.3f}")
+    print("  -> JS near 0 means no drift; larger means the batches have moved apart")
+    print("     (KL and JS are relative-entropy measures between distributions)\n")
+    return reference, drifted, kl, js
+
+
+def main():
+    print("=" * 60)
+    print("DISTRIBUTION DRIFT — KL & JENSEN-SHANNON DIVERGENCE")
+    print("=" * 60)
+    drift_detection()
+
+
+if __name__ == "__main__":
+    main()

tests/test_divergence.py

@@ -0,0 +1,104 @@
+import numpy as np
+import pytest
+
+from entroscope import divergence as dv
+
+
+def test_binned_probs_shared_edges_align_lengths():
+    # Two samples over DIFFERENT ranges must still produce equal-length vectors.
+    p = np.array([0.0, 1.0, 2.0])
+    q = np.array([10.0, 11.0, 12.0])
+    pp, qq = dv._binned_probs(p, q, bins=5)
+    assert len(pp) == len(qq) == 5
+    assert pp.sum() == pytest.approx(1.0)
+    assert qq.sum() == pytest.approx(1.0)
+    # p occupies the low bins, q the high bins -> their mass is in different bins.
+    assert pp[0] > 0 and qq[-1] > 0
+    assert pp[-1] == 0 and qq[0] == 0
+
+
+def test_binned_probs_degenerate_constant_input():
+    p = np.array([5.0, 5.0, 5.0])
+    q = np.array([5.0, 5.0, 5.0])
+    pp, qq = dv._binned_probs(p, q, bins=10)
+    assert pp.sum() == pytest.approx(1.0)
+    assert qq.sum() == pytest.approx(1.0)
+    assert len(pp) == len(qq)
+
+
+def test_kl_identical_is_zero():
+    rng = np.random.RandomState(0)
+    x = rng.randn(5000)
+    assert dv.kl(x, x.copy(), bins=20) == pytest.approx(0.0, abs=1e-6)
+
+
+def test_kl_known_value():
+    # Construct samples that bin into known probabilities over 2 bins on [0,1).
+    # p: 75% in bin0, 25% in bin1 ; q: 50/50.
+    p = np.array([0.1, 0.1, 0.1, 0.6])  # 3 in [0,0.5), 1 in [0.5,1)
+    q = np.array([0.1, 0.6])  # 1 in each
+    # KL = 0.75*log2(0.75/0.5) + 0.25*log2(0.25/0.5) = 0.75*0.585 - 0.25 = 0.1887
+    val = dv.kl(p, q, bins=2)
+    assert val == pytest.approx(0.1887, abs=0.005)
+
+
+def test_kl_finite_when_q_has_empty_bin():
+    # q has no mass where p does -> would be +inf without smoothing.
+    p = np.array([0.0, 0.0, 1.0, 1.0])
+    q = np.array([0.0, 0.0, 0.0, 0.0])
+    val = dv.kl(p, q, bins=4)
+    assert np.isfinite(val)
+
+
+def test_kl_bad_bins_raises():
+    with pytest.raises(ValueError):
+        dv.kl(np.array([1.0, 2.0]), np.array([1.0, 2.0]), bins=0)
+
+
+def test_js_identical_is_zero():
+    rng = np.random.RandomState(1)
+    x = rng.randn(5000)
+    assert dv.js(x, x.copy(), bins=20) == pytest.approx(0.0, abs=1e-6)
+
+
+def test_js_is_symmetric():
+    rng = np.random.RandomState(2)
+    a = rng.randn(3000)
+    b = rng.randn(3000) + 2.0
+    assert dv.js(a, b, bins=20) == pytest.approx(dv.js(b, a, bins=20), abs=1e-9)
+
+
+def test_js_bounded_zero_to_one():
+    rng = np.random.RandomState(3)
+    a = rng.randn(3000)
+    b = rng.randn(3000) + 50.0  # effectively disjoint supports
+    val = dv.js(a, b, bins=20)
+    assert 0.0 <= val <= 1.0
+    assert val > 0.9  # near the upper bound for disjoint distributions
+
+
+def test_js_known_value():
+    # Two 2-bin distributions: p=[1,0], q=[0,1] (disjoint) -> JS = 1 bit exactly.
+    p = np.array([0.1, 0.1])  # both in bin0 of shared range [0.1,0.9]
+    q = np.array([0.9, 0.9])  # both in bin1
+    val = dv.js(p, q, bins=2)
+    assert val == pytest.approx(1.0, abs=1e-3)
+
+
+import matplotlib
+
+
+def test_plot_returns_figure():
+    rng = np.random.RandomState(4)
+    a = rng.randn(1000)
+    b = rng.randn(1000) + 1.0
+    fig = dv.plot(a, b, bins=20)
+    assert isinstance(fig, matplotlib.figure.Figure)
+
+
+def test_divergence_in_public_api():
+    import entroscope
+
+    assert "divergence" in entroscope.__all__
+    assert hasattr(entroscope.divergence, "kl")
+    assert hasattr(entroscope.divergence, "js")

tests/test_examples.py

@@ -25,7 +25,8 @@ def _load(module_name):
 
 
 @pytest.mark.parametrize(
-    "name", ["medical", "business", "correlation_stability", "information_flow"]
+    "name",
+    ["medical", "business", "correlation_stability", "information_flow", "drift"],
 )
 def test_example_runs(name, capsys):
     module = _load(name)
@@ -83,3 +84,13 @@ def test_lead_lag_assets_is_directional_frame():
     assert isinstance(df, pd.DataFrame)
     assert list(df.columns) == ["a", "b"]
     assert len(df) > 0
+
+
+def test_distribution_drift_returns_two_series():
+    import pandas as pd
+
+    data = _load("_synthetic")
+    reference, drifted = data.distribution_drift()
+    assert isinstance(reference, pd.Series)
+    assert isinstance(drifted, pd.Series)
+    assert len(reference) > 0 and len(drifted) > 0