Suggested method of applying a function to an xarray data array

[Michael Tippett]

I want to apply a function to every grid point of a data array and keep the attributes of the coordinates. Is stack and map the right way to do that?

I tried using stack and map but it lost the coordinate attributes. Adding keep_attrs=True to map gives an error too.

import numpy as np

import xarray as xr

from scipy import stats

from matplotlib import pyplot as plt

%config InlineBackend.figure_format = 'retina'

%matplotlib inline

def ks_test_map(x, t_ind1, t_ind2):

    ks_statistic,p_value = stats.ks_2samp(x.isel(T = t_ind1),x.isel(T = t_ind2))

    return xr.DataArray(p_value) 

global_prcp = 'http://iridl.ldeo.columbia.edu/SOURCES/.NOAA/.NCEP/.CPC/.Merged_Analysis/.monthly/.latest/.ver2/.prcp_est/'

t_range = 'T/(Jan 1979)/(Aug 2020)/RANGE/'

JJA = 'T/3/runningAverage/T/(Jun-Aug)VALUES/'

prcp_url = global_prcp + JJA + 'dods'

prcp_ds = xr.open_dataset(prcp_url, decode_times=False)

da = prcp_ds.prcp_est

t_ind1 = np.full_like(prcp_ds.T.values, False, dtype=bool)

t_ind2 = np.full_like(prcp_ds.T.values, False, dtype=bool)

t_ind1[0:10] = True

t_ind2[10:20] = True

stacked = da.stack(allpoints=['X','Y'])  

stacked_p = stacked.groupby('allpoints').map(ks_test_map, args=(t_ind1, t_ind2))

p_map = stacked_p.unstack('allpoints')

p_map.plot()

[Maximilian Roos]

Hi there,

I replied to this on SO: https://stackoverflow.com/questions/63973873/xarray-map-apply-gives-error-when-keep-attrs-is-specified/63998573#63998573

Hope that helps — at least to retain attrs — feel free to respond there if not.

Max

[Mathias Hauser]

> I want to apply a function to every grid point of a data array and keep the attributes of the coordinates. Is stack and map the right way to do that?

I think you should look at xr.apply_gufunc. The function you want probably looks something like this:

import numpy as np
import xarray as xr
from scipy import stats

def ks_test_gufunc(da1, da2, dim):

    dim = [dim] if isinstance(dim, str) else dim

    statistic, pvalue = xr.apply_ufunc(
        stats.ks_2samp,
        da1,
        da2,
        input_core_dims=[dim, dim],
        output_core_dims=[[], []],
        exclude_dims=set(dim),
        vectorize=True,
        keep_attrs=True,
    )

    statistic.name = "statistic"
    pvalue.name = "pvalue"

    result = xr.merge([statistic, pvalue])

    return result


# example
da = xr.tutorial.open_dataset("air_temperature").air

da1 = da.isel(time=slice(None, 50))
da2 = da.isel(time=slice(-50, None))

ks_test_gufunc(da1, da2, dim="time")

# NOTE you need to pass DataArrays and not Datasets

--
You received this message because you are subscribed to the Google Groups "xarray" group.
To unsubscribe from this group and stop receiving emails from it, send an email to xarray+un...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/xarray/f71b2da7-bcc8-499f-922a-13868a4ddd2fn%40googlegroups.com.

[Michael Tippett]

Thanks I will open an issue about the bad error message.

[Michael Tippett]

Thanks, this looks good. I will take some time to digest what is exactly going on.