Mapping a list of subarrays in R.

map_subarrays.r

array_or_vector_p <- function(object) {
  ## Test if object is what we consider an array for the purposes of
  ## the code below.
  ##
  ## Lists are not arrays, but vectors are, even if they don't have a
  ## dim attribute. The following give an error:
  ## - objects which are neither lists, (non-list) vectors, or arrays
  ## - vectors and arrays with zero length
  ## - arrays with rank lower than 1
  if (is.list(object))
    FALSE
  else {
    stopifnot(is.array(object) || is.vector(object))
    stopifnot(length(object) > 0)
    if (is.array(object))
      stopifnot(length(dim(object)) > 0)
    TRUE
  }
}

leading_dimension <- function(array) {
  ## Return the leading dimension of array.
  ##
  ## For vectors, it is the length,
  stopifnot(array_or_vector_p(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 that, otherwise signal an error.
  if (array_or_vector_p(array_list))
    leading_dimension(array_list)
  else {
    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, , ...]
  stopifnot(array_or_vector_p(array))
  if (is.vector(array))
    array[index]
  else
    do.call("[",c(list(array,index),rep(TRUE,length(dim(array))-1)))
}

set_subarray <- function(value, array, index) {
  ## equivalent to array[index, , ...] <- value, except that it returns the new value.
  if (is.vector(array)) {
    array[index] <- value
    array
  } else
    do.call("[<-", c(list(array, index),
                     as.list(rep(TRUE, length(dim(array))-1)), list(value)))
}

subarray_row_list <- function(subarray_list, index) {
  ## Extract a row with given index from a list of arrays or a single array.
  if (array_or_vector_p(subarray_list)) {
    list(subarray(subarray_list,index))
  } else {
    names_ <- names(subarray_list)
    row <- lapply(names_,
                  function(name) subarray(subarray_list[[name]],index))
    names(row) <- names_
    row
  }
}

conforming_result_array <- function(leading_dimension, value) {
  ## Return a vector or array for storing results, conforming in dimension to VALUE.
  stopifnot(array_or_vector_p(value))
  dims <- if(is.vector(value)) {
    length_ <- length(value)
    if (length_ == 1)
      NULL
    else
      length_
  } else {
    dim(value)
  }
  if (is.null(dims))
    rep(NA,leading_dimension)
  else
    array(NA,c(leading_dimension,dims))
}

map_subarrays <- function(subarrays, f) {
  ## FIX: add documentation
  ld <- common_leading_dimension(subarrays)
  result <- NULL
  for (index in 1:ld) {
    row_result <- do.call(f,subarray_row_list(subarrays, index))
    if (is.null(result)) {
      if (array_or_vector_p(row_result))
        result <- conforming_result_array(ld,row_result)
      else {
        result <- lapply(row_result, function(value) conforming_result_array(ld,value))
        names(result) <- names(row_result)
      }
    }
    if (array_or_vector_p(result))
      result <- set_subarray(row_result, result, index)
    else {
      row_names_ <- names(row_result)
      for (row_result_index in seq_along(row_result)) {
        position <- match(row_names_[row_result_index],names(result))
        stopifnot(is.integer(position))
        value <- row_result[[row_result_index]]
        if (is.vector(result[[position]]))
          result[[position]][index] <- value
        else
          result[[position]] <- set_subarray(value, result[[position]], index)
      }
    }
  }
  result
}

### TESTS
## subarrays <- list(a=1:3,b=matrix(4:9,nrow=3))
## map_subarrays(subarrays, function(a,b) { list(d=a*sum(b)) })
## map_subarrays(subarrays, function(a,b) { a*sum(b) })
## map_subarrays(subarrays$b, function(b) list(d=2*b))
## map_subarrays(subarrays$b, function(b) 2*b)
## map_subarrays(1:3, function(b) 2*b)