Dear community,
For fitting DDM I use Wiener First Passage Time Distribution and all works fine in terms of chain diagnostics. However, I cannot run predictive checks since the wiener_rng function is not yet implemented in Stan math library. What can be a current alternative? We tried writing our own wiener_rng, but it didn’t work…
Thank you!
Based on this old thread it seems like it should be possible to implement the wiener_rng function using uniform_rng in pure-Stan. It will likely not be particularly fast, due to the use of a kind of rejection sampling
Thanks! This is what we tried :) didn’t work for us.
I just wrote it up as a Stan program and found that it took an unreasonably long time to produce a result. It compiled and ran, so it “worked” but I doubt it’s usable in a real program.
first of all, thank so much with your help with this! We also wrote something that complies alright, but the predictive checks looked bad. We used the very basic modeling approach, so I think something in the rng was funky.
we also compared results from our custom code and R rwiener code, and we got different outputs.
do you have a wiener model that I could look at? There’s this paper Fast and accurate Monte Carlo sampling of first-passage times from Wiener diffusion models | Scientific Reports which seems promising but there’s some nuance with asymmetric boundaries and I’m not familiar with this density. Having an example in Stan would be helpful.
I think I got the wiener_rng to work. Please test and let me know
functions {
real fs_cdf(real t, real a) {
if (a < 1)
reject("a must be >= 1, found a = ", a);
return erfc(inv_sqrt(2 * a * t));
}
vector make_vars(real mu) {
real mu2 = pow(mu, 2);
real t_tilde = 0.12 + 0.5 * exp(-mu2 / 3);
real a = (3 + sqrt(9 + 4 * mu2)) / 6;
real sqrtamu = sqrt((a - 1) * mu2 / a);
real fourmu2pi = (4 * mu2 + pi() ^ 2) / 8;
real Cf1s = sqrt(a) * exp(-sqrtamu);
real Cf1l = pi() / (4 * fourmu2pi);
real CF1st = Cf1s * fs_cdf(t_tilde | a);
real F1lt = -expm1(-t_tilde * fourmu2pi);
real F1inf = CF1st + Cf1l * (1 - F1lt);
return [mu2, //.......1
t_tilde, //.. 2
a, //.........3
sqrtamu, //...4
fourmu2pi, //.5
Cf1s, //......6
Cf1l, //......7
CF1st, //.....8
F1lt, //......9
F1inf]'; //...10
}
int acceptt_rng(real t_star, real ft, real c) {
if (c <= 0.06385320297074884)
reject("c is ", c);
if (is_nan(c))
reject("c is nan!");
real z = ft * uniform_rng(0, 1);
real b = exp(-c);
int k = 3;
while (1) {
if (z > b)
return 0;
b -= k * exp(-c * k ^ 2);
if (z < b)
return 1;
k += 2;
b += k * exp(-c * k ^ 2);
k += 2;
}
return 0;
}
real sample_small_mu_rng(vector vars) {
real t_star;
real pi_sq = pi() ^ 2;
real mu2 = vars[1];
real a = vars[3];
real sqrtamu = vars[4];
real fourmu2pi = vars[5];
real Cf1s = vars[6];
real Cf1l = vars[7];
real CF1st = vars[8];
real F1lt = vars[9];
real F1inf = vars[10];
int counter_outer = 0;
while (1) {
real p = F1inf * uniform_rng(0, 1);
if (p <= CF1st) {
t_star = 1. / (2 * a * pow(inv_erfc(p / Cf1s), 2));
while (0.5 * t_star <= 0.06385320297074884) {
p = uniform_rng(0.06385320297074884, CF1st);
t_star = 1. / (2 * a * pow(inv_erfc(p / Cf1s), 2));
}
real ft = exp(-1. / (2 * a * t_star) - sqrtamu + mu2 * t_star);
if (acceptt_rng(t_star, ft, 0.5 * t_star) == 1)
return t_star;
} else {
t_star = -log1p(-(p - CF1st) / Cf1l - F1lt) / fourmu2pi;
real pisqt = pi_sq * t_star / 8;
int counter = 0;
while (pisqt <= 0.06385320297074884) {
p = uniform_rng(CF1st, F1inf);
t_star = -log1p(-(p - CF1st) / Cf1l - F1lt) / fourmu2pi;
pisqt = pi_sq * t_star / 8;
}
if (acceptt_rng(t_star, exp(-pisqt), pisqt) == 1)
return t_star;
}
}
return 0;
}
real inverse_gaussian_rng(real mu, real mu_sq) {
real v = pow(std_normal_rng(), 2);
real z = uniform_rng(0, 1);
real x = mu + 0.5 * mu_sq * v
- 0.5 * mu * sqrt(4 * mu * v + mu_sq * v ^ 2);
if (z <= (mu / (mu + x)))
return x;
else
return mu_sq / x;
}
real sample_large_mu_rng(vector vars) {
real mu2 = vars[1];
real t_tilde = vars[2];
real a = vars[3];
real sqrtamu = vars[4];
real fourmu2pi = vars[5];
real Cf1s = vars[6];
real Cf1l = vars[7];
real CF1st = vars[8];
real F1lt = vars[9];
real F1inf = vars[10];
real invabsmu = inv_sqrt(mu2);
int counter_outer = 0;
if (t_tilde >= 0.63662) {
Cf1l = -log(pi() * 0.25) - 0.5 * log(2 * pi());
Cf1s = 0;
} else {
Cf1l = -pi() ^ 2 * t_tilde / 8 + (3. / 2.) * log(t_tilde)
+ 0.5 * inv(t_tilde);
Cf1s = Cf1l + 0.5 * log(2 * pi()) + log(pi() * 0.25);
}
while (1) {
real t_star = inverse_gaussian_rng(invabsmu, inv(mu2));
if (is_nan(t_star))
reject("t_star is nan! ", mu2);
real one2t = 0.5 * inv(t_star);
counter_outer += 1;
if (counter_outer > 10000)
reject("counter_outer_large", invabsmu);
if (t_star <= 2.5) {
real expone2t = exp(Cf1s - one2t);
if (expone2t == 0)
expone2t = 1e-8;
if (acceptt_rng(t_star, expone2t, one2t) == 0 || invabsmu < 0.000666)
return t_star;
} else {
real expone2t = exp(-log(pi() / 4) - 0.5 * log(2 * pi()) - one2t
- (3. / 2.) * log(t_star));
if (acceptt_rng(t_star, expone2t, pi() ^ 2 * t_star / 8) == 0)
return t_star;
}
}
return 0;
}
real fast_pt_rng(real alpha, real tau, real beta, real delta) {
real absmu = abs(delta);
vector[10] vars = make_vars(absmu);
real pt;
if (absmu < 1) {
pt = sample_small_mu_rng(vars);
} else
pt = sample_large_mu_rng(vars);
return beta * alpha * pt;
}
vector wiener_rng(real alpha, real tau, real beta, real delta) {
real t = tau;
real x = beta * alpha;
real mu = abs(delta);
real boundary_cond = delta > 0 ? 0 : 1;
real hit_bound;
vector[2] out;
int counter = 0;
while (1) {
real mutheta;
// lower bound is 0
// upper bound is alpha in stan parmeterization
real xlo = x;
real xhi = alpha - x;
if (alpha < 1e-6) {
mutheta = xhi * mu;
real pt = fast_pt_rng(alpha, tau, beta, xhi * delta);
hit_bound = inv_logit(2 * mutheta * beta * alpha);
real bound = uniform_rng(0, 1) < hit_bound ? boundary_cond : 1 - boundary_cond;
return [t + (xhi ^ 2 * pt), bound]';
} else if (xlo > xhi) {
mutheta = xhi * mu;
t += beta * alpha * (xhi ^ 2 * fast_pt_rng(alpha, tau, beta, xhi * delta));
hit_bound = inv_logit(2 * mutheta * beta * alpha);
if (uniform_rng(0, 1) > hit_bound) {
return [t, boundary_cond]';
}
x -= xhi;
}
mutheta = xlo * mu;
t += beta * alpha * (xlo ^ 2 * fast_pt_rng(alpha, tau, beta, xlo * delta));
hit_bound = inv_logit(2 * mutheta * beta * alpha);
if (uniform_rng(0, 1) < hit_bound) {
out[1] = t;
out[2] = 1 - boundary_cond;
break;
}
x += beta * alpha * xlo;
}
return out;
}
}
data {
int N_lower; // number of trials;
int N_upper;
array[N_lower + N_upper] real rt;
array[N_lower] int x_lower; // correct/error responses per trial;
array[N_upper] int x_upper; // response times per trial
}
parameters {
real tau_raw; // logit of non-decision time
real delta; // drift-rate
real<lower=0> alpha; // boundary separation
real<lower=0, upper=1> beta; // starting point
}
transformed parameters {
real tau = inv_logit(tau_raw) * min(rt); // non-decision time at RT scale
}
model {
delta ~ normal(0, 3);
alpha ~ normal(1, .1);
beta ~ normal(0, 2);
tau_raw ~ normal(-.2, .48);
target += wiener_lpdf(rt[x_upper] | alpha, tau, beta, delta);
target += wiener_lpdf(rt[x_lower] | alpha, tau, 1 - beta, -delta);
}
generated quantities {
array[N_lower + N_upper] vector[2] out;
for (n in 1 : (N_lower + N_upper)) {
out[n] = wiener_rng(alpha, tau, beta, delta);
}
}
edit: fixed an issue with the bounds and not properly reparameterizing for Stan. @nerpa please update to this version.
this is fantastic! thank you so much! We will test it of course and let you know if it worked. :) :)
There’s some parameterization issue that I’m working through. Hold off on using it for now. It’s nearly there though.
This works, don’t ask me how I got it to (it was messy).
functions {
real fs_cdf(real t, real a) {
if (a < 1) {
reject("a must be >= 1, found a = ", a);
}
return erfc(inv_sqrt(2 * a * t));
}
vector make_vars(real mu) {
real mu2 = pow(mu, 2);
real t_tilde = 0.12 + 0.5 * exp(-mu2 / 3);
real a = (3 + sqrt(9 + 4 * mu2)) / 6;
real sqrtamu = sqrt((a - 1) * mu2 / a);
real fourmu2pi = (4 * mu2 + pi() ^ 2) / 8;
real Cf1s = sqrt(a) * exp(-sqrtamu);
real Cf1l = pi() / (4 * fourmu2pi);
real CF1st = Cf1s * fs_cdf(t_tilde | a);
real F1lt = -expm1(-t_tilde * fourmu2pi);
real F1inf = CF1st + Cf1l * (1 - F1lt);
return [mu2, //.......1
t_tilde, //.. 2
a, //.........3
sqrtamu, //...4
fourmu2pi, //.5
Cf1s, //......6
Cf1l, //......7
CF1st, //.....8
F1lt, //......9
F1inf]'; //...10
}
int acceptt_rng(real t_star, real ft, real c) {
if (c <= 0.06385320297074884) {
reject("c is ", c);
}
if (is_nan(c)) {
reject("c is nan!");
}
real z = ft * uniform_rng(0, 1);
real b = exp(-c);
int k = 3;
while (1) {
if (z > b) {
return 0;
}
b -= k * exp(-c * k ^ 2);
if (z < b) {
return 1;
}
k += 2;
b += k * exp(-c * k ^ 2);
k += 2;
}
return 0;
}
real sample_small_mu_rng(vector vars) {
real t_star;
real pi_sq = pi() ^ 2;
real mu2 = vars[1];
real a = vars[3];
real sqrtamu = vars[4];
real fourmu2pi = vars[5];
real Cf1s = vars[6];
real Cf1l = vars[7];
real CF1st = vars[8];
real F1lt = vars[9];
real F1inf = vars[10];
int counter_outer = 0;
while (1) {
real p = F1inf * uniform_rng(0, 1);
if (p <= CF1st) {
t_star = 1. / (2 * a * pow(inv_erfc(p / Cf1s), 2));
while (0.5 * t_star <= 0.06385320297074884) {
p = uniform_rng(0.06385320297074884, CF1st);
t_star = 1. / (2 * a * pow(inv_erfc(p / Cf1s), 2));
}
real ft = exp(-1. / (2 * a * t_star) - sqrtamu + mu2 * t_star);
if (acceptt_rng(t_star, ft, 0.5 * t_star) == 1) {
return t_star;
}
} else {
t_star = -log1p(-(p - CF1st) / Cf1l - F1lt) / fourmu2pi;
real pisqt = pi_sq * t_star / 8;
while (pisqt <= 0.06385320297074884) {
p = uniform_rng(CF1st, F1inf);
t_star = -log1p(-(p - CF1st) / Cf1l - F1lt) / fourmu2pi;
pisqt = pi_sq * t_star / 8;
}
if (acceptt_rng(t_star, exp(-pisqt), pisqt) == 1) {
return t_star;
}
}
}
return 0;
}
real inverse_gaussian_rng(real mu, real mu_sq) {
real v = pow(std_normal_rng(), 2);
real z = uniform_rng(0, 1);
real x = mu + 0.5 * mu_sq * v - 0.5 * mu * sqrt(4 * mu * v + mu_sq * v ^ 2);
if (z <= (mu / (mu + x))) {
return x;
} else {
return mu_sq / x;
}
}
real sample_large_mu_rng(vector vars) {
real mu2 = vars[1];
real t_tilde = vars[2];
real a = vars[3];
real sqrtamu = vars[4];
real fourmu2pi = vars[5];
real Cf1s = vars[6];
real Cf1l = vars[7];
real CF1st = vars[8];
real F1lt = vars[9];
real F1inf = vars[10];
real invabsmu = inv_sqrt(mu2);
if (t_tilde >= 0.63662) {
Cf1l = -log(pi() * 0.25) - 0.5 * log(2 * pi());
Cf1s = 0;
} else {
Cf1l = -pi() ^ 2 * t_tilde / 8 + (3. / 2.) * log(t_tilde) + 0.5 * inv(t_tilde);
Cf1s = Cf1l + 0.5 * log(2 * pi()) + log(pi() * 0.25);
}
while (1) {
real t_star = inverse_gaussian_rng(invabsmu, inv(mu2));
if (is_nan(t_star)) {
reject("t_star is nan! ", mu2);
}
real one2t = 0.5 * inv(t_star);
if (t_star <= 2.5) {
real expone2t = exp(Cf1s - one2t);
if (expone2t == 0) {
expone2t = 1e-8;
}
if (acceptt_rng(t_star, expone2t, one2t) == 0 || invabsmu < 0.000666) {
return t_star;
}
} else {
real expone2t = exp(-log(pi() / 4) - 0.5 * log(2 * pi()) - one2t - (3. / 2.) * log(t_star));
if (acceptt_rng(t_star, expone2t, pi() ^ 2 * t_star / 8) == 0) {
return t_star;
}
}
}
return 0;
}
real fast_pt_rng(real alpha, real tau, real beta, real delta) {
real absmu = abs(delta) ;
vector[10] vars = make_vars(absmu);
real pt;
if (absmu < 1) {
pt = sample_small_mu_rng(vars);
} else {
pt = sample_large_mu_rng(vars);
}
return pt;
}
vector wiener_rng(real alpha, real tau, real beta, real delta) {
real t = 0 ;
real sign_delta = delta > 0 ? 1 : -1;
real x = beta * alpha ;
real mu = abs(delta);
real hit_bound;
vector[2] out;
int counter = 0;
if (beta == 0 || beta == 1) {
return [tau, beta]';
}
while (1) {
real mutheta;
real xlo = x ;
real xhi = alpha - x ;
// lower bound is 0
// upper bound is alpha in stan parmeterization
// symmetric case, [x - xup, x + xup]
if (abs(xlo - xhi) < 1e-6) {
mutheta = xhi * mu;
real pt = fast_pt_rng(alpha, tau, beta, xhi * abs(delta));
hit_bound = sign_delta == 1 ? inv_logit( 2 * mutheta ) : 1 - inv_logit( 2 * mutheta );
real bound = uniform_rng(0, 1) < hit_bound ? 1 : 0;
return [ tau + t + ( square(xhi) * pt), bound]';
// x is closer to upper bound, [x - xup, x + xup]
} else if (xlo > xhi) {
mutheta = xhi * mu;
t += ( square(xhi ) * fast_pt_rng(alpha, tau, beta, xhi* abs(delta))) ;
hit_bound = sign_delta == 1 ? inv_logit( 2 * mutheta ) : 1 - inv_logit( 2 * mutheta );
if (uniform_rng(0, 1) < hit_bound ) {
return [tau + t, 1]';
}
x -= xhi ;
} else {
// x is closer to lower bound, [x - xlo, x + xlo]
mutheta = xlo * mu ;
t += ( square(xlo ) * fast_pt_rng(alpha, tau, beta, xlo* abs(delta) )) ;
hit_bound = sign_delta == 1 ? inv_logit( 2 * mutheta ) : 1 - inv_logit( 2 * mutheta );
if (uniform_rng(0, 1) > hit_bound) {
out[1] = tau + t;
out[2] = 0 ;
break;
}
x += xlo ;
}
}
return out;
}
}
data {
int N;
real<lower=0> alpha_in;
real<lower=0> tau_in;
real<lower=0, upper=1> beta_in;
real delta_in;
}
transformed data {
int N_lower = 0;
int N_upper = 0;
array[N] real rt;
array[N] real idx;
int counter = 0;
for (n in 1 : N) {
vector[2] tmp = wiener_rng(alpha_in, tau_in, beta_in, delta_in);
rt[n] = tmp[1];
idx[n] = tmp[2];
if (tmp[2] == 0) {
N_lower += 1;
} else {
N_upper += 1;
}
}
array[N_lower] int id_lower;
array[N_upper] int id_upper;
array[2] int cnt;
cnt[1] = 0;
cnt[2] = 0;
for (n in 1 : N) {
if (idx[n] == 0) {
cnt[1] += 1;
id_lower[cnt[1]] = n;
} else {
cnt[2] += 1;
id_upper[cnt[2]] = n;
}
}
}
parameters {
real tau_raw; // logit of non-decision time
real delta; // drift-rate
real<lower=0> alpha; // boundary separation
real<lower=0, upper=1> beta; // starting point
}
transformed parameters {
real tau = inv_logit(tau_raw) * min(rt); // non-decision time at RT scale
}
model {
delta ~ normal(0, 3);
alpha ~ normal(1, 2);
beta ~ normal(0, 2);
tau_raw ~ normal(-.2, .48);
target += wiener_lpdf(rt[id_upper] | alpha, tau, beta, delta);
target += wiener_lpdf(rt[id_lower] | alpha, tau, 1 - beta, -delta);
}
In R you can test by
library(cmdstanr)
mod <- cmdstan_model("wiener_rng.stan")
out <- mod$sample(
data = list(
N = 400,
alpha_in = 1.25,
tau_in = 0.2,
beta_in = 0.7,
delta_in = -0.3
),
parallel_chains = 4
)
out
with the following results
> out
variable mean median sd mad q5 q95 rhat ess_bulk ess_tail
lp__ -138.60 -138.29 1.42 1.23 -141.28 -136.94 1.00 1991 2944
tau_raw 2.45 2.45 0.19 0.18 2.15 2.77 1.00 2494 2652
delta -0.36 -0.36 0.11 0.11 -0.54 -0.18 1.00 2315 2585
alpha 1.25 1.25 0.03 0.03 1.20 1.31 1.00 2853 2533
beta 0.72 0.72 0.01 0.02 0.69 0.74 1.00 2279 2817
tau 0.21 0.21 0.00 0.00 0.20 0.21 1.00 2494 2652
edit: found an issue with large delta. Fixed by simplification :)
I will not ask :) and thank you so so much, @spinkney ! We will test it and let you know. This is a tremendous help!
Hey, updated the code a bit. Noticed some issue with large delta that I didn’t test with initially. The fix made the code simpler, please update to the above.
@spinkney , we used your weiner_rng function and thanks so much again for writing it. Overall, the predictive checks are not so bad (especially for the lower boundary condition), but I have a question about how we use your function.
We use it in the generated quantities block, similarly to what you had in your first post. It means that we have multiple draws from the function for every data point (we have 1000 samples in each chain). So we get 1000 predicted RTs (out[1]) and 1000 predicted upper-or-lower boundary responses (out[2]) for each real data sample. What is the right way to deal with this kind of output? What we did, is to average separately the RTs for upper and lower boundary predicted responses, meaning that for every real data sample we get 2 predicted RTs. Does it make sense or should we use all the raw draws from the rng function for estimating the predictive accuracy? I hope my question makes sense.
That’s a great question! I’m not familiar with modeling using this distribution, I put the rng function together to learn more about it. Even with that said, it seems reasonable to average but this could mislead you. You may want to have 4 graphs instead of 2 (or overlay them) and see what the output looks like.
Something a bit more sophisticated that requires a bit of coding is to grab the hit_bound parameter from the rng function and return it as well. You’ll also have to get an estimate of this parameter for the data because it will be the probability of hitting the bound. Then you can do a cool 3d graph. Where the x, y are the same but you have a z axis which goes from 0 to 1. This way you can see how your data distribution overlays with this predicted distribution across the entire domain. If you do this, please, please, please come back and post it. I’m curious what it looks like.
The parameter I’m talking about is
hit_bound = sign_delta == 1 ? inv_logit( 2 * mutheta ) : 1 - inv_logit( 2 * mutheta );
to get this for the data you’ll have to do it in gen quant where you take mutheta = abs(delta) * rt and then calculate as above for hit_bound.
Something like (I haven’t tested this)
array[N] vector[3] out; // return hit_bounds as 3rd parameter
vector[N] probs_dat;
for (n in 1:N){
real mutheta = abs(delta) * rt[n];
probs_dat = sign_delta == 1 ? inv_logit( 2 * mutheta ) : 1 - inv_logit( 2 * mutheta );
// bound == 1 if upper, 0 if lower
if (bound[n]) {
out[n] = wiener_rng(alpha, tau, beta, delta);
} else {
out[n] = wiener_rng(alpha, tau, 1 - beta, -delta);
}
}
@spinkney have you given any thought to what is needed for a stan::math implementation of the rng?
It should be easy. All the functions are available and we don’t need to autodiff through. I just don’t have the time to add this at the moment.
Thank you so much for sharing the code for wiener_rng, @spinkney ! My collaborators and I are working on a project, in which we intend to make predictions from a drift-diffusion model. Therefore, we find this to be extremely helpful!
However, I ran into a few issues when putting the code into my Stan model. In total I ran into three errors and I was unable to proceed at the third one. The first error on the line:
SYNTAX ERROR, MESSAGE(S) FROM PARSER:
error in 'modelf682371fa33_ddm' at line 18, column 30
-------------------------------------------------
16: real Cf1s = sqrt(a) * exp(-sqrtamu);
17: real Cf1l = pi() / (4 * fourmu2pi);
18: real CF1st = Cf1s * fs_cdf(t_tilde | a);
^
19: real F1lt = -expm1(-t_tilde * fourmu2pi);
-------------------------------------------------
PARSER EXPECTED: "("
Error in stanc(file = file, model_code = model_code, model_name = model_name, :
failed to parse Stan model 'ddm' due to the above error.
This error was gone after I changed the | to ,.
The second error was:
SYNTAX ERROR, MESSAGE(S) FROM PARSER:
Variable "real" does not exist.
error in 'modelf6834f4b1b2_ddm' at line 39, column 8
-------------------------------------------------
37: reject("c is ", c);
38: if (is_nan(c))
39: real z = ft * uniform_rng(0, 1);
^
40: real b = exp(-c);
-------------------------------------------------
Error in stanc(file = file, model_code = model_code, model_name = model_name, :
failed to parse Stan model 'ddm' due to the above error.
I managed to solve this by putting the initializations before the if loops.
The last one was:
No matches for:
inv_erfc(real)
Function inv_erfc not found.
error in 'modelf68379d46c4_ddm' at line 76, column 53
-------------------------------------------------
74:
75: if (p <= CF1st) {
76: t_star = 1. / (2 * a * pow(inv_erfc(p / Cf1s), 2));
^
77: while (0.5 * t_star <= 0.06385320297074884) {
-------------------------------------------------
Error in stanc(file = file, model_code = model_code, model_name = model_name, :
failed to parse Stan model 'ddm' due to the above error.
which really made me confused. I was unable to find too much information online, except for the link from Stan Math Library. It seems to me that inv_erfc is a function shipped with Stan so my conjecture of the error is the mismatch of the input type (const vs. real) but I was unable to find more information to solve my confusion.
Therefore, may I ask for your help on this matter? Thank you for the time in advance.
It seems like you are using quite an old version of Stan based on those errors. The code provided requires Stan 2.29 or newer (that is the version which added inv_erfc)
@WardBrian Thanks for your reply! I was using the stable release of rstan. It seems like the latest version of rstan is 2.26.X so I upgraded to that version. I now got a different error saying:
Error in stanc(file = file, model_code = model_code, model_name = model_name, :
parser failed badly
May I ask for your suggestion? Thank you…