processing.py

#!/usr/bin/env python3

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import norm

import rpy2.robjects as ro


def ensure_dir(p: Path) -> None:
    p.mkdir(parents=True, exist_ok=True)


def require_file(p: Path) -> None:
    if not p.exists():
        raise FileNotFoundError(f"Required file not found: {p}")


def save_fig(fig: plt.Figure, out_path: Path, dpi: int) -> None:
    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")


def set_pub_style(base_font: int = 14) -> None:
    plt.rcParams.update(
        {
            "font.size": base_font,
            "axes.labelsize": base_font,
            "axes.titlesize": base_font + 2,
            "axes.titleweight": "bold",
            "axes.labelweight": "bold",
            "legend.fontsize": base_font - 2,
            "xtick.labelsize": base_font - 2,
            "ytick.labelsize": base_font - 2,
            "lines.linewidth": 2.8,
        }
    )


def set_2h_time_axis(ax: plt.Axes) -> None:
    ax.set_xlim(0, 24)
    ticks = np.arange(0, 24, 2)
    ax.set_xticks(ticks)
    ax.set_xticklabels([f"{int(t):02d}:00" for t in ticks])


def set_4h_time_axis(ax: plt.Axes) -> None:
    ax.set_xlim(0, 24)
    ticks = np.arange(0, 24, 4)
    ax.set_xticks(ticks)
    ax.set_xticklabels([f"{int(t):02d}:00" for t in ticks], rotation=45)


def load_Y(processed_dir: Path) -> np.ndarray:
    y_path = processed_dir / "function_on_scalar_regression_steps_time_series.csv"
    require_file(y_path)
    Y = np.loadtxt(y_path, delimiter=",")
    if Y.ndim != 2 or Y.shape[1] != 48:
        raise ValueError(f"Expected Y shape (n,48). Got {Y.shape}")
    return Y


def load_X(processed_dir: Path) -> pd.DataFrame:
    x_path = (
        processed_dir
        / "function_on_scalar_regression_features_mean_EDSS_mean_t25fw.csv"
    )
    require_file(x_path)
    X = pd.read_csv(x_path)
    required_cols = {"mean_EDSS", "mean_t25fw"}
    if not required_cols.issubset(set(X.columns)):
        raise ValueError("Missing required columns in feature file.")
    return X


def read_rds(path: Path):
    require_file(path)
    return ro.r(f'readRDS("{path.as_posix()}")')


def extract_fields(m):
    return {
        "yhat": np.array(m.rx2("yhat")),
        "resid": np.array(m.rx2("resid")),
        "est.func": np.array(m.rx2("est.func")),
        "se.func": np.array(m.rx2("se.func")),
    }


def load_models(processed_dir: Path):
    m_both = read_rds(processed_dir / "fosr_model.rds")
    m_mean_EDSS_excluded = read_rds(processed_dir / "fosr_model_mean_EDSS_excluded.rds")
    m_mean_t25fw_excluded = read_rds(
        processed_dir / "fosr_model_mean_t25fw_excluded.rds"
    )
    return {
        "both": extract_fields(m_both),
        "mean_EDSS_excluded": extract_fields(m_mean_EDSS_excluded),
        "mean_t25fw_excluded": extract_fields(m_mean_t25fw_excluded),
    }


def _extract_bayes_results_listvec(m):
    names = list(getattr(m, "names", []) or [])
    if "beta.hat" in names and "beta.LB" in names and "beta.UB" in names:
        return m
    if "results" in names:
        r = m.rx2("results")
        rnames = list(getattr(r, "names", []) or [])
        if "beta.hat" in rnames and "beta.LB" in rnames and "beta.UB" in rnames:
            return r
    raise KeyError(
        "Could not find beta.hat / beta.LB / beta.UB in bayes_fosr model object."
    )


def load_bayes_fosr_results(processed_dir: Path):
    m = read_rds(processed_dir / "bayes_fosr_model.rds")
    r = _extract_bayes_results_listvec(m)
    beta_hat = np.array(r.rx2("beta.hat"))
    beta_lb = np.array(r.rx2("beta.LB"))
    beta_ub = np.array(r.rx2("beta.UB"))
    return {"beta.hat": beta_hat, "beta.LB": beta_lb, "beta.UB": beta_ub}


