[R] modify subset of array in list in a function

[Tamas Papp]

Hi,

I am trying to implement a function that would allow functional
transformations of posterior simulations (useful in posterior predictive
checks after MCMC, eg in Stan).

A posterior simulation is a list of vectors and arrays. Usually, one
uses loops to transform it, but that's error prone. I was thinking of
ending up with an interface like this (toy example):

--8<---------------cut here---------------start------------->8---
posterior <- list(a=1:3,b=matrix(4:9,nrow=3))

map_posterior(posterior, function(a,b) {
  list(d=a*sum(b))
})        # should be equivalent to list(d=posterior$a + rowSums(posterior$b))
--8<---------------cut here---------------end--------------->8---

I started coding it: I make an arglist for do.call, matching names,
then deconstruct the value and save it where it belongs. But the last
part is giving me trouble, since in R calls are by value, not reference,
so I don't end up modifying the original array in the code below (when
set_subarray is called):

--8<---------------cut here---------------start------------->8---
leading_dimension <- function(array) {
  if (is.vector(array))
    length(array)
  else
    dim(array)[1]
}

common_leading_dimension <- function(array_list) {
  ## if all leading dimensions are the same, return it, otherwise signal an error
  length_ <- length(array_list)
  stopifnot(length_ > 0)
  ld <- leading_dimension(array_list[[1]])
  if (length_ > 1)
    for (i in 2:length_)
      stopifnot(leading_dimension(array_list[[i]])==ld)
  ld
}

subarray <- function(array, index) {
  ## equivalent to array[index, , ...]
  if (is.vector(array))
    array[index]
  else {
    rank_ <- length(dim(array))
    stopifnot(rank_ >= 1)
    do.call("[",c(list(array,index),rep(TRUE,rank_-1)))
  }
}

set_subarray <- function(value, array, index) {
  ## equivalent to array[index, , ...] <- value
  if (is.vector(array))
    array[index] <- value
  else {
    rank_ <- length(dim(array))
    stopifnot(rank_ >= 1)
    do.call("[<-",c(list(array,value,index),rep(TRUE,rank_-1)))
  }
}

map_posterior <- function(posterior,f) {
  names_ <- names(posterior)
  ld <- common_leading_dimension(posterior)
  result <- NULL
  for (index in 1:ld) {
    row_args <- Map(function(name) subarray(posterior[[name]],index),names_)
    names(row_args) <- names_
    row_result <- do.call(f,row_args)
    if (is.null(result)) {
      result <- Map(function(value) {
        dims <- if(is.vector(value)) {
          length_ <- length(value)
          if (length_ == 1)
            NULL
          else
            length_
        } else {
          dim(value)
        }
        array(NA,c(ld,dims))
      },row_result)
      names(result) <- names(row_result)
    }
    Map(function(row_result_name,row_result_value) {
      set_subarray(row_result_value,result[[row_result_name]],index)
    })
  }
  result
}
--8<---------------cut here---------------end--------------->8---

Any help or hints on how to do this would be appreciated, including
alternative approaches of doing/programming the same thing.

Best,

Tamas