notebook.community

Edit and run 





In [4]:



    


%load_ext autoreload
%autoreload 2
%matplotlib inline
import random
random.seed(13847942484)
import survivalstan
import numpy as np
import pandas as pd
from stancache import stancache
from matplotlib import pyplot as plt



    


















    






INFO:stancache.seed:Setting seed to 1245502385










    






The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

















In [5]:



    


print(survivalstan.models.pem_survival_model_timevarying)



    


















    






/*  Variable naming:
 // dimensions
 N          = total number of observations (length of data)
 S          = number of sample ids
 T          = max timepoint (number of timepoint ids)
 M          = number of covariates
 
 // data
 s          = sample id for each obs
 t          = timepoint id for each obs
 event      = integer indicating if there was an event at time t for sample s
 x          = matrix of real-valued covariates at time t for sample n [N, X]
 obs_t      = observed end time for interval for timepoint for that obs
 
*/
// Jacqueline Buros Novik <jackinovik@gmail.com>

functions {
  matrix spline(vector x, int N, int H, vector xi, int P) {
    matrix[N, H + P] b_x;         // expanded predictors
    for (n in 1:N) {
        for (p in 1:P) {
            b_x[n,p] <- pow(x[n],p-1);  // x[n]^(p-1)
        }
        for (h in 1:H)
          b_x[n, h + P] <- fmax(0, pow(x[n] - xi[h],P-1)); 
    }
    return b_x;
  }
}
data {
  // dimensions
  int<lower=1> N;
  int<lower=1> S;
  int<lower=1> T;
  int<lower=0> M;
  
  // data matrix
  int<lower=1, upper=N> s[N];     // sample id
  int<lower=1, upper=T> t[N];     // timepoint id
  int<lower=0, upper=1> event[N]; // 1: event, 0:censor
  matrix[N, M] x;                 // explanatory vars

  // timepoint data
  vector<lower=0>[T] t_obs;
  vector<lower=0>[T] t_dur;
}
transformed data {
  vector[T] log_t_dur;
  int n_trans[S, T];  

  log_t_dur = log(t_obs);
  
  // n_trans used to map each sample*timepoint to n (used in gen quantities)
  // map each patient/timepoint combination to n values
  for (n in 1:N) {
      n_trans[s[n], t[n]] = n;
  }

  // fill in missing values with n for max t for that patient
  // ie assume "last observed" state applies forward (may be problematic for TVC)
  // this allows us to predict failure times >= observed survival times
  for (samp in 1:S) {
      int last_value;
      last_value = 0;
      for (tp in 1:T) {
          // manual says ints are initialized to neg values
          // so <=0 is a shorthand for "unassigned"
          if (n_trans[samp, tp] <= 0 && last_value != 0) {
              n_trans[samp, tp] = last_value;
          } else {
              last_value = n_trans[samp, tp];
          }
      }
  }  
}
parameters {
  vector[T] log_baseline_raw;    // unstructured baseline hazard for each timepoint t
  real<lower=0> baseline_sigma;
  real log_baseline_mu;
  
  vector[M] beta; // beta-intercept
  vector<lower=0>[M] beta_time_sigma;
  vector[T-1] raw_beta_time_deltas[M]; // for each coefficient
                                       // change in coefficient value from previous time
}
transformed parameters {
  vector[N] log_hazard;
  vector[T] log_baseline;
  vector[T] beta_time[M];
  vector[T] beta_time_deltas[M];

  // adjust baseline hazard for duration of each period
  log_baseline = log_baseline_raw + log_t_dur;
  
  // compute timepoint-specific betas 
  // offsets from previous time
  for (coef in 1:M) {
      beta_time_deltas[coef][1] = 0;
      for (time in 2:T) {
          beta_time_deltas[coef][time] = raw_beta_time_deltas[coef][time-1];
      }
  }
  
  // coefficients for each timepoint T
  for (coef in 1:M) {
      beta_time[coef] = beta[coef] + cumulative_sum(beta_time_deltas[coef]);
  }

  // compute log-hazard for each obs
  for (n in 1:N) {
    real log_linpred;
    log_linpred <- 0;
    for (coef in 1:M) {
      // for now, handle each coef separately
      // (to be sure we pull out the "right" beta..)
      log_linpred = log_linpred + x[n, coef] * beta_time[coef][t[n]]; 
    }
    log_hazard[n] = log_baseline_mu + log_baseline[t[n]] + log_linpred;
  }
}
model {
  // priors on time-varying coefficients
  for (m in 1:M) {
    raw_beta_time_deltas[m][1] ~ normal(0, 100);
    for(i in 2:(T-1)){
        raw_beta_time_deltas[m][i] ~ normal(raw_beta_time_deltas[m][i-1], beta_time_sigma[m]);
    }
  }
  beta_time_sigma ~ cauchy(0, 1);
  beta ~ cauchy(0, 1);
  
  // priors on baseline hazard
  log_baseline_mu ~ normal(0, 1);
  baseline_sigma ~ normal(0, 1);
  log_baseline_raw[1] ~ normal(0, 1);
  for (i in 2:T) {
      log_baseline_raw[i] ~ normal(log_baseline_raw[i-1], baseline_sigma);
  }
  
  // model
  event ~ poisson_log(log_hazard);
}
generated quantities {
  real log_lik[N];
  vector[T] baseline;
  int y_hat_mat[S, T];     // ppcheck for each S*T combination
  real y_hat_time[S];      // predicted failure time for each sample
  int y_hat_event[S];      // predicted event (0:censor, 1:event)
  
  // compute raw baseline hazard, for summary/plotting
  baseline = exp(log_baseline_raw);
  
  // log_likelihood for loo-psis
  for (n in 1:N) {
      log_lik[n] <- poisson_log_lpmf(event[n] | log_hazard[n]);
  }
  
  // posterior predicted values
  for (samp in 1:S) {
      int sample_alive;
      sample_alive = 1;
      for (tp in 1:T) {
        if (sample_alive == 1) {
              int n;
              int pred_y;
              real log_linpred;
              real log_haz;
              
              // determine predicted value of y
              n = n_trans[samp, tp];
              
              // (borrow code from above to calc linpred)
              // but use sim tp not t[n] 
              log_linpred = 0;
              for (coef in 1:M) {
                  // for now, handle each coef separately
                  // (to be sure we pull out the "right" beta..)
                  log_linpred = log_linpred + x[n, coef] * beta_time[coef][tp]; 
              }
              log_haz = log_baseline_mu + log_baseline[tp] + log_linpred;
              
              // now, make posterior prediction 
              if (log_haz < log(pow(2, 30))) 
                  pred_y = poisson_log_rng(log_haz);
              else
                  pred_y = 9; 
              
              // mark this patient as ineligible for future tps
              // note: deliberately make 9s ineligible 
              if (pred_y >= 1) {
                  sample_alive = 0;
                  y_hat_time[samp] = t_obs[tp];
                  y_hat_event[samp] = 1;
              }
              
              // save predicted value of y to matrix
              y_hat_mat[samp, tp] = pred_y;
          }
          else if (sample_alive == 0) {
              y_hat_mat[samp, tp] = 9;
          } 
      } // end per-timepoint loop
      
      // if patient still alive at max
      // 
      if (sample_alive == 1) {
          y_hat_time[samp] = t_obs[T];
          y_hat_event[samp] = 0;
      }
  } // end per-sample loop
}


