def add_bins_notebook(
    predictors: pd.DataFrame,
) -> tuple[pd.DataFrame, list[str], list[str]]:
    predictors = predictors.copy()
    predictors["mean_EDSS_rounded"] = np.round(predictors["mean_EDSS"] * 2) / 2
    predictors.loc[predictors["mean_EDSS_rounded"] < 3.0, "mean_EDSS_rounded"] = 3.0

    EDSS_levels = [2.75, 3.5, 4.5, 5.5, 6.25, 6.75, 7.25]
    EDSS_labels = ["3-3.5", "4-4.5", "5-5.5", "6", "6.5", "7"]

    predictors["mean_edss_bin"] = pd.cut(
        predictors["mean_EDSS_rounded"],
        bins=EDSS_levels,
        labels=EDSS_labels,
        right=True,
        include_lowest=True,
    )

    T25FW_edges = [-np.inf, 6.0, 8.0, np.inf]
    T25FW_labels = ["<6", "6-7.99", ">=8"]

    predictors["mean_t25fw_bin"] = pd.cut(
        predictors["mean_t25fw"],
        bins=T25FW_edges,
        labels=T25FW_labels,
        right=False,
        include_lowest=True,
    )

    return predictors, EDSS_labels, T25FW_labels


def make_fig1a(steps_time_series, predictors, yhat, EDSS_labels, out_dir, dpi):
    time_points = np.linspace(0, 1, 48)
    fig, ax = plt.subplots(figsize=(12.5, 6.5))

    for edss in EDSS_labels:
        indices = predictors.index[predictors["mean_edss_bin"] == edss]
        n = len(indices)
        if n == 0:
            continue
        avg_curve = np.expm1(np.mean(yhat[indices], axis=0))
        ax.plot(time_points * 24, avg_curve, label=f"EDSS {edss} (n={n})")

    for i in range(steps_time_series.shape[0]):
        ax.plot(time_points * 24, steps_time_series[i], "k.", markersize=2, alpha=0.1)

    ax.set_xlabel("Time of Day")
    ax.set_ylabel("Mean Step Count per Minute")
    ax.set_title("Mean Step Count Curves by EDSS Bin")
    set_2h_time_axis(ax)
    ax.grid(True)
    ax.legend(
        title="EDSS", frameon=False, loc="center left", bbox_to_anchor=(1.02, 0.5)
    )

    save_fig(fig, out_dir / "Fig_1A.tif", dpi)
    plt.close(fig)


def make_fig1b(steps_time_series, predictors, yhat, T25FW_labels, out_dir, dpi):
    time_points = np.linspace(0, 1, 48)
    fig, ax = plt.subplots(figsize=(12.5, 6.5))

    for t25fw in T25FW_labels:
        indices = predictors.index[predictors["mean_t25fw_bin"] == t25fw]
        n = len(indices)
        if n == 0:
            continue
        avg_curve = np.expm1(np.mean(yhat[indices], axis=0))
        ax.plot(time_points * 24, avg_curve, label=f"T25FW {t25fw} (n={n})")

    for i in range(steps_time_series.shape[0]):
        ax.plot(time_points * 24, steps_time_series[i], "k.", markersize=2, alpha=0.1)

    ax.set_xlabel("Time of Day")
    ax.set_ylabel("Mean Step Count per Minute")
    ax.set_title("Mean Step Count Curves by T25FW Level")
    set_2h_time_axis(ax)
    ax.grid(True)
    ax.legend(
        title="T25FW", frameon=False, loc="center left", bbox_to_anchor=(1.02, 0.5)
    )

    save_fig(fig, out_dir / "Fig_1B.tif", dpi)
    plt.close(fig)


