在 Stan 中对矩阵进行条件索引

问题描述

我有以下代码正确指定了我正在考虑的模型：

for (iT in 1:T)
  for (iN in 1:N)
    if (!is_inf(y[iT,iN]))
      y[iT,iN] ~ normal(Gamma[iN] * F[iT],sigma[iT]);

但是，与我没有检查 y 是否为 inf（我用来表示缺失值）的矢量化版本相比，这似乎有点慢。

我希望能够做类似的事情

for (iT in 1:T) {
  isObserved = !is_inf(y[iT,]);
  y[iT,isObserved] ~ normal(Gamma[isObserved ] * F[iT],sigma[iT]);
}

但似乎不支持这样的逻辑索引。

问题：

是否有我遗漏的逻辑索引语法？
还有其他方法可以矢量化此代码以提高性能吗？

解决方法

您可以使用 transformed data 块创建一个包含 y 的所有非缺失值的向量，然后使用多个索引访问模型中的相应参数。这种方法的优点是对数据的计算只发生一次。这是一个如何工作的例子；为方便起见，我还添加了一个用户定义的函数来计算我们有多少非缺失值。（我假设 y 中的任何单元格都可能丢失 - 但如果有模式，例如整行或整列丢失，则可以简化。）

functions {
  int count_non_missing_values(matrix y) {
    int nm = 0;
    for(i in 1:rows(y)) {
      for(j in 1:cols(y)) {
        if(!is_inf(y[i,j])) {
          nm += 1;
        }
      }
    }
    return(nm);
  }
}

data {
  int T;
  int N;
  matrix[T,N] y;
}

transformed data {
  vector[count_non_missing_values(y)] flat_ys;
  int flat_Ts[count_non_missing_values(y)];
  int flat_Ns[count_non_missing_values(y)];
  {
    int ind = 1;
    for(iT in 1:T) {
      for(iN in 1:N) {
        if(!is_inf(y[iT,iN])) {
          flat_ys[ind] = y[iT,iN];
          flat_Ts[ind] = iT;
          flat_Ns[ind] = iN;
          ind += 1;
        }
      }
    }
  }
}

parameters {
  vector[T] F;
  vector<lower=0>[T] sigma;
  vector[N] Gamma;
}

model {
  flat_ys ~ normal(Gamma[flat_Ns] .* F[flat_Ts],sigma[flat_Ts]);
}

indexing indexing stan vectorization