bitPhile

map_blocks

2 min read

xarray.map_blocks

It applies custom function on each chunk block of the DataArray or Dataset.

def add_wrapper(d):
    print('------------------- START ---------------------')
    print(type(d))
    print(d.shape)
    print(d)
    print('-------------------- END ------------------------')
    return d + 1
 
xda = ds.air[:2, :5, :5].chunk({"time": 1})
 
 xr.map_blocks(add_wrapper, ds.air).compute()

Output

------------------- START ---------------------
<class 'xarray.core.dataarray.DataArray'>
(0, 0, 0)
<xarray.DataArray 'air' (time: 0, lat: 0, lon: 0)> Size: 0B
array([], shape=(0, 0, 0), dtype=float64)
-------------------- END ------------------------
------------------- START ---------------------
<class 'xarray.core.dataarray.DataArray'>
(1, 5, 5)
<xarray.DataArray 'air' (time: 1, lat: 5, lon: 5)> Size: 200B
array([[[241.2 , 242.5 , 243.5 , 244.  , 244.1 ],
        [243.8 , 244.5 , 244.7 , 244.2 , 243.39],
        [250.  , 249.8 , 248.89, 247.5 , 246.  ],
        [266.5 , 267.1 , 267.1 , 266.7 , 265.9 ],
        [274.5 , 274.29, 274.1 , 274.  , 273.79]]])
-------------------- END ------------------------
------------------- START ---------------------
<class 'xarray.core.dataarray.DataArray'>
(1, 5, 5)
<xarray.DataArray 'air' (time: 1, lat: 5, lon: 5)> Size: 200B
array([[[242.1 , 242.7 , 243.1 , 243.39, 243.6 ],
        [243.6 , 244.1 , 244.2 , 244.1 , 243.7 ],
        [253.2 , 252.89, 252.1 , 250.8 , 249.3 ],
        [269.7 , 269.4 , 268.6 , 267.4 , 266.  ],
        [272.5 , 271.5 , 270.4 , 269.4 , 268.5 ]]])
-------------------- END ------------------------
Out[40]:
<xarray.DataArray 'air' (time: 2, lat: 5, lon: 5)> Size: 400B
array([[[242.2 , 243.5 , 244.5 , 245.  , 245.1 ],
        [244.8 , 245.5 , 245.7 , 245.2 , 244.39],
        [251.  , 250.8 , 249.89, 248.5 , 247.  ],
        [267.5 , 268.1 , 268.1 , 267.7 , 266.9 ],
        [275.5 , 275.29, 275.1 , 275.  , 274.79]],

       [[243.1 , 243.7 , 244.1 , 244.39, 244.6 ],
        [244.6 , 245.1 , 245.2 , 245.1 , 244.7 ],
        [254.2 , 253.89, 253.1 , 251.8 , 250.3 ],
        [270.7 , 270.4 , 269.6 , 268.4 , 267.  ],
        [273.5 , 272.5 , 271.4 , 270.4 , 269.5 ]]])

Observations

  1. Function add_wrapper receives loaded DataArray.
  2. map_blocks calls function first time with random data for output type inferencing. If it is not able to determine the type, we need explicitly specify type using template kwarg.
  3. So in the above example, the function is called three times.
    1. One for inference
    2. First chunk
    3. Second chunk

Graph View