run_lesions_model.py

import math
import os
import logging
import warnings
import re  

# HACK: needed for numpyro 0.16.1
os.environ["XLA_FLAGS"] = "--xla_cpu_use_thunk_runtime=false" 


import matplotlib
matplotlib.use('Cairo')

import arviz as az
import seaborn as sns
import matplotlib.pyplot as plt

import pandas as pd
from pandas.api.types import CategoricalDtype
import numpy as np # do we need this?
import jax.numpy as jnp
import jax

from jax import lax, random
import numpyro
from numpyro.infer import MCMC, NUTS, Predictive
from numpyro.diagnostics import hpdi,summary

## for model
import numpyro.distributions as dist

# for storing data
import pyarrow.feather as feather

## colors
import matplotlib as mpl
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

# data handling
from ipmsa.ipmsa_data import *

# Bayesian model
from model.lesion_model import *
from model.posterior_predictive import *


import argparse

def parse_options():
    parser = argparse.ArgumentParser(description='Run bayesian model')

    parser.add_argument("--prefix", default="results_20240612",
                        help="Prefix for files")
    
    parser.add_argument("--target", default="les_t2",
                        choices=['les_t2','les_t1'],
                        help="Prediction target")
    
    parser.add_argument("--fast", action="store_true", default=False,
                        help="Run fast (for testing)")
    
    parser.add_argument("--nodrugs", action="store_true", default=False,
                        help="Replace drug names with category")

    parser.add_argument("--prior", action="store_true", default=False,
                        help="Genereate prior samples")
    
    parser.add_argument("--posterior", action="store_true", default=False,
                        help="Genereate posterior predictive samples")
    
    parser.add_argument("--all", action="store_true", default=False,
                        help="Replace drug names with category")

    parser.add_argument("-j",  default=8,type=int,
                        help="Number of threads to use")


    return parser.parse_args()


