Conditional expectations to fill in missing values

Question

Given a matrix, which represents a multivariate time series, wherein some observations are missing, I want to calculate the most likely observation, given a covariance matrix.

Notation: \$x_{-i}\$ represents some vector with the \$i\$th elements removed.

Let \$m\$ be an index of missing observations, \$X\$ is a vector of observations, \$\Omega\$ is a covariance matrix, then \$X_m=X_{-m} \Omega_{-m,-m}^{-1}\Omega_{-m,m}\$

This calculation is repeated for every row.

#include <RcppArmadillo.h>
using namespace Rcpp;
// [[Rcpp::depends(RcppArmadillo)]]


// [[Rcpp::export]]
arma::mat runLinearModelOverTimeInPlace(arma::mat returnsMat, arma::mat covMat) {
  int nrow = returnsMat.n_rows;

  for (int iRow = 0; iRow < nrow; iRow++){
    arma::mat rowMat = returnsMat.row(iRow);
    arma::uvec missingIndex = find_nonfinite(rowMat);
    if (missingIndex.n_elem > 0) {
      arma::uvec notMissingIndex = find_finite(rowMat);
      arma::mat A = covMat.submat(notMissingIndex, notMissingIndex);
      arma::mat B = covMat.submat(notMissingIndex, missingIndex);

      arma::mat beta = arma::solve(A,B);
      arma::uvec rowIndexVec = arma::linspace<arma::uvec>(iRow, iRow, 1);
      arma::mat res = returnsMat.submat(rowIndexVec, notMissingIndex) * beta;
      returnsMat.submat(rowIndexVec, missingIndex) = res;
    }
  }
  return returnsMat;
}

Here is a test case that runs in R:

runLinearModelOverTime = function(returns_with_holes, cov_returns) {
  N_sample = nrow(returns_with_holes)
  
  sapply(seq(1, N_sample), \(i){
    res = returns_with_holes[i,]
    missingIndex = which(is.na(returns_with_holes[i,]))
    if(length(missingIndex)>0){
      beta = solve(cov_returns[-missingIndex, -missingIndex], cov_returns[-missingIndex, missingIndex])
      res[missingIndex] <- returns_with_holes[i,-missingIndex] %*% beta
      res
    } else {
      res
    }
  }) |> t()
}

set.seed(123)
x = matrix(runif(60), nrow = 10)
# x_ = x + 1
# x_ = x_ - 1
# all.equal(x, x_)
y = cov(x)
x[sample(60, 10)] <- NA

expectedOutput = runLinearModelOverTime(x,y)
actualOutput = runLinearModelOverTimeInPlace(x,y)
all.equal(expectedOutput, actualOutput)

Answer 1

Naming things

The way you name things can be improved. First, don't add the type of a variable to its name. The type is already known to the compiler in another way, and it's probably clear to the reader of the code as well, so it just adds noise.

Second, don't abbreviate things unnecessarily. Instead of cov, write covariance, so I don't have to guess whether it meant something else, like coverage or covalence.

The name missingIndex is also confusing; it's not a single index, but rather a set of indices (I assume, because it's stored in a vector). I recommend you always use the plural for something that is a collection of things. So: missingIndices.

Avoid overly verbose names. While missingIndices describes itself very well, it's quite long, and I think you can just name it missing; the fact that it's the indices is clear from the context in which it is used.

iRow: I would just name this i; i and j are very commonly used names for indices, both in many programming languages and in mathematics. You might want to use i for rows and j for columns, since that's the order in which you pass them to many Armadillo functions.

The parameter returnsMat is very misleading, it doesn't return anything. From your description of the problem, it sounds like it should have been named observations instead. This also brings me to the name of the function itself. It mentions InPlace, but since returnsMat is passed by value, it doesn't do anything in-place, instead a copy will have been made and you are modifying the copy before returning it.

Keep things simple

It took me a while to figure out what this was meant to be:

arma::uvec rowIndexVec = arma::linspace<arma::uvec>(iRow, iRow, 1);

It's just a vector with one element with value iRow. You could have created this much simpler this way:

arma::uvec rowIndexVec = {iRow};

This causes a compiler warning about a narrowing conversion because iRow is a signed integer. Make sure it is arma::uword.

But you could simplify things further. Consider that you already have rowMat. Instead of calculating returnsMat.submat(...), you can write:

arma::mat result = rowMat.cols(notMissingIndex) * beta;

Code after the above changes

arma::mat runLinearModelOverTime(arma::mat observations, arma::mat covariance) {
  for (arma::uword i = 0; i < observations.n_rows; ++i) {
    auto row = observations.row(i);
    auto missing = find_nonfinite(row);

    if (missing.empty()) {
      continue;
    }

    auto notMissing = find_finite(row);
    auto A = covariance.submat(notMissing, notMissing);
    auto B = covariance.submat(notMissing, missing);
    auto beta = arma::solve(A, B);
    row.cols(missing) = row.cols(notMissing) * beta;
  }

  return observations;
}

Possible optimizations

Consider that you can have multiple rows that have exactly the same non-finite columns. In that case, beta will have the same value for all those rows. If you have large matrices and a high chance of this happening, then it would make sense to somehow cache previously calculated values of beta.