"""DSL operators — the small composable language the GP engine will search over. These operators are label-agnostic: `Select` picks columns by whatever labels the matrix carries (real gene symbols or opaque IDs), `Reduce` collapses rows, and so on. The H1 fixture composes them over the NAMED matrix; the engine (next chunk) will compose them over the ANONYMISED matrix via [`airgap`](../airgap). This module has no gene names, no pathway names, no MSI-specific constants. The only convention it knows about is the Cohort schema produced by `Load`: three frames (expression, clinical, labels) sharing a sample-id index. """ from __future__ import annotations import re from dataclasses import dataclass from pathlib import Path from typing import Callable import numpy as np import pandas as pd from sklearn.linear_model import LogisticRegression from sklearn.metrics import balanced_accuracy_score, roc_auc_score from sklearn.model_selection import train_test_split _OPAQUE_ID_RE = re.compile(r"^g\d+$") # --- Cohort ----------------------------------------------------------------- @dataclass class Cohort: """Three aligned frames sharing a sample-id index. - expression : samples x features (features may be gene symbols OR opaque IDs) - clinical : samples x {stage, age, sex, site, os_event, os_months, ...} - labels : samples x {msi_status, tmb} """ expression: pd.DataFrame clinical: pd.DataFrame labels: pd.DataFrame @property def sample_ids(self) -> pd.Index: return self.expression.index def restrict(self, mask: pd.Series) -> "Cohort": mask = mask.reindex(self.sample_ids).fillna(False).astype(bool) idx = self.sample_ids[mask] return Cohort( expression=self.expression.loc[idx], clinical=self.clinical.loc[idx], labels=self.labels.loc[idx], ) # --- Load ------------------------------------------------------------------- def Load(source: str | Path = "processed") -> Cohort: """Read the data-pipeline parquets into a samples-as-rows Cohort. `source="processed"` resolves the path from `data_pipeline.schema`; otherwise `source` is treated as a directory containing clinical.parquet + expression.parquet. """ if isinstance(source, str) and source == "processed": from data_pipeline import schema processed_dir = schema.PROCESSED_DIR else: processed_dir = Path(source) clin = pd.read_parquet(processed_dir / "clinical.parquet").set_index("sample_id") expr = pd.read_parquet(processed_dir / "expression.parquet").T # samples x genes expr.index.name = "sample_id" samples = expr.index.intersection(clin.index) expr = expr.loc[samples] clin = clin.loc[samples] clinical_cols = [ c for c in [ "stage", "age", "sex", "site", "os_event", "os_months", # HNSC-only confounders (filtered out for cohorts that # don't carry them). "race", "ethnicity", "tissue_site", "icd_o_3_site", "is_oropharynx", ] if c in clin.columns ] label_cols = [ c for c in ["msi_status", "tmb", "hpv_status"] if c in clin.columns ] return Cohort( expression=expr, clinical=clin[clinical_cols].copy(), labels=clin[label_cols].copy(), ) # --- Select / Reduce / Split ------------------------------------------------ def Select(matrix: pd.DataFrame, feature_ids: list[str]) -> pd.DataFrame: """Restrict matrix columns to feature_ids, preserving their order.""" missing = [f for f in feature_ids if f not in matrix.columns] if missing: raise KeyError(f"Select: features missing from matrix: {missing}") return matrix.loc[:, list(feature_ids)] _REDUCE_AGGS = {"mean", "median", "max", "min", "var"} def Reduce(matrix: pd.DataFrame, agg: str = "mean") -> pd.Series: """Collapse each row to a single score. Supported aggregators: mean, median, max, min, var.""" if agg not in _REDUCE_AGGS: raise ValueError( f"Reduce: unknown agg={agg!r}; supported: {sorted(_REDUCE_AGGS)}" ) if agg == "mean": return matrix.mean(axis=1) if agg == "median": return matrix.median(axis=1) if agg == "max": return matrix.max(axis=1) if agg == "min": return matrix.min(axis=1) return matrix.var(axis=1) def Split(cohort: Cohort, predicate: Callable[[Cohort], pd.Series]) -> tuple[Cohort, Cohort]: """Split a cohort by a boolean predicate(cohort) -> Series. Returns (subset_true, subset_false). The predicate must return a boolean Series indexed by sample_id; missing values are treated as False. """ mask = predicate(cohort) if not isinstance(mask, pd.Series): raise TypeError("Split: predicate must return a pandas Series of bools") mask = mask.reindex(cohort.sample_ids).fillna(False).astype(bool) return cohort.restrict(mask), cohort.restrict(~mask) # --- Associate / Effect ----------------------------------------------------- def _align(a: pd.Series, b: pd.Series) -> tuple[pd.Series, pd.Series]: df = pd.concat([pd.Series(a).astype(float), pd.Series(b).astype(float)], axis=1).dropna() return df.iloc[:, 0], df.iloc[:, 1] def Associate(a: pd.Series, b: pd.Series, kind: str = "pearson") -> float: """Observational correlation; supports 'pearson' and 'spearman'.""" a, b = _align(a, b) if kind == "pearson": return float(a.corr(b, method="pearson")) if kind == "spearman": return float(a.corr(b, method="spearman")) raise ValueError(f"Associate: unknown kind={kind!r}") def _design(adjust: pd.DataFrame) -> pd.DataFrame: """Build a numeric design matrix: object/bool/category -> one-hot, numeric kept.""" parts: list[pd.DataFrame] = [] for col in adjust.columns: s = adjust[col] if s.dtype.kind in ("O", "b") or isinstance(s.dtype, pd.CategoricalDtype): d = pd.get_dummies(s, prefix=col, drop_first=True, dummy_na=False).astype(float) if d.shape[1] > 0: parts.append(d) else: parts.append(s.astype(float).to_frame()) return pd.concat(parts, axis=1) if parts else pd.DataFrame(index=adjust.index) def _residualise(y: pd.Series, X: pd.DataFrame) -> pd.Series: X1 = np.column_stack([np.ones(len(X)), X.values]) if X.shape[1] else np.ones((len(y), 1)) beta, *_ = np.linalg.lstsq(X1, y.values, rcond=None) return pd.Series(y.values - X1 @ beta, index=y.index) @dataclass class EffectResult: """Adjusted vs unadjusted association under observational backdoor adjustment. `partial_corr` is the pearson correlation of the residuals of cause and effect after each is regressed on the covariates. `unadjusted` is the plain pearson on the SAME rows (so the comparison is apples-to-apples). `n_used` is the number of rows kept after dropping any NA in cause/effect/covariates. """ partial_corr: float unadjusted: float n_used: int note: str = ( "Observational backdoor adjustment: only as good as the measured " "confounders. Unmeasured confounding can still bias the estimate." ) def Effect(cause: pd.Series, effect: pd.Series, adjust: pd.DataFrame) -> EffectResult: """Partial correlation of cause vs effect given `adjust`.""" df = pd.concat([cause.rename("__cause"), effect.rename("__effect"), adjust], axis=1).dropna() n_used = len(df) if n_used < 3: return EffectResult(float("nan"), float("nan"), n_used) X = _design(df.drop(columns=["__cause", "__effect"])) r_cause = _residualise(df["__cause"], X) r_effect = _residualise(df["__effect"], X) partial = float(r_cause.corr(r_effect, method="pearson")) unadj = float(df["__cause"].corr(df["__effect"], method="pearson")) return EffectResult(partial_corr=partial, unadjusted=unadj, n_used=n_used) # --- Search ----------------------------------------------------------------- def Search( matrix: pd.DataFrame, objective: Callable[[pd.Series], float], k: int, ) -> list[str]: """Placeholder baseline: rank single features by `objective(scores)`, top-k. Sees only the airgapped view: column labels MUST match ``^g\\d+$``. The GP engine (next chunk) replaces this with a real search over compositions. The objective is a function of the feature's scores (and any labels the caller has closed over) — never of feature identity. """ cols = list(matrix.columns) bad = [c for c in cols if not _OPAQUE_ID_RE.match(str(c))] if bad: raise ValueError( "Search: matrix columns must be opaque IDs matching ^g\\d+$; " f"got non-conforming columns e.g. {bad[:5]}" ) scored = [(c, float(objective(matrix[c]))) for c in cols] scored.sort(key=lambda kv: kv[1], reverse=True) return [c for c, _ in scored[: int(k)]] # --- Fit / Apply ------------------------------------------------------------ @dataclass class FitResult: """A trained logistic-regression model plus its held-out performance. `auroc` and `balanced_acc` are computed on the held-out test split — we report balanced accuracy rather than raw accuracy because the MSI-H vs MSS classes are imbalanced. `train_index` / `test_index` are the original sample-id Index objects on each side of the split (None if `state` had no index), so downstream code can pick example patients that are genuinely out-of-sample. """ model: LogisticRegression auroc: float balanced_acc: float n_train: int n_test: int feature_names: list[str] train_index: pd.Index | None = None test_index: pd.Index | None = None def _state_matrix(state) -> tuple[np.ndarray, list[str], pd.Index | None]: if isinstance(state, pd.Series): name = state.name if state.name is not None else "feature" return state.values.reshape(-1, 1), [str(name)], state.index if isinstance(state, pd.DataFrame): return state.values, [str(c) for c in state.columns], state.index arr = np.asarray(state) if arr.ndim == 1: arr = arr.reshape(-1, 1) return arr, [f"x{i}" for i in range(arr.shape[1])], None def Fit( state, y, *, test_size: float = 0.3, random_state: int = 0, ) -> FitResult: """Logistic regression from `state` to binary `y`; held-out AUROC + balanced acc. `state` is a Series (1D) or DataFrame (2D); `y` is a binary {0, 1} Series (or any array-like). The caller is responsible for restricting to the usable cohort and binarising the label. """ X, feature_names, idx = _state_matrix(state) y_arr = np.asarray(y).astype(int) if len(X) != len(y_arr): raise ValueError(f"Fit: state ({len(X)}) and y ({len(y_arr)}) length mismatch") pos = np.arange(len(X)) Xtr, Xte, ytr, yte, pos_tr, pos_te = train_test_split( X, y_arr, pos, test_size=test_size, random_state=random_state, stratify=y_arr, ) model = LogisticRegression(max_iter=1000) model.fit(Xtr, ytr) proba_te = model.predict_proba(Xte)[:, 1] pred_te = (proba_te >= 0.5).astype(int) return FitResult( model=model, auroc=float(roc_auc_score(yte, proba_te)), balanced_acc=float(balanced_accuracy_score(yte, pred_te)), n_train=int(len(Xtr)), n_test=int(len(Xte)), feature_names=feature_names, train_index=idx[pos_tr] if idx is not None else None, test_index=idx[pos_te] if idx is not None else None, ) def Apply(fit: FitResult, state) -> pd.Series: """Return P(y=1 | state) from a fitted model, indexed like `state`.""" X, _, idx = _state_matrix(state) proba = fit.model.predict_proba(X)[:, 1] return pd.Series(proba, index=idx if idx is not None else pd.RangeIndex(len(proba)), name="probability")