inference.py

"""Inference script to run on data not included in the manuscript.

This best works if you have a complete cohort and not just a few samples.
"""

import argparse
import logging
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression

from calc_scores import single_slide_grade_features
from grade_pred import load_model, prep_grade_features
from muc_new import single_slide_mucin
from run_eval import log_scale
from src.constants import EXPLANATORY, FEATURE_SUBSET, GRADE_RADII, TO_SCALE
from src.statutils import normalize_percentile

logger = logging.getLogger(__name__)


def normalize_pct(inp_df, outlier_threshold=None):
    voi = []
    tmp_df_ = inp_df.copy()
    for var in EXPLANATORY:
        if tmp_df_[var].dtype is float:
            tmp_df_[var] = normalize_percentile(tmp_df_[var], 5, 95, clip=False)
            voi.append(var)
    if outlier_threshold is not None:
        keep = (
            np.sum(
                [(tmp_df_[r] > outlier_threshold).to_numpy() for r in list(set(voi))],
                axis=0,
            )
            == 0
        )
        return tmp_df_[keep].copy()
    return tmp_df_.copy()


def get_viable_paths(args: dict) -> list[Path]:
    """Find all potential slides to be processed."""
    root = Path(args["nuc_input"])
    poss_paths = list(root.glob("*/tum_inst.tsv"))
    logger.debug("Total potential cases: %i", len(poss_paths))
    final_paths = []
    for p in poss_paths:
        slide_id = p.parent.stem
        # check for necessary files
        pc = Path(root, slide_id, "pred_com.csv")
        srma = Path(args["srma_input"], slide_id, slide_id + "_fullhd.png")
        if pc.exists() & srma.exists():
            final_paths.append((p, pc, srma))
    logger.info("Found %i cases to process.", len(final_paths))
    return final_paths


def write_grade_features(grade_features: list[str], grade_res_path: Path) -> None:
    """Save grade features to file."""
    with grade_res_path.open("w") as gr:
        r0, r1 = GRADE_RADII
        gr.write(
            f"id,G,ndeg_{r0},ndeg_{r1},ndeg_{r0}_weight,ndeg_{r1}_weight,ndeg_{r0}_sort,ndeg_{r1}_sort,katz_mean,katz_std,katz_mean_w,katz_std_w,katz_mean_sort,katz_std_sort,n_crypts\n"
        )
        gr.writelines(grade_features)


def write_muc_features(muc_features: list[str], muc_res_path: Path) -> None:
    """Save mucin amounts to file."""
    with muc_res_path.open("w") as mr:
        mr.write("slide_id,muc_tum_ratio_refined\n")
        mr.writelines(muc_features)
    return pd.DataFrame(muc_features, columns=["slide_id", "muc_tum_ratio_refined"])


def prepare_clinical_data(
    source: Path, grade_features: pd.DataFrame, muc_features: pd.DataFrame
) -> pd.DataFrame:
    """Get clinical datasheet an preprocess."""
    source_df = pd.read_csv(source)
    source_df.columns = source_df.columns.map(str.lower)

    source_df = source_df.merge(grade_features, how="left", left_on="slide_id", right_on="id")
    source_df = source_df.merge(muc_features, how="left", left_on="slide_id", right_on="slide_id")

    if source_df.isna().any():
        source_df = source_df.dropna(axis=0)
        logger.warning("Nan detected in input dataframe, dropping rows with missing data.")

    return normalize_pct(source_df, outlier_threshold=None)


def verify_model_path(path: str) -> Path:
    """Check if model path exists."""
    p = Path(path)
    if p.exists() & p.suffix == ".pkl":
        return p
    logger.warning("Specified MSI prediction model not found, fallback to default.")
    p = Path("models/all_to_one/LogisticRegression_tcga_cross_cohort_lm_all_to_one.pkl")
    if not p.exists():
        msg = "Model files not found, please verify msi prediction models exist."
        logger.error(msg)
        raise FileNotFoundError(msg)
    return p


def main(args: dict) -> None:
    """Generate all necessary features for MSI prediction and predict."""
    paths = get_viable_paths(args)
    output_folder = Path(args["output_folder"])
    grade_out = Path(output_folder, "grade_features.csv")
    muc_out = Path(output_folder, "muc_amount.csv")
    final_out = Path(output_folder, "msi_prediction.csv")

    # Grade and features
    grade_features = []
    muc_features = []
    for p, pc, srma in paths:
        grade_features.append(single_slide_grade_features(p, pc))
        muc_features.append(single_slide_mucin(srma))
    write_grade_features(grade_features, grade_out)
    muc_features = write_muc_features(muc_features, muc_out)

    # predict grade with stored weights.
    model_path = Path("models", "GaussianNB_grade.pkl")
    grade_model = load_model(model_path)
    grade_features = prep_grade_features(grade_out, None, output_folder.stem, normalize=True)
    grade_features["pred"] = grade_model.predict_proba(grade_features[FEATURE_SUBSET])[:, 1]
    grade_features[["id", "pred"]].to_csv(Path(output_folder, "grade_pred.csv"))

    # predict msi
    dataset = prepare_clinical_data(Path(args["clinical_data"]), grade_features, muc_features)
    mod_path = verify_model_path(args["msi_model"])
    msi_model = load_model(mod_path)
    ds_prep = log_scale(dataset, TO_SCALE) if type(msi_model) is LogisticRegression else dataset
    test_ds = ds_prep[EXPLANATORY].copy()
    results = msi_model.predict_proba(test_ds)[:, 1]
    ds_prep["msi_pred"] = results
    ds_prep.to_csv(final_out)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-l",
        "--log_level",
        type=str,
        default="INFO",
        help="Log level, one of DEBUG, INFO, WARNING, ERROR, CRITICAL",
    )
    parser.add_argument(
        "-n",
        "--nuc_input",
        type=str,
        help="HoVer-NeXt results directory. This should be the cohort you want to process.",
    )
    parser.add_argument(
        "-s",
        "--srma_input",
        type=str,
        help="SRMA results directory. Results should be ran for each case included in the cohort.",
    )
    parser.add_argument(
        "-o",
        "--output_folder",
        type=str,
        default="results/test_cohort/",
        help="Output folder where all intermediate and final results are stored.",
    )
    parser.add_argument(
        "-m",
        "--msi_model",
        type=str,
        help="which model to use for prediction. Provide path to model - models/../.",
    )
    parser.add_argument(
        "-c",
        "--clinical_data",
        type=str,
        help=(
            "Path to clinical data sheet including pre-processed scores from CRC-Eos-Lym-IEL. "
            "Needs to include the following variables: slide_id, age, gender, location, n_tls, iel_ratio,"
            "neu_full_r20, fil_full_r20"
        ),
    )
    args = vars(parser.parse_args())

    logging.basicConfig(level=args["log_level"].upper())
    logger.info("Parsed arguments: ", extra=args)

    main(args)