{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e24ffe2b",
   "metadata": {},
   "source": [
    "# Batched Linear Operations\n",
    "\n",
    "Almost all of SciPy's linear algebra functions now support N-dimensional array input. These operations have not been mathematically generalized to higher-order tensors; rather, the indicated operation is performed on a *batch* (or \"stack\") of input scalars, vectors, and/or matrices.\n",
    "\n",
    "Consider the `linalg.det` function, which maps a matrix to a scalar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "156270ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import linalg\n",
    "A = np.eye(3)\n",
    "linalg.det(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5495b039",
   "metadata": {},
   "source": [
    "Sometimes we need the determinant of a batch of matrices of the same dimensionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68724c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = [i*np.eye(3) for i in range(1, 4)]\n",
    "batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0595c099",
   "metadata": {},
   "source": [
    "We could perform the operation for each element of the batch in a loop or list comprehension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf19554",
   "metadata": {},
   "outputs": [],
   "source": [
    "[linalg.det(A) for A in batch]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fc6f88a",
   "metadata": {},
   "source": [
    "However, just as we might use NumPy broadcasting and vectorization rules to create the batch of matrices in the first place:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c175a7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = np.arange(1, 4).reshape(-1, 1, 1)\n",
    "batch = i * np.eye(3)\n",
    "batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f83c1f6",
   "metadata": {},
   "source": [
    "we might also wish to perform the determinant operation on all of the matrices in one function call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b65c539f",
   "metadata": {},
   "outputs": [],
   "source": [
    "linalg.det(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a57dd587",
   "metadata": {},
   "source": [
    "In SciPy, we prefer the term \"batch\" instead of \"stack\" because the idea is generalized to N-dimensional batches. Suppose the input is a 2 x 4 batch of 3 x 3 matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0f1be0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_shape = (2, 4)\n",
    "i = np.arange(np.prod(batch_shape)).reshape(*batch_shape, 1, 1)\n",
    "input = i * np.eye(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c25cb00c",
   "metadata": {},
   "source": [
    "In this case, we say that the *batch shape* is `(2, 4)`, and the *core shape* of the input is `(3, 3)`. The net shape of the input is the sum (concatenation) of the batch shape and core shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "211d7ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "input.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f6e2bb3",
   "metadata": {},
   "source": [
    "Since each 3 x 3 matrix is converted to a zero-dimensional scalar, we say that the core shape of the outuput is `()`. The shape of the output is the sum of the batch shape and core shape, so the result is a 2 x 4 array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b6818ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = linalg.det(input)\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c897b9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "output.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "744b16ed",
   "metadata": {},
   "source": [
    "Not all linear algebra functions map to scalars. For instance, the {func}`scipy.linalg.expm` function maps from a matrix to a matrix with the same shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d71e580",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.eye(3)\n",
    "linalg.expm(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35d44691",
   "metadata": {},
   "source": [
    "In this case, the core shape of the output is `(3, 3)`, so with a batch shape of `(2, 4)`, we expect an output of shape `(2, 4, 3, 3)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db913db9",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = linalg.expm(input)\n",
    "output.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c689168",
   "metadata": {},
   "source": [
    "Generalization of these rules to functions with multiple inputs and outputs is straightforward. For instance, the {func}`scipy.linalg.eig` function produces two outputs by default, a vector and a matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df5be48f",
   "metadata": {},
   "outputs": [],
   "source": [
    "evals, evecs = linalg.eig(A)\n",
    "evals.shape, evecs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd86bba7",
   "metadata": {},
   "source": [
    "In this case, the core shape of the output vector is `(3,)` and the core shape of the output matrix is `(3, 3)`. The shape of each output is the batch shape plus the core shape as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07509b2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "evals, evecs = linalg.eig(input)\n",
    "evals.shape, evecs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "202039a9",
   "metadata": {},
   "source": [
    "When there is more than one input, there is no complication if the input shapes are identical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "307212d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "evals, evecs = linalg.eig(input, b=input)\n",
    "evals.shape, evecs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd1ce4ef",
   "metadata": {},
   "source": [
    "The rules when the shapes are not identical follow logically. Each input can have its own batch shape as long as the shapes are broadcastable according to [NumPy's broadcasting rules](#array-broadcasting-in-numpy). The net batch shape is the broadcasted shape of the individual batch shapes, and the shape of each output is the net batch shape plus its core shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f0a2cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(2859239482)\n",
    "\n",
    "# Define input core shapes\n",
    "m = 3\n",
    "core_shape_a = (m, m)\n",
    "core_shape_b = (m, m)\n",
    "\n",
    "# Define broadcastable batch shapes\n",
    "batch_shape_a = (2, 4)\n",
    "batch_shape_b = (5, 1, 4)\n",
    "\n",
    "# Define output core shapes\n",
    "core_shape_evals = (m,)\n",
    "core_shape_evecs = (m, m)\n",
    "\n",
    "# Predict shapes of outputs: broadcast batch shapes,\n",
    "# and append output core shapes\n",
    "net_batch_shape = np.broadcast_shapes(batch_shape_a, batch_shape_b)\n",
    "output_shape_evals = net_batch_shape + core_shape_evals\n",
    "output_shape_evecs = net_batch_shape + core_shape_evecs\n",
    "output_shape_evals, output_shape_evecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66aae16e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check predictions\n",
    "input_a = rng.random(batch_shape_a + core_shape_a)\n",
    "input_b = rng.random(batch_shape_b + core_shape_b)\n",
    "evals, evecs = linalg.eig(input_a, b=input_b)\n",
    "evals.shape, evecs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9b564c9",
   "metadata": {},
   "source": [
    "There are a few functions for which the core dimensionality (i.e., the length of the core shape) of an argument or output can be either 1 or 2. In these cases, the core dimensionality is taken to be 1 if the array has only one dimension and 2 if the array has two or more dimensions. For instance, consider the following calls to {func}`scipy.linalg.solve`. The simplest case is a single square matrix `A` and a single vector `b`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654ecc8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.eye(5)\n",
    "b = np.arange(5)\n",
    "linalg.solve(A, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faadeead",
   "metadata": {},
   "source": [
    "In this case, the core dimensionality of `A` is 2 (shape `(5, 5)`), the core dimensionality of `b` is 1  (shape `(5,)`), and the core dimensionality of the output is 1  (shape `(5,)`).\n",
    "\n",
    "However, `b` can also be a two-dimensional array in which the *columns* are taken to be one-dimensional vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a21a50e",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = np.empty((5, 2))\n",
    "b[:, 0] = np.arange(5)\n",
    "b[:, 1] = np.arange(5, 10)\n",
    "linalg.solve(A, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fdb928d",
   "metadata": {},
   "outputs": [],
   "source": [
    "b.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4322c53",
   "metadata": {},
   "source": [
    "At first glance, it might seem that the core shape of `b` is still `(5,)`, and we have simply performed the operation with a batch shape of `(2,)`. However, if this were the case, the batch shape of `b` would be *prepended* to the core shape, resulting in `b` and the output having shape `(2, 5)`. Thinking more carefully, it is correct to consider the core dimensionality of both inputs and the output to be 2; the batch shape is `()`.\n",
    "\n",
    "Likewise, whenever `b` has more than two dimensions, the core dimensionality of `b` and the output is considered to be 2. For example, to solve a batch of three entirely separate linear systems, each with only one right hand side, `b` must be provided as a three-dimensional array: one dimensions for the batch shape (`(3,)`) and two for the core shape (`(5, 1)`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "602ea577",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = rng.random((3, 5, 5))\n",
    "b = rng.random((3, 5, 1))  # batch shape (3,), core shape (5, 1)\n",
    "linalg.solve(A, b).shape"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "default_lexer": "ipython3",
   "text_representation": {
    "extension": ".md",
    "format_name": "myst",
    "format_version": 0.13,
    "jupytext_version": "1.16.4"
   }
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
