SomaScan Inflammation aging clock (85 proteins / 125 aptamers, ROSMAP) — TabM student + standalone predict.py

944651d verified 4 months ago

7.74 kB

#!/usr/bin/env python3
"""Standalone inference for the Inflammatory Aging Clock TabM student models.
Runs the distilled TabM student with public dependencies only:
    pip install torch tabm rtdl_num_embeddings numpy pandas
Usage
-----
    python predict.py --input proteins.csv --output ages.csv
Input
-----
A CSV/TSV with one row per sample and one column per protein. Column names must
match the model's feature names (e.g. ``OID00471_CXCL8`` for MARS, the SomaScan
SeqId-style names for ROSMAP — see ``<model_dir>/meta.json`` -> ``feature_name``).
An optional first ``sample_id`` (or ``SAMPLE.ID``) column is carried through to
the output. Missing proteins are an error; NaN cells are median-imputed for
models trained with a median-impute preprocessor (MARS), and rejected otherwise.
Output
------
A CSV with columns ``sample_id, predicted_age``.
"""
import argparse
import json
import os
import sys

import numpy as np
import pandas as pd
import torch


# --- pinball5 reparameterization (numpy mirror of the training code) ---------
PINBALL5_ALPHAS = np.array([0.1, 0.25, 0.5, 0.75, 0.9], dtype=np.float64)


def _reparameterize_pinball5_np(raw_5):
    """raw (..., 5) -> monotone quantiles (q01, q025, q05, q075, q09)."""
    def softplus(x):
        return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
    m = raw_5[..., 0]
    s1 = softplus(raw_5[..., 1]); s2 = softplus(raw_5[..., 2])
    s3 = softplus(raw_5[..., 3]); s4 = softplus(raw_5[..., 4])
    return np.stack([m - s1 - s2, m - s2, m, m + s3, m + s3 + s4], axis=-1)


def _pinball5_point_estimate_np(raw_5, alphas=PINBALL5_ALPHAS):
    """Trapezoidal-mean point estimate from raw (..., 5). Matches the teacher's mean."""
    q = _reparameterize_pinball5_np(raw_5)
    slope_low = (q[..., 1] - q[..., 0]) / (alphas[1] - alphas[0])
    slope_high = (q[..., 4] - q[..., 3]) / (alphas[4] - alphas[3])
    q0 = q[..., 0] - slope_low * alphas[0]
    q1 = q[..., 4] + slope_high * (1.0 - alphas[4])
    q_full = np.concatenate([q0[..., None], q, q1[..., None]], axis=-1)
    a_full = np.array([0.0, *alphas, 1.0])
    trapz = np.trapezoid if hasattr(np, "trapezoid") else np.trapz  # numpy>=2 vs <2
    return trapz(q_full, a_full, axis=-1).astype(np.float32, copy=False)


# --- preprocessing (median-impute + observed-mask channels) ------------------
def _apply_preproc(X, npz_path):
    """Median-impute NaN and append a mask channel for partially observed columns."""
    d = np.load(npz_path)
    medians = d["medians"].astype(np.float32)
    has_nan_cols = (d["has_nan_cols"].astype(bool) if "has_nan_cols" in d.files
                    else np.ones(medians.shape[0], dtype=bool))
    X = np.asarray(X, dtype=np.float32)
    nan = np.isnan(X)
    X_imp = np.where(nan, medians, X).astype(np.float32, copy=False)
    if not has_nan_cols.any():
        return X_imp
    M = (~nan[:, has_nan_cols]).astype(np.float32)
    return np.concatenate([X_imp, M], axis=1)


