Skills › Scientific Agent Skills › statistical-analysis › assumption_checks.py
assumption_checks.py
352 lines · skills/statistical-analysis/scripts/assumption_checks.py at 5f2bc03
1"""Descriptive diagnostics, not an automatic test selector or proof of assumptions.
3NaNs are omitted with counts; infinities, malformed arrays, insufficient samples
4and undefined diagnostics raise ValueError. Independence requires study-design
5review. Legacy ``is_normal``/``is_homogeneous`` fields mean only non-rejection.
6"""
8from typing import Dict, List, Optional, Union
10import matplotlib.pyplot as plt
11import numpy as np
12import pandas as pd
13from scipy import stats
16def _alpha(alpha):
17 if not np.isfinite(alpha) or not 0 < alpha < 1:
18 raise ValueError("alpha must be finite and between 0 and 1")
21def _array(data):
22 if np.iscomplexobj(data) or np.ma.isMaskedArray(data):
23 raise ValueError("data must be real and unmasked; encode missing values as NaN")
24 if isinstance(data, pd.Series):
25 values = data.to_numpy(dtype=float, na_value=np.nan)
26 else:
27 values = np.asarray(data, dtype=float)
28 if values.ndim != 1:
29 raise ValueError("data must be one-dimensional")
30 if np.isinf(values).any():
31 raise ValueError("data must not contain infinite values")
32 return values
35def _clean(data, minimum=1, variable=False):
36 values = _array(data)
37 positions = np.flatnonzero(~np.isnan(values))
38 clean = values[positions]
39 if clean.size < minimum:
40 raise ValueError(f"at least {minimum} non-missing observations are required")
41 if variable and np.ptp(clean) == 0:
42 raise ValueError("constant data do not support this diagnostic")
43 return clean, positions, len(values) - len(clean)
46def _groups(data, value_col, group_col, minimum=1):
47 if data.empty or data[group_col].isna().any():
48 raise ValueError("group labels must be non-missing and data nonempty")
49 groups = []
50 for label, frame in data.groupby(group_col, sort=False, observed=True):
51 clean, positions, omitted = _clean(frame[value_col], minimum=minimum)
52 groups.append((label, clean, omitted, frame, positions))
53 return groups
56def _finish_plot():
57 plt.tight_layout()
58 plt.show()
59 plt.close()
62def check_normality(
63 data: Union[np.ndarray, pd.Series, List], name: str = "data",
64 alpha: float = 0.05, plot: bool = True
65) -> Dict:
66 """Shapiro-Wilk screen with Q-Q/histogram; n>5000 has no binary verdict.
68 ``is_normal`` is a legacy alias for ``normality_not_rejected`` and never
69 establishes normality. NaN omissions are returned as ``n_missing``.
70 """
71 _alpha(alpha)
72 clean, _, missing = _clean(data, minimum=3, variable=True)
73 statistic, p_value = stats.shapiro(clean)
74 if not np.isfinite([statistic, p_value]).all():
75 raise ValueError("Shapiro-Wilk returned a non-finite result")
76 reliable = len(clean) <= 5000
77 decision = bool(p_value > alpha) if reliable else None
78 finding = ("not rejected" if decision else "rejected") if reliable else "not assessed by p-value (n > 5000)"
79 interpretation = f"Normality {finding} (W = {statistic:.3f}, p = {p_value:.3g}); this does not verify assumptions."
80 if plot:
81 _, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
82 stats.probplot(clean, dist="norm", plot=ax1)
83 ax1.set_title(f"Q-Q Plot: {name}")
84 ax1.grid(alpha=0.3)
85 ax2.hist(clean, bins='auto', density=True, alpha=0.7, color='steelblue', edgecolor='black')
86 grid = np.linspace(clean.min(), clean.max(), 100)
87 ax2.plot(grid, stats.norm.pdf(grid, clean.mean(), clean.std()), 'r-', label='Normal curve')
88 ax2.set(xlabel='Value', ylabel='Density', title=f'Histogram: {name}')
89 ax2.legend()
90 _finish_plot()
91 return {
92 'test': 'Shapiro-Wilk', 'statistic': statistic, 'p_value': p_value,
93 'is_normal': decision, 'normality_not_rejected': decision,
94 'p_value_reliable': reliable, 'interpretation': interpretation,
95 'n': len(clean), 'n_missing': missing,
96 'recommendation': "Review the Q-Q plot, tails, design and estimand. Neither parametric nor non-parametric inference is selected by this p-value.",
97 }
100def check_normality_per_group(data: pd.DataFrame, value_col: str, group_col: str,
101 alpha: float = 0.05, plot: bool = True) -> pd.DataFrame:
102 """One Shapiro-Wilk screen per observed group; 'Normal' is a legacy label.
104 Yes means not rejected, No means rejected, and Unassessed means n>5000.
105 Missing group labels are rejected; unused categorical levels are ignored.
106 """
107 _alpha(alpha)
108 groups = _groups(data, value_col, group_col, minimum=3)
109 results = []
110 if plot:
111 _, axes = plt.subplots(1, len(groups), figsize=(5 * len(groups), 4), squeeze=False)
112 for i, (label, clean, missing, _, _) in enumerate(groups):
113 result = check_normality(clean, alpha=alpha, plot=False)
114 decision = result['normality_not_rejected']
115 results.append({'Group': label, 'N': len(clean), 'N_missing': missing,
116 'W': result['statistic'], 'p-value': result['p_value'],
117 'Normal': 'Unassessed' if decision is None else ('Yes' if decision else 'No')})
118 if plot:
119 stats.probplot(clean, dist="norm", plot=axes[0, i])
120 axes[0, i].set_title(f"Q-Q Plot: {label}")
121 if plot:
122 _finish_plot()
123 return pd.DataFrame(results)
126def check_homogeneity_of_variance(data: pd.DataFrame, value_col: str, group_col: str,
127 alpha: float = 0.05, plot: bool = True) -> Dict:
128 """Median-centered Levene (Brown-Forsythe) screen; does not select a test."""
129 _alpha(alpha)
130 groups = _groups(data, value_col, group_col, minimum=2)
131 if len(groups) < 2:
132 raise ValueError("at least two observed groups are required")
133 arrays = [g[1] for g in groups]
134 # All within-group absolute deviations constant makes Levene undefined.
135 deviations = [np.abs(g - np.median(g)) for g in arrays]
136 if all(np.ptp(z) == 0 for z in deviations):
137 raise ValueError("Levene diagnostic is undefined for these within-group deviations")
138 statistic, p_value = stats.levene(*arrays, center='median')
139 if not np.isfinite([statistic, p_value]).all():
140 raise ValueError("Levene returned a non-finite result")
141 variances = np.array([np.var(g, ddof=1) for g in arrays])
142 if not np.isfinite(variances).all():
143 raise ValueError("variances overflowed; rescale values before diagnostics")
144 ratio = float(variances.max() / variances.min()) if variances.min() else np.inf
145 decision = bool(p_value > alpha)
146 if plot:
147 _, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
148 labels = [str(g[0]) for g in groups]
149 ax1.boxplot(arrays, tick_labels=labels)
150 ax1.set(xlabel=group_col, ylabel=value_col, title='Box Plots by Group')
151 ax2.bar(np.arange(len(groups)), variances, color='steelblue')
152 ax2.set_xticks(np.arange(len(groups)), labels, rotation=45)
153 ax2.set(ylabel='Variance', title='Variance by Group')
154 _finish_plot()
155 return {
156 'test': 'Levene', 'center': 'median', 'statistic': statistic,
157 'p_value': p_value, 'is_homogeneous': decision,
158 'variance_equality_not_rejected': decision, 'variance_ratio': ratio,
159 'groups': [g[0] for g in groups], 'variances': variances,
160 'n_missing': sum(g[2] for g in groups),
161 'interpretation': f"Equal variances {'not rejected' if decision else 'rejected'} (F = {statistic:.3f}, p = {p_value:.3g}); non-rejection is not evidence of equality.",
162 'recommendation': "For independent mean comparisons, prespecify Welch inference rather than choosing pooled versus Welch by a variance pretest.",
163 }
166def check_linearity(x: Union[np.ndarray, pd.Series], y: Union[np.ndarray, pd.Series],
167 x_name: str = "X", y_name: str = "Y", plot: bool = True) -> Dict:
168 """Simple-regression scatter/residual plots; Pearson r does not test linearity.
170 Rows missing either variable are omitted together, preserving pairing.
171 """
172 x, y = _array(x), _array(y)
173 if len(x) != len(y):
174 raise ValueError("x and y must have the same length and aligned rows")
175 keep = ~np.isnan(x) & ~np.isnan(y)
176 missing = int((~keep).sum())
177 x, y = x[keep], y[keep]
178 _clean(x, minimum=3, variable=True)
179 _clean(y, minimum=3, variable=True)
180 slope, intercept, r_value, _, _ = stats.linregress(x, y)
181 y_pred = intercept + slope * x
182 residuals = y - y_pred
183 if not np.isfinite([slope, intercept, r_value]).all():
184 raise ValueError("linear fit returned non-finite results; review/rescale the data")
185 if plot:
186 _, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
187 ax1.scatter(x, y, alpha=0.6)
188 order = np.argsort(x)
189 ax1.plot(x[order], y_pred[order], 'r-', label=f'y = {intercept:.2f} + {slope:.2f}x')
190 ax1.set(xlabel=x_name, ylabel=y_name, title='Scatter Plot with Regression Line')
191 ax1.legend()
192 ax2.scatter(y_pred, residuals, alpha=0.6)
193 ax2.axhline(0, color='r', linestyle='--')
194 ax2.set(xlabel='Fitted values', ylabel='Residuals', title='Residuals vs Fitted Values')
195 _finish_plot()
196 return {
197 'r': r_value, 'r_squared': r_value ** 2, 'n': len(x), 'n_missing': missing,
198 'interpretation': "Examine residual plots for curvature and changing spread; correlation magnitude does not establish linearity.",
199 'recommendation': "Review the conditional mean model and design; consider justified nonlinear terms or variance/dependence models.",
200 }
203def detect_outliers(data: Union[np.ndarray, pd.Series, List], name: str = "data",
204 method: str = "iqr", threshold: Optional[float] = None,
205 plot: bool = True) -> Dict:
206 """Flag values for review; return original zero-based positions and labels.
208 Defaults: IQR multiplier 1.5; population-SD z-score threshold 3. A flag is
209 not permission to exclude/winsorize a legitimate observation.
210 """
211 if method not in ('iqr', 'zscore'):
212 raise ValueError("method must be 'iqr' or 'zscore'")
213 threshold = (1.5 if method == 'iqr' else 3.0) if threshold is None else threshold
214 if not np.isfinite(threshold) or threshold <= 0:
215 raise ValueError("threshold must be finite and positive")
216 clean, positions, missing = _clean(data)
217 if method == 'iqr':
218 q1, q3 = np.percentile(clean, [25, 75])
219 lower, upper = q1 - threshold * (q3 - q1), q3 + threshold * (q3 - q1)
220 else:
221 mean, sd = clean.mean(), clean.std()
222 lower, upper = mean - threshold * sd, mean + threshold * sd
223 if not np.isfinite([lower, upper]).all():
224 raise ValueError("outlier bounds overflowed; rescale values")
225 mask = (clean < lower) | (clean > upper)
226 indices = positions[mask]
227 labels = np.asarray(data.index)[indices] if isinstance(data, pd.Series) else indices.copy()
228 if plot:
229 _, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
230 ax1.boxplot(clean, orientation='vertical', patch_artist=True)
231 ax1.set(ylabel='Value', title=f'Box Plot: {name}')
232 ax2.scatter(positions[~mask], clean[~mask], label='Unflagged', alpha=0.6)
233 ax2.scatter(indices, clean[mask], label='Flagged', marker='D', color='red')
234 ax2.axhline(lower, linestyle='--', color='orange', label='Bounds')
235 ax2.axhline(upper, linestyle='--', color='orange')
236 ax2.set(xlabel='Original row position', ylabel='Value', title=f'Outlier Screen: {name}')
237 ax2.legend()
238 _finish_plot()
239 count = int(mask.sum())
240 pct = 100 * count / len(clean)
241 return {
242 'method': method, 'threshold': threshold, 'n': len(clean), 'n_missing': missing,
243 'n_outliers': count, 'pct_outliers': pct, 'outlier_indices': indices,
244 'outlier_labels': labels, 'outlier_values': clean[mask],
245 'lower_bound': lower, 'upper_bound': upper,
246 'interpretation': f"Flagged {count} observations ({pct:.1f}% of non-missing data)",
247 'recommendation': "Investigate provenance and data-entry errors. Retain valid observations; justify exclusions and report prespecified sensitivity/robust analyses.",
248 }
251def check_regression_diagnostics(model, alpha: float = 0.05, plot: bool = True,
252 ordered: bool = False) -> Dict:
253 """Diagnostics for full-rank OLS with intercept, positive residual df/spread.
255 Durbin-Watson is reported only for meaningfully ordered rows (ordered=True)
256 and has no universal pass threshold. Dependence/causal design is not tested.
257 """
258 from statsmodels.regression.linear_model import OLS
259 from statsmodels.stats.diagnostic import het_breuschpagan
260 from statsmodels.stats.stattools import durbin_watson
261 from statsmodels.stats.outliers_influence import variance_inflation_factor
263 _alpha(alpha)
264 if not isinstance(model.model, OLS):
265 raise ValueError("this helper requires fitted OLS results")
266 residuals = _array(model.resid)
267 fitted = _array(model.fittedvalues)
268 exog = np.asarray(model.model.exog, dtype=float)
269 if not all(np.isfinite(x).all() for x in (residuals, fitted, exog)):
270 raise ValueError("fitted data must be finite")
271 constants = np.flatnonzero((np.ptp(exog, axis=0) == 0) & (exog[0] != 0))
272 if len(constants) != 1 or exog.shape[1] < 2:
273 raise ValueError("OLS must have an explicit nonzero intercept and at least one predictor")
274 if np.linalg.matrix_rank(exog) != exog.shape[1] or model.df_resid <= 0:
275 raise ValueError("OLS design must be full rank with positive residual degrees of freedom")
276 if np.linalg.norm(residuals) <= np.finfo(float).eps * max(1, np.linalg.norm(model.model.endog)) * len(residuals):
277 raise ValueError("near-perfect fit does not support residual diagnostics")
278 normality = check_normality(residuals, alpha=alpha, plot=False)
279 normality['ok'] = normality['normality_not_rejected'] # legacy non-rejection alias
280 lm, lm_p, f_value, f_p = het_breuschpagan(residuals, exog, robust=True)
281 if not np.isfinite([lm, lm_p, f_value, f_p]).all():
282 raise ValueError("heteroscedasticity diagnostic is undefined")
283 dw = float(durbin_watson(residuals)) if ordered else None
284 vif_table = pd.DataFrame([
285 {'Variable': name, 'VIF': variance_inflation_factor(exog, i)}
286 for i, name in enumerate(model.model.exog_names) if i not in constants
287 ])
288 max_vif = float(vif_table['VIF'].max())
289 if plot:
290 _, axes = plt.subplots(2, 2, figsize=(12, 10))
291 axes[0, 0].scatter(fitted, residuals, alpha=0.6)
292 axes[0, 0].axhline(0, color='r', linestyle='--')
293 axes[0, 0].set(xlabel='Fitted values', ylabel='Residuals', title='Residuals vs Fitted')
294 stats.probplot(residuals, dist='norm', plot=axes[0, 1])
295 axes[0, 1].set_title('Normal Q-Q')
296 standardized = model.get_influence().resid_studentized_internal
297 axes[1, 0].scatter(fitted, np.sqrt(np.abs(standardized)), alpha=0.6)
298 axes[1, 0].set(xlabel='Fitted values', ylabel='sqrt(|Studentized residual|)', title='Scale-Location')
299 axes[1, 1].hist(residuals, bins='auto', edgecolor='black')
300 axes[1, 1].set(xlabel='Residuals', ylabel='Frequency', title='Residual Histogram')
301 _finish_plot()
302 issues = []
303 if normality['normality_not_rejected'] is False:
304 issues.append('residual normality rejected; review tails and inference robustness')
305 if normality['normality_not_rejected'] is None:
306 issues.append('Shapiro p-value unreliable above 5000 observations; inspect residual plots')
307 if lm_p <= alpha:
308 issues.append('heteroscedasticity screen rejected; consider justified HC3 or variance modeling')
309 if not np.isfinite(max_vif) or max_vif > 5:
310 issues.append('large VIF; examine identifiability and precision without automatically removing adjustment covariates')
311 return {
312 'residual_normality': normality,
313 'heteroscedasticity': {'test': 'Koenker Breusch-Pagan', 'statistic': lm,
314 'p_value': lm_p, 'f_statistic': f_value, 'f_p_value': f_p,
315 'ok': bool(lm_p > alpha)},
316 'autocorrelation': {'test': 'Durbin-Watson', 'statistic': dw, 'ok': None,
317 'interpretation': 'Requires meaningful row order and design-specific inference; cannot establish independence.'},
318 'vif': vif_table, 'max_vif': max_vif, 'issues': issues,
319 'interpretation': ('Screen flags: ' + '; '.join(issues) if issues else 'No flags in these numerical screens.') + ' Review plots, sampling units, dependence, exogeneity and model specification separately.',
320 }
323def comprehensive_assumption_check(data: pd.DataFrame, value_col: str,
324 group_col: Optional[str] = None,
325 alpha: float = 0.05, plot: bool = True) -> Dict:
326 """Grouped descriptive screens; does not handle pairing/clustering itself.
328 For paired tests pass aligned difference scores without group_col. Review
329 regression residuals using check_regression_diagnostics instead.
330 """
331 _alpha(alpha)
332 if group_col is not None:
333 groups = _groups(data, value_col, group_col, minimum=3)
334 results = {
335 'outliers_per_group': {label: detect_outliers(frame[value_col], name=str(label), plot=plot)
336 for label, _, _, frame, _ in groups},
337 'normality_per_group': check_normality_per_group(data, value_col, group_col, alpha, plot),
338 'homogeneity': check_homogeneity_of_variance(data, value_col, group_col, alpha, plot),
339 }
340 else:
341 results = {'outliers': detect_outliers(data[value_col], name=value_col, plot=plot),
342 'normality': check_normality(data[value_col], name=value_col, alpha=alpha, plot=plot)}
343 print('[OK] Diagnostic calculations completed. Missing-value counts are in each result.')
344 print('-> These screens do not establish assumptions or select a test. Review design, estimand and plots.')
345 return results
348if __name__ == '__main__':
349 rng = np.random.default_rng(42)
350 df = pd.DataFrame({'score': np.r_[rng.normal(75, 8, 50), rng.normal(68, 10, 50)],
351 'group': ['A'] * 50 + ['B'] * 50})
352 results = comprehensive_assumption_check(df, 'score', 'group')