From b687a08473a74741620c065caff11249b21cb785 Mon Sep 17 00:00:00 2001 From: lukezhang01 Date: Sat, 28 Feb 2026 22:17:34 -0500 Subject: [PATCH 1/3] data preprocessing --- cais/discovery/data_preprocess.py | 569 ++++++++++++++++++++++++++++++ data/profile_data.py | 266 ++++++++++++++ 2 files changed, 835 insertions(+) create mode 100644 cais/discovery/data_preprocess.py create mode 100644 data/profile_data.py diff --git a/cais/discovery/data_preprocess.py b/cais/discovery/data_preprocess.py new file mode 100644 index 0000000..0dabd97 --- /dev/null +++ b/cais/discovery/data_preprocess.py @@ -0,0 +1,569 @@ +from __future__ import annotations +import re +from dataclasses import dataclass, field + +import numpy as np +import pandas as pd + + +# ────────────────────────────────────────────────────────────────────── +# Data structures +# ────────────────────────────────────────────────────────────────────── + +COLUMN_TYPES = [ + "constant", + "id", + "temporal", + "continuous", + "binary_numeric", + "binary_string", + "nominal", + "low_cardinality_string", + "likely_categorical_numeric", + "ordinal", + "log_derived", + "polynomial_derived", + "cumulative_derived", + "duplicate", + "count", + "rate_or_proportion", +] + + +@dataclass +class ClassificationResult: + # col_name -> {col_type, is_string, n_unique, n_missing, note} + metadata: dict[str, dict] = field(default_factory=dict) + cols_to_drop: list[str] = field(default_factory=list) + cols_for_discovery: list[str] = field(default_factory=list) + discrete_flags: list[str] = field(default_factory=list) + temporal_var: str | None = None + warnings: list[str] = field(default_factory=list) + +_ID_PATTERNS = re.compile( + r"^(id|_id|row_?id|index|obs|observation|unit_?id|" + r"subject_?id|participant_?id|record_?id|serial|" + r"respondent_?id|sample_?id|case_?id)$", + re.IGNORECASE, +) +_ID_SUFFIX = re.compile(r"_id$", re.IGNORECASE) + +_TEMPORAL_PATTERNS = re.compile( + r"^(year|yr|month|mon|date|day|time|period|quarter|qtr|" + r"week|t|trend|wave|round|session|epoch)$", + re.IGNORECASE, +) + +_LOG_PREFIX = re.compile(r"^(ln|log|lg|l_)", re.IGNORECASE) + +_RATE_PATTERNS = re.compile( + r"(rate|ratio|proportion|pct|percent|share|frac|frequency|prevalence)", + re.IGNORECASE, +) + +_COUNT_PATTERNS = re.compile( + r"(count|total|tot|num_|number|n_|pop|cases|deaths|births|incidents|events)", + re.IGNORECASE, +) + +_CUMULATIVE_PATTERNS = re.compile( + r"(cumul|cum_|acc|running_total|cumsum)", re.IGNORECASE +) + +_ORDINAL_PATTERNS = re.compile( + r"(level|grade|stage|rank|score|scale|rating|class|tier|severity|priority|order)", + re.IGNORECASE, +) + + +def _safe_corr(a: np.ndarray, b: np.ndarray) -> float: + """Pearson correlation that returns 0.0 when either input is constant.""" + if a.std() < 1e-12 or b.std() < 1e-12: + return 0.0 + return np.corrcoef(a, b)[0, 1] + + +def _is_integer_valued(s: pd.Series) -> bool: + """True if a numeric series contains only integer values (ignoring NaN).""" + s = s.dropna() + if len(s) == 0: + return False + return np.allclose(s, s.round(), equal_nan=True) + + +def _looks_sequential(s: pd.Series) -> bool: + """True if sorted unique values are roughly 1,2,3,...,n (an ID or time index).""" + u = np.sort(s.dropna().unique()) + if len(u) < 5: + return False + ideal = np.arange(1, len(u) + 1, dtype=float) + # Allow for small scaling/offset + if np.allclose(u, ideal, atol=1) or np.allclose(u - u[0], np.arange(len(u)), atol=1): + return True + return False + + +def _corr_with_log(df: pd.DataFrame, col: str) -> tuple[str, float] | None: + """ + If `col` looks like a log-transform of another column, return (parent, corr). + Checks: is there a column X such that log(X) ≈ col? + """ + if not _LOG_PREFIX.match(col): + return None + s = df[col].dropna() + if len(s) < 10: + return None + # Strip prefix to guess parent name + stem = _LOG_PREFIX.sub("", col) + # Try exact stem, stem with various cases + candidates = [c for c in df.columns if c != col and c.lower().startswith(stem.lower())] + # Also try all other numeric columns + if not candidates: + candidates = [c for c in df.select_dtypes(include="number").columns if c != col] + for cand in candidates: + other = df[cand].dropna() + if (other <= 0).any(): + continue + common = s.index.intersection(other.index) + if len(common) < 10: + continue + log_other = np.log(other.loc[common]) + corr = _safe_corr(s.loc[common].values, log_other.values) + if abs(corr) > 0.999: + return (cand, corr) + return None + + +def _find_polynomial_parent(df: pd.DataFrame, col: str) -> tuple[str, int, float] | None: + """ + Check if col ≈ parent^k for k in {2, 3}. + Returns (parent_name, degree, corr) or None. + """ + s = df[col].dropna() + if len(s) < 10: + return None + for cand in df.select_dtypes(include="number").columns: + if cand == col: + continue + other = df[cand].dropna() + common = s.index.intersection(other.index) + if len(common) < 10: + continue + for k in [2, 3]: + powered = other.loc[common] ** k + corr = _safe_corr(s.loc[common].values, powered.values) + if abs(corr) > 0.999: + return (cand, k, corr) + return None + + +def _find_cumulative_parent(df: pd.DataFrame, col: str, group_col: str | None = None) -> str | None: + """ + Check if col ≈ cumsum of another column (optionally within groups). + """ + if not _CUMULATIVE_PATTERNS.search(col): + return None + s = df[col].dropna() + if len(s) < 10: + return None + for cand in df.select_dtypes(include="number").columns: + if cand == col: + continue + if group_col and group_col in df.columns: + cumulated = df.groupby(group_col)[cand].cumsum() + else: + cumulated = df[cand].cumsum() + common = s.index.intersection(cumulated.dropna().index) + if len(common) < 10: + continue + corr = _safe_corr(s.loc[common].values, cumulated.loc[common].values) + if abs(corr) > 0.999: + return cand + return None + + +def _find_duplicate(df: pd.DataFrame, col: str, already_seen: list[str]) -> str | None: + """Check if col is essentially a duplicate of a previously-seen column.""" + s = df[col].dropna() + if len(s) < 5: + return None + for prev in already_seen: + other = df[prev].dropna() + common = s.index.intersection(other.index) + if len(common) < 5: + continue + # Exact or near-exact match + if np.allclose(s.loc[common], other.loc[common], rtol=1e-4, atol=1e-8, equal_nan=True): + return prev + # Or perfect correlation with same scale + if s.loc[common].std() > 1e-12 and other.loc[common].std() > 1e-12: + corr = _safe_corr(s.loc[common].values, other.loc[common].values) + if abs(corr) > 0.9999: + return prev + return None + + +# ────────────────────────────────────────────────────────────────────── +# Main classifier +# ────────────────────────────────────────────────────────────────────── + +def classify_columns( + df: pd.DataFrame, + *, + max_onehot_levels: int = 5, + id_uniqueness_threshold: float = 0.95, + categorical_nunique_threshold: int = 15, +) -> ClassificationResult: + """ + Classify each column in `df` and return a ClassificationResult. + + Parameters + ---------- + df : pd.DataFrame + Raw dataset. + max_onehot_levels : int + Max unique values for one-hot encoding nominal columns. + id_uniqueness_threshold : float + If nunique/nrows > this AND column looks like an ID, classify as ID. + categorical_nunique_threshold : int + Numeric columns with nunique <= this that don't match other patterns + are flagged as likely_categorical_numeric. + """ + result = ClassificationResult() + n_rows = len(df) + seen_numeric: list[str] = [] # for duplicate detection + # Track which cols were classified as "id" for cumulative parent detection + id_cols: list[str] = [] + + for col in df.columns: + s = df[col] + is_numeric = pd.api.types.is_numeric_dtype(s) + is_string = pd.api.types.is_string_dtype(s) or pd.api.types.is_object_dtype(s) + n_unique = s.nunique() + n_missing = int(s.isna().sum()) + + col_type = None + note = "" + classified = False + + # ── 1. CONSTANT ────────────────────────────────────────────── + if n_unique <= 1: + col_type = "constant" + note = f"Only {n_unique} unique value(s)" + classified = True + + # ── 2. ID detection ────────────────────────────────────────── + if not classified: + is_id = False + if _ID_PATTERNS.match(col) or _ID_SUFFIX.search(col): + is_id = True + if is_string and n_unique / n_rows > id_uniqueness_threshold: + is_id = True + if is_numeric and n_unique / n_rows > id_uniqueness_threshold: + if _is_integer_valued(s) and _looks_sequential(s): + is_id = True + if is_numeric and _is_integer_valued(s): + low_name = col.lower() + if any(tag in low_name for tag in ["fip", "fips", "state_id", "county", "country_code", "entity", "unit"]): + is_id = True + + if is_id: + col_type = "id" + note = "Detected as identifier" + classified = True + + # ── 3. TEMPORAL detection ──────────────────────────────────── + if not classified and n_unique > 2: + is_temporal = False + low = col.lower().strip() + if _TEMPORAL_PATTERNS.match(low): + is_temporal = True + if is_numeric and ("year" in low or "yr" in low): + vals = s.dropna() + if len(vals) > 0 and vals.min() >= 1900 and vals.max() <= 2100: + is_temporal = True + if is_numeric and "date" in low: + is_temporal = True + + if is_temporal: + col_type = "temporal" + note = "Time / period variable" + classified = True + + # ── 4. DERIVED: log-transform ──────────────────────────────── + if not classified and is_numeric: + log_result = _corr_with_log(df, col) + if log_result is not None: + parent, corr = log_result + col_type = "log_derived" + note = f"log({parent}), corr={corr:.6f}" + classified = True + + # ── 5. DERIVED: polynomial (tsq = t^2) ────────────────────── + if not classified and is_numeric: + poly_result = _find_polynomial_parent(df, col) + if poly_result is not None: + parent, degree, corr = poly_result + col_type = "polynomial_derived" + note = f"{parent}^{degree}, corr={corr:.6f}" + classified = True + + # ── 6. DERIVED: cumulative ─────────────────────────────────── + if not classified and is_numeric: + group_col = id_cols[0] if id_cols else None + cum_parent = _find_cumulative_parent(df, col, group_col) + if cum_parent is not None: + col_type = "cumulative_derived" + note = f"Cumulative sum of {cum_parent}" + classified = True + + # ── 7. DUPLICATE of already-seen column ────────────────────── + if not classified and is_numeric: + dup = _find_duplicate(df, col, seen_numeric) + if dup is not None: + col_type = "duplicate" + note = f"Near-duplicate of {dup}" + classified = True + + # ── 8. BINARY NUMERIC (exactly 2 unique numeric values) ────── + if not classified and is_numeric and n_unique == 2: + col_type = "binary_numeric" + note = f"Binary: values {sorted(s.dropna().unique().tolist())}" + classified = True + + # ── 9. BINARY STRING ───────────────────────────────────────── + if not classified and is_string and n_unique == 2: + col_type = "binary_string" + note = f"Binary string: values {sorted(s.dropna().unique().tolist())}" + classified = True + + # ── 10. STRING with low cardinality → nominal ──────────────── + if not classified and is_string: + if n_unique <= max_onehot_levels: + col_type = "nominal" + note = f"Nominal string, {n_unique} levels" + elif n_unique <= categorical_nunique_threshold: + col_type = "low_cardinality_string" + note = f"Low-cardinality string, {n_unique} levels" + else: + col_type = "id" + note = f"High-cardinality string ({n_unique} unique), treating as ID" + classified = True + + # ── 11. NUMERIC: rate / proportion ─────────────────────────── + if not classified and is_numeric: + vals = s.dropna() + is_rate = False + if _RATE_PATTERNS.search(col): + is_rate = True + if len(vals) > 0 and vals.min() >= 0 and vals.max() <= 1.0 and not _is_integer_valued(s): + is_rate = True + if is_rate: + col_type = "rate_or_proportion" + note = "Rate/proportion — treat as continuous" + classified = True + + # ── 12. NUMERIC: count data ────────────────────────────────── + if not classified and is_numeric: + vals = s.dropna() + if _COUNT_PATTERNS.search(col) and _is_integer_valued(s) and vals.min() >= 0: + col_type = "count" + note = "Count data — treat as continuous" + classified = True + + # ── 13. NUMERIC: likely categorical ────────────────────────── + if not classified and is_numeric: + if n_unique <= categorical_nunique_threshold and _is_integer_valued(s): + unique_ratio = n_unique / n_rows + if unique_ratio < 0.05 or n_unique <= 10: + col_type = "likely_categorical_numeric" + note = f"{n_unique} unique integer values, ratio={unique_ratio:.3f}" + classified = True + + # ── 14. ORDINAL (name-based heuristic) ─────────────────────── + if col_type == "likely_categorical_numeric" and _ORDINAL_PATTERNS.search(col): + col_type = "ordinal" + note += " — name suggests ordinal" + + # ── 15. FALLBACK: continuous ───────────────────────────────── + if not classified: + if is_numeric: + col_type = "continuous" + note = "Numeric, high cardinality" + else: + col_type = "nominal" + note = "Fallback: unclassified string column" + + # Store metadata + result.metadata[col] = { + "col_type": col_type, + "is_string": is_string, + "n_unique": n_unique, + "n_missing": n_missing, + "note": note, + } + + if col_type == "id": + id_cols.append(col) + if is_numeric: + seen_numeric.append(col) + + # ── Build action lists ─────────────────────────────────────────── + for col, meta in result.metadata.items(): + ct = meta["col_type"] + match ct: + case "constant" | "id" | "temporal": + result.cols_to_drop.append(col) + if ct == "temporal": + result.temporal_var = col + + case "log_derived" | "polynomial_derived" | "cumulative_derived" | "duplicate": + result.cols_to_drop.append(col) + + case "continuous" | "rate_or_proportion" | "count": + result.cols_for_discovery.append(col) + + case "binary_numeric" | "binary_string": + result.cols_for_discovery.append(col) + result.discrete_flags.append(col) + + case "nominal" | "low_cardinality_string": + if meta["n_unique"] <= max_onehot_levels: + result.cols_for_discovery.append(col) + result.discrete_flags.append(col) + else: + result.cols_to_drop.append(col) + result.warnings.append( + f"Dropping {col}: {meta['n_unique']} levels, too many for one-hot" + ) + + case "likely_categorical_numeric" | "ordinal": + result.cols_for_discovery.append(col) + result.discrete_flags.append(col) + + return result + + +# ────────────────────────────────────────────────────────────────────── +# Clean data for causal discovery +# ────────────────────────────────────────────────────────────────────── + +def prepare_data( + df: pd.DataFrame, + result: ClassificationResult, +) -> pd.DataFrame: + """ + Given raw data and its ClassificationResult, return a cleaned numeric + DataFrame ready for causal-learn. + + Steps + ----- + 1. Keep only cols_for_discovery (drop IDs, constants, derived, etc.). + 2. Handle missing values (drop rows with any NaN). + 3. Encode string columns: + - binary_string → 0/1 + - nominal / low_cardinality_string → one-hot (drop_first) + 4. Standardise binary numerics whose values aren't {0, 1} → remap to 0/1. + 5. Return a float64 ndarray-backed DataFrame (causal-learn expects this). + """ + kept = [c for c in result.cols_for_discovery if c in df.columns] + out = df[kept].copy() + + # ── Drop rows with missing values ──────────────────────────────── + out = out.dropna() + + # ── Per-column transforms based on col_type ────────────────────── + cols_to_onehot: list[str] = [] + + for col in list(out.columns): + meta = result.metadata[col] + ct = meta["col_type"] + + if ct == "binary_string": + # Map the two unique string values to 0 / 1 alphabetically + vals = sorted(out[col].unique()) + out[col] = out[col].map({vals[0]: 0, vals[1]: 1}).astype(np.float64) + + elif ct in ("nominal", "low_cardinality_string"): + cols_to_onehot.append(col) + + elif ct == "binary_numeric": + # Ensure values are exactly {0, 1} + vals = sorted(out[col].dropna().unique()) + if vals != [0, 1]: + out[col] = out[col].map({vals[0]: 0, vals[1]: 1}).astype(np.float64) + else: + out[col] = out[col].astype(np.float64) + + elif ct == "likely_categorical_numeric" or ct == "ordinal": + out[col] = out[col].astype(np.float64) + + else: + # continuous, rate_or_proportion, count — already numeric + out[col] = out[col].astype(np.float64) + + # ── One-hot encode nominal / low-cardinality string columns ────── + if cols_to_onehot: + out = pd.get_dummies(out, columns=cols_to_onehot, drop_first=True, dtype=np.float64) + + out = out.reset_index(drop=True) + return out + + +# ────────────────────────────────────────────────────────────────────── +# Pretty printer +# ────────────────────────────────────────────────────────────────────── + +def print_classification(res: ClassificationResult) -> None: + print(f"{'Column':<25s} {'Type':<28s} {'Action':<8s} {'Note'}") + print("─" * 100) + for col, meta in res.metadata.items(): + action = "DROP" if col in res.cols_to_drop else "KEEP" + disc = " [D]" if col in res.discrete_flags else "" + print(f"{col:<25s} {meta['col_type']:<28s} {action:<8s} {meta['note']}{disc}") + + print(f"\n── Summary ──") + print(f" Keep for discovery: {len(res.cols_for_discovery)} columns") + print(f" Drop: {len(res.cols_to_drop)} columns") + print(f" Discrete flags: {len(res.discrete_flags)}") + if res.temporal_var: + print(f" Temporal variable: {res.temporal_var}") + for w in res.warnings: + print(f" ⚠ {w}") + + +# ────────────────────────────────────────────────────────────────────── +# Test on the four datasets +# ────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + import sys + from pathlib import Path + + DATA_DIR = Path(__file__).resolve().parents[2] / "data" / "all_data" + + # The four original datasets + named_files = [ + DATA_DIR / "abortion_bf15.csv", + DATA_DIR / "drinking.csv", + DATA_DIR / "fulton.csv", + DATA_DIR / "ihdp_0.csv", + ] + + # Pass --all to run on every CSV in the data folder instead + if "--all" in sys.argv: + files = sorted(DATA_DIR.glob("*.csv")) + else: + files = named_files + + for path in files: + if not path.exists(): + print(f" ⚠ Skipping {path.name}: file not found") + continue + df = pd.read_csv(path) + print(f"\n{'='*100}") + print(f" {path.name} ({df.shape[0]} rows × {df.shape[1]} cols)") + print(f"{'='*100}") + res = classify_columns(df) + print(res) + print_classification(res) diff --git a/data/profile_data.py b/data/profile_data.py new file mode 100644 index 0000000..6ebbab7 --- /dev/null +++ b/data/profile_data.py @@ -0,0 +1,266 @@ +""" +Profiles every CSV in data/all_data and data/synthetic_data. +Per file: total variable count + named lists for each semantic category. +Goal: see what's directly usable by causal-learn (needs continuous numeric) +and what needs cleaning (and what kind). +""" + +import pandas as pd +import numpy as np +from pathlib import Path +from collections import defaultdict +import re +import json +import warnings +warnings.filterwarnings("ignore") + + +# ── keyword / pattern banks for heuristic classification ───────────────────── + +# Patterns matched case-insensitively against column names. +# The distinction between ordinal-int and nominal-int is inherently ambiguous +# from data alone — the heuristic here leans on naming conventions seen in +# this specific corpus (severity levels, scales, grade_level, etc.). + +ID_PATTERNS = re.compile( + r"^(unnamed|row|index|Unnamed)" + r"|_?id$|^id$|^id_" + r"|^fip$|^fips$|^sid$" + r"|^unit$|^cluster$" + r"|^hh_id$|^store_id$|^factory_id$" + r"|^household_id$|^student_id$|^school_id$" + r"|^region_id$|^club_id$|^village$|^villnum$" + r"|^statedisdec$", + re.IGNORECASE, +) + +YEAR_TIME_PATTERNS = re.compile( + r"^year$|^yr$|^yob$|^yod$" + r"|year_of_birth|quarter_of_birth" + r"|^quarter$|^quarter_num$" + r"|^date\d*$|^date2$" + r"|^academic_year$|^measurement_year$" + r"|_year$|^t$|^time$", + re.IGNORECASE, +) + +LOG_PREFIXES = re.compile(r"^l_|^ln|^log_|^lh", re.IGNORECASE) + +COUNT_PATTERNS = re.compile( + r"^tot|total|^num_|^n_|_count$|^count" + r"|^ncalls|^nregs|^popwt$|^population$|^police$", + re.IGNORECASE, +) + +RATE_PATTERNS = re.compile( + r"rate$|^pct|^perc|percent|proportion|_mean$|_pct$|share$|^p\d{4}", + re.IGNORECASE, +) + +# Names that suggest an ordinal scale (severity, grade_level, satisfaction, etc.) +ORDINAL_HINTS = re.compile( + r"severity|grade_level|satisfaction|education|experience|income_level" + r"|quality|score|level|scale|rank|priority", + re.IGNORECASE, +) + + +def profile_column(col_name: str, series: pd.Series, all_col_names: list[str]) -> str: + s = series.dropna() + n = len(series) + + if s.nunique() <= 1: + return "constant_or_near_constant" + + # ── ID / index ─────────────────────────────────────────────────────── + if ID_PATTERNS.search(col_name): + return "id_index" + numeric = pd.to_numeric(s, errors="coerce") + if numeric.notna().all() and s.nunique() / max(n, 1) > 0.90: + if numeric.min() in (0, 1) and (numeric.diff().dropna() == 1).mean() > 0.8: + return "id_index" + + # ── string-typed columns ───────────────────────────────────────────── + if series.dtype == object or s.dtype == object: + s_str = s.astype(str).str.strip() + nuniq = s_str.nunique() + avg_len = s_str.str.len().mean() + + lower_vals = set(s_str.str.lower().unique()) + if lower_vals <= {"true", "false", "yes", "no", "t", "f", "y", "n"}: + return "binary" + + if avg_len > 40 or nuniq > 100: + return "natural_language" + + return "nominal_categorical_str" + + # ── numeric from here ──────────────────────────────────────────────── + numeric = pd.to_numeric(s, errors="coerce") + if numeric.isna().any(): + return "nominal_categorical_str" + + vals = set(numeric.dropna().unique()) + nuniq = len(vals) + + if vals <= {0, 1, 0.0, 1.0}: + return "binary" + + if YEAR_TIME_PATTERNS.search(col_name): + return "year_time" + if nuniq > 2 and numeric.min() >= 1900 and numeric.max() <= 2030 and (numeric == numeric.astype(int)).all(): + return "year_time" + + # log-transformed: check if a plausible raw counterpart exists in the same file + if LOG_PREFIXES.search(col_name): + raw_candidates = [ + re.sub(r"^l_", "", col_name, flags=re.I), + re.sub(r"^ln", "", col_name, flags=re.I), + re.sub(r"^log_", "", col_name, flags=re.I), + re.sub(r"^lh", "", col_name, flags=re.I), + ] + has_raw = any(rc in all_col_names for rc in raw_candidates if rc and rc != col_name) + return "log_transformed_has_raw" if has_raw else "log_transformed" + + if COUNT_PATTERNS.search(col_name) and numeric.min() >= 0 and (numeric == numeric.round(0)).all(): + return "count" + + if RATE_PATTERNS.search(col_name): + return "bounded_rate_proportion" + if 0 <= numeric.min() and numeric.max() <= 1 and nuniq > 2: + return "bounded_rate_proportion" + + # ordinal vs nominal for low-cardinality integers + is_int = (numeric == numeric.astype(int)).all() + if is_int and 2 < nuniq <= 10: + if ORDINAL_HINTS.search(col_name): + return "ordinal_int" + return "nominal_int" + if is_int and 10 < nuniq <= 20: + return "nominal_int" + + return "continuous" + + +def profile_file(path: Path) -> dict | None: + try: + df = pd.read_csv(path, low_memory=False, nrows=50_000) + except Exception as e: + return {"error": str(e)} + + all_col_names = list(df.columns) + buckets = defaultdict(list) + missing_cols = {} + + for col in df.columns: + tag = profile_column(col, df[col], all_col_names) + buckets[tag].append(col) + miss = df[col].isna().mean() * 100 + if miss > 0: + missing_cols[col] = round(miss, 1) + + return { + "rows": df.shape[0], + "cols": df.shape[1], + "buckets": dict(buckets), + "missing": missing_cols, + } + + +DISPLAY_ORDER = [ + ("continuous", "Continuous"), + ("binary", "Binary indicators"), + ("ordinal_int", "Ordinal categoricals (int-encoded)"), + ("nominal_int", "Nominal categoricals (int-encoded)"), + ("nominal_categorical_str", "Nominal categoricals (string)"), + ("year_time", "Year / time columns"), + ("log_transformed_has_raw", "Log-transformed (raw version present)"), + ("log_transformed", "Log-transformed (standalone)"), + ("count", "Count data"), + ("bounded_rate_proportion", "Bounded rates / proportions"), + ("id_index", "IDs / indices"), + ("constant_or_near_constant", "Constant or near-constant"), + ("natural_language", "Natural language descriptions"), +] + +READY_TAGS = {"continuous", "binary"} + + +def main(): + base = Path(__file__).parent + folders = [base / "all_data", base / "synthetic_data"] + + all_profiles = {} + for folder in folders: + if not folder.exists(): + continue + for csv_path in sorted(folder.glob("*.csv")): + key = f"{folder.name}/{csv_path.name}" + all_profiles[key] = profile_file(csv_path) + + agg_type_counts = defaultdict(int) + agg_file_ready = [] + agg_file_needs_work = [] + + for fname, prof in sorted(all_profiles.items()): + print(f"\n{'━' * 80}") + print(f" {fname} ({prof['rows']} rows × {prof['cols']} cols)") + print(f"{'━' * 80}") + + if "error" in prof: + print(f" ERROR: {prof['error']}") + continue + + buckets = prof["buckets"] + file_clean = True + + for tag, label in DISPLAY_ORDER: + cols = buckets.get(tag, []) + if not cols: + continue + agg_type_counts[tag] += len(cols) + if tag not in READY_TAGS: + file_clean = False + marker = "✅" if tag in READY_TAGS else "🔧" + print(f" {marker} {label} ({len(cols)}): {', '.join(cols)}") + + if prof["missing"]: + items = [f"{c} ({v}%)" for c, v in prof["missing"].items()] + print(f" ⚠️ Missing values: {', '.join(items)}") + + if file_clean: + agg_file_ready.append(fname) + else: + agg_file_needs_work.append(fname) + + total_cols = sum(agg_type_counts.values()) + ready_cols = agg_type_counts.get("continuous", 0) + agg_type_counts.get("binary", 0) + + print(f"\n{'═' * 80}") + print(" AGGREGATE SUMMARY") + print(f"{'═' * 80}") + print(f" Files scanned: {len(all_profiles)}") + print(f" Ready for causal-learn (all continuous/binary): {len(agg_file_ready)}") + print(f" Need some cleaning: {len(agg_file_needs_work)}") + print(f" Total columns: {total_cols} (ready: {ready_cols}, need work: {total_cols - ready_cols})") + + print(f"\n Column type breakdown:") + for tag, label in DISPLAY_ORDER: + n = agg_type_counts.get(tag, 0) + if n == 0: + continue + marker = "✅" if tag in READY_TAGS else "🔧" + print(f" {marker} {label:48s} {n:4d}") + + print(f"\n Files ready as-is:") + for f in sorted(agg_file_ready): + print(f" ✅ {f}") + + out_path = base / "data_profile.json" + with open(out_path, "w") as fh: + json.dump(all_profiles, fh, indent=2, default=str) + print(f"\n 💾 Full profile → {out_path}") + + +if __name__ == "__main__": + main() From 3aeaea690c95d8b5193a5dee8b079056b0442e0a Mon Sep 17 00:00:00 2001 From: lukezhang01 Date: Wed, 24 Jun 2026 09:43:26 -0400 Subject: [PATCH 2/3] fci pc validation --- cais/discovery/causal_validation.py | 855 ++++++++++++++++++++++++++++ cais/discovery/data_preprocess.py | 18 +- 2 files changed, 866 insertions(+), 7 deletions(-) create mode 100644 cais/discovery/causal_validation.py diff --git a/cais/discovery/causal_validation.py b/cais/discovery/causal_validation.py new file mode 100644 index 0000000..a4cd054 --- /dev/null +++ b/cais/discovery/causal_validation.py @@ -0,0 +1,855 @@ +""" +Causal discovery validation pipeline. + +Validates CAIS outputs by learning causal graphs from data using causal-learn +and checking whether the assumptions behind CAIS's chosen method are consistent +with the discovered graph structure. +""" +from __future__ import annotations + +import json +import logging +import os +import signal +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import networkx as nx +import numpy as np +import pandas as pd + +from causallearn.search.ConstraintBased.FCI import fci +from causallearn.search.ConstraintBased.PC import pc +from causallearn.search.ScoreBased.GES import ges + +from cais.discovery.data_preprocess import ClassificationResult, classify_columns, prepare_data + +logger = logging.getLogger(__name__) + +# Fraction of discrete columns above which we fall back to KCI +CATEGORICAL_FRACTION_THRESHOLD = 0.5 + +# Subsample large datasets to keep discovery tractable +MAX_ROWS_FOR_DISCOVERY = 5000 + +# Skip all discovery when column count exceeds this (PC/GES blow up too) +MAX_COLS_FOR_DISCOVERY = 50 + +# Skip FCI when column count exceeds this (FCI is exponential in depth) +MAX_COLS_FOR_FCI = 25 + +# Minimum rows-per-column ratio; below this CI tests are underpowered +MIN_ROWS_PER_COL = 10 + +# Per-algorithm timeout in seconds +ALGORITHM_TIMEOUT = 120 + +# Map CAIS method names to canonical families +_OLS_FAMILY = {"ols", "linear_regression", "backdoor_adjustment",} +_IV_FAMILY = {"iv", "instrumental_variable", "2sls", "tsls"} +_DID_FAMILY = {"did", "did_canonical", "difference_in_differences"} + + +# ────────────────────────────────────────────────────────────────────── +# Data structures +# ────────────────────────────────────────────────────────────────────── + +@dataclass +class GraphResult: + """Holds a learned causal graph and helper lookups.""" + graph_matrix: np.ndarray # shape (p, p) + node_names: list[str] + algorithm: str # "pc", "ges", "fci" + + def _idx(self, name: str) -> int | None: + try: + return self.node_names.index(name) + except ValueError: + return None + + # ── edge queries ──────────────────────────────────────────────── + def has_directed_edge(self, src: str, dst: str) -> bool: + """True if src -> dst (arrow at dst, tail at src).""" + i, j = self._idx(src), self._idx(dst) + if i is None or j is None: + return False + return self.graph_matrix[j, i] == -1 and self.graph_matrix[i, j] == 1 + + def has_any_edge(self, a: str, b: str) -> bool: + i, j = self._idx(a), self._idx(b) + if i is None or j is None: + return False + return self.graph_matrix[i, j] != 0 or self.graph_matrix[j, i] != 0 + + def has_bidirected_edge(self, a: str, b: str) -> bool: + """True if a <-> b (arrow at both ends — FCI latent confounder).""" + i, j = self._idx(a), self._idx(b) + if i is None or j is None: + return False + return self.graph_matrix[i, j] == 1 and self.graph_matrix[j, i] == 1 + + def parents(self, node: str) -> list[str]: + """Return nodes with a directed edge into `node`.""" + j = self._idx(node) + if j is None: + return [] + out = [] + for i, name in enumerate(self.node_names): + # arrow at j from i: graph[j, i] == -1 and graph[i, j] == 1 + if self.graph_matrix[j, i] == -1 and self.graph_matrix[i, j] == 1: + out.append(name) + return out + + def children(self, node: str) -> list[str]: + """Return nodes that `node` has a directed edge to.""" + i = self._idx(node) + if i is None: + return [] + out = [] + for j, name in enumerate(self.node_names): + if self.graph_matrix[j, i] == -1 and self.graph_matrix[i, j] == 1: + out.append(name) + return out + + def descendants(self, node: str) -> set[str]: + """All descendants of `node` via directed edges (BFS).""" + visited: set[str] = set() + queue = self.children(node) + while queue: + cur = queue.pop(0) + if cur not in visited: + visited.add(cur) + queue.extend(self.children(cur)) + return visited + + def adjacent(self, node: str) -> list[str]: + """All nodes connected to `node` by any edge.""" + i = self._idx(node) + if i is None: + return [] + out = [] + for j, name in enumerate(self.node_names): + if self.graph_matrix[i, j] != 0 or self.graph_matrix[j, i] != 0: + out.append(name) + return out + + def to_nx_digraph(self) -> nx.DiGraph: + """Convert directed edges to a networkx DiGraph (ignores undirected/bidirected).""" + G = nx.DiGraph() + G.add_nodes_from(self.node_names) + for i, src in enumerate(self.node_names): + for j, dst in enumerate(self.node_names): + # src -> dst: graph[j, i] == -1 and graph[i, j] == 1 + if self.graph_matrix[j, i] == -1 and self.graph_matrix[i, j] == 1: + G.add_edge(src, dst) + return G + + +@dataclass +class ValidationResult: + """Collects validation flags and messages for a single CAIS output.""" + query_index: str + method_family: str + flags: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + info: list[str] = field(default_factory=list) + graphs_used: list[str] = field(default_factory=list) + skipped: bool = False + skip_reason: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "query_index": self.query_index, + "method_family": self.method_family, + "flags": self.flags, + "warnings": self.warnings, + "info": self.info, + "graphs_used": self.graphs_used, + "skipped": self.skipped, + "skip_reason": self.skip_reason, + } + + +# ────────────────────────────────────────────────────────────────────── +# Step 1: preprocess +# ────────────────────────────────────────────────────────────────────── + +def get_data_preprocess(dataset_path: str) -> tuple[pd.DataFrame, ClassificationResult]: + """Load CSV, classify columns, and return cleaned data + classification.""" + df = pd.read_csv(dataset_path) + classification = classify_columns(df) + cleaned = prepare_data(df, classification) + # Subsample if too large for discovery algorithms + if len(cleaned) > MAX_ROWS_FOR_DISCOVERY: + logger.info( + "Subsampling %d -> %d rows for discovery", len(cleaned), MAX_ROWS_FOR_DISCOVERY + ) + cleaned = cleaned.sample(n=MAX_ROWS_FOR_DISCOVERY, random_state=42).reset_index(drop=True) + return cleaned, classification + + +# ────────────────────────────────────────────────────────────────────── +# Step 2: learn graphs +# ────────────────────────────────────────────────────────────────────── + +def _pick_ci_test(classification: ClassificationResult) -> str: + """Choose conditional-independence test based on data types.""" + n_discovery = len(classification.cols_for_discovery) + if n_discovery == 0: + return "fisherz" + n_discrete = len(classification.discrete_flags) + frac = n_discrete / n_discovery + if frac > CATEGORICAL_FRACTION_THRESHOLD: + return "chisq" + return "fisherz" + + +class _Timeout: + """Context manager that raises TimeoutError after `seconds` using SIGALRM.""" + + def __init__(self, seconds: int, label: str = ""): + self.seconds = seconds + self.label = label + + def _handler(self, signum, frame): + raise TimeoutError(f"{self.label} timed out after {self.seconds}s") + + def __enter__(self): + self._old = signal.signal(signal.SIGALRM, self._handler) + signal.alarm(self.seconds) + return self + + def __exit__(self, *args): + signal.alarm(0) + signal.signal(signal.SIGALRM, self._old) + + +def get_causal_learn_graphs( + cleaned_df: pd.DataFrame, + classification: ClassificationResult, + alpha: float = 0.05, +) -> dict[str, GraphResult]: + """ + Run PC, GES, and FCI on the cleaned data. Returns a dict keyed by + algorithm name. Returns empty dict if the data is too wide or too + underpowered for reliable structure learning. + """ + data = cleaned_df.values.astype(np.float64) + col_names = list(cleaned_df.columns) + n_rows, n_cols = data.shape + graphs: dict[str, GraphResult] = {} + + if n_cols > MAX_COLS_FOR_DISCOVERY: + logger.info( + "Skipping all discovery: %d columns exceeds limit of %d", + n_cols, MAX_COLS_FOR_DISCOVERY, + ) + return graphs + + rows_per_col = n_rows / n_cols if n_cols > 0 else 0 + if rows_per_col < MIN_ROWS_PER_COL: + logger.info( + "Skipping all discovery: rows/columns ratio %.1f < %d " + "(n=%d, p=%d) — CI tests would be underpowered", + rows_per_col, MIN_ROWS_PER_COL, n_rows, n_cols, + ) + return graphs + + ci_test = _pick_ci_test(classification) + logger.info("CI test selected: %s (shape: %d x %d)", ci_test, n_rows, n_cols) + + # --- PC --- + logger.info(" Running PC...") + t0 = time.time() + try: + with _Timeout(ALGORITHM_TIMEOUT, "PC"): + cg = pc(data, alpha=alpha, indep_test=ci_test, show_progress=False) + graphs["pc"] = GraphResult( + graph_matrix=cg.G.graph.copy(), + node_names=col_names, + algorithm="pc", + ) + logger.info(" PC finished in %.1fs", time.time() - t0) + except Exception as e: + logger.warning(" PC failed after %.1fs: %s", time.time() - t0, e) + + # --- GES --- + logger.info(" Running GES...") + t0 = time.time() + try: + with _Timeout(ALGORITHM_TIMEOUT, "GES"): + record = ges(data, score_func="local_score_BIC", node_names=col_names) + graphs["ges"] = GraphResult( + graph_matrix=record["G"].graph.copy(), + node_names=col_names, + algorithm="ges", + ) + logger.info(" GES finished in %.1fs", time.time() - t0) + except Exception as e: + logger.warning(" GES failed after %.1fs: %s", time.time() - t0, e) + + # --- FCI (skip if too many columns) --- + if n_cols > MAX_COLS_FOR_FCI: + logger.info("Skipping FCI: %d columns exceeds limit of %d", n_cols, MAX_COLS_FOR_FCI) + else: + logger.info(" Running FCI...") + t0 = time.time() + try: + with _Timeout(ALGORITHM_TIMEOUT, "FCI"): + g, _ = fci(data, independence_test_method=ci_test, alpha=alpha, show_progress=False) + graphs["fci"] = GraphResult( + graph_matrix=g.graph.copy(), + node_names=col_names, + algorithm="fci", + ) + logger.info(" FCI finished in %.1fs", time.time() - t0) + except Exception as e: + logger.warning(" FCI failed after %.1fs: %s", time.time() - t0, e) + + return graphs + + +# ────────────────────────────────────────────────────────────────────── +# Step 3: validation checks +# ────────────────────────────────────────────────────────────────────── + +def _classify_method(method_str: str) -> str: + """Return 'ols', 'iv', 'did', or 'other'.""" + m = method_str.lower().strip() + if m in _OLS_FAMILY: + return "ols" + if m in _IV_FAMILY: + return "iv" + if m in _DID_FAMILY: + return "did" + return "other" + + +def _resolve_column_name(var_name: str, available: list[str]) -> str | None: + """ + Try to find `var_name` in the available column list, accounting for + one-hot encoding (e.g. "treatment" might become "treatment_Hawthorne"). + Returns exact match first, then prefix match, else None. + """ + if var_name in available: + return var_name + # Check for one-hot encoded variants + prefix = var_name + "_" + matches = [c for c in available if c.startswith(prefix)] + if matches: + return matches[0] # return first variant + return None + + +def _validate_ols( + final_result: dict, + graphs: dict[str, GraphResult], + vr: ValidationResult, +) -> None: + """Checks for OLS / linear regression / backdoor methods.""" + treatment = final_result.get("treatment_variable") + outcome = final_result.get("outcome_variable") + covariates = final_result.get("covariates") or [] + + if not treatment or not outcome: + vr.warnings.append("Missing treatment or outcome variable — cannot validate.") + return + + # --- PC + GES checks --- + for alg_name in ("pc", "ges"): + g = graphs.get(alg_name) + if g is None: + continue + vr.graphs_used.append(alg_name) + + t_col = _resolve_column_name(treatment, g.node_names) + y_col = _resolve_column_name(outcome, g.node_names) + + if t_col is None or y_col is None: + vr.info.append( + f"[{alg_name}] Treatment or outcome not found in graph columns " + f"(T={treatment}, Y={outcome}). Possibly dropped during preprocessing." + ) + continue + + # Check covariates: should be parents of T, Y, or both + t_descendants = g.descendants(t_col) + t_parents = set(g.parents(t_col)) + y_parents = set(g.parents(y_col)) + adjustment_set = t_parents | y_parents + + for cov in covariates: + cov_col = _resolve_column_name(cov, g.node_names) + if cov_col is None: + continue + if cov_col in t_descendants: + vr.flags.append( + f"[{alg_name}] Covariate '{cov}' is a descendant of treatment " + f"'{treatment}' — adjusting for it may induce bias." + ) + elif cov_col not in adjustment_set and not g.has_any_edge(cov_col, t_col) and not g.has_any_edge(cov_col, y_col): + vr.info.append( + f"[{alg_name}] Covariate '{cov}' has no edge to treatment or outcome " + f"in the graph — may be unnecessary but not harmful." + ) + + # D-separation check: does the covariate set block all backdoor paths? + try: + nx_g = g.to_nx_digraph() + if t_col in nx_g and y_col in nx_g: + # Mutilated graph: remove all edges out of treatment + mutilated = nx_g.copy() + mutilated.remove_edges_from(list(nx_g.out_edges(t_col))) + cov_set = set() + for cov in covariates: + c = _resolve_column_name(cov, g.node_names) + if c is not None: + cov_set.add(c) + if not nx.d_separated(mutilated, {t_col}, {y_col}, cov_set): + vr.flags.append( + f"[{alg_name}] Covariates do not block all backdoor paths from " + f"'{treatment}' to '{outcome}' — backdoor criterion may not be satisfied." + ) + else: + vr.info.append( + f"[{alg_name}] Covariates satisfy the backdoor criterion " + f"(d-separation holds in the mutilated graph)." + ) + except Exception as e: + vr.warnings.append(f"[{alg_name}] D-separation check failed: {e}") + + # --- FCI check for latent confounding --- + fci_g = graphs.get("fci") + if fci_g is not None: + vr.graphs_used.append("fci") + t_col = _resolve_column_name(treatment, fci_g.node_names) + y_col = _resolve_column_name(outcome, fci_g.node_names) + + if t_col and y_col and fci_g.has_bidirected_edge(t_col, y_col): + vr.flags.append( + f"[fci] Bidirected edge {treatment} <-> {outcome} detected — " + f"suggests latent confounding. Conditional ignorability may be " + f"violated and the OLS effect estimate may be biased." + ) + elif t_col and y_col: + vr.info.append( + f"[fci] No bidirected edge between {treatment} and {outcome} — " + f"no evidence of latent confounding from FCI." + ) + + +def _validate_iv( + final_result: dict, + graphs: dict[str, GraphResult], + vr: ValidationResult, +) -> None: + """Checks for instrumental variable methods.""" + treatment = final_result.get("treatment_variable") + outcome = final_result.get("outcome_variable") + instrument = final_result.get("instrument_variable") + + if not treatment or not outcome: + vr.warnings.append("Missing treatment or outcome variable — cannot validate.") + return + if not instrument: + vr.warnings.append("No instrument variable specified in CAIS output — cannot validate IV assumptions.") + return + + # --- FCI is the primary tool for IV validation --- + fci_g = graphs.get("fci") + if fci_g is not None: + vr.graphs_used.append("fci") + t_col = _resolve_column_name(treatment, fci_g.node_names) + y_col = _resolve_column_name(outcome, fci_g.node_names) + z_col = _resolve_column_name(instrument, fci_g.node_names) + + if not all([t_col, y_col, z_col]): + vr.info.append( + f"[fci] One or more IV variables not found in graph columns " + f"(Z={instrument}, T={treatment}, Y={outcome})." + ) + else: + # Exclusion restriction: Z should NOT directly cause Y + if fci_g.has_directed_edge(z_col, y_col): + vr.flags.append( + f"[fci] Directed edge {instrument} -> {outcome} — " + f"exclusion restriction may be violated." + ) + # Reverse: Y causing Z means instrument is endogenous + if fci_g.has_directed_edge(y_col, z_col): + vr.flags.append( + f"[fci] Directed edge {outcome} -> {instrument} — " + f"instrument may be endogenous (caused by the outcome)." + ) + # Bidirected Z <-> Y: latent common cause of instrument and outcome + if fci_g.has_bidirected_edge(z_col, y_col): + vr.flags.append( + f"[fci] Bidirected edge {instrument} <-> {outcome} — " + f"latent common cause of instrument and outcome, " + f"violating instrument exogeneity." + ) + # Other edge types (undirected/circle) + if (not fci_g.has_directed_edge(z_col, y_col) + and not fci_g.has_directed_edge(y_col, z_col) + and not fci_g.has_bidirected_edge(z_col, y_col)): + if fci_g.has_any_edge(z_col, y_col): + vr.warnings.append( + f"[fci] Undirected/circle edge between {instrument} and " + f"{outcome} — exclusion restriction uncertain." + ) + else: + vr.info.append( + f"[fci] No direct edge from {instrument} to {outcome} — " + f"exclusion restriction appears satisfied." + ) + + # Relevance: Z should affect T + if fci_g.has_directed_edge(z_col, t_col): + vr.info.append( + f"[fci] Directed edge {instrument} -> {treatment} — " + f"instrument relevance condition supported." + ) + elif fci_g.has_any_edge(z_col, t_col): + vr.info.append( + f"[fci] Edge between {instrument} and {treatment} — " + f"instrument relevance plausible but direction uncertain." + ) + else: + vr.flags.append( + f"[fci] No edge from {instrument} to {treatment} — " + f"instrument relevance condition may be violated." + ) + + # T <-> Y bidirected is actually expected for IV (unmeasured confounding) + if fci_g.has_bidirected_edge(t_col, y_col): + vr.info.append( + f"[fci] Bidirected edge {treatment} <-> {outcome} detected — " + f"consistent with IV rationale (unmeasured confounding present)." + ) + else: + vr.warnings.append( + f"[fci] No bidirected edge between {treatment} and {outcome} — " + f"if no latent confounding, IV may be unnecessary; " + f"OLS/backdoor adjustment might suffice." + ) + + # --- PC/GES cross-check --- + for alg_name in ("pc", "ges"): + g = graphs.get(alg_name) + if g is None: + continue + vr.graphs_used.append(alg_name) + + t_col = _resolve_column_name(treatment, g.node_names) + y_col = _resolve_column_name(outcome, g.node_names) + z_col = _resolve_column_name(instrument, g.node_names) + + if not all([t_col, y_col, z_col]): + continue + + # Exclusion restriction + if g.has_directed_edge(z_col, y_col): + vr.flags.append( + f"[{alg_name}] Direct edge {instrument} -> {outcome} — " + f"exclusion restriction may be violated." + ) + + # Relevance: Z should affect T + if not g.has_directed_edge(z_col, t_col) and not g.has_any_edge(z_col, t_col): + vr.flags.append( + f"[{alg_name}] No edge from {instrument} to {treatment} — " + f"instrument relevance condition may be violated." + ) + + +def _validate_did( + final_result: dict, + graphs: dict[str, GraphResult], + vr: ValidationResult, +) -> None: + """Checks for difference-in-differences methods.""" + treatment = final_result.get("treatment_variable") + outcome = final_result.get("outcome_variable") + covariates = final_result.get("covariates") or [] + + if not treatment or not outcome: + vr.warnings.append("Missing treatment or outcome variable — cannot validate.") + return + + # --- PC/GES: check for post-treatment bias in covariates --- + for alg_name in ("pc", "ges"): + g = graphs.get(alg_name) + if g is None: + continue + vr.graphs_used.append(alg_name) + + t_col = _resolve_column_name(treatment, g.node_names) + if t_col is None: + continue + + t_descendants = g.descendants(t_col) + + for cov in covariates: + cov_col = _resolve_column_name(cov, g.node_names) + if cov_col is None: + continue + if cov_col in t_descendants: + vr.flags.append( + f"[{alg_name}] Covariate '{cov}' is a descendant of treatment " + f"'{treatment}' — conditioning on post-treatment variables " + f"may bias the DiD estimate." + ) + + # --- FCI: flag latent time-varying confounders --- + fci_g = graphs.get("fci") + if fci_g is not None: + vr.graphs_used.append("fci") + t_col = _resolve_column_name(treatment, fci_g.node_names) + y_col = _resolve_column_name(outcome, fci_g.node_names) + + if t_col and y_col and fci_g.has_bidirected_edge(t_col, y_col): + vr.warnings.append( + f"[fci] Bidirected edge {treatment} <-> {outcome} — " + f"latent confounders detected. Parallel trends assumption " + f"may not handle time-varying unobserved confounding." + ) + elif t_col and y_col: + vr.info.append( + f"[fci] No bidirected edge between {treatment} and {outcome} — " + f"no evidence of latent confounders from FCI." + ) + + +def get_validation( + cais_output: dict, + graphs: dict[str, GraphResult], + query_index: str, +) -> ValidationResult: + """ + Validate a single CAIS output entry against discovered graphs. + + Parameters + ---------- + cais_output : dict + A single entry from the CAIS output JSON (contains 'final_result', etc.). + graphs : dict[str, GraphResult] + Graphs learned from the same dataset. + query_index : str + The key/index of this entry in the output file. + """ + final_result = cais_output.get("final_result", {}) + method_raw = final_result.get("method", "unknown") + method_family = _classify_method(method_raw) + + vr = ValidationResult(query_index=query_index, method_family=method_family) + + if not graphs: + vr.skipped = True + vr.skip_reason = "No graphs were learned (all algorithms failed)." + return vr + + if method_family == "ols": + _validate_ols(final_result, graphs, vr) + elif method_family == "iv": + _validate_iv(final_result, graphs, vr) + elif method_family == "did": + _validate_did(final_result, graphs, vr) + else: + vr.info.append( + f"Method '{method_raw}' (family='{method_family}') — " + f"no specific validation checks implemented yet." + ) + + # Deduplicate graphs_used + vr.graphs_used = list(dict.fromkeys(vr.graphs_used)) + return vr + + +# ────────────────────────────────────────────────────────────────────── +# Main pipeline +# ────────────────────────────────────────────────────────────────────── + +def run_validation_pipeline( + cais_outputs_dir: str, + base_data_dir: str | None = None, + output_path: str | None = None, + alpha: float = 0.05, + datasets: set[str] | None = None, +) -> dict[str, list[dict]]: + """ + For every JSON file in `cais_outputs_dir`, load each query entry, + preprocess the dataset, learn graphs, and validate. + + Parameters + ---------- + cais_outputs_dir : str + Directory containing CAIS output JSON files. + base_data_dir : str or None + If dataset_path in the JSON is relative, resolve it relative to this. + Defaults to the project root (two levels up from this file). + output_path : str or None + If provided, write all results to this JSON file. + alpha : float + Significance level for CI tests. + datasets : set of str or None + If provided, only process entries whose dataset filename (without + extension) is in this set. E.g. {"women", "smoking2", "rct_data_0"}. + + Returns + ------- + dict mapping filename -> list of ValidationResult dicts. + """ + if base_data_dir is None: + base_data_dir = str(Path(__file__).resolve().parents[2]) + + outputs_dir = Path(cais_outputs_dir) + all_results: dict[str, list[dict]] = {} + + # Cache preprocessed data + graphs per dataset path + _graph_cache: dict[str, dict[str, GraphResult]] = {} + _preprocess_cache: dict[str, tuple[pd.DataFrame, ClassificationResult]] = {} + + json_files = sorted(outputs_dir.glob("*.json")) + logger.info("Found %d JSON files in %s", len(json_files), cais_outputs_dir) + + for json_path in json_files: + filename = json_path.name + logger.info("Processing %s", filename) + print(f"\n{'='*60}") + print(f" {filename}") + print(f"{'='*60}") + + with open(json_path) as f: + data = json.load(f) + + file_results: list[dict] = [] + + for idx, entry in data.items(): + if isinstance(entry, str): + # Error entries are stored as plain strings + file_results.append({ + "query_index": idx, + "skipped": True, + "skip_reason": f"CAIS output was an error string: {entry[:200]}", + }) + continue + + dataset_path_raw = entry.get("dataset_path", "") + + # Filter by dataset name if requested + if datasets is not None: + dataset_stem = Path(dataset_path_raw).stem + if dataset_stem not in datasets: + continue + + dataset_path = dataset_path_raw + if not os.path.isabs(dataset_path): + dataset_path = os.path.join(base_data_dir, dataset_path) + + if not os.path.exists(dataset_path): + file_results.append({ + "query_index": idx, + "skipped": True, + "skip_reason": f"Dataset not found: {dataset_path}", + }) + continue + + # Preprocess + learn graphs (cached per dataset) + if dataset_path not in _graph_cache: + dataset_stem = Path(dataset_path).stem + try: + logger.info("Preprocessing dataset '%s'...", dataset_stem) + cleaned_df, classification = get_data_preprocess(dataset_path) + _preprocess_cache[dataset_path] = (cleaned_df, classification) + logger.info( + "Learning graphs for dataset '%s' (shape: %d x %d)", + dataset_stem, cleaned_df.shape[0], cleaned_df.shape[1], + ) + graphs = get_causal_learn_graphs(cleaned_df, classification, alpha=alpha) + _graph_cache[dataset_path] = graphs + logger.info("Done with dataset '%s' — %d graphs learned", dataset_stem, len(graphs)) + except Exception as e: + logger.error("Failed to process dataset %s: %s", dataset_path, e) + _graph_cache[dataset_path] = {} + + graphs = _graph_cache[dataset_path] + + vr = get_validation(entry, graphs, query_index=idx) + result_dict = vr.to_dict() + result_dict["query"] = entry.get("query", "") + result_dict["dataset_path"] = dataset_path_raw + result_dict["cais_method"] = entry.get("final_result", {}).get("method", "unknown") + file_results.append(result_dict) + + # Print summary + status = "SKIP" if vr.skipped else ("FLAG" if vr.flags else "OK") + print(f" [{idx}] {status} | {vr.method_family} | flags={len(vr.flags)} warns={len(vr.warnings)}") + for flag in vr.flags: + print(f" FLAG: {flag}") + for warn in vr.warnings: + print(f" WARN: {warn}") + + all_results[filename] = file_results + + if output_path: + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + with open(output_path, "w") as f: + json.dump(all_results, f, indent=2) + print(f"\nResults written to {output_path}") + + return all_results + + +# ────────────────────────────────────────────────────────────────────── +# CLI entry point +# ────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Validate CAIS outputs with causal discovery.") + parser.add_argument( + "--cais_outputs_dir", + type=str, + default="cais_outputs", + help="Directory containing CAIS output JSON files.", + ) + parser.add_argument( + "--base_data_dir", + type=str, + default=None, + help="Base directory for resolving relative dataset paths.", + ) + parser.add_argument( + "--output", + type=str, + default="validation_results.json", + help="Path to write validation results JSON.", + ) + parser.add_argument( + "--alpha", + type=float, + default=0.05, + help="Significance level for CI tests.", + ) + parser.add_argument( + "--datasets", + type=str, + default=None, + help="Comma-separated list of dataset names (without extension) to validate. " + "E.g. 'women,smoking2,rct_data_0'. If omitted, all datasets are processed.", + ) + args = parser.parse_args() + + ds_filter = set(args.datasets.split(",")) if args.datasets else None + + logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") + run_validation_pipeline( + cais_outputs_dir=args.cais_outputs_dir, + base_data_dir=args.base_data_dir, + output_path=args.output, + alpha=args.alpha, + datasets=ds_filter, + ) diff --git a/cais/discovery/data_preprocess.py b/cais/discovery/data_preprocess.py index 0dabd97..3de60c3 100644 --- a/cais/discovery/data_preprocess.py +++ b/cais/discovery/data_preprocess.py @@ -262,7 +262,7 @@ def classify_columns( if is_numeric and n_unique / n_rows > id_uniqueness_threshold: if _is_integer_valued(s) and _looks_sequential(s): is_id = True - if is_numeric and _is_integer_valued(s): + if is_numeric and _is_integer_valued(s) and _looks_sequential(s): # the if statements underneath, such as searching for "fip" is based on the specific dataset i gave claude (fip is taken from abortion dataset) low_name = col.lower() if any(tag in low_name for tag in ["fip", "fips", "state_id", "county", "country_code", "entity", "unit"]): is_id = True @@ -278,12 +278,10 @@ def classify_columns( low = col.lower().strip() if _TEMPORAL_PATTERNS.match(low): is_temporal = True - if is_numeric and ("year" in low or "yr" in low): + if is_numeric and ("year" in low or "yr" in low or "date" in low): vals = s.dropna() if len(vals) > 0 and vals.min() >= 1900 and vals.max() <= 2100: is_temporal = True - if is_numeric and "date" in low: - is_temporal = True if is_temporal: col_type = "temporal" @@ -392,7 +390,7 @@ def classify_columns( note = "Numeric, high cardinality" else: col_type = "nominal" - note = "Fallback: unclassified string column" + note = "Unclassified string column" # Store metadata result.metadata[col] = { @@ -480,16 +478,22 @@ def prepare_data( ct = meta["col_type"] if ct == "binary_string": - # Map the two unique string values to 0 / 1 alphabetically vals = sorted(out[col].unique()) + if len(vals) < 2: + # Collapsed to a single value after dropna — constant column, drop it + out = out.drop(columns=[col]) + continue out[col] = out[col].map({vals[0]: 0, vals[1]: 1}).astype(np.float64) elif ct in ("nominal", "low_cardinality_string"): cols_to_onehot.append(col) elif ct == "binary_numeric": - # Ensure values are exactly {0, 1} vals = sorted(out[col].dropna().unique()) + if len(vals) < 2: + # Collapsed to a single value after dropna — constant column, drop it + out = out.drop(columns=[col]) + continue if vals != [0, 1]: out[col] = out[col].map({vals[0]: 0, vals[1]: 1}).astype(np.float64) else: From c2a25c5552c2d45598beaa5cc1d511baf5a6cd48 Mon Sep 17 00:00:00 2001 From: lukezhang01 Date: Wed, 24 Jun 2026 09:49:59 -0400 Subject: [PATCH 3/3] integration changes --- analyze_methods.py | 84 ++++++++++++++++++++++++++++ cais/agent.py | 8 ++- cais/components/query_interpreter.py | 12 ++-- 3 files changed, 96 insertions(+), 8 deletions(-) create mode 100644 analyze_methods.py diff --git a/analyze_methods.py b/analyze_methods.py new file mode 100644 index 0000000..35cfae0 --- /dev/null +++ b/analyze_methods.py @@ -0,0 +1,84 @@ +""" +Analyze cais_outputs JSON files to summarize method choices per model and data type. +""" + +import json +import os +import pandas as pd +from collections import Counter +from pathlib import Path + +OUTPUT_DIR = "cais_outputs" + + +def parse_filename(fname): + """Extract model name and data type from filename like 'gpt-4o_real.json'.""" + stem = Path(fname).stem # e.g. 'gpt-4o_real' + # Split on last underscore to separate model from data type + parts = stem.rsplit("_", 1) + if len(parts) == 2: + return parts[0], parts[1] + return stem, "unknown" + + +def get_data_shape(dataset_path): + """Read csv and return (rows, cols) or None if file not found.""" + try: + df = pd.read_csv(dataset_path) + return df.shape + except Exception: + return None + + +def main(): + json_files = sorted( + f for f in os.listdir(OUTPUT_DIR) if f.endswith(".json") + ) + + for fname in json_files: + model, data_type = parse_filename(fname) + filepath = os.path.join(OUTPUT_DIR, fname) + + with open(filepath) as f: + data = json.load(f) + + # Group queries by dataset to show shape info + dataset_queries = {} # dataset_path -> list of chosen methods + for key, entry in data.items(): + dataset_path = entry.get("dataset_path", "unknown") + chosen_method = entry.get("final_result", {}).get("method", "unknown") + dataset_queries.setdefault(dataset_path, []).append(chosen_method) + + # Print header + print("=" * 70) + print(f"Model: {model} | Data Type: {data_type}") + print(f"Total queries: {len(data)}") + print("-" * 70) + + # Overall method counts + all_methods = [] + for methods in dataset_queries.values(): + all_methods.extend(methods) + method_counts = Counter(all_methods) + + print(f"{'Method':<30} {'Count':>6}") + print(f"{'------':<30} {'-----':>6}") + for method, count in sorted(method_counts.items(), key=lambda x: -x[1]): + print(f"{method:<30} {count:>6}") + print() + + # Per-dataset breakdown + for dpath, methods in sorted(dataset_queries.items()): + shape = get_data_shape(dpath) + shape_str = f"{shape[0]} x {shape[1]}" if shape else "N/A" + dataset_name = Path(dpath).stem + print(f" Dataset: {dataset_name:<35} Shape: {shape_str}") + per_ds_counts = Counter(methods) + for method, count in sorted(per_ds_counts.items(), key=lambda x: -x[1]): + print(f" {method:<28} {count:>4}") + + print() + + +if __name__ == "__main__": + main() diff --git a/cais/agent.py b/cais/agent.py index 6651f0a..0a37909 100644 --- a/cais/agent.py +++ b/cais/agent.py @@ -345,6 +345,10 @@ def run_causal_analysis(query: str, dataset_path: str, original_query = input_parsing_result["original_query"], excluded_methods=None) + # Check for errors from method selection before accessing 'method_info' + if "error" in method_selector_output and "method_info" not in method_selector_output: + raise ValueError(f"Method selection failed: {method_selector_output['error']}") + # NEW: Select control variables based on chosen method method_info = MethodInfo( **method_selector_output['method_info'] @@ -425,8 +429,8 @@ def run_causal_analysis(query: str, dataset_path: str, logger.info("Causal analysis run finished.") # Remove the cleaned csv - logger.info("Removing cleaned csv.") - os.remove(cleaned_path) + # logger.info("Removing cleaned csv.") + # os.remove(cleaned_path) # Ensure result is a dict and extract the 'output' part if isinstance(result, dict): diff --git a/cais/components/query_interpreter.py b/cais/components/query_interpreter.py index de661c9..0f22df6 100644 --- a/cais/components/query_interpreter.py +++ b/cais/components/query_interpreter.py @@ -220,7 +220,7 @@ def interpret_query(query_info: Dict[str, Any], dataset_analysis: Dict[str, Any] # --- Identify Treatment --- - treatment_hints = query_info.get("potential_treatments", []) + treatment_hints = query_info.get("potential_treatments") or [] dataset_treatments = dataset_analysis.get("potential_treatments", []) treatment_variable = _identify_variable_hybrid(role="treatment", query_hints=treatment_hints, dataset_suggestions=dataset_treatments, columns=columns, @@ -232,7 +232,7 @@ def interpret_query(query_info: Dict[str, Any], dataset_analysis: Dict[str, Any] # --- Identify Outcome --- - outcome_hints = query_info.get("outcome_hints", []) + outcome_hints = query_info.get("potential_outcomes") or [] dataset_outcomes = dataset_analysis.get("potential_outcomes", []) outcome_variable = _identify_variable_hybrid(role="outcome", query_hints=outcome_hints, dataset_suggestions=dataset_outcomes, columns=columns, column_categories=column_categories, @@ -242,14 +242,14 @@ def interpret_query(query_info: Dict[str, Any], dataset_analysis: Dict[str, Any] logger.info(f"Identified Outcome: {outcome_variable}") # --- Identify Covariates --- - covariate_hints = query_info.get("covariates_hints", []) + covariate_hints = query_info.get("covariates_hints") or [] covariates = _identify_covariates_hybrid("covars", treatment_variable=treatment_variable, outcome_variable=outcome_variable, columns=columns, column_categories=column_categories, query_hints=covariate_hints, query_text=query_text, dataset_description=dataset_description, llm=llm) logger.info(f"Identified Covariates: {covariates}") # --- Identify Confounders --- - confounder_hints = query_info.get("covariates_hints", []) + confounder_hints = query_info.get("covariates_hints") or [] confounders = _identify_covariates_hybrid("confounders", treatment_variable=treatment_variable, outcome_variable=outcome_variable, columns=columns, column_categories=column_categories, query_hints=confounder_hints, query_text=query_text, dataset_description=dataset_description, llm=llm) @@ -456,11 +456,11 @@ def _identify_variable_hybrid(role: str, query_hints: List[str], dataset_suggest if not available_columns: return None # 1. Exact matches from hints - for hint in query_hints: + for hint in (query_hints or []): if hint in available_columns: candidates.add(hint) # 2. Add dataset suggestions - for sugg in dataset_suggestions: + for sugg in (dataset_suggestions or []): if sugg in available_columns: candidates.add(sugg)