Issues and fixes related to the handling of the regridding weights as a sparse.COO array and/or an xarray.DataArray backed by a sparse.COO array

[lemieuxbenedicte]

Quick description :

1. Issue 9 and FIX10 :

   MemoryError: Unable to allocate 832. GiB for an array with shape
   (334167, 334000) and data type float64
   

   This problem is encountered with the xesmf package while
   encapsulating the sparse sparse.COO matrix of the regridding weights
   under an xarray.DataArray structure.
   xarray fails interpreting the sparse.COO object hidden behind the
   numpy.ndarray type. It reserves Memory for the full size array instead.
   This can cause MemoryError failures and/or result in the use of an
   unecessary large amount of Memory.

2. Issue 1 and 2 and FIX2 :

   AttributeError: 'numpy.ndarray' object has no attribute
   'to_scipy_sparse'
   

   Issue 1 arises when the regridding involves a source grid with masked data.

   AttributeError: 'numpy.ndarray' object has no attribute 'coords'
   

   Issue 2 is related to the writing of the
   regridding weights to a netCDF file with the dedicated Class
   BaseRegridder method.

Issue 1, 2 and 9 are all connected to the handling of the
regridding weights as a sparse.COO array and/or an xarray.DataArray
backed by a sparse.COO array.

We propose fixes and tests under the github in a forked repository. In the master branch you will find the sources with fixes as well as two jupyter-notebooks in the dedicated directory xESMF/xesmf/doc/notebooks with full details :

1. notebook with issue9 and FIX10

2. notebook with issue1 and 2 and FIX2

More Details :

The xesmf package was installed with conda using the conda-forge channel. The installed version of the conda binaries is xesmf-0.6.2 but we did the appropriate checks confirming that the latest conda binaries version xesmf-0.6.3 were not solving the issues documented hereafter. The underlying package versions are therefore the following:

• esmf-8.0.1, esmpy-8.0.1
• xesmf-0.6.2

Issues were experimented with real data but they are reproduced in the notebook with synthetic grids prepared with xesmf.util.grid_2d.
Tests were conducted on HPC facilities without applying any type of parallelization (frontend).

TODO

Urgent : test the fix FIX10 with dask.

[raphaeldussin]

Hi @lemieuxbenedicte

First of all, thanks a lot for the tremendous work you've done here! I've gone trough the notebooks
and found them very comprehensive and well written. I think the fixes you propose are solving important
problems and I'd very much like to help you adding them to the code base.

Unfortunately I was not able to access the branches where the individual fixes live and I could only see your master
but this one has a bit too many changes to be pulled as is in the repo. Ideally we would like to see a pull request per
issue accompanied by the unit test (under xesmf/tests/). I think your tests are well written and should be added into
our test suite. One way to proceed could be:

• create a new branch based on our master
• add tests for small size problem that demonstrate the issues with current code
• add your fix for small size problems and pull request to pangeo-xesmf/master

repeat for large size problem. This test may or may not be able to run on github CI but we could still host it.

Let me know if this sounds good to you. It's a bit more work but I'd really like for you to get full credit for this work.
We can have a video/phone call if that helps. Both I, as well as @huard and @aulemahal speak french as well.

[lemieuxbenedicte]

Hi @raphaeldussin ,

Thanks for your reply. I did create one branch per proposed fix and made a pull request for each branch (xesmf_fix2, xesmf_fix10) but it didn't pass the tests (made a stupid error omitting import sparse)... But being given @aulemahal's explanation about sparse package, I must investigate and revise the changes. Is-it a problem proposing new merge requests even with small modifications ? Thanks again.

[raphaeldussin]

not a problem at all! and no contribution is too small!

[aulemahal]

@lemieuxbenedicte I am deeply sorry for this bug. It's my fault because I'm the one that coded sparse into xESMF but I overlooked some aspects, the major one being the dependencies and versions.

I believe all your issues can be boiled down to #158 (and maybe #170, which might be related). In both your notebooks, you are using sparse 0.3.1 . This version is extremely old and I didn't mean to have xESMF compatible with it : I simply forgot to pin it in the 0.6.2 release. There's something happening with the numba releases on conda-forge that almost systematically gives broken environments. Usually, this resulted in the installation of an incompatible numpy version and one would just be unable to import xESMF. However, I suspect that in your case is resulted in the installation of an ancient version of sparse...

I have run your notebooks with sparse 0.13, xESMF 0.6.2 and 0.6.3, xarray 2022.6.0 and scipy 1.9.1 and did not get the errors you show.
In theory, any sparse version above 0.8 should work, but our test CI is not extensive as to test all version combinations.

So désolé again that you spent so much time detailing and fixing the bug, but I think updating sparse will simply do it.

[lemieuxbenedicte]

Thanks for your explanations. I will try to test updated packages as soon as possible. Cheers !