def make_fig2a(steps_time_series, models, out_dir, dpi):
    Y_mean = np.mean(steps_time_series, axis=0)
    SS_tot = np.sum((steps_time_series - Y_mean) ** 2, axis=0)

    SS_res_both = np.sum(models["both"]["resid"] ** 2, axis=0)
    SS_res_edss = np.sum(models["mean_t25fw_excluded"]["resid"] ** 2, axis=0)
    SS_res_t25fw = np.sum(models["mean_EDSS_excluded"]["resid"] ** 2, axis=0)

    R2_both = 1 - (SS_res_both / SS_tot)
    R2_edss = 1 - (SS_res_edss / SS_tot)
    R2_t25fw = 1 - (SS_res_t25fw / SS_tot)

    time_of_day = np.linspace(0, 24, 48)

    fig, ax = plt.subplots(figsize=(10.5, 6.5))
    ax.plot(time_of_day, R2_both, label="Both")
    ax.plot(time_of_day, R2_edss, label="EDSS")
    ax.plot(time_of_day, R2_t25fw, label="T25FW")

    ax.set_xlabel("Time of Day")
    ax.set_ylabel("$R^2$")
    ax.set_title("$R^2$ over Time of Day")
    ax.legend(
        title="Model", frameon=False, loc="center left", bbox_to_anchor=(1.02, 0.5)
    )
    ax.grid(True, linestyle="--", alpha=0.7)
    set_2h_time_axis(ax)

    save_fig(fig, out_dir / "Fig_2A.tif", dpi)
    plt.close(fig)


def make_fig2b(models_both, out_dir, dpi):
    coef_est = models_both["est.func"]
    se_est = models_both["se.func"]

    time_points = np.arange(0, 24, 0.5)

    alpha = 0.05
    z_score = norm.ppf(1 - alpha / 2)
    conf_int_lower = coef_est - z_score * se_est
    conf_int_upper = coef_est + z_score * se_est

    bonferroni_alpha = alpha / coef_est.shape[0]
    z_score_bonferroni = norm.ppf(1 - bonferroni_alpha / 2)
    conf_int_lower_bonferroni = coef_est - z_score_bonferroni * se_est
    conf_int_upper_bonferroni = coef_est + z_score_bonferroni * se_est

    significant = (conf_int_lower_bonferroni > 0) | (conf_int_upper_bonferroni < 0)

    labels = ["Sex", "Age", "BMI", "Mean EDSS", "Mean T25FW", "Intercept"]

    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    axes = axes.flatten()

    for i, label in enumerate(labels):
        ax = axes[i]
        ax.plot(time_points, coef_est[:, i], color="blue")
        ax.fill_between(
            time_points,
            conf_int_lower[:, i],
            conf_int_upper[:, i],
            color="lightblue",
            alpha=0.5,
        )
        ax.fill_between(
            time_points,
            conf_int_lower[:, i],
            conf_int_upper[:, i],
            where=significant[:, i],
            color="orange",
            alpha=0.5,
        )
        ax.axhline(0, color="black", linestyle="--")
        ax.set_xlabel("Time of Day")
        ax.set_ylabel("Estimate")
        ax.set_title(label)
        set_4h_time_axis(ax)

    fig.suptitle(
        "Model Coefficient Estimate Throughout the Day", fontsize=20, fontweight="bold"
    )
    plt.tight_layout(rect=[0, 0, 1, 0.96])

    save_fig(fig, out_dir / "Fig_2B.tif", dpi)
    plt.close(fig)


