Yogesh18018/drift-detection

1"""
2Statistical Drift Detection Module
3===================================
4Implements multiple statistical tests for detecting data drift between
5reference (training) and production (live) distributions.
6
7Supported Tests:
8    - Kolmogorov-Smirnov (KS) Test
9    - Population Stability Index (PSI)
10    - Jensen-Shannon (JS) Divergence
11"""
12
13import numpy as np
14import pandas as pd
15from scipy import stats
16from scipy.spatial.distance import jensenshannon
17
18
19class DriftDetector:
20    """Detects distribution drift between reference and production datasets."""
21
22    # Thresholds 
23    KS_P_THRESHOLD = 0.05          # p-value below this → drift detected
24    PSI_THRESHOLD = 0.2            # PSI above this → significant drift
25    JS_THRESHOLD = 0.1             # JS divergence above this → drift
26
27    def __init__(self, n_bins: int = 20):
28        self.n_bins = n_bins
29
30    # Kolmogorov-Smirnov Test 
31    def ks_test(
32        self,
33        reference: pd.DataFrame,
34        production: pd.DataFrame,
35    ) -> dict:
36        """
37        Run a two-sample KS test on every shared numeric feature.
38
39        Returns
40        -------
41        dict  {feature: {"statistic", "p_value", "drift"}}
42        """
43        results = {}
44        features = reference.select_dtypes(include=[np.number]).columns
45        for feat in features:
46            stat, p_val = stats.ks_2samp(
47                reference[feat].dropna(),
48                production[feat].dropna(),
49            )
50            results[feat] = {
51                "statistic": round(float(stat), 6),
52                "p_value": round(float(p_val), 6),
53                "drift": p_val < self.KS_P_THRESHOLD,
54            }
55        return results
56
57    # Population Stability Index 
58    def psi(
59        self,
60        reference: pd.DataFrame,
61        production: pd.DataFrame,
62    ) -> dict:
63        """
64        Compute PSI for each numeric feature.
65
66        Returns
67        -------
68        dict  {feature: {"psi_value", "drift"}}
69        """
70        results = {}
71        features = reference.select_dtypes(include=[np.number]).columns
72        for feat in features:
73            psi_val = self._compute_psi(
74                reference[feat].dropna().values,
75                production[feat].dropna().values,
76            )
77            results[feat] = {
78                "psi_value": round(float(psi_val), 6),
79                "drift": psi_val > self.PSI_THRESHOLD,
80            }
81        return results
82
83    def _compute_psi(self, reference: np.ndarray, production: np.ndarray) -> float:
84        """Calculate PSI between two 1-D arrays."""
85        eps = 1e-4
86        # Use reference quantiles so bins are consistent
87        breakpoints = np.linspace(0, 100, self.n_bins + 1)
88        edges = np.percentile(reference, breakpoints)
89        edges[0] = -np.inf
90        edges[-1] = np.inf
91
92        ref_counts = np.histogram(reference, bins=edges)[0].astype(float)
93        prod_counts = np.histogram(production, bins=edges)[0].astype(float)
94
95        ref_pct = ref_counts / ref_counts.sum() + eps
96        prod_pct = prod_counts / prod_counts.sum() + eps
97
98        psi_value = np.sum((prod_pct - ref_pct) * np.log(prod_pct / ref_pct))
99        return psi_value
100
101    # Jensen-Shannon Divergence 
102    def js_divergence(
103        self,
104        reference: pd.DataFrame,
105        production: pd.DataFrame,
106    ) -> dict:
107        """
108        Compute JS divergence for each numeric feature.
109
110        Returns
111        -------
112        dict  {feature: {"js_value", "drift"}}
113        """
114        results = {}
115        features = reference.select_dtypes(include=[np.number]).columns
116        for feat in features:
117            js_val = self._compute_js(
118                reference[feat].dropna().values,
119                production[feat].dropna().values,
120            )
121            results[feat] = {
122                "js_value": round(float(js_val), 6),
123                "drift": js_val > self.JS_THRESHOLD,
124            }
125        return results
126
127    def _compute_js(self, reference: np.ndarray, production: np.ndarray) -> float:
128        """Calculate JS divergence between two 1-D arrays."""
129        eps = 1e-10
130        all_vals = np.concatenate([reference, production])
131        edges = np.linspace(all_vals.min(), all_vals.max(), self.n_bins + 1)
132
133        ref_hist = np.histogram(reference, bins=edges)[0].astype(float) + eps
134        prod_hist = np.histogram(production, bins=edges)[0].astype(float) + eps
135
136        ref_hist /= ref_hist.sum()
137        prod_hist /= prod_hist.sum()
138
139        return float(jensenshannon(ref_hist, prod_hist) ** 2)
140
141    # Aggregate 
142    def detect_all(
143        self,
144        ref_df: pd.DataFrame,
145        prod_df: pd.DataFrame,
146    ) -> dict:
147        """
148        Run every drift test and consolidate results.
149
150        Returns
151        -------
152        dict  {feature: {"ks": {...}, "psi": {...}, "js": {...}, "overall_drift": bool}}
153        """
154        ks_results = self.ks_test(ref_df, prod_df)
155        psi_results = self.psi(ref_df, prod_df)
156        js_results = self.js_divergence(ref_df, prod_df)
157
158        combined: dict = {}
159        for feat in ks_results:
160            combined[feat] = {
161                "ks": ks_results[feat],
162                "psi": psi_results[feat],
163                "js": js_results[feat],
164                "overall_drift": (
165                    ks_results[feat]["drift"]
166                    or psi_results[feat]["drift"]
167                    or js_results[feat]["drift"]
168                ),
169            }
170        return combined
171