dsl/operators.py

"""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")