## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
    collapse = TRUE,
    comment = "#>"
)

if (!requireNamespace("imager", quietly = TRUE)) {
    cat("The suggested dependency `imager` is not installed, skipping vignette build.")
    knitr::knit_exit()
}

# figure dimensions
pane_length <- function(p) 0.5 * p
SINGLE_PANE <- "7.142857%"

## ----setup--------------------------------------------------------------------
library(einops)

## ----load_einops_image--------------------------------------------------------
options(einops_row_major = TRUE)
ims <- get(data("einops_image"))
print(dim(ims))

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# display the first image (whole 4d tensor can't be rendered)
ims[1, , , ]

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# second image in a batch
ims[2, , , ]

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# rearrange, as the name suggests, rearranges elements
# below we swapped height and width.
# In other words, transposed first two axes (dimensions)
rearrange(ims[1, , , ], "h w c -> w h c")

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# we could use more verbose names for axes, and result is the same:
rearrange(ims[1, , , ], "height width color -> width height color")
# when you operate on same set of axes many times,
# you usually come up with short names.
# That's what we do throughout this vignette - we'll use b (for batch), h, w, and c

## ----fig.width = pane_length(1)-----------------------------------------------
# einops allows seamlessly composing batch and height to a new height dimension
# We just rendered all images by collapsing to 3d tensor!
rearrange(ims, "b h w c -> (b h) w c")

## ----fig.height = pane_length(1)----------------------------------------------
# or compose a new dimension of batch and width
rearrange(ims, "b h w c -> h (b w) c")

## -----------------------------------------------------------------------------
# resulting dimensions are computed very simply
# length of newly composed axis is a product of components
# [6, 96, 96, 3] -> [96, (6 * 96), 3]
dim(rearrange(ims, "b h w c -> h (b w) c"))

## -----------------------------------------------------------------------------
# we can compose more than two axes.
# let's flatten 4d array into 1d, resulting array has as many elements as the original
dim(rearrange(ims, "b h w c -> (b h w c)"))

## -----------------------------------------------------------------------------
# decomposition is the inverse process - represent an axis as a combination of new axes
# several decompositions possible, so b1=2 is to decompose 6 to b1=2 and b2=3
dim(rearrange(ims, "(b1 b2) h w c -> b1 b2 h w c ", b1 = 2))

## ----fig.width = pane_length(3), fig.height = pane_length(2)------------------
# finally, combine composition and decomposition:
rearrange(ims, "(b1 b2) h w c -> (b1 h) (b2 w) c ", b1 = 2)

## ----fig.width = pane_length(2), fig.height = pane_length(3)------------------
# slightly different composition: b1 is merged with width, b2 with height
# ... so letters are ordered by w then by h
rearrange(ims, "(b1 b2) h w c -> (b2 h) (b1 w) c ", b1 = 2)

## ----fig.height = pane_length(4)----------------------------------------------
# move part of width dimension to height.
# we should call this width-to-height as image width shrunk by 2 and height doubled.
# but all pixels are the same!
# Can you write reverse operation (height-to-width)?
rearrange(ims, "b h (w w2) c -> (h w2) (b w) c", w2 = 2)

## ----fig.height = pane_length(1)----------------------------------------------
# compare with the next example
rearrange(ims, "b h w c -> h (b w) c")

## ----fig.height = pane_length(1)----------------------------------------------
# order of axes in composition is different
# rule is just as for digits in the number: leftmost digit is the most significant,
# while neighboring numbers differ in the rightmost axis.

# you can also think of this as lexicographic sort
rearrange(ims, "b h w c -> h (w b) c")

## ----fig.height = pane_length(1)----------------------------------------------
# what if b1 and b2 are reordered before composing to width?
rearrange(ims, "(b1 b2) h w c -> h (b1 b2 w) c ", b1 = 2)  # produces 'einops'

## ----fig.height = pane_length(1)----------------------------------------------
rearrange(ims, "(b1 b2) h w c -> h (b2 b1 w) c ", b1 = 2)  # produces 'eoipns'

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# average over batch
reduce(ims, "b h w c -> h w c", "mean")

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# the previous is identical to familiar:
as_image_tensor(apply(ims, c(2, 3, 4), mean))
# but is so much more readable

## ----fig.width = pane_length(1), fig.height = pane_length(1)------------------
# Example of reducing of several axes
# besides mean, there are also min, max, sum, prod
reduce(ims, "b h w c -> h w", "min")

## ----fig.height = pane_length(1)----------------------------------------------
# this is mean-pooling with 2x2 kernel
# image is split into 2x2 patches, each patch is averaged
reduce(ims, "b (h h2) (w w2) c -> h (b w) c", "mean", h2 = 2, w2 = 2)

## ----fig.height = pane_length(1)----------------------------------------------
# max-pooling is similar
# result is not as smooth as for mean-pooling
reduce(ims, "b (h h2) (w w2) c -> h (b w) c", "max", h2 = 2, w2 = 2)

## ----fig.width = pane_length(2), fig.height = pane_length(3)------------------
# yet another example. Can you compute the resulting shape?
reduce(ims, "(b1 b2) h w c -> (b2 h) (b1 w)", "mean", b1 = 2)

## -----------------------------------------------------------------------------
# rearrange can also take care of lists of arrays with the same shape
x <- lapply(1:6, function(i) ims[i, , , ])
cat("list with", length(x), "tensors of shape", paste(dim(x[[1]]), collapse=" "))
# that's how we can stack inputs
# "list axis" becomes first ("b" in this case), and we left it there
dim(rearrange(x, "b h w c -> b h w c"))

