-
Notifications
You must be signed in to change notification settings - Fork 297
Using XArray and dask in satpy
DEPRECATED: See https://satpy.readthedocs.io/en/latest/dev_guide/xarray_migration.html for the newest version.
import xarray as xrXArray's DataArray is now the standard data structure for arrays in satpy. They allow the array to define dimensions, coordinates, and attributes (that we use for the metadata).
To create such an array, you can do for example
my_dataarray = xr.DataArray(my_data, dims=['y', 'x'],
coords={'x': np.arange(...)},
attrs={'sensor': 'olci'})my_data can be a regular numpy array, a numpy memmap, or, if you want to keep things lazy, a dask array (more on dask later).
In satpy, the dimensions of the arrays should include
-
xfor the x or column or pixel dimension -
yfor the y or row or line dimension -
bandsfor composites -
timecan also be provided, but we have limited support for it at the moment. Use metadata for common cases (start_time,end_time)
Dimensions are accessible through my_dataarray.dims. To get the size of a given dimension, use sizes:
my_dataarray.sizes['x']Coordinates can be defined for those dimensions when it makes sense:
-
xandy: they are usually defined when the data's area is an AreaDefinition, and the contain the projection coordinates in x and y. -
bands: they contain the letter of the color they represent, eg['R', 'G', 'B']for an RGB composite.
This allows then to select for example a single band like this:
red = my_composite.sel(bands='R')or even multiple bands:
red_and_blue = my_composite.sel(bands=['R', 'B'])To access the coordinates of the data array, use the following syntax:
x_coords = my_dataarray['x']
my_dataarray['y'] = np.arange(...)Most of the time, satpy will fill the coordinates for you, so you just need to provide the dimension names.
To save metadata, we use the .attrs dictionary.
my_dataarray.attrs['platform_name'] = 'Sentinel-3A'Some metadata that should always be present in our dataarrays:
-
areathe area of the dataset. This should be handled in the reader. -
start_time,end_time sensor
DataArrays work with regular arithmetic operation as one would expect of eg numpy arrays, with the exception that using an operator on two DataArrays requires both arrays to share the same dimensions, and coordinates if those are defined.
For mathematical functions like cos or log, use the ufuncs module:
import xarray.ufuncs as xu
cos_zen = xu.cos(zen_xarray)Note that the xu.something function also work on numpy arrays.
In DataArrays, masked data is represented with NaN values. Hence the default type is float64, but float32 works also in this case. XArray can't handle masked data for integer data, but in satpy we try to use the special _FillValue attribute (in .attrs) to handle this case. If you come across a case where this isn't handled properly, contact us.
Masking data from a condition can be done with:
result = my_dataarray.where(my_dataarray > 5)Result is then analogous to my_dataarray, with values lower or equal to 5 replaced by NaNs.
http://xarray.pydata.org/en/stable/generated/xarray.DataArray.html#xarray.DataArray
import dask.array as daThe data part of the DataArrays we use in satpy are mostly dask Arrays. That allows lazy and chunked operations for efficient processing.
To create a dask array from a numpy array, one can call the from_array function:
darr = da.from_array(my_numpy_array, chunks=4096)The chunks keyword tells dask the size of a chunk of data. If the numpy array is 3-dimensional, the chunk size provide above means that one chunk will be 4096x4096x4096 elements. To prevent this, one can provide a tuple:
darr = da.from_array(my_numpy_array, chunks=(4096, 1024, 2))meaning a chunk will be 4096x1024x2 elements in size.
Even more detailed sizes for the chunks can be provided if needed, see the dask documentation.
To avoid loading the data into memory when creating a dask array, other kinds of arrays can be passed to from_array. For example, a numpy memmap allows dask to know where the data is, and will only be loaded when the actual values need to be computed. Another example is a hdf5 variable read with h5py.
Certain operations are easiest to perform on dask arrays by themselves, especially when certain functions are only available from the dask library. In these cases you can operate on the dask array beneath the DataArray and create a new DataArray when done. Note dask arrays do not support in-place operations. In-place operations on xarray DataArrays will reassign the dask array automatically.
dask_arr = my_dataarray.data
dask_arr = dask_arr + 1
# ... other non-xarray operations ...
new_dataarr = xr.DataArray(dask_arr, dims=my_dataarray.dims, attrs=my_dataarray.attrs.copy())Or if the operation should be assigned back to the original DataArray (if and only if the data is the same size):
my_dataarray.data = dask_arrSome procedural generation function are available in dask, eg meshgrid, arange, or random.random.
Regular arithmetic operations are provided, and generate another dask array.
arr1 = da.random.uniform(0, 1000, size=(1000, 1000), chunks=100)
arr2 = da.random.uniform(0, 1000, size=(1000, 1000), chunks=100)
arr1 + arr2
# -> dask.array<add, shape=(1000, 1000), dtype=float64, chunksize=(100, 100)>In order to compute the actual data during testing, use the compute method. In normal SatPy operations you will want the data to be evaluated as late as possible to improve performance so compute should only be used when needed.
(arr1 + arr2).compute()
->
array([[ 898.08811639, 1236.96107629, 1154.40255292, ...,
1537.50752674, 1563.89278664, 433.92598566],
[ 1657.43843608, 1063.82390257, 1265.08687916, ...,
1103.90421234, 1721.73564104, 1276.5424228 ],
[ 1620.11393216, 212.45816261, 771.99348555, ...,
1675.6561068 , 585.89123159, 935.04366354],
...,
[ 1533.93265862, 1103.33725432, 191.30794159, ...,
520.00434673, 426.49238283, 1090.61323471],
[ 816.6108554 , 1526.36292498, 412.91953023, ...,
982.71285721, 699.087645 , 1511.67447362],
[ 1354.6127365 , 1671.24591983, 1144.64848757, ...,
1247.37586051, 1656.50487092, 978.28184726]])Dask also provides cos, log and other mathematical function, that you can use with da.cos, da.log. However, since satpy uses xarrays as standard data structure, prefer the xarray functions when possible (they call in turn the dask counterparts when possible).
Some operations are not supported by dask yet or are difficult to convert to take full advantage of dask's multithreaded operations. In these cases you can wrap a function to run on an entire dask array when it is being computed and pass on the result. Note that this requires fully computing all of the dask inputs to the function and are passed as a numpy array or in the case of an XArray DataArray they will be a DataArray with a numpy array underneath. You should NOT use dask functions inside the delayed function.
import dask
import dask.array as da
def _complex_operation(my_arr1, my_arr2):
return my_arr1 + my_arr2
delayed_result = dask.delayed(_complex_operation)(my_dask_arr1, my_dask_arr2)
# to create a dask array to use in the future
my_new_arr = da.from_delayed(delayed_result, dtype=my_dask_arr1.dtype, shape=my_dask_arr1.shape)
Dask Delayed objects can also be computed delayed_result.compute() if the array is not needed or if the function doesn't return an array.
http://dask.pydata.org/en/latest/array-api.html#dask.array.from_delayed
If the complicated operation you need to perform can be vectorized and does not need the entire data array to do its operations you can use da.map_blocks to get better performance than creating a delayed function. Similar to delayed functions the inputs to the function are fully computed DataArrays or numpy arrays, but only the individual chunks of the dask array at a time. Note that map_blocks must be provided dask arrays and won't function properly on XArray DataArrays.
my_new_arr = da.map_blocks(_complex_operation, my_dask_arr1, my_dask_arr2, dtype=my_dask_arr1.dtype)
http://dask.pydata.org/en/latest/array-api.html#dask.array.core.map_blocks
Helpful functions:
map_blocksmap_overlapatopstoretokenizecomputedelayedrechunkvindex