search

ESMValGroup / ESMValTool

ESMValTool: A community diagnostic and performance metrics tool for routine evaluation of Earth system models in CMIP
https://www.esmvaltool.org
Apache License 2.0
228 stars 128 forks source link

Speed up curvilinear grid regridding #2661

Open dhohn opened 2 years ago

dhohn commented 2 years ago

Hi all,

while looking at several ocean variables from several different models, I had the need to regrid from the curvilinear grids to regular ones. The built-in regrid preprocessor works well for that use case. However, I runs very slowly. I had a deeper look and discovered that the new(-ish) xesmf package also provides the required regridding functiontionality (curvilinear -> regular) but is also ~2 orders of magnitude faster. I made this notebook to illustrate my findings.

The sticking point of course being that xesmf uses xarray data, so its not quite plug and play for use in esmvaltool. I see the following options:

  1. Convert to xarray, regrid, convert back to iris:
  1. Convert to xarray, get regrid weights, apply to iris numpy data, adjust DimCoords:
  1. Replicate xesmf idea (sparse matrix weights multiplication with scipy) in esmvaltool/iris-esmf-regrid

Please let me know what you think. Maybe I overlooked something?

Thanks, David

zklaus commented 2 years ago

You may be interested in giving the generic regridding interface a try, that we added in 2.5.0. Together with https://github.com/SciTools-incubator/iris-esmf-regrid it should provide what you want, and more.

dhohn commented 2 years ago

Thanks @zklaus. I forgot to say I found that package and had a look before posting. I couldn't recognise a scheme that works for curvilinear -> rectilinear. Did I overlook it?

zklaus commented 2 years ago

I think the example in the documentation should work for this case. Reproducing it here:

preprocessors:
  regrid_preprocessor:
    regrid:
      target_grid: 2.5x2.5
      scheme:
        reference: esmf_regrid.schemes:ESMFAreaWeighted
        mdtol: 0.7
dhohn commented 2 years ago

Wow indeed! And its fast too! Thanks a lot @zklaus. Although I wish I had seen it earlier... (I was thrown off by some functions being called something rectilinear_to_rectilinear)

zklaus commented 2 years ago

It's still in development, so sometimes things change names, but I think it already shows promise.

bouweandela commented 2 years ago

@pepcos This may also be interesting for you.

pepcos commented 2 years ago

@zklaus are there other regriding schemes provided by the esmf_regrid other than the area weighted?

I'd like to use a bilinear scheme and the old scheme: linear has gotten slower in version 2.5. That's why I was wondering if there is a faster linear scheme using esmf_regrid, but didn't find it. Did I overlook it?

zklaus commented 2 years ago

@pepcos, I am not sure. But I am sure that it is reasonably easy to add if it's not there yet.

sloosvel commented 2 years ago

@pepcos what type of performance issues have you been having?

pepcos commented 2 years ago

Using the same regriding preprocessor in v2.4 and v2.5 the performance varies quite a bit for tos datasets in the native grid. The preprocessor is the following:

preprocessors:
  general_preproc:
    regrid:
      target_grid: 3x3
      scheme: area_weighted

v2.4 takes about 20s to regrid 5 years of monthly data from ACCESS-CM2:

2022-06-21 10:16:39,295 UTC [754432] DEBUG   esmvalcore.preprocessor:293 Running preprocessor step regrid
2022-06-21 10:16:39,301 UTC [754432] DEBUG   esmvalcore.preprocessor:283 Running regrid(sea_surface_temperature / (degC)          (time: 72; cell index along second dimension: 300; cell index along first dimension: 360)
    Dimension coordinates:
        time                                   x                                      -                                      -
        cell index along second dimension      -                                      x                                      -
        cell index along first dimension       -                                      -                                      x
    Auxiliary coordinates:
        latitude                               -                                      x                                      x
        longitude                              -                                      x                                      x
    Cell methods:
        mean where sea                    area
        mean                              time
    Attributes:
        Conventions                       CF-1.7 CMIP-6.2
        NCO                               netCDF Operators version 4.9.2 (Homepage = http://nco.sf.net, Code = h...
        activity_id                       CMIP
        branch_method                     standard
        branch_time_in_child              0.0
        branch_time_in_parent             -328713.0
        cmor_version                      3.4.0
        data_specs_version                01.00.30
        experiment                        all-forcing simulation of the recent past
        experiment_id                     historical
        external_variables                areacello
        forcing_index                     1
        frequency                         mon
        further_info_url                  https://furtherinfo.es-doc.org/CMIP6.CSIRO-ARCCSS.ACCESS-CM2.historica...
        grid                              native atmosphere N96 grid (144x192 latxlon)
        grid_label                        gn
        initialization_index              1
        institution                       CSIRO (Commonwealth Scientific and Industrial Research Organisation, Aspendale,...
        institution_id                    CSIRO-ARCCSS
        license                           CMIP6 model data produced by CSIRO is licensed under a Creative Commons...
        mip_era                           CMIP6
        nominal_resolution                250 km
        notes                             Exp: CM2-historical; Local ID: bj594; Variable: tos (['sst'])
        parent_activity_id                CMIP
        parent_experiment_id              piControl
        parent_mip_era                    CMIP6
        parent_source_id                  ACCESS-CM2
        parent_time_units                 days since 1850-1-1 00:00:00
        parent_variant_label              r1i1p1f1
        physics_index                     1
        product                           model-output
        realization_index                 1
        realm                             ocean
        run_variant                       forcing: GHG, Oz, SA, Sl, Vl, BC, OC, (GHG = CO2, N2O, CH4, CFC11, CFC12,...
        source                            'ACCESS-CM2 (2019): \naerosol: UKCA-GLOMAP-mode\natmos: MetUM-HadGEM3-GA7.1...
        source_id                         ACCESS-CM2
        source_type                       AOGCM
        sub_experiment                    none
        sub_experiment_id                 none
        table_id                          Omon
        table_info                        Creation Date:(30 April 2019) MD5:e14f55f257cceafb2523e41244962371
        title                             ACCESS-CM2 output prepared for CMIP6
        variable_id                       tos
        variant_label                     r1i1p1f1
        version                           v20191108, {'target_grid': '3x3', 'scheme': 'area_weighted'})
2022-06-21 10:16:59,654 UTC [754432] DEBUG   esmvalcore.preprocessor:293 Running preprocessor step remove_fx_variables

v2.5 takes about 56s to regrid the same dataset:

2022-06-21 10:03:51,905 UTC [748849] DEBUG   esmvalcore.preprocessor:336 Running preprocessor step regrid
2022-06-21 10:03:51,909 UTC [748849] DEBUG   esmvalcore.preprocessor:297 Running preprocessor function 'regrid' on the data
<iris 'Cube' of sea_surface_temperature / (degC) (time: 72; cell index along second dimension: 300; cell index along first dimension: 360)>
loaded from original input file(s)
['/esarchive/exp/CMIP6/historical/ACCESS-CM2/CMIP/CSIRO-ARCCSS/ACCESS-CM2/historical/r1i1p1f1/Omon/tos/gn/v20191108/tos_Omon_ACCESS-CM2_historical_r1i1p1f1_gn_185001-201412.nc']
with function argument(s)
target_grid = '3x3',
scheme = 'area_weighted'
2022-06-21 10:04:47,899 UTC [748849] DEBUG   esmvalcore.preprocessor:336 Running preprocessor step remove_fx_variables
2022-06-21 10:04:47,902 UTC [748849] DEBUG   esmvalcore.preprocessor:297 Running preprocessor function 'remove_fx_variables' on the data
<iris 'Cube' of sea_surface_temperature / (degC) (time: 72; latitude: 60; longitude: 120)>
loaded from original input file(s)
['/esarchive/exp/CMIP6/historical/ACCESS-CM2/CMIP/CSIRO-ARCCSS/ACCESS-CM2/historical/r1i1p1f1/Omon/tos/gn/v20191108/tos_Omon_ACCESS-CM2_historical_r1i1p1f1_gn_185001-201412.nc']
with function argument(s)

2022-06-21 10:04:47,902 UTC [748849] DEBUG   esmvalcore.preprocessor:336 Running preprocessor step save

The time difference escalates with the amount of years taken from the dataset, and I have been able to reproduce this with many other CMIP6 datasets.

If I use the 2.5 with the esmf area weighted scheme it turns out to be quite sweet as the regriding time goes down to just 3s:

2022-06-21 13:48:21,219 UTC [883633] DEBUG   esmvalcore.preprocessor:336 Running preprocessor step regrid
2022-06-21 13:48:21,224 UTC [883633] DEBUG   esmvalcore.preprocessor:297 Running preprocessor function 'regrid' on the data
<iris 'Cube' of sea_surface_temperature / (degC) (time: 72; cell index along second dimension: 300; cell index along first dimension: 360)>
loaded from original input file(s)
['/esarchive/exp/CMIP6/historical/ACCESS-CM2/CMIP/CSIRO-ARCCSS/ACCESS-CM2/historical/r1i1p1f1/Omon/tos/gn/v20191108/tos_Omon_ACCESS-CM2_historical_r1i1p1f1_gn_185001-201412.nc']
with function argument(s)
target_grid = '3x3',
scheme = {'reference': 'esmf_regrid.schemes:ESMFAreaWeighted'}
2022-06-21 13:48:24,212 UTC [883633] DEBUG   esmvalcore.preprocessor:336 Running preprocessor step remove_fx_variables

Nevertheless, many datasets which I could regrid with the old area_weighted scheme, now break during the regriding process with the esmf scheme:

Traceback (most recent call last):
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_main.py", line 499, in run
    fire.Fire(ESMValTool())
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/fire/core.py", line 466, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/fire/core.py", line 681, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_main.py", line 443, in run
    process_recipe(recipe_file=recipe, config_user=cfg)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_main.py", line 124, in process_recipe
    recipe.run()
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_recipe.py", line 1882, in run
    self.tasks.run(max_parallel_tasks=self._cfg['max_parallel_tasks'])
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_task.py", line 725, in run
    self._run_parallel(max_parallel_tasks)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_task.py", line 768, in _run_parallel
    _copy_results(task, running[task])
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_task.py", line 791, in _copy_results
    task.output_files, task.products = future.get()
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/multiprocessing/pool.py", line 771, in get
    raise self._value
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_task.py", line 796, in _run_task
    output_files = task.run()
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/_task.py", line 263, in run
    self.output_files = self._run(input_files)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/preprocessor/__init__.py", line 599, in _run
    product.apply(step, self.debug)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/preprocessor/__init__.py", line 410, in apply
    self.cubes = preprocess(self.cubes, step,
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/preprocessor/__init__.py", line 346, in preprocess
    result.append(_run_preproc_function(function, item, settings,
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/preprocessor/__init__.py", line 302, in _run_preproc_function
    return function(items, **kwargs)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmvalcore/preprocessor/_regrid.py", line 628, in regrid
    cube = cube.regrid(target_grid, loaded_scheme)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/iris/cube.py", line 4177, in regrid
    regridder = scheme.regridder(self, grid)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmf_regrid/schemes.py", line 317, in regridder
    return ESMFAreaWeightedRegridder(src_grid, tgt_grid, mdtol=self.mdtol)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmf_regrid/schemes.py", line 346, in __init__
    regrid_info = _regrid_rectilinear_to_rectilinear__prepare(src_grid, tgt_grid)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmf_regrid/schemes.py", line 173, in _regrid_rectilinear_to_rectilinear__prepare
    srcinfo = _cube_to_GridInfo(src_grid_cube)
  File "/gpfs/projects/bsc32/software/suselinux/11/software/ESMValTool/2.5.0-exp/lib/python3.8/site-packages/esmf_regrid/schemes.py", line 44, in _cube_to_GridInfo
    assert lon.is_contiguous()
AssertionError
2022-06-21 10:50:18,830 UTC [781294] INFO    esmvalcore._main:515

So, using esmf solves the time problem but generates a new one for datasets which could previously be regrided.

zklaus commented 2 years ago

@pepcos, I encourage you to take up issues with esmf_regrid at its own issue tracker here.

dhohn commented 2 years ago

Ive seen the same assert lon.is_contiguous() AssertionError and looked at the failing datasets a bit. Indeed they dont have contiguous bounds, so it seems to be a dataset issue.

Having said that: Commenting out the particular line in esmf_regrid and does this check allows the regridding to finish, and I didnt see any side effects. So not sure how strict that requirement really needs to be!