# --- model construction ------------------------------------------------------
def _build_model(meta, weights_path):
    import tabm as _tabm
    state = torch.load(weights_path, map_location="cpu", weights_only=False)
    if isinstance(state, dict) and "state_dict" in state:
        sd, saved_bins = state["state_dict"], state.get("bins")
    else:
        sd, saved_bins = state, None
    arch = meta["arch"]
    em = arch["embed_meta"]
    if em.get("kind") == "piecewise_linear":
        import rtdl_num_embeddings as _rne
        if saved_bins is None:
            raise ValueError("Model declares piecewise_linear embeddings but .pt has no 'bins'.")
        emb = _rne.PiecewiseLinearEmbeddings(
            saved_bins, d_embedding=int(em["d_embedding"]),
            activation=bool(em.get("activation", False)), version=em.get("version", "B"),
        )
    elif em.get("kind") in (None, "none"):
        emb = None
    else:
        raise ValueError(f"Unsupported embed_meta.kind={em.get('kind')!r}")
    model = _tabm.TabM(
        n_num_features=int(arch["d_in"]), cat_cardinalities=[],
        d_out=int(arch["n_outputs"]), num_embeddings=emb, **arch["tabm_kwargs"],
    )
    model.load_state_dict(sd)
    model.eval()
    return model


def _predict_scaled(model, X, distill_loss, batch_size=512):
    """Forward TabM -> (n,) scaled-y prediction, averaged over the k ensemble members."""
    out_chunks = []
    with torch.no_grad():
        for i in range(0, len(X), batch_size):
            xb = torch.from_numpy(X[i:i + batch_size].astype(np.float32, copy=False))
            out = model(x_num=xb, x_cat=None).float().cpu().numpy()  # (B, k, d_out)
            if distill_loss == "pinball5":
                point_bk = _pinball5_point_estimate_np(out)          # (B, k)
                out_chunks.append(point_bk.mean(axis=1))
            else:
                out_chunks.append(out.squeeze(-1).mean(axis=1))
    return np.concatenate(out_chunks, axis=0)


