issue_comments

I have created a library for easier linear algebra (plus others) with xarray: https://xarray-einstats.readthedocs.io/en/latest/.

It currently has: * wrappers for many numpy.linalg functions * wrappers for scipy.stats distributions plus a few functions * wrappers for einops functions * a numba.guvectorize-decorated version of numpy.histogram for dataarrays

I have mostly added wrappers for things I personally use, also trying to not overlap with xr-scipy nor xarray-extras, so it might not make much sense as a group, but the modules are completely independent between them and could be reorganized into independent packages or merged into existing ones.

Feedback very welcome!

Great, thanks for the offer :smile: I think I'll start with a minimal repo back at ArviZ and see how much of that can be used more generally. I probably should have named the example above arviz_dot instead. I definitely want to be able to do arviz_dot(a, b) and have it work automatically because in most ArviZ cases we do have all the information we need for this to be possible. We'll therefore definitely need a higher ArviZ layer, but I think an xarray-linalg would be a great base on which to build and we can keep both in different files/modules so it can be split easily.

I also commented here as I thought it was the most related issue, many of the functions I have in mind are related to linalg, but we'll probably have some other functions too.

Would there be any interest in an "xarray linalg library"? Some kind of library of that sort would be really useful to ArviZ and our users, I'll probably write some very basic functions myself on a minimal library or add them to ArviZ directly, but if there are other people interested we can try to find synergies and get something more general.

In our case, 99% of the time we want to "batch" over the chain and draw dimensions. So things as simple as:

def xarray_dot(a, b, dim=None): if dim is None: a_dims = set(a.dims) b_dims = set(b.dims) dim_set = a_dims.intersection(b_dims) - {"chain", "draw"} if len(dim_set) == 1: dim = dim_set.pop() else: ValueError if not isinstance(dim, str): raise ValueError return xr.apply_ufunc( np.einsum, "...i,...i", a, b, input_core_dims=[[], [dim], [dim]] )

are already extremely useful.

Being able to invert, cholesky... without using apply_ufunc would be convenient, even more so if we define some conventions. i.e. I generally use dim and dim bis for covariance matrices so they have shape chain, draw, ..., dim, dim bis, so I can invert those independently of their number and order of dimensions without even needing to say which are the "matrix" dimensions.

In our case, being able to ds.transpose("dim1", "dim2") and have it ignore chain, draw would also be convenient, which is probably a very ArviZ specific situation, but maybe there are other libraries/people that commonly have some batch dimensions, even if they use different names for them or if the user chooses those names freely.