1#' Coerce array to list
2#'
3#' `array_branch()` and `array_tree()` enable arrays to be
4#' used with purrr's functionals by turning them into lists. The
5#' details of the coercion are controlled by the `margin`
6#' argument. `array_tree()` creates an hierarchical list (a tree)
7#' that has as many levels as dimensions specified in `margin`,
8#' while `array_branch()` creates a flat list (by analogy, a
9#' branch) along all mentioned dimensions.
10#'
11#' When no margin is specified, all dimensions are used by
12#' default. When `margin` is a numeric vector of length zero, the
13#' whole array is wrapped in a list.
14#' @param array An array to coerce into a list.
15#' @param margin A numeric vector indicating the positions of the
16#'   indices to be to be enlisted. If `NULL`, a full margin is
17#'   used. If `numeric(0)`, the array as a whole is wrapped in a
18#'   list.
19#' @name array-coercion
20#' @export
21#' @examples
22#' # We create an array with 3 dimensions
23#' x <- array(1:12, c(2, 2, 3))
24#'
25#' # A full margin for such an array would be the vector 1:3. This is
26#' # the default if you don't specify a margin
27#'
28#' # Creating a branch along the full margin is equivalent to
29#' # as.list(array) and produces a list of size length(x):
30#' array_branch(x) %>% str()
31#'
32#' # A branch along the first dimension yields a list of length 2
33#' # with each element containing a 2x3 array:
34#' array_branch(x, 1) %>% str()
35#'
36#' # A branch along the first and third dimensions yields a list of
37#' # length 2x3 whose elements contain a vector of length 2:
38#' array_branch(x, c(1, 3)) %>% str()
39#'
40#' # Creating a tree from the full margin creates a list of lists of
41#' # lists:
42#' array_tree(x) %>% str()
43#'
44#' # The ordering and the depth of the tree are controlled by the
45#' # margin argument:
46#' array_tree(x, c(3, 1)) %>% str()
47array_branch <- function(array, margin = NULL) {
48  dims <- dim(array) %||% length(array)
49  margin <- margin %||% seq_along(dims)
50
51  if (length(margin) == 0) {
52    list(array)
53  } else if (is.null(dim(array))) {
54    if (!identical(as.integer(margin), 1L)) {
55      abort(sprintf(
56        "`margin` must be `NULL` or `1` with 1D arrays, not `%s`",
57        toString(margin)
58      ))
59    }
60    as.list(array)
61  } else {
62    flatten(apply(array, margin, list))
63  }
64}
65
66#' @rdname array-coercion
67#' @export
68array_tree <- function(array, margin = NULL) {
69  dims <- dim(array) %||% length(array)
70  margin <- margin %||% seq_along(dims)
71
72  if (length(margin) > 1) {
73    new_margin <- ifelse(margin[-1] > margin[[1]], margin[-1] - 1, margin[-1])
74    apply(array, margin[[1]], array_tree, new_margin)
75  } else {
76    array_branch(array, margin)
77  }
78}
79