# --- I/O ---------------------------------------------------------------------
def _read_table(path):
    sep = "\t" if path.lower().endswith((".tsv", ".txt", ".tsv.gz", ".txt.gz")) else ","
    return pd.read_csv(path, sep=sep)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model_dir", default=".",
                    help="Folder containing meta.json + models/ (default: current directory).")
    ap.add_argument("--input", required=True, help="CSV/TSV of proteins (samples x features).")
    ap.add_argument("--output", required=True, help="Where to write predictions CSV.")
    args = ap.parse_args()

    meta = json.load(open(os.path.join(args.model_dir, "meta.json")))
    feats = meta["feature_name"]

    df = _read_table(args.input)
    # Carry a sample id column through if present.
    id_col = next((c for c in ("sample_id", "SAMPLE.ID", "SampleID", "ID") if c in df.columns), None)
    sample_ids = df[id_col].astype(str).values if id_col else np.arange(len(df)).astype(str)

    # The model needs EVERY trained protein present, by exact name. Column order
    # does not matter (selected/reordered below). Hard-fail on any missing one.
    feat_set = set(feats)
    missing = [f for f in feats if f not in df.columns]
    if missing:
        sys.exit(f"ERROR: input is missing {len(missing)}/{len(feats)} required protein "
                 f"column(s), e.g. {missing[:5]}. Column names must match "
                 f"{args.model_dir}/meta.json -> feature_name exactly.")
    # Warn about input columns the model does not use (often a naming mismatch).
    extra = [c for c in df.columns if c not in feat_set and c != id_col]
    if extra:
        print(f"WARNING: ignoring {len(extra)} input column(s) not used by the model, "
              f"e.g. {extra[:5]}.", file=sys.stderr)
    print(f"Matched all {len(feats)} required feature columns.", file=sys.stderr)
    X = df[feats].to_numpy(dtype=np.float32)  # (n, n_features_raw), column order = feats

    if meta.get("preproc"):
        X = _apply_preproc(X, os.path.join(args.model_dir, meta["preproc"]["file"]))
    elif np.isnan(X).any():
        sys.exit("ERROR: input contains NaN but this model has no imputer. "
                 "Provide complete protein values.")

    model = _build_model(meta, os.path.join(args.model_dir, meta["student_weights"]))
    scaled = _predict_scaled(model, X, meta["regression_distill_loss"])
    age = scaled * meta["y_scaler"]["scale"] + meta["y_scaler"]["mean"]

    out = pd.DataFrame({"sample_id": sample_ids, "predicted_age": age})
    out.to_csv(args.output, index=False)
    print(f"Wrote {len(out)} predictions to {args.output}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Standalone inference for the Inflammatory Aging Clock TabM student models.

	Runs the distilled TabM student with public dependencies only:

	pip install torch tabm rtdl_num_embeddings numpy pandas

	Usage
	-----
	python predict.py --input proteins.csv --output ages.csv

	Input
	-----
	A CSV/TSV with one row per sample and one column per protein. Column names must
	match the model's feature names (e.g. ``OID00471_CXCL8`` for MARS, the SomaScan
	SeqId-style names for ROSMAP — see ``<model_dir>/meta.json`` -> ``feature_name``).
	An optional first ``sample_id`` (or ``SAMPLE.ID``) column is carried through to
	the output. Missing proteins are an error; NaN cells are median-imputed for
	models trained with a median-impute preprocessor (MARS), and rejected otherwise.

	Output
	------
	A CSV with columns ``sample_id, predicted_age``.
	"""
	import argparse
	import json
	import os
	import sys

	import numpy as np
	import pandas as pd
	import torch


	# --- pinball5 reparameterization (numpy mirror of the training code) ---------
	PINBALL5_ALPHAS = np.array([0.1, 0.25, 0.5, 0.75, 0.9], dtype=np.float64)


	def _reparameterize_pinball5_np(raw_5):
	"""raw (..., 5) -> monotone quantiles (q01, q025, q05, q075, q09)."""
	def softplus(x):
	return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
	m = raw_5[..., 0]
	s1 = softplus(raw_5[..., 1]); s2 = softplus(raw_5[..., 2])
	s3 = softplus(raw_5[..., 3]); s4 = softplus(raw_5[..., 4])
	return np.stack([m - s1 - s2, m - s2, m, m + s3, m + s3 + s4], axis=-1)


	def _pinball5_point_estimate_np(raw_5, alphas=PINBALL5_ALPHAS):
	"""Trapezoidal-mean point estimate from raw (..., 5). Matches the teacher's mean."""
	q = _reparameterize_pinball5_np(raw_5)
	slope_low = (q[..., 1] - q[..., 0]) / (alphas[1] - alphas[0])
	slope_high = (q[..., 4] - q[..., 3]) / (alphas[4] - alphas[3])
	q0 = q[..., 0] - slope_low * alphas[0]
	q1 = q[..., 4] + slope_high * (1.0 - alphas[4])
	q_full = np.concatenate([q0[..., None], q, q1[..., None]], axis=-1)
	a_full = np.array([0.0, *alphas, 1.0])
	trapz = np.trapezoid if hasattr(np, "trapezoid") else np.trapz # numpy>=2 vs <2
	return trapz(q_full, a_full, axis=-1).astype(np.float32, copy=False)


	# --- preprocessing (median-impute + observed-mask channels) ------------------
	def _apply_preproc(X, npz_path):
	"""Median-impute NaN and append a mask channel for partially observed columns."""
	d = np.load(npz_path)
	medians = d["medians"].astype(np.float32)
	has_nan_cols = (d["has_nan_cols"].astype(bool) if "has_nan_cols" in d.files
	else np.ones(medians.shape[0], dtype=bool))
	X = np.asarray(X, dtype=np.float32)
	nan = np.isnan(X)
	X_imp = np.where(nan, medians, X).astype(np.float32, copy=False)
	if not has_nan_cols.any():
	return X_imp
	M = (~nan[:, has_nan_cols]).astype(np.float32)
	return np.concatenate([X_imp, M], axis=1)


	# --- model construction ------------------------------------------------------
	def _build_model(meta, weights_path):
	import tabm as _tabm
	state = torch.load(weights_path, map_location="cpu", weights_only=False)
	if isinstance(state, dict) and "state_dict" in state:
	sd, saved_bins = state["state_dict"], state.get("bins")
	else:
	sd, saved_bins = state, None
	arch = meta["arch"]
	em = arch["embed_meta"]
	if em.get("kind") == "piecewise_linear":
	import rtdl_num_embeddings as _rne
	if saved_bins is None:
	raise ValueError("Model declares piecewise_linear embeddings but .pt has no 'bins'.")
	emb = _rne.PiecewiseLinearEmbeddings(
	saved_bins, d_embedding=int(em["d_embedding"]),
	activation=bool(em.get("activation", False)), version=em.get("version", "B"),
	)
	elif em.get("kind") in (None, "none"):
	emb = None
	else:
	raise ValueError(f"Unsupported embed_meta.kind={em.get('kind')!r}")
	model = _tabm.TabM(
	n_num_features=int(arch["d_in"]), cat_cardinalities=[],
	d_out=int(arch["n_outputs"]), num_embeddings=emb, **arch["tabm_kwargs"],
	)
	model.load_state_dict(sd)
	model.eval()
	return model


	def _predict_scaled(model, X, distill_loss, batch_size=512):
	"""Forward TabM -> (n,) scaled-y prediction, averaged over the k ensemble members."""
	out_chunks = []
	with torch.no_grad():
	for i in range(0, len(X), batch_size):
	xb = torch.from_numpy(X[i:i + batch_size].astype(np.float32, copy=False))
	out = model(x_num=xb, x_cat=None).float().cpu().numpy() # (B, k, d_out)
	if distill_loss == "pinball5":
	point_bk = _pinball5_point_estimate_np(out) # (B, k)
	out_chunks.append(point_bk.mean(axis=1))
	else:
	out_chunks.append(out.squeeze(-1).mean(axis=1))
	return np.concatenate(out_chunks, axis=0)


	# --- I/O ---------------------------------------------------------------------
	def _read_table(path):
	sep = "\t" if path.lower().endswith((".tsv", ".txt", ".tsv.gz", ".txt.gz")) else ","
	return pd.read_csv(path, sep=sep)


	def main():
	ap = argparse.ArgumentParser(description=__doc__,
	formatter_class=argparse.RawDescriptionHelpFormatter)
	ap.add_argument("--model_dir", default=".",
	help="Folder containing meta.json + models/ (default: current directory).")
	ap.add_argument("--input", required=True, help="CSV/TSV of proteins (samples x features).")
	ap.add_argument("--output", required=True, help="Where to write predictions CSV.")
	args = ap.parse_args()

	meta = json.load(open(os.path.join(args.model_dir, "meta.json")))
	feats = meta["feature_name"]

	df = _read_table(args.input)
	# Carry a sample id column through if present.
	id_col = next((c for c in ("sample_id", "SAMPLE.ID", "SampleID", "ID") if c in df.columns), None)
	sample_ids = df[id_col].astype(str).values if id_col else np.arange(len(df)).astype(str)

	# The model needs EVERY trained protein present, by exact name. Column order
	# does not matter (selected/reordered below). Hard-fail on any missing one.
	feat_set = set(feats)
	missing = [f for f in feats if f not in df.columns]
	if missing:
	sys.exit(f"ERROR: input is missing {len(missing)}/{len(feats)} required protein "
	f"column(s), e.g. {missing[:5]}. Column names must match "
	f"{args.model_dir}/meta.json -> feature_name exactly.")
	# Warn about input columns the model does not use (often a naming mismatch).
	extra = [c for c in df.columns if c not in feat_set and c != id_col]
	if extra:
	print(f"WARNING: ignoring {len(extra)} input column(s) not used by the model, "
	f"e.g. {extra[:5]}.", file=sys.stderr)
	print(f"Matched all {len(feats)} required feature columns.", file=sys.stderr)
	X = df[feats].to_numpy(dtype=np.float32) # (n, n_features_raw), column order = feats

	if meta.get("preproc"):
	X = _apply_preproc(X, os.path.join(args.model_dir, meta["preproc"]["file"]))
	elif np.isnan(X).any():
	sys.exit("ERROR: input contains NaN but this model has no imputer. "
	"Provide complete protein values.")

	model = _build_model(meta, os.path.join(args.model_dir, meta["student_weights"]))
	scaled = _predict_scaled(model, X, meta["regression_distill_loss"])
	age = scaled * meta["y_scaler"]["scale"] + meta["y_scaler"]["mean"]

	out = pd.DataFrame({"sample_id": sample_ids, "predicted_age": age})
	out.to_csv(args.output, index=False)
	print(f"Wrote {len(out)} predictions to {args.output}")


	if __name__ == "__main__":
	main()