def make_fig3(bayes_results, out_dir, dpi):
    beta_hat = bayes_results["beta.hat"]
    beta_lb = bayes_results["beta.LB"]
    beta_ub = bayes_results["beta.UB"]

    n_resp = beta_hat.shape[1]
    time_bins = np.arange(0, 24, 24 / n_resp)

    alpha = 0.05
    bonferroni_alpha = alpha / len(time_bins)
    z_score_corrected = norm.ppf(1 - bonferroni_alpha / 2)
    z_score_original = norm.ppf(0.975)

    titles = ["Intercept", "Sex", "Age", "BMI", "Close EDSS", "Close T25FW"]
    idxs = [0, 1, 2, 3, 4, 5]

    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    axes = axes.flatten()

    for ax, title, idx in zip(axes, titles, idxs):
        coef = beta_hat[idx, :]
        lb95 = beta_lb[idx, :]
        ub95 = beta_ub[idx, :]

        se_est = (ub95 - lb95) / (2 * z_score_original)
        lb_b = coef - z_score_corrected * se_est
        ub_b = coef + z_score_corrected * se_est
        sig = (lb_b > 0) | (ub_b < 0)

        ax.plot(time_bins, coef, color="blue")
        ax.fill_between(time_bins, lb95, ub95, color="lightblue", alpha=0.5)
        ax.fill_between(time_bins, lb95, ub95, where=sig, color="orange", alpha=0.5)
        ax.axhline(0, color="black", linestyle="--")

        ax.set_xlabel("Time of Day")
        ax.set_ylabel("Estimate")
        ax.set_title(title)
        set_4h_time_axis(ax)

    fig.suptitle(
        "Model Coefficient Estimate Throughout the Day", fontsize=20, fontweight="bold"
    )
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    save_fig(fig, out_dir / "Fig_3.tif", dpi)
    plt.close(fig)


def load_kfs_Y(processed_dir: Path) -> np.ndarray:
    y_path = processed_dir / "function_on_scalar_regression_steps_time_series_kfs.csv"
    require_file(y_path)
    Y = np.loadtxt(y_path, delimiter=",")
    if Y.ndim != 2 or Y.shape[1] != 48:
        raise ValueError(f"Expected KFS Y shape (n,48). Got {Y.shape}")
    return Y


def load_kfs_models(processed_dir: Path):
    files = {
        "full": "fosr_model_full.rds",
        "Cerebellar": "fosr_model_Cerebellar.rds",
        "Pyramidal_motor_function": "fosr_model_Pyramidal_motor_function.rds",
        "Sensory": "fosr_model_Sensory.rds",
        "Bowel_and_Bladder": "fosr_model_Bowel_and_Bladder.rds",
        "Cerebral_mental_function": "fosr_model_Cerebral_mental_function.rds",
        "Optic___Visual": "fosr_model_Optic___Visual.rds",
        "Ambulation": "fosr_model_Ambulation.rds",
        "Brainstem": "fosr_model_Brainstem.rds",
    }
    out = {}
    for k, fname in files.items():
        out[k] = extract_fields(read_rds(processed_dir / fname))
    return out


def make_fig4_kfs_r2(
    Y_kfs: np.ndarray, kfs_models: dict, out_dir: Path, dpi: int
) -> None:
    Y_mean = np.mean(Y_kfs, axis=0)
    SS_tot = np.sum((Y_kfs - Y_mean) ** 2, axis=0)

    time_of_day = np.linspace(0, 24, 48)

    label_map = {
        "full": "Full Model",
        "Cerebellar": "Cerebellar",
        "Pyramidal_motor_function": "Pyramidal",
        "Sensory": "Sensory",
        "Bowel_and_Bladder": "Bowel/Bladder",
        "Cerebral_mental_function": "Cerebral",
        "Optic___Visual": "Optic/Visual",
        "Ambulation": "Ambulation",
        "Brainstem": "Brainstem",
    }

    fig, ax = plt.subplots(figsize=(12, 7))

    for key in kfs_models.keys():
        resid = kfs_models[key]["resid"]
        SS_res = np.sum(resid**2, axis=0)
        r2 = 1 - (SS_res / SS_tot)
        ax.plot(time_of_day, r2, label=label_map[key], linewidth=2)

    ax.set_xlabel("Time of Day")
    ax.set_ylabel("$R^2$")
    ax.set_title("$R^2$ over Time of Day")
    ax.legend(title="Model", bbox_to_anchor=(1.05, 1), loc="upper left", frameon=False)
    ax.grid(True, linestyle="--", alpha=0.7)
    set_2h_time_axis(ax)

    plt.tight_layout()
    save_fig(fig, out_dir / "Fig_4.tif", dpi)
    plt.close(fig)


