Is xarray the right tool for my use case?

[im44pos]

I was pleased to recently discover xarray, Thanks for all videos and tutorials.

I'm not sure if I already understand xarray's capabilities well enough.

Currently using numpy, starting with (a small example) of ndarrays with shapes (5520, 3) and (6667, 3) with float64 data that fit easily in 16 Gb RAM memory. For my use case the numpy calculations intermediate come together in a ndarray of shape (k, m, n) or (5520, 5520, 6667) with float64 data that would require 1.5 Tb of RAM. Fortunately I don't require all these data points, so I wrote a for loop and use numpy's masking and fancy indexing to filter only the required data points. In this for loop I used two python lists to create a list of lists. Then converting this list of list into a numpy ndarray for further processing. The shape of this numpy ndarray is ((k), (m*, n*)) or ((5520), (m*, n*)). Note that (k) is separated from (m, n). So for each element of k there is a (m*, n*) ndarray where m* and n* vary in size (with m* = 1 to 5520 and n* = 0 - 6667). This makes it hard to analyse memory consumption (I wrote a for loop to figure out there is still 7.2 Gb required). So swapping to disk causes the for loop to be slowwww.
Can this be solved with xarray? Would it require to store intermediate data on disk?

[keewis]

Just a reminder: for numerical computations python loops (even those in comprehensions) are very slow, so you'd best avoid those whenever possible, or use numba / cython otherwise.

If you truly want to stay with this kind of ragged array, I'd recommend looking at awkward. However, if the only reason you chose that is that your data might be very sparse, I think you can use a combination of dask and the sparse library (not scipy.sparse), and put that into xarray.

[im44pos]

Thank you for quick reply.
Yes python loops are very slow. Thanks to Jake VanderPlas video "Loosing your Loops Fast" I could speedup a lot of my calculations. Nevertheless not all can be done with numpy's ufuncs, broadcasting, slicing, fancy indexing etc. I have some use cases where a for loop outperforms numpy.
Innitially I assumed that the for loop made the computation slow. Therefore I tried cython, couldn't get it to work with this ((k), (m*, n*)) ndarray. After investigating the memory usage it turned out that the swapping to disc is the current bottleneck.

Do you refer with "ragged array" to this ((k), (m*, n*)) ndarray? I don't have any preference to it. I'll take a look at awkward.
I created the for loop because my computer crashed due to memory shorthage on the numpy calculations.

I filtered the data because I only need some of it, not because it is (or could be) sparse, and so reduce memory requirement by numpy.
The many xarray examples and tutorials I have seen start with a big dataset from disc. Some refere to lazy computation. How does xarray handle intermediate arrays? Numpy creates them in memory, and there is my bottleneck.

Yes for future scale up I have dask on my to do list.