Need help optimizing estimation of biomathematical model of fatigue

I am attempting to estimate parameters of the unified model of fatigue (https://www.sciencedirect.com/science/article/pii/S0022519313001811), a dynamic model that calculates predicted fatigue based on sleep history and time of day.

I have a working example of how this could work, with the fundamentals of the model all correct, but I am not confident this is done in the best way in terms of using Stan’s full functionality. I am hoping somebody good with Stan could eyeball what I’ve got and let me know if there are any ways to speed up the sampling or improve the code .I did my best to explain the whole current approach below. This dropbox link also has a working example with Stan code, R code, data and notes:

The equations for the model are nicely laid out in Table 2 of the paper below (ignore the caffeine component as I have not implemented that):

The model assumes that fatigue is the sum of two time-varying processes, a homeostatic process S and a circadian component C. The S parameter is estimated continuously using the previous value, time since previous value and activity (sleep or wake). The C parameter is estimated using time of day. Both are governed also by other parameters (e.g., asymptote and rate parameters for S, circadian phase and amplitude for C).

The Stan script starts by defining functions that calculate the dynamic parameters- S, C, and L which is a governing parameter for S.


functions{
  //calculates S during wake
  real Sfun(real sw, //sw = S upon waking
             real taw, //taw = time awake
             real tau_r, //tau_d = controls rate of rise in S during wake
             real U0
             ){ 
    
    real r=exp(-taw/tau_r);        
    real S=U0-r*(U0-sw);
    return S;
  }
  
  //calculates S during sleep
  real Spfun(real ss, //ss = S upon falling asleep
             real tas, //tas = time asleep
             real tau_s, //tau_d = controls rate of decay in S during sleep
             real U0,     //upper assymptote
             real tau_la, //rate of change in lower assymptote
             real ls     //lower assymptote at sleep onset
             ){ 
    
    real term1 = ss*exp(-tas/tau_s);
    real term2 = -2*U0*(1-exp(-tas/tau_s));
    real term3 = (((ls + 2*U0)*tau_la)/(tau_la-tau_s)) * (exp(-tas/tau_la)-exp(-tas/tau_s));            
    real Sp=term1+term2+term3;
    return Sp;
  }
  
 //calculates L during wake
  real Lfun(real lw, //lower assymptote upon falling asleep
           real taw,    //time asleep
           real tau_la, //rate of change in lower assymptote
           real U0      //upper assymptote
          ){
    real term1 = lw*exp(-taw/tau_la);
    real term2 = U0*(1-exp(-taw/tau_la));
    real L=term1+term2;
    return L;
  }
  
  
  
  //calculates L during sleep
  real Lpfun(real ls, //lower assymptote upon falling asleep
           real tas,    //time asleep
           real tau_la, //rate of change in lower assymptote
           real U0      //upper assymptote
          ){
    real term1 = ls*exp(-tas/tau_la);
    real term2 = -2*U0*(1-exp(-tas/tau_la));
    real L=term1+term2;
    return L;
  }
  
  
  //calculates C (circadian process)
  real Cfun(real tod, //tod = time of day (in decimal hours)
            real phi,  //phi = phase at beginning of the simulation (I think this should be 0 if t = tod)
            real tau, //tau = period of C process
            real A //amplitute of process
            ){
    
    real omega = 2*pi()/tau;
    real term1 = 0.97*sin(omega*(tod+phi));
    real term2 = 0.22*sin(2*omega*(tod+phi));
    real term3 = 0.07*sin(3*omega*(tod+phi));
    real term4 = 0.03*sin(4*omega*(tod+phi));
    real term5 = 0.0001*sin(5*omega*(tod+phi));
    real C = A*(term1+term2+term3+term4+term5);
    return C;
  }
  
}

In the dropbox link I have an R script View_data.R that shows the type of data frame that is likely in sleep settings. Basically there are rows indicating when people went to sleep, when they woke up, and when their level of fatigue was observed (e.g., with a subjective sleepiness rating). I have converted this into a data list that currently works with Stan, details below:

Nsubj - number of participants.

Ntotal - length of total data frame. This includes points where individuals slept,
woke up or an observation of their fatigue was made

subject - subject number.Currently parameters are estimated as if everybody
is one participant, plan to move to a hierarchical model later (will be slow to test)

event_number - event number for each subject

previous_episode_type - 1 means sleep, 2 means wake

time_since_previous - how long since last event

timeofday - time of day

valid - the trial numbers which correspond to an observation (e.g., somebody performing
a task or rating their sleepiness)

Nvalid - how many trials are valid

This format could be changed if that would help optimize the model fitting.

data {
  int<lower=0> Nsubj;
  int<lower=0> Ntotal;
  int<lower=0> subject[Ntotal];
  int<lower=0> event_number[Ntotal];
  int<lower=0> previous_episode_type[Ntotal];
  real<lower=0> time_since_previous[Ntotal];
  vector<lower=0>[Ntotal] timeofday;
  vector[Ntotal] fatigue;
  int<lower=0> Nvalid;
  int<lower=0> valid[Nvalid];
  
  // have fixed tau_la because the model doesn't recover otherwise
  real<lower=0> tau_la;
  
}  

Below are the governing model parameters to be estimated

 
parameters {

  //parameters
  real<lower=0> U0;
  real<lower=0,upper=1> S0_raw;
  real<lower=-0.11,upper=1> L0_raw;
  real phi;
  real<lower=0> kappa;
  real<lower=0> tau_d; //typically fixed to 4.2  h
  real<lower=0> tau_r; //typically fixed to 18.2 h
  real<lower=0> sigma;
  
  
}

To calculate dynamic parameters, the script goes through each event and calculates S/L based on last event as well as C


transformed parameters {

  //Fixed parameters
  real tau = 24;
  real A = 1;
  real timeawake;
  
  //Calculate level of processes for each observation
  real s_prev;  //level of homeostatic process at previous event
  real l_prev;  //level of lower assymptote
  vector[Ntotal] S; //level of homeostatic process
  vector[Ntotal] C; //level of 24-hour circadian process
  vector[Ntotal] L; //lower assumptote of homeostatic process
  
  //Centre parameters
  real L0 = L0_raw*U0;
  real S0 = L0 + (U0-L0)*S0_raw; 

  //Calculate S and C for each event
  for(i in 1:Ntotal){
    //if it is subject's first event, assign S to be S0;
    if(event_number[i] == 1){
      S[i] = S0;
      L[i] = L0;
      s_prev = S[i];
      l_prev = L[i];
      timeawake = 0;
    }
    //if it is not subject's first event, update S based on S at previous event
    if(event_number[i] > 1){
      //if most recent episode was sleep (1)
      if(previous_episode_type[i] == 1){
        S[i] = Spfun(s_prev,time_since_previous[i],tau_d,U0,tau_la,l_prev);
        L[i] = Lpfun(l_prev,time_since_previous[i],tau_la,U0);
        s_prev = S[i];
        l_prev = L[i]; //L process only updates when sleeps
        timeawake = 0;
      }
      //if most recent episode was wake (2)
      if(previous_episode_type[i] == 2){
        S[i] = Sfun(s_prev,time_since_previous[i],tau_r,U0);
         L[i] = Lfun(l_prev,time_since_previous[i],tau_la,U0);
        s_prev = S[i];
        l_prev = L[i];
        timeawake += time_since_previous[i];
      }
    }  
    
    C[i]=Cfun(timeofday[i],phi,tau,A);
    
  }
}

After calculating all the dynamic parameters the script loops through the data points that are “valid” - that is the fatigue observations - and calculates a likelihood using a normal distribution. This normal distribution corresponds to other work in the sleep literature - e.g., here they estimated model parameters with an extended Kalman filter:

https://onlinelibrary.wiley.com/doi/pdf/10.1111/jsr.12535

In practice these models are often fit to data like sleepiness ratings that can’t be negative so I will probably bound the distribution at 0 later on.

model {
  
  //parameters
  U0 ~ normal(20,5); //centred on midpoint of scale
  phi ~ normal(0,10); //centred on zero
  kappa ~ normal(0,5);
  tau_d ~ normal(0,5); //typically fixed to 4.2  h
  tau_r ~ normal(40,20); //typically fixed to 18.2 h
//tau_la has to be fixed or the model doesn't recover it seems  
//  tau_la ~ normal(80,20); 
  sigma ~ normal(0,3);
  
   for(i in 1:Nvalid){
    fatigue[valid[i]] ~ normal(S[valid[i]]+kappa*C[valid[i]],sigma);
  } 
}

Below are generated likelihoods and posterior predictives

generated quantities{
  real pp[Ntotal];
  vector[Nvalid] log_lik;
  for(i in 1:Ntotal){
    pp[i] = normal_rng(S[i]+kappa*C[i],sigma);
  }
  for(j in 1:Nvalid){
    log_lik[j] = normal_lpdf(fatigue[valid[j]] | S[valid[j]]+kappa*C[valid[j]],sigma);
  }
}

Thanks for any help and consideration you may offer.

1 Like

Instead of:

  for(i in 1:Ntotal){
    pp[i] = normal_rng(S[i]+kappa*C[i],sigma);
  }

I think you may be able to write:

pp=normal_rng(S+kappa*C,sigma);

You may have to use rep_vector(sigma,Ntotal), instead of sigma, I’m not sure.

Your model would take me a bit too much time to understand in order to comment on but here are some generic points that might help:

  • HMC benefits from smoothness and continuity. If the dynamic parts of your model cause discontinuities; then the sampler may struggle.

  • If it’s possible to remove loops by using vectorised versions of things; that’s better. I’m no longer sure whether vectorisation benefits from SIMD in Stan but vectorised functions are usually faster for whatever reason.

  • If once you’ve tried everything else it still doesn’t work fast enough, you can lean on map_rect for linear speed-ups if you have cores to spare. Benchmarks differ but in my own experience doubling the number of cores has lead to about 1.5-1.8 times lower runtime.

2 Likes

Great, thanks for that!

It’s going to be hard to do anything other than looping for the transformed parameters block, but I will try storing off a vector with all the valid observations as that loops so that I can subsequently do a vectorized normal() and normal_lpdf()

1 Like