def main():
    params = parse_options()
   
    model_name='lesion_model'
    output_pfx=params.prefix + "/" + model_name + "/" + params.target

    os.makedirs(output_pfx,exist_ok=True)

    use_cpu=True

    if use_cpu:
        numpyro.set_platform("cpu")
        numpyro.enable_x64()
        numpyro.set_host_device_count(params.j)

    az.rcParams["stats.hdi_prob"] = 0.95

    print("Local devices:",jax.local_devices())
    print("count:",jax.local_device_count())
    print("Using HDI:",az.rcParams["stats.hdi_prob"])

    rng_key = random.PRNGKey(0)
    rng_key, rng_key_,rng_key_2,rng_key_3 = random.split(rng_key,num=4)
    logging.basicConfig(format='%(message)s', level=logging.INFO)

    if params.all:
        mstype = None
        clin_col   = [ "age_bl", 'dd_bl', 'les_t2_diff', 'les_t2_bl'] 
    else:
        mstype = 'RRMS'
        clin_col   = [ "age_bl", 'dd_bl', 'les_t1_diff', 'les_t2_diff', 'les_t1_bl', 'les_t2_bl'] # 'edss_bl','age_bl','dd_bl'  'les_t1_bl', 'les_t2_bl'

    # setup data
    # outcome
    output_col = [ params.target,]

    #### DEBUG
    #trials=['OPERA1', 'OPERA2', 'BRAVO', 'ADVANCE', 'DEFINE'] # REMOVING CONFIRM
    ####

    db = IPMSA_db(clin_col, output_col, mstype=mstype, 
                  conf='conf152', ver='20240515',
                  meas="SynthSeg", 
                  #trials=trials,
                  reduced_set=True,filter_bad=True)
    
    features_bl_y = [ "age_bl" , "dd_bl"] 

    #features_bl_y_n = [ i+'_n' for i in features_bl_y  ]
    output_col_n    = [ i+'_n' for i in output_col ]
    norm_t =2.0 # maximum FU is 2yrs

    if params.nodrugs:
        ss = (db.ipmsa_r
                .query("fu_ref>=0.0 and fu_ref<=@norm_t") # remove far away points
                .assign(gid=lambda x: x.gid.cat.remove_unused_categories()) 
                .assign(trial_arm=lambda x: x.trt.cat.remove_unused_categories()) # replace trial_arm with drug efficency category
            )
    else:
        # use drug names
        ss = (db.ipmsa_r
                .query("fu_ref>=0.0 and fu_ref<=@norm_t") # remove far away points
                .assign(gid=lambda x: x.gid.cat.remove_unused_categories()) 
                .assign(trial_arm=lambda x: x.drug.cat.remove_unused_categories()) # replace trial_arm with drug
            )
    
    if params.all:
        ss = (ss.assign(trial_arm = lambda x: x.trial_arm.astype('str') + " " + x.mstype.astype('str'))
                .assign(trial_arm = lambda x: x.trial_arm.astype('category')))
    
    # baseline measurements
    first_ss = ss.groupby('gid',observed=False).first().reset_index()
    n_gid = first_ss.gid.cat.categories.size

    # debug
    print(first_ss.fu_ref) # should print zeros
    # model data, convert to numpyro tensors
    dat={   
        # number of groups
        "n_id":             n_gid,
        "n_trt":            first_ss.trial_arm.cat.categories.size,
        "n_trial":          first_ss.trial.cat.categories.size,
        "n_sex":            first_ss.sex.cat.categories.size,

        # baseline data
        "bl_sex"  :jnp.array(first_ss.sex.cat.codes),
        "bl_trt"  :jnp.array(first_ss.trial_arm.cat.codes),
        "bl_trial":jnp.array(first_ss.trial.cat.codes),

        "bl_age":jnp.array(first_ss.age_bl_n),
        "bl_dd":jnp.array(first_ss.dd_bl_n),

        "t":  jnp.array(ss.fu_ref),
        "id": jnp.array(ss.gid.cat.codes),
        "y":  jnp.array(ss[output_col_n[0]]),
    }

    # run inference
    if params.fast:
        n_samp=200
        n_warm=200
        thinning=1
        n_ch=4
    else:
        n_samp=5000
        n_warm=5000
        thinning=5
        n_ch=4
    
    acc_rate=0.9
    chain_method="parallel" if n_ch>1 and use_cpu else "sequential"
    progress_bar=True # (n_ch==1)

    mcmc = MCMC(NUTS(lesion_model,target_accept_prob=acc_rate), 
        num_warmup=n_warm, num_samples=n_samp, num_chains=n_ch, 
        chain_method=chain_method, thinning=thinning,
        progress_bar=progress_bar) 

    mcmc.run(rng_key,**dat)

    print("Number of divergences: ", mcmc.get_extra_fields()["diverging"].sum())

    # summary
    post = mcmc.get_samples()

    if False: # diagnostics
        sum = summary(mcmc.get_samples(True))    
        print(sum['lambda0']['r_hat']   )

    if params.prior:
        prior = Predictive(lesion_model, num_samples=1000)( rng_key_3, **dat)
    else:
        prior=None

    if params.posterior:
        posterior_predictive = Predictive(lesion_model, post)( rng_key_2, **dat)
    else:
        posterior_predictive = None

    # convert to ARVIZ data
    az_mcmc=az.from_numpyro(mcmc,
            prior=prior,
            posterior_predictive=posterior_predictive,
                           
            coords={
                    "trt":first_ss.trial_arm.cat.categories,
                    "trial":first_ss.trial.cat.categories,
                    "feature" :['intercept','age','dd'],
                    "sex" :first_ss.sex.cat.categories,
                    },
            dims=  {
                    "beta_trt_slope":  ["trt"],
                    #"beta_trial_sex":  ["sex","trial"],
                    "beta_age_dd_sex_trial": ["feature", "sex", "trial"],
                   }
            )    

    output_var_names=["beta_trt_slope", 
                      "beta_intercept",
                      "beta_age", 
                      "beta_dd"]
    
    output_sigma_names=["rand_sigma",
                        "sigma_trt_slope",
                        "sigma_age_dd_sex_trial"]
    
    # show diagnostic plots
    g=az.plot_trace(az_mcmc, var_names=output_var_names, compact=False, figsize=(20,8)) 
    plt.savefig(f"{output_pfx}/diag_trace_main_{params.target}.png", bbox_inches='tight')

    g=az.plot_trace(az_mcmc, var_names=output_sigma_names,
                    compact=False, figsize=(20,8))
    plt.savefig(f"{output_pfx}/diag_trace_output_sigmas_{params.target}.png", bbox_inches='tight')

    g=az.plot_trace(az_mcmc, var_names=["beta_age_dd_sex_trial_z"], compact=False, figsize=(20,6)) # "beta_trt_z"
    plt.savefig(f"{output_pfx}/diag_trace_age_dd_sex_trial_{params.target}.png", bbox_inches='tight')

    # show sigmas
    g=az.plot_forest(az_mcmc,
                var_names=output_sigma_names,
                    ess=False,combined=True,figsize=(8,8))
    g[0].vlines(0,-1,100,linestyles='dotted')
    plt.savefig(f"{output_pfx}/regression_sigmas_{params.target}.png", bbox_inches='tight')

    #### show regression results plots
    g=az.plot_forest(az_mcmc,
                    var_names=output_var_names,
                    ess=False,combined=True,figsize=(8,12),transform=lambda x: x*db.norm_std[output_col[0]])
    g[0].vlines(0,-1,100,linestyles='dotted')
    plt.savefig(f"{output_pfx}/regression_results_{params.target}.png", bbox_inches='tight')

    #### calculate summary statistics
    regression_summary=az.summary(az_mcmc, 
                        var_names=output_var_names,
                        fmt='wide',kind='stats',extend=False,
                        stat_funcs={
                            "mean": np.mean,
                            "std": np.std,
                            "2.5%": lambda x: np.percentile(x, 2.5),
                            "median": lambda x: np.percentile(x, 50),
                            "97.5%": lambda x: np.percentile(x, 97.5),
                        })
                        
    print(regression_summary)
    regression_summary.to_csv(f"{output_pfx}/regression_summary_{params.target}.csv")

    # posterior predictive
    intepolate_t = jnp.linspace(0.0, norm_t, 100)

    sim_dat={
        "n_id": ss.gid.cat.categories.size,
        "n_trt":ss.trial_arm.cat.categories.size,
        "n_sex":ss.sex.cat.categories.size,
        "n_trial": ss.trial.cat.categories.size,

        # time course
        "t": intepolate_t,
    }

    # plot group trajectories
    fig, ax = plt.subplots(
        ncols=1 ,squeeze=False,
        nrows=1, sharex=True, sharey=True, figsize=(8, 8))

    cmap = mpl.colormaps['Set1'].resampled(ss.trial_arm.cat.categories.size)
    samples = mcmc.get_samples()
    print(f"{ax.shape=}")

    for s in [0]: # range(ss.sex.cat.categories.size):
        print(f"{s=}")
        for trt in range(first_ss.trial_arm.cat.categories.size):
            # simulate posterior predictive model with fixed effects only
            sim_dat.update({ 
                "bl_trt":   jnp.full(sim_dat["t"].shape,trt,dtype=jnp.int32),
                "bl_age":   jnp.zeros(sim_dat["t"].shape),
                "bl_dd":    jnp.zeros(sim_dat["t"].shape),
            })

            mu = jnp.squeeze(Predictive(mcmc.sampler.model, samples)(rng_key_, **sim_dat)["mu"]) \
                    * db.norm_std[params.target] + db.norm_mean[params.target]

            az.plot_hdi(intepolate_t*norm_t,  mu, 
                ax=ax[0,s], smooth=False, color=cmap(trt),
                fill_kwargs={'alpha': 0.05}) # 
            
            ax[0,s].plot(intepolate_t*norm_t, 
                    mu.mean(axis=0), 
                    color=cmap(trt), 
                    label=f"{ss.trial_arm.cat.categories[trt]}")
        #
        ax[0,s].set_title(f"{ss.sex.cat.categories[s]} lesion volume")
        ax[0,s].set_xlim(0, norm_t)
        ax[0,s].set_xlabel("Years of follow-up")
        
    ax[0,0].legend(title='Treatment')
    ax[0,0].set_ylabel(params.target)

    plt.suptitle("Posterior Predictive model, mean expected group value")
    fig.tight_layout()
    
    plt.savefig(f"{output_pfx}/posterior_predictive_{params.target}_total.png", bbox_inches='tight')

    ##### calculate R2 values
    r2_median=posterior_predictive_residuals(ss, params.target+"_n", mcmc.sampler.model, mcmc.get_samples(),
                                      output_fig=f"{output_pfx}/residuals_{params.target}.png",
                                      rng_key=rng_key_2,
                                      dat=dat,
                                      title=params.target)
    
    print(f"{model_name} {params.target} R2: {r2_median:.4f}")
    q=pd.DataFrame(data=dict(model=[model_name],r2_median=[r2_median],var=params.target))
    q.to_csv(f"{output_pfx}/r2_summary_{params.target}.csv")

    ### save results for comparing models
    if True:
        az.to_netcdf(az_mcmc, f"{output_pfx}/model_{params.target}.ncdf")


if __name__ == "__main__":
    main()