arima_kernel.stan

Go to the documentation of this file.

 
vector pacf_to_phi(vector r) {
  int p = num_elements(r);
  if (p == 0) return r;
  vector[p] phi = r;
  vector[p] work = r;
  for (k in 2:p) {
    for (j in 1:(k - 1)) {
      work[j] = phi[j] - r[k] * phi[k - j];
    }
    for (j in 1:(k - 1)) {
      phi[j] = work[j];
    }
    phi[k] = r[k];
  }
  return phi;
}
 
vector arma_impulse(vector phi, vector theta, int T) {
  int p = num_elements(phi);
  int q = num_elements(theta);
  vector[T] psi = rep_vector(0.0, T);
  psi[1] = 1.0;
  for (t in 2:T) {
    real s = 0.0;
    int kmax = min(p, t - 1);
    for (k in 1:kmax) s += phi[k] * psi[t - k];
    if (t - 1 <= q) s += theta[t - 1];
    psi[t] = s;
  }
  return psi;
}
 
matrix lower_toeplitz(vector psi) {
  int T = num_elements(psi);
  matrix[T, T] K = rep_matrix(0.0, T, T);
  for (s in 1:T) {
    K[s:T, s] = head(psi, T - s + 1);
  }
  return K;
}
 
matrix cumulative_op(int T, int d) {
  if (d == 0) return diag_matrix(rep_vector(1.0, T));
  matrix[T, T] D1 = rep_matrix(0.0, T, T);
  for (s in 1:T) {
    D1[s:T, s] = rep_vector(1.0, T - s + 1);
  }
  matrix[T, T] D = D1;
  for (i in 2:d) D = D * D1;
  return D;
}
 
matrix arima_kernel_matrix(vector phi, vector theta, int d, int T) {
  vector[T] psi = arma_impulse(phi, theta, T);
  matrix[T, T] K = lower_toeplitz(psi);
  for (i in 1:d) {
    for (col in 1:T) {
      K[, col] = cumulative_sum(K[, col]);
    }
  }
  return K;
}
 
matrix arima_filter(matrix Z, vector phi, vector theta, int d) {
  int T = rows(Z);
  int G = cols(Z);
  if (num_elements(phi) == 0 && num_elements(theta) == 0) {
    if (d == 0) return Z;
    matrix[T, G] result = Z;
    for (i in 1:d) {
      for (g in 1:G) {
        result[, g] = cumulative_sum(result[, g]);
      }
    }
    return result;
  }
  return arima_kernel_matrix(phi, theta, d, T) * Z;
}
 
real arima_latent_mean_offset(int present, int centre, int T, int G,
                              int p, int d, int q,
                              matrix z, vector pacf, vector theta,
                              array[] real sigma) {
  if (!present || !centre || d < 1) return 0.0;
  vector[p] phi = pacf_to_phi(pacf);
  matrix[T, G] eps = sigma[1] * arima_filter(z, phi, theta, d);
  return mean(to_vector(eps));
}
 
vector apply_arima_residual(vector base, int n_obs,
                            int present, int centre, int T, int G,
                            int p, int d, int q,
                            matrix z, vector pacf, vector theta,
                            array[] real sigma,
                            array[] int flat_idx) {
  if (!present) return base;
  vector[p] phi = pacf_to_phi(pacf);
  matrix[T, G] eps = sigma[1] * arima_filter(z, phi, theta, d);
  // Integration (d >= 1) gives the residual a free overall level: the
  // first shock propagates through the cumulative sum and shifts the
  // whole path, so the level of the residual trades off against the
  // intercept (a ridge in the joint posterior that slows HMC). When an
  // intercept is present to carry the level, remove the grand mean (over
  // time and groups) so the intercept owns it. This mirrors brms's
  // design-matrix centring and EpiNow2's `gp -= mean(gp)`.
  //
  // Only the grand mean is removed -- the single level direction that
  // competes with the shared intercept. Each group keeps its own level
  // relative to that grand mean, so this is exact for any number of groups
  // and a grouped latent still gives each group its own level and drift.
  if (centre && d >= 1) {
    eps = eps - mean(to_vector(eps));
  }
  // Vectorised gather: flat_idx is precomputed in transformed data
  // as (group_idx - 1) * T + time_idx, so to_vector(eps) (column-
  // major flatten) [flat_idx] returns the per-observation residuals
  // in one shot, replacing the previous element-by-element loop.
  return base + to_vector(eps)[flat_idx];
}