## -----------------------------------------------------------------------------
# but new axis can appear in the other place:
dim(rearrange(x, "b h w c -> h w c b"))

## -----------------------------------------------------------------------------
# that's equivalent to array stacking, but written more explicitly
all(rearrange(x, "b h w c -> h w c b") ==
    aperm(array(unlist(x), dim = c(dim(x[[1]]), length(x))), 1:4)
)

## -----------------------------------------------------------------------------
# ... or we can concatenate along axes
dim(rearrange(x, "b h w c -> h (b w) c"))

## -----------------------------------------------------------------------------
# which is equivalent to concatenation
all(rearrange(x, "b h w c -> h (b w) c") ==
    abind::abind(x, along = 2))

## -----------------------------------------------------------------------------
x <- rearrange(ims, "b h w c -> b 1 h w 1 c")  # functionality of array expansion
print(dim(x))
print(dim(rearrange(x, "b 1 h w 1 c -> b h w c")))  # functionality of array squeeze

## ----fig.height = pane_length(1)----------------------------------------------
# compute max in each image individually, then show a difference
x <- reduce(ims, "b h w c -> b () () c", "max")
x <- `repeat`(x, "b 1 1 c -> b 96 96 c") - ims # we will learn about this in a sec
rearrange(x, "b h w c -> h (b w) c")

## -----------------------------------------------------------------------------
# repeat along a new axis. New axis can be placed anywhere
dim(einops.repeat(ims[1, , , ], "h w c -> h new_axis w c", new_axis = 5))

## -----------------------------------------------------------------------------
# shortcut
dim(`repeat`(ims[1, , , ], "h w c -> h 5 w c"))

## ----fig.height = pane_length(1), fig.width = pane_length(3)------------------
# repeat along w (existing axis)
`repeat`(ims[1, , , ], "h w c -> h (repeated w) c", repeated = 3)

## ----fig.width = pane_length(2), fig.height = pane_length(2)------------------
# repeat along two existing axes
`repeat`(ims[1, , , ], "h w c -> (2 h) (2 w) c")

## ----fig.width = pane_length(3), fig.height = pane_length(1)------------------
# order of axes matters as usual - you can repeat each element (pixel) 3 times
# by changing order in parenthesis
`repeat`(ims[1, , , ], "h w c -> h (w repeated) c", repeated = 3)

## -----------------------------------------------------------------------------
repeated <- `repeat`(ims, "b h w c -> b h new_axis w c", new_axis = 2)
reduced <- reduce(repeated, "b h new_axis w c -> b h w c", "min")
all(ims == reduced)

## -----------------------------------------------------------------------------
# interweaving pixels of different pictures
# all letters are observable
rearrange(ims, "(b1 b2) h w c -> (h b1) (w b2) c ", b1 = 2)

## -----------------------------------------------------------------------------
# interweaving along vertical for couples of images
rearrange(ims, "(b1 b2) h w c -> (h b1) (b2 w) c", b1 = 2)

## -----------------------------------------------------------------------------
# interweaving lines for couples of images
# exercise: achieve the same result without einops in your favourite framework
reduce(ims, "(b1 b2) h w c -> h (b2 w) c", "max", b1 = 2)

## -----------------------------------------------------------------------------
# color can be also composed into dimension
# ... while image is downsampled
reduce(ims, "b (h 2) (w 2) c -> (c h) (b w)", "mean")

## ----fig.width = pane_length(5/3), fig.height = pane_length(1/4)--------------
# disproportionate resize
reduce(ims, "b (h 4) (w 3) c -> (h) (b w)", "mean")

## ----fig.height = pane_length(1)----------------------------------------------
# split each image in two halves, compute mean of the two
reduce(ims, "b (h1 h2) w c -> h2 (b w)", "mean", h1 = 2)

## ----fig.height = pane_length(1)----------------------------------------------
# split in small patches and transpose each patch
rearrange(ims, "b (h1 h2) (w1 w2) c -> (h1 w2) (b w1 h2) c", h2 = 8, w2 = 8)

## ----fig.height = pane_length(1)----------------------------------------------
# stop me someone!
rearrange(ims, "b (h1 h2 h3) (w1 w2 w3) c -> (h1 w2 h3) (b w1 h2 w3) c", h2 = 2, w2 = 2, w3 = 2, h3 = 2)

## -----------------------------------------------------------------------------
rearrange(ims, "(b1 b2) (h1 h2) (w1 w2) c -> (h1 b1 h2) (w1 b2 w2) c", h1 = 3, w1 = 3, b2 = 3)

## ----fig.height = pane_length(1)----------------------------------------------
# patterns can be arbitrarily complicated
reduce(ims, "(b1 b2) (h1 h2 h3) (w1 w2 w3) c -> (h1 w1 h3) (b1 w2 h2 w3 b2) c", "mean", h2 = 2, w1 = 2, w3 = 2, h3 = 2, b2 = 2)

## ----fig.height = pane_length(1)----------------------------------------------
# pixelate: first downscale by averaging, then upscale back using the same pattern
averaged <- reduce(ims, "b (h h2) (w w2) c -> b h w c", "mean", h2 = 6, w2 = 8)
`repeat`(averaged, "b h w c -> (h h2) (b w w2) c", h2 = 6, w2 = 8)

## ----fig.height = pane_length(1)----------------------------------------------
rearrange(ims, "b h w c -> w (b h) c")

## ----fig.height = pane_length(1)----------------------------------------------
# let's bring color dimension as part of horizontal axis
# at the same time horizontal axis is downsampled by 2x
reduce(ims, "b (h h2) (w w2) c -> (h w2) (b w c)", "mean", h2 = 3, w2 = 3)