def make_fig5_kfs_coefficients(kfs_full_model: dict, out_dir: Path, dpi: int) -> None:
    coef_est = kfs_full_model["est.func"]
    se_est = kfs_full_model["se.func"]

    time_points = np.arange(0, 24, 0.5)

    alpha = 0.05
    z_score = norm.ppf(1 - alpha / 2)
    conf_int_lower = coef_est - z_score * se_est
    conf_int_upper = coef_est + z_score * se_est

    bonferroni_alpha = alpha / coef_est.shape[1]
    z_score_bonferroni = norm.ppf(1 - bonferroni_alpha / 2)
    conf_int_lower_bonferroni = coef_est - z_score_bonferroni * se_est
    conf_int_upper_bonferroni = coef_est + z_score_bonferroni * se_est

    significant = (conf_int_lower_bonferroni > 0) | (conf_int_upper_bonferroni < 0)

    labels = [
        "Sex",
        "Age",
        "BMI",
        "Cerebellar",
        "Pyramidal (motor function)",
        "Sensory",
        "Bowel and Bladder",
        "Cerebral (mental function)",
        "Optic / Visual",
        "Ambulation",
        "Brainstem",
        "Intercept",
    ]

    n_labels = len(labels)
    n_cols = 3
    n_rows = int(np.ceil(n_labels / n_cols))

    fig, axes = plt.subplots(n_rows, n_cols, figsize=(6 * n_cols, 4.5 * n_rows))
    axes = axes.flatten()

    for i, label in enumerate(labels):
        ax = axes[i]
        ax.plot(time_points, coef_est[:, i], color="blue")
        ax.fill_between(
            time_points,
            conf_int_lower[:, i],
            conf_int_upper[:, i],
            color="lightblue",
            alpha=0.5,
        )
        ax.fill_between(
            time_points,
            conf_int_lower[:, i],
            conf_int_upper[:, i],
            where=significant[:, i],
            color="orange",
            alpha=0.5,
        )
        ax.axhline(0, color="black", linestyle="--")
        ax.set_xlabel("Time of Day")
        ax.set_ylabel("Estimate")
        ax.set_title(label)
        set_4h_time_axis(ax)

    for j in range(i + 1, len(axes)):
        fig.delaxes(axes[j])

    fig.suptitle("Model Coefficient Estimate Throughout the Day", fontsize=20)
    plt.tight_layout(rect=[0, 0, 1, 0.96])

    save_fig(fig, out_dir / "Fig_5.tif", dpi)
    plt.close(fig)


def main():
    repo_root = Path(__file__).resolve().parents[1]
    processed_dir = repo_root / "data" / "processed"
    out_dir = repo_root / "figures"

    ap = argparse.ArgumentParser()
    ap.add_argument("--dpi", type=int, default=300)
    ap.add_argument("--base_font", type=int, default=14)
    args = ap.parse_args()

    set_pub_style(args.base_font)
    ensure_dir(out_dir)

    steps_time_series = load_Y(processed_dir)
    predictors_raw = load_X(processed_dir)
    predictors, EDSS_labels, T25FW_labels = add_bins_notebook(predictors_raw)

    models = load_models(processed_dir)
    yhat = models["both"]["yhat"]

    if yhat.shape != steps_time_series.shape:
        raise ValueError("Shape mismatch between model yhat and steps_time_series.")
    if len(predictors) != steps_time_series.shape[0]:
        raise ValueError("Row mismatch between predictors and steps_time_series.")

    make_fig1a(steps_time_series, predictors, yhat, EDSS_labels, out_dir, args.dpi)
    make_fig1b(steps_time_series, predictors, yhat, T25FW_labels, out_dir, args.dpi)
    make_fig2a(steps_time_series, models, out_dir, args.dpi)
    make_fig2b(models["both"], out_dir, args.dpi)

    bayes_results = load_bayes_fosr_results(processed_dir)
    make_fig3(bayes_results, out_dir, args.dpi)

    Y_kfs = load_kfs_Y(processed_dir)
    kfs_models = load_kfs_models(processed_dir)
    make_fig4_kfs_r2(Y_kfs, kfs_models, out_dir, args.dpi)
    make_fig5_kfs_coefficients(kfs_models["full"], out_dir, args.dpi)

    print(out_dir.resolve())


if __name__ == "__main__":
    main()