In [6]:



    


d = stancache.cached(
    survivalstan.sim.sim_data_exp_correlated,
    N=100,
    censor_time=20,
    rate_form='1 + sex',
    rate_coefs=[-3, 0.5],
)
d['age_centered'] = d['age'] - d['age'].mean()
d.head()



    


















    






INFO:stancache.stancache:sim_data_exp_correlated: cache_filename set to sim_data_exp_correlated.cached.N_100.censor_time_20.rate_coefs_54462717316.rate_form_1 + sex.pkl
INFO:stancache.stancache:sim_data_exp_correlated: Loading result from cache










    
Out[6]:










  
    

      
      age
      sex
      rate
      true_t
      t
      event
      index
      age_centered
    

  
  
    

      0
      59
      male
      0.082085
      20.948771
      20.000000
      False
      0
      4.18
    

    

      1
      58
      male
      0.082085
      12.827519
      12.827519
      True
      1
      3.18
    

    

      2
      61
      female
      0.049787
      27.018886
      20.000000
      False
      2
      6.18
    

    

      3
      57
      female
      0.049787
      62.220296
      20.000000
      False
      3
      2.18
    

    

      4
      55
      male
      0.082085
      10.462045
      10.462045
      True
      4
      0.18
    

  




















In [7]:



    


