DeConveil 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

deconveil/nb_regression_fit.py

@@ -0,0 +1,313 @@
+from cmdstanpy import CmdStanModel
+import numpy as np
+import pandas as pd
+import scipy.stats as st
+
+
+def fit_one_gene(
+    gene_df: pd.DataFrame,
+    model: CmdStanModel,
+    gene: str | None = None,      # optional convenience
+    cna: str = "all",             # "amp" | "del" | "all"
+    et: float = 0.15,
+    min_aneup: int = 5,
+    min_unique_counts: int = 5,
+    min_cn_abs_sum: float = 1.0,  # identifiability filter for cna="all"
+    chains: int = 4,
+    iter_warmup: int = 1000,
+    iter_sampling: int = 1000,
+    seed: int = 1,
+    show_progress: bool = False,
+    adapt_delta: float = 0.99,
+    max_treedepth: int = 15,
+):
+    """
+    Fit Stan NB regression model for a single gene.
+    Expects gene_df to be filtered to one gene OR pass gene and full df.
+
+    Required columns in gene_df:
+      gene, expr, copies, purity, stroma, sf, eup_dev_cancer, eup_equiv_cancer
+    Optional columns:
+      covar (if missing -> set to 'ALL')
+    """
+    # subset to one gene if gene provided and gene_df contains multiple genes
+    df = gene_df.copy()
+    if gene is not None and "gene" in df.columns and df["gene"].nunique() > 1:
+        df = df.loc[df["gene"] == gene].copy()
+    if gene is None and "gene" in df.columns and df["gene"].nunique() == 1:
+        gene = str(df["gene"].iloc[0])
+
+    if df.empty:
+        return {"status": "skipped", "gene": gene, "reason": "no_rows_for_gene"}
+
+    required = {"expr","copies","purity","stroma","sf","eup_dev_cancer","eup_equiv_cancer"}
+    missing = required - set(df.columns)
+    if missing:
+        return {"status": "error", "gene": gene, "reason": f"missing_columns: {sorted(missing)}"}
+
+    # CNA subset (optional)
+    if cna == "amp":
+        df = df[df["copies"] > (2 - et)]
+    elif cna == "del":
+        df = df[df["copies"] < (2 + et)]
+    elif cna == "all":
+        pass
+    else:
+        raise ValueError("cna must be 'amp', 'del', or 'all'")
+
+    if df.empty:
+        return {"status": "skipped", "gene": gene, "reason": "no_samples_after_cna_filter"}
+
+    # basic data QC
+    df = df.dropna(subset=["expr","sf","purity","stroma","eup_dev_cancer","eup_equiv_cancer"])
+    if df.empty:
+        return {"status": "skipped", "gene": gene, "reason": "all_na_after_dropna"}
+
+    if (df["expr"] < 0).any():
+        return {"status": "error", "gene": gene, "reason": "negative_counts"}
+
+    if not df["purity"].between(0, 1).all():
+        return {"status": "error", "gene": gene, "reason": "purity_out_of_bounds"}
+
+    if not (df["sf"] > 0).all():
+        return {"status": "error", "gene": gene, "reason": "nonpositive_sf"}
+
+    if not (df["eup_equiv_cancer"] > 0).all():
+        return {"status": "error", "gene": gene, "reason": "nonpositive_eup_equiv_cancer"}
+
+    # skip degenerate genes early (prevents phi->inf + treedepth explosions)
+    if df["expr"].nunique() < min_unique_counts:
+        return {"status": "skipped", "gene": gene, "reason": "too_few_unique_counts"}
+
+    # aneuploid count check (use copies as cancer_copies analogue)
+    n_aneup = int((np.abs(df["copies"].astype(float) - 2.0) > (1.0 - et)).sum())
+    if n_aneup < min_aneup or (df["expr"] == 0).all():
+        return {"status": "skipped", "gene": gene, "n_aneup": n_aneup, "reason": "low_aneup_or_all_zero"}
+
+    # identifiability check for cna="all": need some CN deviation mass
+    if cna == "all" and df["eup_dev_cancer"].abs().sum() < min_cn_abs_sum:
+        return {"status": "skipped", "gene": gene, "n_aneup": n_aneup, "reason": "too_little_cn_variation"}
+
+    # within one cohort: enforce K=1, covar='ALL'
+    df["covar"] = "ALL"
+    covar_levels = ["ALL"]
+    covar_idx = np.ones(len(df), dtype=int)
+
+    # per-gene sf scaling 
+    mean_expr = float(df["expr"].mean())
+    df["sf_scaled"] = df["sf"].astype(float) * mean_expr
+
+    stan_data = {
+        "N": int(len(df)),
+        "y": df["expr"].astype(int).to_numpy(),
+        "K": 1,
+        "covar": covar_idx,
+        "sf": df["sf_scaled"].to_numpy(dtype=float),
+        "purity": df["purity"].to_numpy(dtype=float),
+        "stroma": df["stroma"].to_numpy(dtype=float),
+        "eup_equiv_cancer": df["eup_equiv_cancer"].to_numpy(dtype=float),
+        "eup_dev_cancer": df["eup_dev_cancer"].to_numpy(dtype=float),
+    }
+
+    rng = np.random.default_rng(seed)
+    init = {
+        "b_scaling": rng.uniform(0.5, 1.5, size=1).tolist(),
+        "b_noncancer": rng.uniform(0.5, 1.5, size=1).tolist(),
+        "b_deviation": 0.0,
+        "phi": 1.0,
+    }
+
+    fit = model.sample(
+        data=stan_data,
+        chains=chains,
+        iter_warmup=iter_warmup,
+        iter_sampling=iter_sampling,
+        seed=seed,
+        inits=init,
+        show_progress=show_progress,
+        adapt_delta=adapt_delta,
+        max_treedepth=max_treedepth,
+    )
+
+    draws = fit.draws_pd()
+
+    # extract posterior summaries
+   
+    scaling_col = next((c for c in draws.columns if c.startswith("b_scaling")), None)
+    if scaling_col is None:
+        return {"status": "error", "gene": gene, "reason": "missing_b_scaling_draws"}
+
+    if "b_deviation" not in draws.columns:
+        return {"status": "error", "gene": gene, "reason": "missing_b_deviation_draws"}
+
+    phi_col = "phi" if "phi" in draws.columns else None
+
+    b_scaling = draws[scaling_col].to_numpy()
+    b_dev = draws["b_deviation"].to_numpy()
+
+    # z and p 
+    z_comp = float(b_dev.mean() / b_dev.std(ddof=1))
+    p_value = float(2.0 * (1.0 - st.norm.cdf(abs(z_comp))))
+
+    summ = fit.summary()
+
+    # robust extraction of Rhat / ESS
+    rhat_col = next((c for c in ["R_hat", "Rhat"] if c in summ.columns), None)
+    ess_col  = next((c for c in ["Ess_bulk", "ESS_bulk", "N_Eff", "Ess"] if c in summ.columns), None)
+
+    rhat_dev = float(summ.loc["b_deviation", rhat_col]) if (rhat_col and "b_deviation" in summ.index) else np.nan
+    ess_dev  = float(summ.loc["b_deviation", ess_col]) if (ess_col and "b_deviation" in summ.index) else np.nan
+
+    # mark borderline fits as warn (you can filter later)
+    status = "ok"
+    if (not np.isnan(rhat_dev) and rhat_dev > 1.05) or (not np.isnan(ess_dev) and ess_dev < 200):
+        status = "warn"
+
+    out = {
+        "status": status,
+        "gene": gene,
+        "N": int(len(df)),
+        "n_aneup": n_aneup,
+        "cna": cna,
+        # posterior summaries
+        "mean_b_scaling": float(b_scaling.mean()),
+        "sd_b_scaling": float(b_scaling.std(ddof=1)),
+        "mean_b_deviation": float(b_dev.mean()),
+        "sd_b_deviation": float(b_dev.std(ddof=1)),
+        "z_comp": z_comp,
+        "p_value": p_value,
+        # optional
+        "mean_phi": float(draws[phi_col].mean()) if phi_col else np.nan,
+        "Rhat_b_deviation": rhat_dev,
+        "ess_b_deviation": ess_dev,
+        "covar_levels": covar_levels,
+    }
+    return out
+
+
+def postprocess_nb_results(
+    res_df: pd.DataFrame,
+    alpha: float = 0.05,
+    comp_thr: float = 0.5,
+    min_aneup_frac: float = 0.30,   # >=30% aneuploid samples
+    min_scaling: float = 1e-6,
+    require_scaling_for_dsg: bool = True,
+    warn_suffix: str | None = None,  # e.g. "_LOWCONF" to tag warn calls
+) -> pd.DataFrame:
+    """
+    Post-process Stan NB regression results.
+
+    Expected columns in `res_df`:
+      - status, gene, N, n_aneup, cna
+      - mean_b_scaling, sd_b_scaling
+      - mean_b_deviation, sd_b_deviation
+      - z_comp, p_value, mean_phi
+      - Rhat_b_deviation, ess_b_deviation, covar_levels
+
+    Adds:
+      - aneup_frac  : n_aneup / N (CN informativeness proxy)
+      - comp_score  : normalized deviation  = mean_b_deviation / mean_b_scaling
+      - signed_comp : sign-normalized comp_score (flip for 'del', keep for 'amp'/'all')
+      - shrunk_comp : comp_score shrunk toward 0 by (1 - p_value)  [orthogonal score]
+      - label_nb    : {'DSG','DCG','HYPER','OTHER','SKIP','ERROR'}
+
+    Notes:
+      - Without a per-gene CN direction summary, signed_comp is only partially
+        identifiable for cna='all'; here we flip sign for 'del' fits and leave
+        'amp'/'all' unchanged.
+      - DCG/HYPER require sufficient magnitude |signed_comp| >= comp_thr
+      - DSG small magnitude |signed_comp| <= comp_thr, with sufficient CN signal.
+    """
+
+    df = res_df.copy()
+
+    # required columns check 
+    required = {
+        "status", "gene", "N", "n_aneup", "cna", "p_value",
+        "mean_b_scaling", "mean_b_deviation"
+    }
+    missing = required - set(df.columns)
+    if missing:
+        raise ValueError(f"Missing required columns: {sorted(missing)}")
+
+    # usable rows: include ok + warn 
+    status_lower = df["status"].astype(str).str.lower()
+    usable = status_lower.isin(["ok", "warn"]) & df["p_value"].notna()
+
+    # default labels 
+    df["label_nb"] = "OTHER"
+    df.loc[status_lower.isin(["skip", "skipped"]), "label_nb"] = "SKIP"
+    df.loc[status_lower.isin(["error", "failed", "fail"]), "label_nb"] = "ERROR"
+
+
+    # CN informativeness proxy (fraction aneuploid) 
+    df["aneup_frac"] = np.nan
+    df.loc[usable, "aneup_frac"] = (
+        df.loc[usable, "n_aneup"] / df.loc[usable, "N"]
+    ).astype(float)
+
+    # compute compensation scores 
+    df["comp_score"] = np.nan
+    df["signed_comp"] = np.nan
+    df["shrunk_comp"] = np.nan  # normalized + shrunk score (orthogonal summary)
+
+    ok2 = (
+        usable
+        & df["mean_b_scaling"].notna()
+        & df["mean_b_deviation"].notna()
+        & (df["mean_b_scaling"].abs() > min_scaling)
+        & df["aneup_frac"].notna()
+        & (df["aneup_frac"] >= min_aneup_frac)
+    )
+
+    # Normalized deviation: beta2* = mean_b_deviation / mean_b_scaling
+    df.loc[ok2, "comp_score"] = (
+        df.loc[ok2, "mean_b_deviation"] / df.loc[ok2, "mean_b_scaling"]
+    )
+
+    # Sign-normalize by CNA type if available: del -> flip sign
+    cna_vals = df.loc[ok2, "cna"].astype(str).str.lower()
+    sign = np.ones(len(cna_vals), dtype=float)
+    sign[cna_vals == "del"] = -1.0    # treat deletions as negative direction
+    # amp/all remain +1
+    df.loc[ok2, "signed_comp"] = (
+        df.loc[ok2, "comp_score"].to_numpy(dtype=float) * sign
+    )
+
+    # Shrink normalized deviation by pseudo-p (here: raw p_value)
+    # s_g ≈ beta2*_g * (1 - p_g)
+    pseudo_p = df.loc[ok2, "p_value"].astype(float).clip(0.0, 1.0)
+    df.loc[ok2, "shrunk_comp"] = df.loc[ok2, "comp_score"] * (1.0 - pseudo_p)
+
+    # classification gates 
+    # Significant deviation (
+    is_sig = ok2 & (df["p_value"] <= alpha)
+    is_nonsig = ok2 & (df["p_value"] > alpha)
+
+    # Optional: require non-trivial scaling to call DSG
+    scaling_ok = ok2
+    if require_scaling_for_dsg:
+        scaling_ok = ok2 & (df["mean_b_scaling"] > 0.3)
+
+    # DCG: significant negative deviation (compensation)
+    #df.loc[is_sig & (df["shrunk_comp"] <= -comp_thr), "label_nb"] = "DCG"
+    df.loc[df["shrunk_comp"] <= -comp_thr, "label_nb"] = "DCG"
+
+    # HYPER: significant positive deviation
+    #df.loc[is_sig & (df["signed_comp"] >= comp_thr), "label_nb"] = "HYPER"
+    df.loc[df["shrunk_comp"] >= comp_thr, "label_nb"] = "HYPER"
+
+    # DSG: non-significant deviation AND near CN-proportional expectation
+    df.loc[
+        #is_nonsig & scaling_ok & (df["shrunk_comp"].abs() <= comp_thr),
+        scaling_ok & (df["shrunk_comp"].abs() <= comp_thr),
+        "label_nb"
+    ] = "DSG"
+
+    # optional: tag warn calls as low confidence 
+    if warn_suffix:
+        is_warn = status_lower.eq("warn")
+        mask = is_warn & df["label_nb"].isin(["DSG", "DCG", "HYPER"])
+        df.loc[mask, "label_nb"] = df.loc[mask, "label_nb"] + warn_suffix
+
+    return df