Lazy evaluation on Dask arrays¶

If you are unfamiliar with Dask, read Parallel computing with Dask in Xarray documentation first. The current version only supports dask arrays on a single machine. Support of Dask.distributed is in roadmap.

xESMF’s Dask support is mostly for lazy evaluation and out-of-core computing, to allow processing large volumes of data with limited memory. You might also get moderate speed-up on a multi-core machine by choosing proper chunk sizes, but that generally won’t help your entire pipeline too much, because the read-regrid-write pipeline is severely I/O limited (see this issue for more discussions). On a single machine, the disk bandwidth is typically limited to ~500 MB/s, and you cannot process data faster than such rate. If you need much faster data processing rate, you should resort to parallel file systems on HPC clusters or distributed storage on public cloud platforms. Please refer to the Pangeo project for more information.

[1]:

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import dask.array as da  # need to have dask.array installed, although not directly using it here.
import xarray as xr
import xesmf as xe

A simple example¶

Prepare input data¶

[2]:

ds = xr.tutorial.open_dataset("air_temperature", chunks={"time": 500})
ds

[2]:

xarray.Dataset

Dimensions:
- lat: 25
- lon: 53
- time: 2920

Coordinates: (3)

lat

(lat)

float32

75.0 72.5 70.0 ... 20.0 17.5 15.0

standard_name :: latitude
long_name :: Latitude
units :: degrees_north
axis :: Y

array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
       45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
       15. ], dtype=float32)

lon

(lon)

float32

200.0 202.5 205.0 ... 327.5 330.0

standard_name :: longitude
long_name :: Longitude
units :: degrees_east
axis :: X

array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
       225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
       250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
       275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
       300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
       325. , 327.5, 330. ], dtype=float32)

time

(time)

datetime64[ns]

2013-01-01 ... 2014-12-31T18:00:00

standard_name :: time
long_name :: Time

array(['2013-01-01T00:00:00.000000000', '2013-01-01T06:00:00.000000000',
       '2013-01-01T12:00:00.000000000', ..., '2014-12-31T06:00:00.000000000',
       '2014-12-31T12:00:00.000000000', '2014-12-31T18:00:00.000000000'],
      dtype='datetime64[ns]')

Data variables: (1)

air

(time, lat, lon)

float32

dask.array<chunksize=(500, 25, 53), meta=np.ndarray>

long_name :: 4xDaily Air temperature at sigma level 995
units :: degK
precision :: 2
GRIB_id :: 11
GRIB_name :: TMP
var_desc :: Air temperature
dataset :: NMC Reanalysis
level_desc :: Surface
statistic :: Individual Obs
parent_stat :: Other
actual_range :: [185.16 322.1 ]





  
      Array  Chunk 
  
  
     Bytes  15.48 MB   2.65 MB 
     Shape  (2920, 25, 53)   (500, 25, 53) 
     Count  7 Tasks  6 Chunks 
     Type  float32  numpy.ndarray

Attributes: (5)
Conventions :
COARDS
title :
4x daily NMC reanalysis (1948)
description :
Data is from NMC initialized reanalysis (4x/day). These are the 0.9950 sigma level values.
platform :
Model
references :
http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis.html

[3]:

ds.chunks

[3]:

Frozen(SortedKeysDict({'time': (500, 500, 500, 500, 500, 420), 'lat': (25,), 'lon': (53,)}))

[4]:

ds["air"].data

[4]:

	Array	Chunk
Bytes	15.48 MB	2.65 MB
Shape	(2920, 25, 53)	(500, 25, 53)
Count	7 Tasks	6 Chunks
Type	float32	numpy.ndarray

Build regridder¶

[5]:

ds_out = xr.Dataset(
    {
        "lat": (["lat"], np.arange(16, 75, 1.0)),
        "lon": (["lon"], np.arange(200, 330, 1.5)),
    }
)

regridder = xe.Regridder(ds, ds_out, "bilinear")
regridder

[5]:

xESMF Regridder
Regridding algorithm:       bilinear
Input grid shape:           (25, 53)
Output grid shape:          (59, 87)
Output grid dimension name: ('lat', 'lon')
Periodic in longitude?      False

Apply to xarray Dataset/DataArray¶

[6]:

# only build the dask graph; actual computation happens later when calling compute()
%time ds_out = regridder(ds)
ds_out

using dimensions ('lat', 'lon') from data variable air as the horizontal dimensions for this dataset.
CPU times: user 2.09 ms, sys: 0 ns, total: 2.09 ms
Wall time: 2.06 ms

[6]:

xarray.Dataset

Dimensions:
- lat: 59
- lon: 87
- time: 2920

Coordinates: (3)

time

(time)

datetime64[ns]

2013-01-01 ... 2014-12-31T18:00:00

standard_name :: time
long_name :: Time

array(['2013-01-01T00:00:00.000000000', '2013-01-01T06:00:00.000000000',
       '2013-01-01T12:00:00.000000000', ..., '2014-12-31T06:00:00.000000000',
       '2014-12-31T12:00:00.000000000', '2014-12-31T18:00:00.000000000'],
      dtype='datetime64[ns]')

lon

(lon)

float64

200.0 201.5 203.0 ... 327.5 329.0