survivalstan.utils.plot_observed_survival(df=d[d['sex']=='female'], event_col='event', time_col='t', label='female')
survivalstan.utils.plot_observed_survival(df=d[d['sex']=='male'], event_col='event', time_col='t', label='male')
plt.legend()



    


















    
Out[7]:








<matplotlib.legend.Legend at 0x7f6f245cc550>










    
























In [8]:



    


dlong = stancache.cached(
    survivalstan.prep_data_long_surv,
    df=d, event_col='event', time_col='t'
)
dlong.sort_values(['index', 'end_time'], inplace=True)



    


















    






INFO:stancache.stancache:prep_data_long_surv: cache_filename set to prep_data_long_surv.cached.df_33772694934.event_col_event.time_col_t.pkl
INFO:stancache.stancache:prep_data_long_surv: Loading result from cache

















In [9]:



    


dlong.head()



    


















    
Out[9]:










  
    

      
      age
      sex
      rate
      true_t
      t
      event
      index
      age_centered
      key
      end_time
      end_failure
    

  
  
    

      62
      59
      male
      0.082085
      20.948771
      20.0
      False
      0
      4.18
      1
      0.118611
      False
    

    

      3
      59
      male
      0.082085
      20.948771
      20.0
      False
      0
      4.18
      1
      0.196923
      False
    

    

      61
      59
      male
      0.082085
      20.948771
      20.0
      False
      0
      4.18
      1
      0.262114
      False
    

    

      71
      59
      male
      0.082085
      20.948771
      20.0
      False
      0
      4.18
      1
      0.641174
      False
    

    

      26
      59
      male
      0.082085
      20.948771
      20.0
      False
      0
      4.18
      1
      0.944220
      False
    

  




















In [10]:



    


testfit = survivalstan.fit_stan_survival_model(
    model_cohort = 'test model',
    model_code = survivalstan.models.pem_survival_model_timevarying,
    df = dlong,
    sample_col = 'index',
    timepoint_end_col = 'end_time',
    event_col = 'end_failure',
    formula = '~ age_centered + sex',
    iter = 10000,
    chains = 4,
    seed = 9001,
    FIT_FUN = stancache.cached_stan_fit,
    )



    


















    






INFO:stancache.stancache:Step 1: Get compiled model code, possibly from cache
INFO:stancache.stancache:StanModel: cache_filename set to anon_model.cython_0_25_1.model_code_9304163442804524267.pystan_2_12_0_0.stanmodel.pkl
INFO:stancache.stancache:StanModel: Loading result from cache
INFO:stancache.stancache:Step 2: Get posterior draws from model, possibly from cache
INFO:stancache.stancache:sampling: cache_filename set to anon_model.cython_0_25_1.model_code_9304163442804524267.pystan_2_12_0_0.stanfit.chains_4.data_75232070308.iter_10000.seed_9001.pkl
INFO:stancache.stancache:sampling: Loading result from cache
/home/jacquelineburos/miniconda3/envs/python3/lib/python3.5/site-packages/stanity/psis.py:228: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison
  elif sort == 'in-place':
/home/jacquelineburos/miniconda3/envs/python3/lib/python3.5/site-packages/stanity/psis.py:246: VisibleDeprecationWarning: using a non-integer number instead of an integer will result in an error in the future
  bs /= 3 * x[sort[np.floor(n/4 + 0.5) - 1]]
/home/jacquelineburos/miniconda3/envs/python3/lib/python3.5/site-packages/stanity/psis.py:262: RuntimeWarning: overflow encountered in exp
  np.exp(temp, out=temp)

