Issue with xr.apply_ufunc

[pecos27]

Hello!

I am trying to find the correlation coefficients of single grid point of (time, lat, lon) data in relation to the entire data field.
Using the rasm air temperature data as an example, I first compute the anomalies (subtracting the time mean). And then I pick one grid point and I am trying to apply scipy's pearson r correlation coefficient function of the grid point to the entire air temp field.
However I am getting an error which I can't seem to figure out:

ds = xr.tutorial.open_dataset('rasm').load()

Tair = ds.Tair

Tair_anom = Tair - Tair.mean(dim='time') # obtaining anomaly by subtracting the time mean

Tair_grid = Tair_anom[:,25,170] # taking one spatial grid as example

plt.pcolormesh(Tair_grid[2])
plt.plot(170,25,'ro',label='gridpoint')
plt.legend()
plt.colorbar();

# defining the function
def pearson_r(x, y):
    pearsonr = stats.pearsonr(x, y)
    return pearsonr

pearsonr = xr.apply_ufunc(pearson_r, ma.masked_invalid(Tair_grid), ma.masked_invalid(Tair_anom), 
                     vectorize = True, dask = 'parallelized', 
                     input_core_dims = [['time'], ['time']], output_dtypes = ['d']).compute()


---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-462-51d86f869b29> in <module>
----> 1 pearsonr = xr.apply_ufunc(pearson_r, ma.masked_invalid(Tair_grid), ma.masked_invalid(Tair_anom), 
      2                      vectorize = True, dask = 'parallelized',
      3                      input_core_dims = [['time'], ['time']], output_dtypes = ['d']).compute()

~/miniconda3/envs/py3_std_maps/lib/python3.8/site-packages/xarray/core/computation.py in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
   1139     else:
   1140         # feed anything else through apply_array_ufunc
-> 1141         return apply_array_ufunc(func, *args, dask=dask)
   1142 
   1143 

~/miniconda3/envs/py3_std_maps/lib/python3.8/site-packages/xarray/core/computation.py in apply_array_ufunc(func, dask, *args)
    789         else:
    790             raise ValueError(f"unknown setting for dask array handling: {dask}")
--> 791     return func(*args)
    792 
    793 

<ipython-input-444-2798ce5ac981> in pearson_r(x, y)
      1 def pearson_r(x, y):
----> 2     pearsonr = stats.pearsonr(x, y)
      3     return pearsonr

~/miniconda3/envs/py3_std_maps/lib/python3.8/site-packages/scipy/stats/stats.py in pearsonr(x, y)
   3854         return dtype(np.sign(x[1] - x[0])*np.sign(y[1] - y[0])), 1.0
   3855 
-> 3856     xmean = x.mean(dtype=dtype)
   3857     ymean = y.mean(dtype=dtype)
   3858 

~/miniconda3/envs/py3_std_maps/lib/python3.8/site-packages/numpy/core/_methods.py in _mean(a, axis, dtype, out, keepdims)
    158             is_float16_result = True
    159 
--> 160     ret = umr_sum(arr, axis, dtype, out, keepdims)
    161     if isinstance(ret, mu.ndarray):
    162         ret = um.true_divide(

TypeError: No loop matching the specified signature and casting was found for ufunc add

I was hoping to output a spatial map showing all the correlation coefficients of the spatial field based on that single grid point. What might be wrong with the function?
I'm still quite new to dask and parallelization, so any tips would be great. thank you!

[max-sixty]

Does xr.corr suffice?

Assuming not, I'm guessing the issue is that ma.masked_invalid is no longer a DataArray. What happens if you do the coercion in the pearson_r function?

[pecos27]

Thank you! Yes, xr.corr appears to be work instead!

[raybellwaves]

You could also take a look at https://xskillscore.readthedocs.io/en/stable/api/xskillscore.pearson_r.html#xskillscore.pearson_r