xarray_regress_two_fields.py

xarray_regress_two_fields.py

# Andrew Williams - 2020
# Written in response to: https://stackoverflow.com/questions/38960903

import numpy as np
import xarray as xr

# Define a linear trend function to use
def linear_trend(x, y):
    pf = np.polyfit(x, y, 1)
    return xr.DataArray(pf[0])

# Generate example DataArrays
da = xr.DataArray(np.random.randn(20, 20, 180, 360),
                  dims=('plev', 'time', 'lat', 'lon'),
                  coords={'plev': np.linspace(0,19, 20), 
                  'time': np.linspace(0,19, 20), 
                  'lat': np.linspace(-90,90,180), 
                  'lon': np.linspace(0,359, 360)})

# Use to create example dataset w two variables
ds = xr.Dataset({'Temp': da1, 'Height': da1})

# Chunk our datarrays so we can use 'dask=parallelized' in xr.apply_ufunc()
dask_height = ds.Height.chunk({'lat':10, 'lon':10, 'time': 10})
dask_temp   = ds.Temp.chunk({'lat':10, 'lon':10, 'time': 10})

# Apply ufunc !
slopes_dask = xr.apply_ufunc(linear_trend,
                             dask_height, dask_temp,
                             vectorize=True,
                             dask='parallelized',
                             input_core_dims=[['plev'], ['plev']], # reduce along 'plev'
                             output_dtypes=['d'],
                             )