issue_comments

Closing in favour of xarray-einstats

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.

"xarray linalg library"?

Thanks @OriolAbril . This sounds like a nice xarray-contrib repo if it's general enough. Note that we do have xarray.dot https://xarray.pydata.org/en/stable/generated/xarray.dot.html

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.

I think the way to write something like an LU decomposition would be to use xarray's apply_ufunc, which in theory could wrap scipy.linalg.lu_factor in one line.

I don't know if this could/should live in xarray proper. We tend to have pretty high standards, which slows development.

Hi @weipeng1999 , could you link the reference implementation in numpy/scipy?

I think this would be niche-ish. I would personally try to keep xarray free of functionality that only a tiny fraction of the users actually use - particularly when such functionality can be implemented with a trivial wrapper by the users themselves. e.g. at the moment we have exactly one scipy function being wrapped, and that's linear interpolation which is useful to a lot of people.

I think this falls into a more general discussion on how niche a function must be in order to be excluded from the library - @shoyer what's your opinion?

Regardless, I would like to point you to https://xarray-extras.readthedocs.io which is a module that I created exactly for this kind of cases (PRs are welcome).

I realize that I am a totally green finger here are my trial implement I think I have long way to make it commitable

qr.txt

Somewhat related: https://github.com/pydata/xarray/pull/2766

Hi @weipeng1999 , could you link the reference implementation in numpy/scipy?

I think this would be niche-ish. I would personally try to keep xarray free of functionality that only a tiny fraction of the users actually use - particularly when such functionality can be implemented with a trivial wrapper by the users themselves. e.g. at the moment we have exactly one scipy function being wrapped, and that's linear interpolation which is useful to a lot of people.

I think this falls into a more general discussion on how niche a function must be in order to be excluded from the library - @shoyer what's your opinion?

Regardless, I would like to point you to https://xarray-extras.readthedocs.io which is a module that I created exactly for this kind of cases (PRs are welcome).