superficial check of convergence





In [11]:



    


survivalstan.utils.print_stan_summary([testfit], pars='lp__')



    


















    






            mean    se_mean          sd        2.5%         50%       97.5%      Rhat
lp__  407.328012  15.738683  101.998321  214.786192  407.532163  606.393469  1.059665

















In [12]:



    


survivalstan.utils.plot_stan_summary([testfit], pars='log_baseline')

summarize coefficient estimates





In [13]:



    


survivalstan.utils.plot_coefs([testfit], element='baseline')



    


















    
























In [14]:



    


survivalstan.utils.plot_coefs([testfit])

posterior-predictive checks





In [15]:



    


survivalstan.utils.plot_pp_survival([testfit], fill=False)
survivalstan.utils.plot_observed_survival(df=d, event_col='event', time_col='t', color='green', label='observed')
plt.legend()



    


















    
Out[15]:








<matplotlib.legend.Legend at 0x7f6f0407b630>










    
























In [16]:



    


survivalstan.utils.plot_pp_survival([testfit], by='sex')



    


















    
























In [23]:



    


survivalstan.utils.plot_pp_survival([testfit], by='sex', pal=['red', 'blue'])

summarize time-varying effect of sex on survival

Standard behavior is to plot estimated betas at each timepoint, for each coefficient in the model.





In [31]:



    


survivalstan.utils.plot_coefs([testfit], element='beta_time', ylim=[-1, 2.5])

accessing lower-level functions for plotting effects over time





In [26]:



    


survivalstan.utils.plot_time_betas(models=[testfit], by=['coef'], y='exp(beta)', ylim=[0, 10])

Alternatively, you can extract the beta-estimates for each timepoint & plot them yourself.





In [27]:



    


testfit['time_beta'] = survivalstan.utils.extract_time_betas([testfit])
testfit['time_beta'].head()



    


















    
Out[27]:










  
    

      
      timepoint_id
      beta
      coef
      end_time
      exp(beta)
      model_cohort
    

  
  
    

      0
      1
      0.281855
      sex[T.male]
      0.118611
      1.325587
      test model
    

    

      1
      1
      1.263018
      sex[T.male]
      0.118611
      3.536076
      test model
    

    

      2
      1
      0.875265
      sex[T.male]
      0.118611
      2.399512
      test model
    

    

      3
      1
      0.128097
      sex[T.male]
      0.118611
      1.136664
      test model
    

    

      4
      1
      0.185878
      sex[T.male]
      0.118611
      1.204276
      test model

You can also extract and/or plot data for single coefficients of interest at a time.





In [28]:



    


first_beta = survivalstan.utils.extract_time_betas([testfit], coefs=['sex[T.male]'])
first_beta.head()



    


















    
Out[28]:










  
    

      
      timepoint_id
      beta
      coef
      end_time
      exp(beta)
      model_cohort
    

  
  
    

      0
      1
      0.281855
      sex[T.male]
      0.118611
      1.325587
      test model
    

    

      1
      1
      1.263018
      sex[T.male]
      0.118611
      3.536076
      test model
    

    

      2
      1
      0.875265
      sex[T.male]
      0.118611
      2.399512
      test model
    

    

      3
      1
      0.128097
      sex[T.male]
      0.118611
      1.136664
      test model
    

    

      4
      1
      0.185878
      sex[T.male]
      0.118611
      1.204276
      test model
    

  




















In [30]:



    


import seaborn as sns
sns.boxplot(data=first_beta, x='timepoint_id', y='beta')



    


















    
Out[30]:








<matplotlib.axes._subplots.AxesSubplot at 0x7f6f0d2f3518>










    
























In [32]:



    


survivalstan.utils.plot_time_betas(models=[testfit], y='beta', x='end_time', coefs=['sex[T.male]'])

Note that this same plot can be produced by passing data to plot_time_betas directly.





In [33]:



    


survivalstan.utils.plot_time_betas(df=first_beta, by=['coef'], y='beta', x='end_time')



    


















    
























In [ ]:

Content source: jburos/survivalstan

Similar notebooks:

notebook.community | gallery | about