Intro to the linear algebra module¶
Most of the linear algebra module are wrappers with very few lines and an API nearly equal to their numpy counterpart. In general, the only thing you need to do is pass the input DataArray and indicate which dimensions correspond to the matrices. There are only a couple exceptions which have their own section.
import xarray_einstats
from xarray_einstats.tutorial import generate_matrices_dataarray
We start by generating syntetic data to work with:
da = generate_matrices_dataarray(7)
da
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB 0.7075 1.025 0.5685 0.8951 0.2065 3.384 ... 1.239 0.4527 0.5749 0.4766 0.859 Dimensions without coordinates: batch, experiment, dim, dim2
The data represents a collection of matrices. dim
and dim2
indicate the matrix dimensions, the whole array is 4d, with 30 matrices in total from 10 batches and 3 experiments.
General linalg functions¶
You can get the trace of all 30 matrices in a single line, you only need the input DataArray and the dimensions corresponding to the matrices:
da.linalg.trace(dims=["dim", "dim2"])
<xarray.DataArray (batch: 10, experiment: 3)> Size: 240B 4.854 4.74 4.457 2.637 2.79 3.163 1.998 ... 2.804 4.58 2.888 4.936 5.983 4.07 Dimensions without coordinates: batch, experiment
The main feature of the wrappers is that they know what is the expected shape of the output, you don’t need to take care of it. See how the inverse which doesn’t reduce the matrix dimension can be called with the exact same arguments.
da.linalg.inv(dims=["dim", "dim2"])
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB 11.26 -2.363 -10.84 -0.2744 10.99 -2.017 ... -3.444 0.7703 0.316 0.01949 -1.162 Dimensions without coordinates: batch, experiment, dim, dim2
Even a qr decomposition which returns multiple matrices (which could even have different shapes) needs only these two arguments to work. (batched qr decomposition requires numpy>=1.22)
q, r = da.linalg.qr(dims=["dim", "dim2"])
q
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB -0.5452 0.01652 -0.5624 -0.6214 -0.1592 ... -0.3322 -0.4013 0.2607 0.8128 Dimensions without coordinates: batch, experiment, dim, dim2
r
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB -1.298 -1.975 -1.858 -1.228 0.0 -3.137 ... -0.4307 1.052 0.0 0.0 0.0 -0.6995 Dimensions without coordinates: batch, experiment, dim, dim2
Tip
Do you always follow the same convention to name your matrix dimensions and feel that even having to repeat that is
unnecessary? Take a look at xarray_einstats.linalg.get_default_dims
to see how to modify the default dims used by the linalg wrappers
matmul: 1st exception¶
The general representation of a matrix multiplication is:
There are conceptually 3 dimensions involved in the operation because the 2nd dimension of \(\mathcal{M}_1\) needs to be the same as the 1st dimension of \(\mathcal{M}_2\). Moreover, when working with square matrices, \(N==M==K\) and there is only 1 dimension.
When working with xarray however, there can’t be repeated dimension names, so as we have already seen, conceptually equivalent dimensions will have potentially different names, i.e. dim
and dim2
.
Taking all of this into account, matmul
’s dims
argument supports indicating the dimensions in 3 different ways. The following table summarizes the inputs dims
accepts and how they are interpreted:
|
dim_a1 |
dim_a2 |
dim_b1 |
dim_b2 |
---|---|---|---|---|
|
dim1 |
dim2 |
dim1 |
dim2 |
|
dim1 |
dim2 |
dim2 |
dim3 |
|
dim_a1 |
dim_a2 |
dim_b1 |
dim_b2 |
where dim_a1, dim_a2
are the matrix dimensions of the first matrix, and dim_b#
are the matrix dimensions
of the 2nd matrix. Like in (1), the dimensions present in the output are dim_a1, dim_b2
.
List of two elements¶
This first example uses square matrices, so when doing a matrix multiplication, the two dimensions are common in both inputs. You only need a list with two strings to indicate how to perform the multiplication:
from xarray_einstats import linalg
linalg.matmul(da, da, dims=["dim", "dim2"])
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB 1.845 5.326 2.407 3.89 3.378 14.68 5.449 ... 5.586 6.55 1.279 1.373 1.791 2.658 Dimensions without coordinates: batch, experiment, dim, dim2
List of three elements¶
However, the input matrices for matrix multiplication might not be square or might not have the exact same dimension names. As we have seen, what is necessary if for the 2nd dimension of the 1st matrix to match with the 1st dimension of the 2nd matrix. This 3 element list of dimensions is arguable the most common way to specify matrix multiplications.
You could interpret the DataArray as a collection of matrices of dimension batch, experiment
, or with experiment, dim2
indicating the matrices. Those two collections of matrices are valid inputs for matrix multiplication.
As there is still one that need to match, matmul
can also take a list of 3 dimensions:
linalg.matmul(da, da, dims=["batch", "experiment", "dim2"], out_append="_bis")
<xarray.DataArray (dim: 4, dim2_bis: 4, batch_bis: 10, batch: 10, dim2: 4)> Size: 51kB 10.79 3.926 1.503 3.986 0.1886 0.1844 ... 1.289 4.187 5.251 3.372 2.81 13.1 Dimensions without coordinates: dim, dim2_bis, batch_bis, batch, dim2
Here, batch
and dim2
were matrix dimensions in one of the matrices and batch dimensions in the other. While this
might not be very common, xarray-einstats
check for dimensions that would end up being duplicated in the output and renames them if necessary using out_append
to avoid collisions.
A similar thing happens when both dim1 and dim3 have the same name:
linalg.matmul(da, da, dims=["batch", "experiment", "batch"])
<xarray.DataArray (dim: 4, dim2: 4, batch: 10, batch2: 10)> Size: 13kB 10.79 0.1886 5.402 1.471 1.243 5.348 2.639 ... 3.462 3.618 11.21 9.47 4.187 13.1 Dimensions without coordinates: dim, dim2, batch, batch2
List of 2 element lists¶
The 3rd option is the more verbose and explicit, but still necessary to avoid the need for manual renamings before being able to multiply some matrices.
To see how it works, you’ll need a db
object, with the same shape but different dimension names:
db = da.rename(dim="different_dim", dim2="different_dim2")
db
<xarray.DataArray (batch: 10, experiment: 3, different_dim: 4, different_dim2: 4)> Size: 4kB 0.7075 1.025 0.5685 0.8951 0.2065 3.384 ... 1.239 0.4527 0.5749 0.4766 0.859 Dimensions without coordinates: batch, experiment, different_dim, different_dim2
Now da
and db
are compatible and you might want to multiply them, after all, it’s the same operation we did in the first matmul
example (you can check the result if running the notebook). But given the name mismatch it wasn’t possible to use the first nor second option:
linalg.matmul(da, db, dims=[["dim", "dim2"], ["different_dim", "different_dim2"]])
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, different_dim2: 4)> Size: 4kB 1.845 5.326 2.407 3.89 3.378 14.68 5.449 ... 5.586 6.55 1.279 1.373 1.791 2.658 Dimensions without coordinates: batch, experiment, dim, different_dim2
Whenever the dimension being multiplied/reduced doesn’t have the same name in both matrices, you’ll need to use this 2+2 dims specification. Like in the list of 3 elements case, matmul
avoids name clashes:
dc = da.rename(batch="batch_bis")
linalg.matmul(da, dc, dims=[["experiment", "batch"], ["batch_bis", "experiment"]])
<xarray.DataArray (dim: 4, dim2: 4, experiment: 3, experiment2: 3)> Size: 1kB 9.727 6.68 3.595 6.68 18.66 6.065 3.595 ... 10.81 36.08 8.181 3.233 8.181 14.77 Dimensions without coordinates: dim, dim2, experiment, experiment2
einsum: 2nd and most notable exception¶
einsum
is a such a flexible function that it can even be intimidating. It can cover from sum
operations, to xarray.dot
reductions and obviously some operations similar to einops
which after all is inspired in einsum.
The goal of this page is not to be an extensive nor in depth guide on einsum but to act as a small ladder from simple operations that you can do without einsum until reaching operations that are only possible with einsum. This will give you a good look into xarray_einstats
unique version of einsum
that works with named dimensions, you’ll see how most einsum operations translate to our syntax.
If you want to master einsum however, we direct you to numpy.einsum
documentation and the einops package. To ease a little bit your ability to follow the tutorial without needing to understand einsum beforehand, we provide the equivalent in non-einsum functions (which is often multiple operations) inside of toggle-able note boxes. But keep in mind that the goal of this section is not teaching how to use einsum
but showing how to use
xarray_einstats
to perform einsum operations with named dimension names.
from xarray_einstats import einsum
import xarray as xr
Start reducing the experiment
dimension. Any ellipsis, broadcasting and transposition is handled by xarray and xarray-einstats. You only need to care about the dimensions you want to operate on. Use []
to indicate you want to reduce the dimension (or ->
in string syntax):
einsum([["experiment"], []], da)
einsum("experiment->", da)
<xarray.DataArray (batch: 10, dim: 4, dim2: 4)> Size: 1kB 4.487 3.158 0.9252 2.683 0.5319 3.799 ... 3.387 1.796 2.601 2.455 1.538 5.402 Dimensions without coordinates: batch, dim, dim2
The same can be dome with multiple dimensions.
einsum([["batch", "experiment"], []], da)
einsum("batch experiment->", da)
<xarray.DataArray (dim: 4, dim2: 4)> Size: 128B 22.27 32.55 29.06 40.96 23.96 33.48 ... 25.27 29.59 34.97 20.57 34.89 30.26 Dimensions without coordinates: dim, dim2
Note
These two calls are respectively equivalent to
da.sum("experiment")
da.sum(("batch", "experiment"))
einsum
also takes multiple outputs. In those cases, if there are repeated dimensions in the expressions
corresponding to different inputs and we want to reduce all of them, the output expression can be omitted, just like
you’d do with numpy.einsum
.
einsum([["experiment"], ["experiment"]], da, da)
einsum("experiment,experiment", da, da)
<xarray.DataArray (batch: 10, dim: 4, dim2: 4)> Size: 1kB 10.79 3.543 0.4447 2.399 0.111 11.58 10.95 ... 5.104 1.799 2.513 3.052 0.79 13.1 Dimensions without coordinates: batch, dim, dim2
Note
This call combines a product and a summation, and has two equivalents. One using xarray.dot
(also quite einsum-like), another in simple mathematical operations:
xr.dot(da, da, dims"experiment")
(da * da).sum("experiment")
When there are no repeated indexes between inputs, then the results of implicit and explicit mode are different, again, just like in numpy.einsum
. After all, xarray_einstats.einsum
is an interface to it that uses
dimension names and needs no ellipsis.
Implicit mode:
einsum([["experiment"], ["batch"]], da, da)
einsum("experiment,batch", da, da)
<xarray.DataArray (dim: 4, dim2: 4, batch: 10, experiment: 3)> Size: 4kB 33.15 44.26 22.52 1.318 1.76 0.8951 ... 19.52 36.93 18.42 42.62 80.64 40.23 Dimensions without coordinates: dim, dim2, batch, experiment
This call no longer has a single operation equivalent. Here we are performing multiple summations
and multiplications. And also reordering the dimensions. The first time they are encountered, dimensions are mapped to a single letter (the input accepted by einsum) in reverse alphabetical order. After that, the saved mapping is used. Therefore, following the xarray.apply_ufunc
convention, the default order of the dimensions is the following:
All the ommitted dimensions in the order they appear in the inputs
All dimensions present in the expressions in the inverse order they appear in the expression for the first time
Thus, the output has dim
and dim2
first as they are not present in the expressions, then comes batch
and finally experiment
in the exact inverted order they appear in the expression.
Note
Even though this computation no longer has a single function/method equivalent, it does have a multiple operation equivalent:
(da.sum("experiment") * da.sum("batch")).transpose(..., "batch", "experiment")
Explicit mode:
einsum([["experiment"], ["batch"], []], da, da)
einsum("experiment,batch->", da, da)
<xarray.DataArray (dim: 4, dim2: 4)> Size: 128B 496.0 1.06e+03 844.2 1.678e+03 573.9 ... 875.4 1.223e+03 423.1 1.218e+03 915.8 Dimensions without coordinates: dim, dim2
Note
Which again has no single operation equivalent but a multiple operation one:
(da.sum("experiment") * da.sum("batch")).sum(("batch", "experiment"))
Relation to xarray.dot
xarray.dot
is also a wrapper on numpy.einsum
, but it takes a single list of dimensions to operate on. This means that none of the two computations above can be reproduced with it. See for yourself:
xr.dot(da, da, dims=["experiment", "batch"])
<xarray.DataArray (dim: 4, dim2: 4)> Size: 128B 32.03 68.57 42.6 101.0 40.06 76.88 59.44 ... 33.78 83.88 72.41 32.43 76.33 60.63 Dimensions without coordinates: dim, dim2
xarray.dot
is operating on both dimensions for both outputs. Therefore, to reproduce its results we need to tell einsum
to operate on both dimensions for both inputs. This is similar to what we did a couple of examples back but now with two dimensions.
einsum([["batch", "experiment"], ["batch", "experiment"], []], da, da)
einsum("batch experiment,batch experiment->", da, da)
<xarray.DataArray (dim: 4, dim2: 4)> Size: 128B 32.03 68.57 42.6 101.0 40.06 76.88 59.44 ... 33.78 83.88 72.41 32.43 76.33 60.63 Dimensions without coordinates: dim, dim2
Note
Similarly to the example using dot
before, the equivalent here is a product followed by a sum over the two provided axis.
(da * da).sum(("batch", "experiment"))
keep_dims
argument
einsum
also has an argument to indicate dimensions that are present in multiple inputs but should be “kept”. That is, instead of treating the dimension as the same for all inputs, its occurrence in multiple inputs should be preserved, as if they were actually different dimensions. einsum
with then rename the repeated dimension names
using the out_append
argument.
Back to using our DataArray as a collection of matrices, we might want to do a matrix multiplication. That would mean reducing the dim2
dimensions and keeping both dim
to get the resulting collection of matrices. You can see how the default einsum
behaviour only keeps one occurence of dim
:
einsum([["dim2"], ["dim2"]], da, da)
einsum("dim2,dim2", da, da)
<xarray.DataArray (batch: 10, experiment: 3, dim: 4)> Size: 960B 2.676 19.38 0.8116 5.562 11.33 2.104 ... 6.259 12.24 6.737 0.5945 7.355 1.5 Dimensions without coordinates: batch, experiment, dim
We need to use the keep_dims
argument to keep the dim
dimension of the first DataArray as dim
and dim
of the 2nd DataArray as the new dim2
:
einsum([["dim2"], ["dim2"]], da, da, keep_dims={"dim"}, out_append="_auto{i}")
einsum("dim2,dim2", da, da, keep_dims={"dim"}, out_append="_auto{i}")
<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim_auto2: 4)> Size: 4kB 2.676 6.135 1.302 3.007 6.135 19.38 2.018 ... 7.355 2.884 2.942 0.8866 2.884 1.5 Dimensions without coordinates: batch, experiment, dim, dim_auto2
Note that here we had dim, dim2
and dim, dim2
and we have reduced dim2
in both arguments. Therefore, we haven’t done the matrix multiplication between the two arguments, but the matrix multiplication between the first argument and the transpose of the second. For seamless matrix multiplication, use xarray_einstats.matmul
.
Important
xarray_einstats.einsum
does not support combining dimensions of different names.
The keep_dims
argument can also be used to perform outer products. In a pure outer product, we don’t want to reduce any dimension, so we give empty lists as input “expression” and we pass the dimension we want to perform the outer product on as keep_dims
:
einsum([[], []], da, da, keep_dims={"batch"})
The history saving thread hit an unexpected error (OperationalError('attempt to write a readonly database')).History will not be written to the database.
<xarray.DataArray (experiment: 3, dim: 4, dim2: 4, batch: 10, batch2: 10)> Size: 38kB 0.5006 0.09001 0.1315 0.3874 0.5949 0.6645 ... 2.931 0.2908 0.5802 0.4342 0.7379 Dimensions without coordinates: experiment, dim, dim2, batch, batch2
%load_ext watermark
%watermark -n -u -v -iv -w -p numpy
The watermark extension is already loaded. To reload it, use:
%reload_ext watermark
Last updated: Thu Sep 19 2024
Python implementation: CPython
Python version : 3.11.8
IPython version : 8.18.1
numpy: 1.26.4
xarray_einstats: 0.8.0
sys : 3.11.8 | packaged by conda-forge | (main, Feb 16 2024, 20:53:32) [GCC 12.3.0]
xarray : 2024.9.0
Watermark: 2.4.3