Remembering weights between consecutive regridding operations

[ThibHlln]

Given two fields fld1 and fld2, if I use for example regrids twice, it seems to take the same amount of time the first and the second time.

fld3 = fld1.regrids(fld2, 'conservative')
fld4 = fld1.regrids(fld2, 'conservative')

So I am wondering if cf is able to remember the weights from the first time it did the regridding in order to make the second time faster?

[sadielbartholomew]

Hi @ThibHlln, thanks for the very good question as usual!

I will double check in the API and code (if I don't comment further assume it is indeed the case) but I suspect that, at the moment, there is not a way to do so. However, caching the weights is certainly a good idea (if, as I assume in practice, you mean a non-trivial case on , not just repeating the exact same regridding operation twice i.e. maybe the metadata differs?) and I think it would not be difficult to implement something that could do that.

We can mark it as a (performance) feature for the next version, as improving performance is one of the core tasks we are working on at the moment and going forward.

[davidhassell]

Hi - it certainly was the case that caching the weights was not possible in ESMPy (up to <= 8.0, I think). However, there was a new release a few weeks back, and I must confess that I only skim-read the release notes ... I'll go back and check them more closely for gold. We have engaged with the ESMF team in the past and requested this feature, so they do know that it is desired, but I don't know how hard it is for them to implement.

Thanks,
David

[ThibHlln]

Hi Sadie and David,

@sadielbartholomew Yes, you are right, my case was rather trivial. Just to give some context, I need to regrid repeatedly (i.e. iterating over time) from one modelling component resolution to another (specified as cf.Field). So, only the data would change between iterations, the metadata would remain unchanged. So more something like this:

# first iteration
fld1.set_data(...)
fld3 = fld1.regrids(fld2, 'conservative')
arr3 = fld3.array

# second iteration
fld1.set_data(...)
fld4 = fld1.regrids(fld2, 'conservative')
arr4 = fld4.array

# etc.

But the weights may remain constant only for some regridding methods though? 🤔 I don't know all methods well enough to answer that.

@davidhassell I was going through their documentation yesterday, and so far, the only solution I found for the caching is through a file http://earthsystemmodeling.org/esmpy_doc/release/ESMF_8_1_0/html/examples.html#read-and-write-a-weight-file where you need to specify the filename parameter in ESMF.api.regrid.Regrid (but I get a segfault on a laptop with macOS 11.3 with esmf==8.0.1, esmpy==8.0.1 if I use this parameter). Also, I think this approach requires parallel-enabled ESMF, so this could be limiting if used in cf.

ESMF.api.regrid.Regrid also has a get_weights_dict method which could be useful, but I haven't been able to find a way for them to be reused later on yet (this one does not yield a segfault for me).

Thank you for considering this for your next release.
Thibault

[davidhassell]

Hi Thibault,

Thanks for looking into the docs. I couldn't see any note about reusing that returned by get_weights_dict, either. I'll have a play with the file option, and also re-ask the qusestion to the ESMPy team.

David

[ThibHlln]

I've done a bit of experimenting with ESMF regridding through files so I am sharing my experience here hoping it may help you.

As mentioned before, I am getting a segfault problem when using Regrid with a value for parameter filename:

regrid_src2dst = ESMF.Regrid(
    src_field, dst_field, filename=filename
    ...
)

I tried to bypass this problem by creating the weights file myself from the dictionary returned by get_weights_dict as follows:

regrid_src2dst = ESMF.Regrid(
    src_field, dst_field, factors=True
    ...
)

d = regrid_src2dst.get_weights_dict()

with netCDF4.Dataset(filename, 'w') as ds:
    ds.createDimension('n_s', len(d['weights']))
    s = ds.createVariable('S', np.float64, ('n_s',))
    s[:] = d['weights']
    col = ds.createVariable('col', np.int32, ('n_s',))
    col[:] = d['col_src']
    row = ds.createVariable('row', np.int32, ('n_s',))
    row[:] = d['row_dst']

regrid_src2dst = ESMF.RegridFromFile(src_field, dst_field, filename=filename)

But I get the same segfault problem on the last line.

However, if I use the "RouteHandle" file instead (i.e. rh_filename parameter), I don't get any problem. For example:

regrid_src2dst = ESMF.Regrid(
    src_field, dst_field, rh_filename=rh_filename
    ...
)

regrid_src2dst = ESMF.RegridFromFile(src_field, dst_field, rh_filename=rh_filename)

I still need to understand why I am having this segmentation error (which may be very specific to my environment), but if you have to choose between using weight files or RouteHandle files, you may want to prefer the latter.

EDIT:

Actually, if one is not interested in the file option, one could just hold on to the ESMF.Regrid instance in memory. Then, it would "just" be a matter of keeping a list of destination fields with their associated regrid instance already used for remapping, and check in this list if there is already what's needed for regrids/regridc before attempting the expensive initialisation of Regrid. Such list could grow to be quite big though, so the file option may be preferable to avoid clogging up the memory with regrid instances which may never be reused. 😕

[davidhassell]

Hi Thibault - thanks very much for sharing this. I shall copy this in my environment and see how I get on. You have given me some good ideas on possible ways forward within cf-python.

[davidhassell]

Hi Thibault

I can do for you something like:

# Calculate and save the weights once ...
regrid = fld1.regrids(fld2, method='conservative', return_regrid=True)
type(regrid)  # <class 'cf.regrid.regrid.Regrid'>, a wrapper to ESMF.Regrid

# ... use them multiple times to do the regridding
fld1.set_data(...)
fld3 = fld1.regrids(regrid)
arr3 = fld3.array

# second iteration
fld1.set_data(...)
fld4 = fld1.regrids(regrid)
arr4 = fld4.array

Note that if the missing data distribution in fld1 changes then when you set_data, then the regrid will still work but will be slower, as the weights will have to be recalculated on-the-fly for the new mask. This is an ESMPy limitation. To be able to use "dynamic" masking (i.e. being able to impose different masks over the same set of weights) is something that we've talked about with the ESMPy developers over many years, but it hasn't yet come to pass (at least I think that's the current situation - I'll double check ...).

I'm currently putting together a PR that enables this.

[ThibHlln]

Hi David,

Thank you very much for already having a solution to this! This looks like a great compromise. 🙂

I am just wondering if this is not going to be a problem for you not to call the destroy() methods of all the ESMF objects just for the eventuality they may be re-used? Is Python going to garbage collect them if what is returned by regrids/regridc is not bound to a name?

For my use case, the mask is not changing overtime, so that is not a problem. It would be great if ESMpy could allow for this at some stage though. 🤞 But I understand that in their workflow, one needs to set the mask before you call the Regrid instance, so this must be why. Although, I thought cf-python was using ESMF Regrid instance without a mask, and then applying one as post-processing?

Once you are happy with your PR, if you need someone to test this, I am of course volunteering.

[davidhassell]

Hi Thibault,

Thank you very much for already having a solution to this! This looks like a great compromise. slightly_smiling_face

You're welcome. If you have any ideas about the API, or any extra functionality that you may need, just let us know.

I am just wondering if this is not going to be a problem for you not to call the destroy() methods of all the ESMF objects just for the eventuality they may be re-used? Is Python going to garbage collect them if what is returned by regrids/regridc is not bound to a name?

Hmm. It may be ... The cf.Regrid instance will call all of the required destroys in its __del__ method, but I'm not sure what will happen if there is no variable assignment (fld3 = fld1.regrids(fld1.regrids(fld2, method='conservative', return_regrid=True)), although I'm not sure that there's any benefit to doing this. What would happens if you assign to _? I presume that next time _ is set the original gets garbage collected as usual?

Although, I thought cf-python was using ESMF Regrid instance without a mask, and then applying one as post-processing?

It's somewhat better than that - the source and (optionally) the destination masks are incorporated into the ESMF.Grid instances and the interplay of these masks is used during the weights calculations. This means that areas adjacent to coastlines are dealt with sensibly.

Once you are happy with your PR, if you need someone to test this, I am of course volunteering.

You're on!

[ThibHlln]

You're welcome. If you have any ideas about the API, or any extra functionality that you may need, just let us know.

The API looks intuitive to me, and no different to what I am used to with cf-python. I cannot think of any additional functionality at this stage.

Hmm. It may be ... The cf.Regrid instance will call all of the required destroys in its del method, but I'm not sure what will happen if there is no variable assignment (fld3 = fld1.regrids(fld1.regrids(fld2, method='conservative', return_regrid=True)), although I'm not sure that there's any benefit to doing this. What would happens if you assign to _? I presume that next time _ is set the original gets garbage collected as usual?

Yes, I think you are right, if all the destroys() are called in __del__ this should be just fine.

[davidhassell]

Hi Thibault - I'd be great if you could test what I have done so far.

You can get my regrid branch with:

$ # Check that this goes where you expect it to!
$ pip install git+https://github.com/davidhassell/cf-python@regrid --user

Here's how I've been testing it (which should be reproducible):

In [19]: import cf

In [20]: assert cf.__version__ == '3.10.0+regrid'

In [21]: f = cf.read('regrid_file1.nc')[0]  # find the file in the `cf/test` directory

In [22]: g = cf.read('regrid_file2.nc')[0]  # find the file in the `cf/test` directory

In [23]: f, g
Out[23]: 
(<CF Field: air_temperature(time(1), latitude(64), longitude(128)) K>,
 <CF Field: air_temperature(time(1), latitude(73), longitude(96)) K>)

In [24]: h = f.regrids(g, 'conservative')

In [25]: r = f.regrids(g, 'conservative', return_operator=True)

In [26]: r
Out[26]: <cf.regrid.regridoperator.RegridOperator at 0x7f5d989d1850>

In [27]: i = f.regrids(r, 'conservative')

In [28]: i.equals(h)
Out[28]: True

In [29]: %timeit h = f.regrids(g, 'conservative')
690 ms ± 7.93 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [30]: %timeit r = f.regrids(g, 'conservative', return_operator=True)
683 ms ± 9.35 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [31]: %timeit i = f.regrids(r, 'conservative')
31.5 ms ± 265 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) 

I'm still working on some comprehensive unit tests.

Thanks.

[ThibHlln]

Hi David,

This works like a charm. 🐍
I have also tried with my own (larger) fields, and it works effortlessly too.

I am just wondering if this is intentional that the method is not optional when using a RegridOperator?

>>> r = f.regrids(g, 'conservative', return_operator=True)
>>> i = f.regrids(r)
ValueError: Can't regrid: Must select a regridding method

Also, I am sure you were planning on it, but if one uses a regridding method that is different from the one in the RegridOperator, the exception raised would benefit from being more informative:

>>> r = f.regrids(g, 'conservative', return_operator=True)
>>> i = f.regrids(r, 'linear')
TypeError: exceptions must derive from BaseException

Thanks a lot for working on this David.