Auxiliary coordinates created by HybridPressureFactory have huge chunks

[bouweandela]

When loading a file that contains an auxiliary coordinate that can be computed using a formula term, the auxiliary coordinate ends up having huge chunks. This leads to memory issues when trying to use such coordinates, as the Dask workers will run out of memory and get killed.

Example

Open the file clw_Amon_FGOALS-f3-L_historical_r1i1p1f1_gr_196001-196912.nc and list the chunks of the computed coordinate:

In [1]: import iris

In [2]: iris.__version__
Out[2]: '3.7.0'

In [3]: cube = iris.load_cube("clw_Amon_FGOALS-f3-L_historical_r1i1p1f1_gr_196001-196912.nc", "mass_fraction_of_cloud_liquid_water_in_air")
/home/bandela/mambaforge/envs/iris/lib/python3.11/site-packages/iris/fileformats/cf.py:859: UserWarning: Missing CF-netCDF measure variable 'areacella', referenced by netCDF variable 'clw'
  warnings.warn(
/home/bandela/mambaforge/envs/iris/lib/python3.11/site-packages/iris/fileformats/cf.py:1264: UserWarning: Ignoring formula terms variable 'ps' referenced by data variable 'b_bnds' via variable 'lev': Dimensions ('time', 'lat', 'lon') do not span ('lev', 'bnds')
  warnings.warn(msg)
/home/bandela/mambaforge/envs/iris/lib/python3.11/site-packages/iris/fileformats/cf.py:1264: UserWarning: Ignoring formula terms variable 'ps' referenced by data variable 'a_bnds' via variable 'lev': Dimensions ('time', 'lat', 'lon') do not span ('lev', 'bnds')
  warnings.warn(msg)
/home/bandela/mambaforge/envs/iris/lib/python3.11/site-packages/iris/fileformats/netcdf/loader.py:353: UserWarning: Unable to find coordinate for variable 'ps'
  warnings.warn(
/home/bandela/mambaforge/envs/iris/lib/python3.11/site-packages/iris/fileformats/netcdf/loader.py:353: UserWarning: Unable to find coordinate for variable 'ps'
  warnings.warn(

In [4]: cube
Out[4]: <iris 'Cube' of mass_fraction_of_cloud_liquid_water_in_air / (kg kg-1) (time: 120; atmosphere_hybrid_sigma_pressure_coordinate: 32; latitude: 180; longitude: 288)>

In [5]: cube.coords()
Out[5]: 
[<DimCoord: time / (days since 0001-01-01 00:00:00)  [...]+bounds  shape(120,)>,
 <DimCoord: atmosphere_hybrid_sigma_pressure_coordinate / (1)  [0.996, ...]+bounds  shape(32,)>,
 <DimCoord: latitude / (degrees)  [-89.5, -88.5, ..., 88.5, 89.5]+bounds  shape(180,)>,
 <DimCoord: longitude / (degrees)  [ 0.625, 1.875, ..., 358.125, 359.375]+bounds  shape(288,)>,
 <DimCoord: vertical coordinate formula term: reference pressure / (Pa)  [...]>,
 <AuxCoord: Surface Air Pressure / (Pa)  <lazy>  shape(120, 180, 288)>,
 <AuxCoord: vertical coordinate formula term: a(k) / (1)  [0., ...]  shape(32,)>,
 <AuxCoord: vertical coordinate formula term: b(k) / (1)  [0.996, ...]  shape(32,)>,
 <AuxCoord: vertical pressure / (Pa)  [ 0. , 316.248, ..., 609.301, 250. ]  shape(32,)>,
 <AuxCoord: air_pressure / (Pa)  <lazy>  shape(120, 32, 180, 288)>]

In [6]: cube.core_data().chunks
Out[6]: ((20, 20, 20, 20, 20, 20), (32,), (180,), (288,))

In [7]: cube.coord('air_pressure').core_points().chunks
Out[7]: ((120,), (32,), (180,), (288,))

i.e. no chunking is applied along the time dimension at all and the 'air_pressure' coordinate has a chunk size of 1.5 GB. For performance it would be best if the coordinate had the same chunks as the data of the cube.

ncdump -hs of the file:

netcdf clw_Amon_FGOALS-f3-L_historical_r1i1p1f1_gr_196001-196912 {
dimensions:
	time = UNLIMITED ; // (120 currently)
	lev = 32 ;
	lat = 180 ;
	lon = 288 ;
	bnds = 2 ;
variables:
	double time(time) ;
		time:bounds = "time_bnds" ;
		time:units = "days since 0001-01-01 00:00:00" ;
		time:calendar = "365_day" ;
		time:axis = "T" ;
		time:long_name = "time" ;
		time:standard_name = "time" ;
	double time_bnds(time, bnds) ;
	double lev(lev) ;
		lev:bounds = "lev_bnds" ;
		lev:units = "1" ;
		lev:axis = "Z" ;
		lev:positive = "down" ;
		lev:long_name = "hybrid sigma pressure coordinate" ;
		lev:standard_name = "atmosphere_hybrid_sigma_pressure_coordinate" ;
		lev:formula = "p = a*p0 + b*ps" ;
		lev:formula_terms = "p0: p0 a: a b: b ps: ps" ;
	double lev_bnds(lev, bnds) ;
		lev_bnds:formula = "p = a*p0 + b*ps" ;
		lev_bnds:standard_name = "atmosphere_hybrid_sigma_pressure_coordinate" ;
		lev_bnds:units = "1" ;
		lev_bnds:formula_terms = "p0: p0 a: a_bnds b: b_bnds ps: ps" ;
	double p0 ;
		p0:long_name = "vertical coordinate formula term: reference pressure" ;
		p0:units = "Pa" ;
	double a(lev) ;
		a:long_name = "vertical coordinate formula term: a(k)" ;
	double b(lev) ;
		b:long_name = "vertical coordinate formula term: b(k)" ;
	float ps(time, lat, lon) ;
		ps:long_name = "Surface Air Pressure" ;
		ps:units = "Pa" ;
	double a_bnds(lev, bnds) ;
		a_bnds:long_name = "vertical coordinate formula term: a(k+1/2)" ;
	double b_bnds(lev, bnds) ;
		b_bnds:long_name = "vertical coordinate formula term: b(k+1/2)" ;
	double lat(lat) ;
		lat:bounds = "lat_bnds" ;
		lat:units = "degrees_north" ;
		lat:axis = "Y" ;
		lat:long_name = "Latitude" ;
		lat:standard_name = "latitude" ;
	double lat_bnds(lat, bnds) ;
	double lon(lon) ;
		lon:bounds = "lon_bnds" ;
		lon:units = "degrees_east" ;
		lon:axis = "X" ;
		lon:long_name = "Longitude" ;
		lon:standard_name = "longitude" ;
	double lon_bnds(lon, bnds) ;
	float clw(time, lev, lat, lon) ;
		clw:standard_name = "mass_fraction_of_cloud_liquid_water_in_air" ;
		clw:long_name = "Mass Fraction of Cloud Liquid Water" ;
		clw:comment = "Includes both large-scale and convective cloud. Calculate as the mass of cloud liquid water in the grid cell divided by the mass of air (including the water in all phases) in the grid cells. Precipitating hydrometeors are included ONLY if the precipitating hydrometeors affect the calculation of radiative transfer in model." ;
		clw:units = "kg kg-1" ;
		clw:cell_methods = "area: time: mean" ;
		clw:cell_measures = "area: areacella" ;
		clw:missing_value = 1.e+20f ;
		clw:_FillValue = 1.e+20f ;
		clw:history = "2019-09-27T10:58:20Z altered by CMOR: Inverted axis: lev." ;

// global attributes:
		:Conventions = "CF-1.7 CMIP-6.2" ;
		:activity_id = "CMIP" ;
		:branch_method = "Spin-up documentation" ;
		:branch_time_in_child = 2310. ;
		:branch_time_in_parent = 12345. ;
		:contact = "Yongqiang Yu (yyq@lasg.iap.ac.cn);Bian He(heb@lasg.iap.ac.cn)" ;
		:creation_date = "2019-09-27T10:58:20Z" ;
		:data_specs_version = "01.00.30" ;
		:experiment = "all-forcing simulation of the recent past" ;
		:experiment_id = "historical" ;
		:external_variables = "areacella" ;
		:forcing_index = 1 ;
		:frequency = "mon" ;
		:further_info_url = "https://furtherinfo.es-doc.org/CMIP6.CAS.FGOALS-f3-L.historical.none.r1i1p1f1" ;
		:grid = "gs1x1" ;
		:grid_label = "gr" ;
		:history = "2019-09-27T10:58:20Z ;rewrote data to be consistent with CMIP for variable clw found in table Amon." ;
		:initialization_index = 1 ;
		:institution = "Chinese Academy of Sciences, Beijing 100029, China" ;
		:institution_id = "CAS" ;
		:mip_era = "CMIP6" ;
		:nominal_resolution = "100 km" ;
		:parent_activity_id = "CMIP" ;
		:parent_experiment_id = "piControl" ;
		:parent_mip_era = "CMIP6" ;
		:parent_source_id = "FGOALS-f3-L" ;
		:parent_time_units = "days since 0001-01-01" ;
		:parent_variant_label = "r1i1p1f1" ;
		:physics_index = 1 ;
		:product = "model-output" ;
		:realization_index = 1 ;
		:realm = "atmos" ;
		:references = "He et al.,2019:CAS FGOALS-f3-L Model datasets for CMIP6 Historical Atmospheric Model Intercomparison Project Simulation. Adv. Atmo. Sci. doi:10.1007/s00376-019-9027-8; Bao, Q et al (2019). Outlook for El Ni?o and the Indian Ocean Dipole in autumn-winter 2018-2019. Chinese Science Bulletin, 64(1), 73-78, DOI: 10.1360/N972018-00913; Li, J., et al (2019). Evaluation of FAMIL2 in Simulating the Climatology and Seasonal-to-Interannual Variability of Tropical Cyclone Characteristics. Journal of Advances in Modeling Earth Systems, 11. https://doi.org/10.1029/2018MS001506." ;
		:source = "FGOALS-f3-L (2017): \n",
			"aerosol: none\n",
			"atmos: FAMIL2.2 (Cubed-sphere, c96; 360 x 180 longitude/latitude; 32 levels; top level 2.16 hPa)\n",
			"atmosChem: none\n",
			"land: CLM4.0\n",
			"landIce: none\n",
			"ocean: LICOM3.0 (LICOM3.0, tripolar primarily 1deg; 360 x 218 longitude/latitude; 30 levels; top grid cell 0-10 m)\n",
			"ocnBgchem: none\n",
			"seaIce: CICE4.0" ;
		:source_id = "FGOALS-f3-L" ;
		:source_type = "AOGCM" ;
		:sub_experiment = "none" ;
		:sub_experiment_id = "none" ;
		:table_id = "Amon" ;
		:table_info = "Creation Date:(09 May 2019) MD5:cde930676e68ac6780d5e4c62d3898f6" ;
		:title = "FGOALS-f3-L output prepared for CMIP6" ;
		:tracking_id = "hdl:21.14100/a2f36f63-6d1a-4f9b-880d-536b462b3149" ;
		:variable_id = "clw" ;
		:variant_label = "r1i1p1f1" ;
		:license = "CMIP6 model data produced by Lawrence Livermore PCMDI is licensed under a Creative Commons Attribution ShareAlike 4.0 International License (https://creativecommons.org/licenses). Consult https://pcmdi.llnl.gov/CMIP6/TermsOfUse for terms of use governing CMIP6 output, including citation requirements and proper acknowledgment. Further information about this data, including some limitations, can be found via the further_info_url (recorded as a global attribute in this file) and at https:///pcmdi.llnl.gov/. The data producers and data providers make no warranty, either express or implied, including, but not limited to, warranties of merchantability and fitness for a particular purpose. All liabilities arising from the supply of the information (including any liability arising in negligence) are excluded to the fullest extent permitted by law." ;
		:cmor_version = "3.5.0" ;
		:_Format = "classic" ;
}

[bouweandela]

I looked into fixing this, but would like to get some input on what the desired chunks are. Input anyone? @SciTools/iris-devs

Some options:

1. Iris default chunks (example in Use reasonable chunk sizes in derived coordinates #5712)
2. Dask default chunks
3. The chunks of the cube data variable
4. Offer some control to the user over what the chunks are, e.g. through the (currently NetCDF specific) iris.fileformats.netcdf.loader.CHUNK_CONTROL

[ESadek-MO]

@SciTools/peloton

We're concerned that this change might detriment other workflows which were otherwise fine. 3.8, which is due to be released soon (#5363), includes a CHUNK_CONTROL context manager (https://scitools-iris.readthedocs.io/en/latest/further_topics/netcdf_io.html#chunk-control) which, when used to chunk the original coordinates, should help avoid this issue.

[bouweandela]

Are you referring to the example implementation in #5712?

An alternative could be to assign chunks to the input coordinates of the derivation at load time, such that the derived variable ends up with reasonably sized chunks. The input coordinates then have rather small chunks, which will be inconvenient for anyone who wants to work with those directly, but maybe that is not a very common scenario so not really an issue.

I prefer to fix this issue on the Iris side and not leave it to the user, as the current behaviour is unlikely to produce working results.