fix(ehgf): invalid eHGF update code

Author: bugoverdose

commons/stan_files/ehgf_ibrb.stan

@@ -218,6 +218,7 @@ model {
   for (i in 1:N) {
     row_vector[L-1] mu = mu0[i];        // prediction (2 ~ L)
     real sa[L] = sigma_base;            // uncertainty of prediction (2 ~ L)
+    real sa_prev[L];                    // uncertainty of prediction from previous trial (2 ~ L)
     real mu_hat[L] = rep_array(0.0, L); // prior prediction (1 ~ L)
     real sa_hat[L] = rep_array(0.0, L); // prior uncertainty of prediction (1 ~ L)
     real m = -1;                        // predictive probability that the next response will be 1 (0~1)
@@ -244,10 +245,12 @@ model {
       for (l in 2:(L-1)) {
         real ka = kappa[i,l-1];
         real om = omega[i,l-1];
-        sa_hat[l] = sa[l] + exp(ka * mu[l] + om);
+        sa_hat[l] = sa[l] + exp(ka * mu[(l+1)-1] + om);
       }
       sa_hat[L] = sa[L] + exp(omega[i,L-1]);
-  
+
+      sa_prev = sa;
+
       // Level 2
       mu_prev = mu_hat[2];
       pe = u[i,t] - mu_hat[1];                     // prediction error
@@ -258,26 +261,26 @@ model {
       for (l in 3:L) {
         real ka = kappa[i,(l-1)-1];
         real om = omega[i,(l-1)-1];
-        real mu_lower = mu[l-1-1];
+        real mu_lower = mu[(l-1)-1];
         real mu_prev_ = mu_hat[l];
         real mu_prev_lower = mu_hat[l-1];
         real sa_lower = sa[l-1];
-        real sa_prev_lower = sa_hat[l-1];
+        real sa_prev_lower = sa_prev[l-1];
         real vv, pimhat, ww, rr, dd;
 
         real v = exp(ka * mu_prev_ + om); // volatility
-        real w = v / sa_prev_lower;       // weighting factor (level: l-1)
+        real w = v / sa_hat[l-1];         // weighting factor (level: l-1)
 
-        real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_prev_lower) - 1; // volatility prediction error
+        real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_hat[l-1]) - 1.0; // volatility prediction error
         real lr = 0.5 * sa_hat[l] * ka * w; // learning rate
         real pwpe = lr * vpe;             // precision-weighted prediction error
         mu[l-1] = mu_prev_ + pwpe;        // posterior prediction
 
-        vv = exp(ka * mu[l-1] +om);
-        pimhat = 1.0 / (sa_prev_lower +vv);
+        vv = exp(ka * mu[l-1] + om);
+        pimhat = 1.0 / (sa_prev_lower + vv);
         ww = vv * pimhat;
         rr = (vv - sa_prev_lower) * pimhat;
-        dd = (sa_lower + square(mu_lower - mu_prev_lower)) *pimhat - 1;
+        dd = (sa_lower + square(mu_lower - mu_prev_lower)) * pimhat - 1.0;
 
         sa[l] = 1.0/((1.0/sa_hat[l]) + fmax(0.0, 0.5 * square(ka) * ww * (ww + rr * dd)));
       }
@@ -302,11 +305,12 @@ generated quantities {
   vector[L-2] mu_kappa;
   vector[L-1] mu_omega;
   real<lower=0> mu_zeta;
-  
-    // Subject loop
+
+  // Subject loop
   for (i in 1:N) {
     row_vector[L-1] mu = mu0[i];        // prediction (2 ~ L)
     real sa[L] = sigma_base;            // uncertainty of prediction (2 ~ L)
+    real sa_prev[L];                    // uncertainty of prediction from previous trial (2 ~ L)
     real mu_hat[L] = rep_array(0.0, L); // prior prediction (1 ~ L)
     real sa_hat[L] = rep_array(0.0, L); // prior uncertainty of prediction (1 ~ L)
     real m = -1;                        // predictive probability that the next response will be 1 (0~1)
@@ -333,10 +337,12 @@ generated quantities {
       for (l in 2:(L-1)) {
         real ka = kappa[i,l-1];
         real om = omega[i,l-1];
-        sa_hat[l] = sa[l] + exp(ka * mu[l+1-1] + om);
+        sa_hat[l] = sa[l] + exp(ka * mu[(l+1)-1] + om);
       }
       sa_hat[L] = sa[L] + exp(omega[i,L-1]);
-  
+
+      sa_prev = sa;
+
       // Level 2
       mu_prev = mu_hat[2];
       pe = u[i,t] - mu_hat[1];                     // prediction error
@@ -347,26 +353,26 @@ generated quantities {
       for (l in 3:L) {
         real ka = kappa[i,(l-1)-1];
         real om = omega[i,(l-1)-1];
-        real mu_lower = mu[l-1-1];
+        real mu_lower = mu[(l-1)-1];
         real mu_prev_ = mu_hat[l];
         real mu_prev_lower = mu_hat[l-1];
         real sa_lower = sa[l-1];
-        real sa_prev_lower = sa_hat[l-1];
+        real sa_prev_lower = sa_prev[l-1];
         real vv, pimhat, ww, rr, dd;
 
-        real v = exp(ka * mu_prev_ + om);// volatility
-        real w = v / sa_prev_lower;      // weighting factor (level: l-1)
+        real v = exp(ka * mu_prev_ + om); // volatility
+        real w = v / sa_hat[l-1];         // weighting factor (level: l-1)
 
-        real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_prev_lower) - 1; // volatility prediction error
+        real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_hat[l-1]) - 1.0; // volatility prediction error
         real lr = 0.5 * sa_hat[l] * ka * w; // learning rate
         real pwpe = lr * vpe;             // precision-weighted prediction error
         mu[l-1] = mu_prev_ + pwpe;        // posterior prediction
 
-        vv = exp(ka * mu[l-1] +om);
-        pimhat = 1.0 / (sa_prev_lower +vv);
+        vv = exp(ka * mu[l-1] + om);
+        pimhat = 1.0 / (sa_prev_lower + vv);
         ww = vv * pimhat;
         rr = (vv - sa_prev_lower) * pimhat;
-        dd = (sa_lower + square(mu_lower - mu_prev_lower)) *pimhat - 1;
+        dd = (sa_lower + square(mu_lower - mu_prev_lower)) * pimhat - 1.0;
 
         sa[l] = 1.0/((1.0/sa_hat[l]) + fmax(0.0, 0.5 * square(ka) * ww * (ww + rr * dd)));
       }

commons/stan_files/ehgf_ibrb_multipleB.stan

@@ -225,6 +225,7 @@ model {
     for (bIdx in 1:Bsubj[i]) {
       row_vector[L-1] mu = mu0[i];        // prediction (2 ~ L)
       real sa[L] = sigma_base;            // uncertainty of prediction (2 ~ L)
+      real sa_prev[L];
       real mu_hat[L] = rep_array(0.0, L); // prior prediction (1 ~ L)
       real sa_hat[L] = rep_array(0.0, L); // prior uncertainty of prediction (1 ~ L)
       real m = -1;                        // predictive probability that the next response will be 1 (0~1)
@@ -251,40 +252,42 @@ model {
         for (l in 2:(L-1)) {
           real ka = kappa[i,l-1];
           real om = omega[i,l-1];
-          sa_hat[l] = sa[l] + exp(ka * mu[l] + om);
+          sa_hat[l] = sa[l] + exp(ka * mu[(l+1)-1] + om);
         }
         sa_hat[L] = sa[L] + exp(omega[i,L-1]);
-    
+
+        sa_prev = sa;
+
         // Level 2
         mu_prev = mu_hat[2];
-        pe = u[i,bIdx,t] - mu_hat[1];                     // prediction error
+        pe = u[i,bIdx,t] - mu_hat[1];                // prediction error
         sa[2] = 1.0 / ((1.0/sa_hat[2]) + sa_hat[1]); // learning rate
         mu[2-1] = mu_prev + sa[2] * pe;              // posterior prediction
 
         // Level 3 ~ L
         for (l in 3:L) {
           real ka = kappa[i,(l-1)-1];
           real om = omega[i,(l-1)-1];
-          real mu_lower = mu[l-1-1];
+          real mu_lower = mu[(l-1)-1];
           real mu_prev_ = mu_hat[l];
           real mu_prev_lower = mu_hat[l-1];
           real sa_lower = sa[l-1];
-          real sa_prev_lower = sa_hat[l-1];
+          real sa_prev_lower = sa_prev[l-1];
           real vv, pimhat, ww, rr, dd;
 
           real v = exp(ka * mu_prev_ + om); // volatility
-          real w = v / sa_prev_lower;       // weighting factor (level: l-1)
+          real w = v / sa_hat[l-1];         // weighting factor (level: l-1)
 
-          real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_prev_lower) - 1; // volatility prediction error
+          real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_hat[l-1]) - 1.0; // volatility prediction error
           real lr = 0.5 * sa_hat[l] * ka * w; // learning rate
           real pwpe = lr * vpe;             // precision-weighted prediction error
           mu[l-1] = mu_prev_ + pwpe;        // posterior prediction
 
-          vv = exp(ka * mu[l-1] +om);
-          pimhat = 1.0 / (sa_prev_lower +vv);
+          vv = exp(ka * mu[l-1] + om);
+          pimhat = 1.0 / (sa_prev_lower + vv);
           ww = vv * pimhat;
           rr = (vv - sa_prev_lower) * pimhat;
-          dd = (sa_lower + square(mu_lower - mu_prev_lower)) *pimhat - 1;
+          dd = (sa_lower + square(mu_lower - mu_prev_lower)) * pimhat - 1.0;
 
           sa[l] = 1.0/((1.0/sa_hat[l]) + fmax(0.0, 0.5 * square(ka) * ww * (ww + rr * dd)));
         }
@@ -311,11 +314,12 @@ generated quantities {
   vector[L-1] mu_omega;
   real<lower=0> mu_zeta;
 
-    // Subject loop
+  // Subject loop
   for (i in 1:N) {
     for (bIdx in 1:Bsubj[i]) {
       row_vector[L-1] mu = mu0[i];        // prediction (2 ~ L)
       real sa[L] = sigma_base;            // uncertainty of prediction (2 ~ L)
+      real sa_prev[L];
       real mu_hat[L] = rep_array(0.0, L); // prior prediction (1 ~ L)
       real sa_hat[L] = rep_array(0.0, L); // prior uncertainty of prediction (1 ~ L)
       real m = -1;                        // predictive probability that the next response will be 1 (0~1)
@@ -342,40 +346,42 @@ generated quantities {
         for (l in 2:(L-1)) {
           real ka = kappa[i,l-1];
           real om = omega[i,l-1];
-          sa_hat[l] = sa[l] + exp(ka * mu[l+1-1] + om);
+          sa_hat[l] = sa[l] + exp(ka * mu[(l+1)-1] + om);
         }
         sa_hat[L] = sa[L] + exp(omega[i,L-1]);
-    
+
+        sa_prev = sa;
+
         // Level 2
         mu_prev = mu_hat[2];
-        pe = u[i,bIdx,t] - mu_hat[1];                     // prediction error
+        pe = u[i,bIdx,t] - mu_hat[1];                // prediction error
         sa[2] = 1.0 / ((1.0/sa_hat[2]) + sa_hat[1]); // learning rate
         mu[2-1] = mu_prev + sa[2] * pe;              // posterior prediction
 
         // Level 3 ~ L
         for (l in 3:L) {
           real ka = kappa[i,(l-1)-1];
           real om = omega[i,(l-1)-1];
-          real mu_lower = mu[l-1-1];
+          real mu_lower = mu[(l-1)-1];
           real mu_prev_ = mu_hat[l];
           real mu_prev_lower = mu_hat[l-1];
           real sa_lower = sa[l-1];
-          real sa_prev_lower = sa_hat[l-1];
+          real sa_prev_lower = sa_prev[l-1];
           real vv, pimhat, ww, rr, dd;
 
-          real v = exp(ka * mu_prev_ + om);// volatility
-          real w = v / sa_prev_lower;      // weighting factor (level: l-1)
-
-          real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_prev_lower) - 1; // volatility prediction error
+          real v = exp(ka * mu_prev_ + om); // volatility
+          real w = v / sa_hat[l-1];         // weighting factor (level: l-1)
+        
+          real vpe = ((sa_lower + pow(mu_lower - mu_prev_lower, 2)) / sa_hat[l-1]) - 1.0; // volatility prediction error
           real lr = 0.5 * sa_hat[l] * ka * w; // learning rate
           real pwpe = lr * vpe;             // precision-weighted prediction error
           mu[l-1] = mu_prev_ + pwpe;        // posterior prediction
 
-          vv = exp(ka * mu[l-1] +om);
-          pimhat = 1.0 / (sa_prev_lower +vv);
+          vv = exp(ka * mu[l-1] + om);
+          pimhat = 1.0 / (sa_prev_lower + vv);
           ww = vv * pimhat;
           rr = (vv - sa_prev_lower) * pimhat;
-          dd = (sa_lower + square(mu_lower - mu_prev_lower)) *pimhat - 1;
+          dd = (sa_lower + square(mu_lower - mu_prev_lower)) * pimhat - 1.0;
 
           sa[l] = 1.0/((1.0/sa_hat[l]) + fmax(0.0, 0.5 * square(ka) * ww * (ww + rr * dd)));
         }