array([200. , 201.5, 203. , 204.5, 206. , 207.5, 209. , 210.5, 212. , 213.5,
       215. , 216.5, 218. , 219.5, 221. , 222.5, 224. , 225.5, 227. , 228.5,
       230. , 231.5, 233. , 234.5, 236. , 237.5, 239. , 240.5, 242. , 243.5,
       245. , 246.5, 248. , 249.5, 251. , 252.5, 254. , 255.5, 257. , 258.5,
       260. , 261.5, 263. , 264.5, 266. , 267.5, 269. , 270.5, 272. , 273.5,
       275. , 276.5, 278. , 279.5, 281. , 282.5, 284. , 285.5, 287. , 288.5,
       290. , 291.5, 293. , 294.5, 296. , 297.5, 299. , 300.5, 302. , 303.5,
       305. , 306.5, 308. , 309.5, 311. , 312.5, 314. , 315.5, 317. , 318.5,
       320. , 321.5, 323. , 324.5, 326. , 327.5, 329. ])

lat

(lat)

float64

16.0 17.0 18.0 ... 72.0 73.0 74.0

array([16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27., 28., 29.,
       30., 31., 32., 33., 34., 35., 36., 37., 38., 39., 40., 41., 42., 43.,
       44., 45., 46., 47., 48., 49., 50., 51., 52., 53., 54., 55., 56., 57.,
       58., 59., 60., 61., 62., 63., 64., 65., 66., 67., 68., 69., 70., 71.,
       72., 73., 74.])

Data variables: (1)

air

(time, lat, lon)

float64

dask.array<chunksize=(500, 59, 87), meta=np.ndarray>





  
      Array  Chunk 
  
  
     Bytes  119.91 MB   20.53 MB 
     Shape  (2920, 59, 87)   (500, 59, 87) 
     Count  13 Tasks  6 Chunks 
     Type  float64  numpy.ndarray

Attributes: (1)
regrid_method :
bilinear

[7]:

ds_out["air"].data  # chunks are preserved

[7]:

	Array	Chunk
Bytes	119.91 MB	20.53 MB
Shape	(2920, 59, 87)	(500, 59, 87)
Count	13 Tasks	6 Chunks
Type	float64	numpy.ndarray

[8]:

%time result = ds_out['air'].compute()  # actually applies regridding

CPU times: user 298 ms, sys: 212 ms, total: 510 ms
Wall time: 154 ms

[9]:

type(result.data), result.data.shape

[9]:

(numpy.ndarray, (2920, 59, 87))

Invalid chunk sizes to avoid¶

Dask support relies on xarray.apply_ufunc, which requires only chunking over extra/broadcasting dimensions (like time and lev), not horizontal/core dimensions (lat, lon, or x, y).

[10]:

# xESMF doesn't like chunking over horizontal dimensions
ds_bad = ds.chunk({"lat": 10, "lon": 10, "time": None})
ds_bad

[10]:

xarray.Dataset

Dimensions:
- lat: 25
- lon: 53
- time: 2920

Coordinates: (3)

lat

(lat)

float32

75.0 72.5 70.0 ... 20.0 17.5 15.0

standard_name :: latitude
long_name :: Latitude
units :: degrees_north
axis :: Y

array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
       45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
       15. ], dtype=float32)

lon

(lon)

float32

200.0 202.5 205.0 ... 327.5 330.0

standard_name :: longitude
long_name :: Longitude
units :: degrees_east
axis :: X

array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
       225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
       250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
       275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
       300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
       325. , 327.5, 330. ], dtype=float32)

time

(time)

datetime64[ns]

2013-01-01 ... 2014-12-31T18:00:00

standard_name :: time
long_name :: Time

array(['2013-01-01T00:00:00.000000000', '2013-01-01T06:00:00.000000000',
       '2013-01-01T12:00:00.000000000', ..., '2014-12-31T06:00:00.000000000',
       '2014-12-31T12:00:00.000000000', '2014-12-31T18:00:00.000000000'],
      dtype='datetime64[ns]')

Data variables: (1)

air

(time, lat, lon)

float32

dask.array<chunksize=(2920, 10, 10), meta=np.ndarray>

long_name :: 4xDaily Air temperature at sigma level 995
units :: degK
precision :: 2
GRIB_id :: 11
GRIB_name :: TMP
var_desc :: Air temperature
dataset :: NMC Reanalysis
level_desc :: Surface
statistic :: Individual Obs
parent_stat :: Other
actual_range :: [185.16 322.1 ]





  
      Array  Chunk 
  
  
     Bytes  15.48 MB   1.17 MB 
     Shape  (2920, 25, 53)   (2920, 10, 10) 
     Count  44 Tasks  18 Chunks 
     Type  float32  numpy.ndarray

Attributes: (5)
Conventions :
COARDS
title :
4x daily NMC reanalysis (1948)
description :
Data is from NMC initialized reanalysis (4x/day). These are the 0.9950 sigma level values.
platform :
Model
references :
http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis.html

[11]:

# regridder(ds_bad)  # uncomment this line to see the error message

[12]:

# besides rechunking data properly, another simple fix is to convert to numpy array without chunking
regridder(ds_bad.load())

using dimensions ('lat', 'lon') from data variable air as the horizontal dimensions for this dataset.