| """analyze_correlation — correlation among numeric columns (KM-608). |
| |
| An analytical "family" tool: in ONE call it measures how strongly numeric |
| columns move together. Returns the full correlation matrix plus a list of |
| column pairs ranked by strength. Answers questions like "does price relate to |
| units sold?". |
| |
| STATUS: compute layer only — the function takes an already-materialized |
| DataFrame. The wrapper layer (fetching data from the catalog via source_id, |
| the ToolOutput envelope, ToolSpec registration) is added once the Planner |
| seam (KM-418) is settled. Keeping compute separate from data-fetching makes |
| this function easy to unit-test in isolation and stable when wrapped. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Any |
|
|
| import pandas as pd |
| from pydantic import BaseModel, Field |
|
|
| from src.tools.analytics.descriptive import ColumnNotFoundError |
|
|
| |
| SUPPORTED_METHODS = ("pearson", "spearman", "kendall") |
|
|
| |
| |
| _WEAK_MIN_N = 10 |
| _ALPHA = 0.05 |
|
|
|
|
| class InvalidMethodError(ValueError): |
| """The requested method is not supported (maps to error_code INVALID_METHOD).""" |
|
|
|
|
| class NonNumericColumnError(ValueError): |
| """A requested column is not numeric (maps to error_code NON_NUMERIC_COLUMN).""" |
|
|
|
|
| class NotEnoughColumnsError(ValueError): |
| """Correlation needs at least two numeric columns (maps to NOT_ENOUGH_COLUMNS).""" |
|
|
|
|
| def _clean(value: object) -> float | None: |
| """Cast to plain float; NaN (e.g. a zero-variance column) -> None.""" |
| if value is None: |
| return None |
| f = float(value) |
| return None if math.isnan(f) else f |
|
|
|
|
| def _p_value(a: pd.Series, b: pd.Series, method: str) -> float | None: |
| """Two-sided p-value of the pair's correlation on its complete rows; None if < 3.""" |
| from scipy import stats |
|
|
| if a.shape[0] < 3: |
| return None |
| test = {"pearson": stats.pearsonr, "spearman": stats.spearmanr, |
| "kendall": stats.kendalltau}[method] |
| return _clean(test(a, b).pvalue) |
|
|
|
|
| def _is_weak(n: int, p_value: float | None) -> bool: |
| return n < _WEAK_MIN_N or p_value is None or p_value >= _ALPHA |
|
|
|
|
| def _quality(top: dict[str, Any] | None) -> dict[str, object]: |
| """Reserved `quality` block (conventions §3) for the strongest pair; CK7 surfaces it.""" |
| if top is None: |
| return {"metric": "p_value", "value": None, "floor": _ALPHA, "weak": True, |
| "note": "no pair had a defined correlation"} |
| n, p = top["n"], top["p_value"] |
| reasons = [] |
| if n < _WEAK_MIN_N: |
| reasons.append(f"only {n} rows") |
| if p is None or p >= _ALPHA: |
| reasons.append("not significant at alpha=0.05") |
| return { |
| "metric": "p_value", |
| "value": None if p is None else round(p, 4), |
| "floor": _ALPHA, |
| "weak": top["weak"], |
| "note": ( |
| f"strongest pair {top['a']}–{top['b']} (r={top['corr']:.3f}, n={n}): " |
| + ("; ".join(reasons) if reasons else f"significant at alpha={_ALPHA}") |
| ), |
| } |
|
|
|
|
| |
| |
|
|
| class CorrelationInput(BaseModel): |
| """Planner-facing inputs for `analyze_correlation` (spec_builder contract).""" |
|
|
| data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).") |
| column_ids: list[str] | None = Field( |
| None, |
| description="NUMERIC columns to correlate pairwise; omit to use every numeric column.", |
| ) |
| method: str = Field( |
| "pearson", |
| description="Correlation method: pearson (linear), spearman or kendall (rank-based).", |
| ) |
|
|
| DESCRIPTION = """\ |
| Summary: Pairwise correlation across numeric columns (pearson, spearman, or \ |
| kendall). Returns a correlation matrix plus the strongest pairs ranked by \ |
| absolute strength. |
| |
| USE WHEN the question is about relationship or association between numeric \ |
| variables. Trigger words: "correlation" (korelasi), "related/relationship" \ |
| (hubungan/keterkaitan), "does X affect Y", "move together". |
| |
| DON'T USE WHEN: |
| - it implies causation — correlation is not causality; stay descriptive |
| - it compares two groups of one metric -> analyze_ab_test |
| - it summarizes a single column -> analyze_descriptive |
| |
| Example questions: |
| - "is there a correlation between price and quantity sold?" |
| - "which variables are most related to revenue?" |
| - "do age and spending move together?" |
| - "show the correlation matrix for the numeric columns" |
| """ |
|
|
|
|
| def analyze_correlation( |
| df: pd.DataFrame, |
| column_ids: list[str] | None = None, |
| method: str = "pearson", |
| ) -> dict[str, object]: |
| """Pairwise correlation across numeric columns. |
| |
| Args: |
| df: already-materialized data (in the real system the wrapper fetches |
| this from a source_id). |
| column_ids: numeric columns to correlate. If None, every numeric |
| column in df is used except `n_*` row-count columns. |
| method: "pearson" (linear), "spearman" (rank), or "kendall". |
| |
| Returns: |
| dict with: |
| method — echo of the chosen method |
| columns — the numeric columns actually correlated |
| matrix — { col: { col: corr|None } } full square matrix |
| pairs — [{"a", "b", "corr", "n", "p_value", "weak"}] unique pairs, |
| strongest |corr| first; n = rows complete for both columns, |
| weak = n < 10 or p >= 0.05 |
| quality — reserved block for the strongest pair (CK7 surfaces `weak`) |
| caveats — a note when the strongest pair is weak evidence |
| |
| Raises: |
| InvalidMethodError: if method is unknown. |
| ColumnNotFoundError: if an explicit column is absent. |
| NonNumericColumnError: if an explicit column is not numeric. |
| NotEnoughColumnsError: if fewer than two numeric columns remain. |
| """ |
| if method not in SUPPORTED_METHODS: |
| raise InvalidMethodError( |
| f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}" |
| ) |
|
|
| if column_ids is None: |
| |
| |
| cols = [ |
| c for c in df.columns |
| if pd.api.types.is_numeric_dtype(df[c]) and not str(c).startswith("n_") |
| ] |
| else: |
| missing = [c for c in column_ids if c not in df.columns] |
| if missing: |
| raise ColumnNotFoundError(f"columns not found: {missing}") |
| non_numeric = [ |
| c for c in column_ids if not pd.api.types.is_numeric_dtype(df[c]) |
| ] |
| if non_numeric: |
| raise NonNumericColumnError(f"columns are not numeric: {non_numeric}") |
| cols = list(column_ids) |
|
|
| if len(cols) < 2: |
| raise NotEnoughColumnsError( |
| f"need >= 2 numeric columns, got {len(cols)}: {cols}" |
| ) |
|
|
| corr = df[cols].corr(method=method) |
| matrix = {a: {b: _clean(corr.loc[a, b]) for b in cols} for a in cols} |
|
|
| pairs: list[dict[str, Any]] = [] |
| for i in range(len(cols)): |
| for j in range(i + 1, len(cols)): |
| val = _clean(corr.iloc[i, j]) |
| if val is not None: |
| both = df[[cols[i], cols[j]]].dropna() |
| n = int(both.shape[0]) |
| p = _p_value(both.iloc[:, 0], both.iloc[:, 1], method) |
| pairs.append({"a": cols[i], "b": cols[j], "corr": val, "n": n, |
| "p_value": p, "weak": _is_weak(n, p)}) |
| pairs.sort(key=lambda p: abs(p["corr"]), reverse=True) |
|
|
| top = pairs[0] if pairs else None |
| caveats: list[str] = [] |
| if top is not None and top["weak"]: |
| if top["n"] < _WEAK_MIN_N: |
| caveats.append( |
| f"The strongest correlation ({top['a']}–{top['b']}, r={top['corr']:.3f}) " |
| f"rests on {top['n']} rows — on so few points a single value can drive r, " |
| "so treat it as a lead to check, not an established relationship." |
| ) |
| else: |
| caveats.append( |
| f"The strongest correlation ({top['a']}–{top['b']}, r={top['corr']:.3f}) " |
| f"is not significant (p >= {_ALPHA}) on {top['n']} rows — it may be noise." |
| ) |
|
|
| return { |
| "method": method, |
| "columns": cols, |
| "matrix": matrix, |
| "pairs": pairs, |
| "quality": _quality(top), |
| "caveats": caveats, |
| } |
|
|
