diff --git a/entroscope/__init__.py b/entroscope/__init__.py index 19c5370..704afa6 100644 --- a/entroscope/__init__.py +++ b/entroscope/__init__.py @@ -9,6 +9,7 @@ differential, multiscale, transfer, + divergence, ) from .utils import plot @@ -22,5 +23,6 @@ "differential", "multiscale", "transfer", + "divergence", "plot", ] diff --git a/entroscope/divergence.py b/entroscope/divergence.py new file mode 100644 index 0000000..3c34246 --- /dev/null +++ b/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 diff --git a/examples/_synthetic.py b/examples/_synthetic.py index 4585678..9106f64 100644 --- a/examples/_synthetic.py +++ b/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 diff --git a/examples/drift.py b/examples/drift.py new file mode 100644 index 0000000..6cb6558 --- /dev/null +++ b/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() diff --git a/tests/test_divergence.py b/tests/test_divergence.py new file mode 100644 index 0000000..2082ff5 --- /dev/null +++ b/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") diff --git a/tests/test_examples.py b/tests/test_examples.py index 4ad7041..49955a3 100644 --- a/tests/test_examples.py +++ b/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