0

Today I ran into the problem of correctly applying xr.apply_ufunc on a 3d xarray dataarray.

I want to filter the array along the 'time' dimension. So the product in the end should have the same dimensions as the input array and should only contain the filtered data in place of the monthly data. To accomplish this I wrote the two following functions:

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

def butter_filt(x,filt_year,fs,order_butter):

    #filt_year = 1 #1 year
    #fs = 12 #monthly data
    #fn = fs/2; # Nyquist Frequency
    fc = (1/filt_year)/2 # cut off frequency 1sample/ 1year = (1/1)/2 equals 1 year filter (two half cycles/sample)
    #fc = (1/2)/2 # cut off frequency 1sample/ 2year = (1/1)/2 equals 2 year filter (two half cycles/sample)
    #fc = (1/4)/2 # cut off frequency 1sample/ 4year = (1/1)/2 equals 4 year filter (two half cycles/sample)
    b, a = signal.butter(order_butter, fc, 'low', fs=fs, output='ba')

    # Check NA values
    co = np.count_nonzero(~np.isnan(x))
    if co < 4: # If fewer than 4 observations return -9999
        return np.empty(x.shape)
    else:
        return signal.filtfilt(b, a, x)

def filtfilt_butter(x,filt_year,fs,order_butter,dim='time'):
    # x ...... xr data array
    # dims .... dimension along which to apply function    
    filt= xr.apply_ufunc(butter_filt, x,filt_year,fs,order_butter,
                     input_core_dims=[[dim], [], [], []],
                       dask='parallelized')

    return filt

x_uv = filtfilt_butter(ds.x,
                   filt_year=1,
                   fs=12,
                   order_butter=2,
                   dim='time')

I tried the butter_filt function which works fine on its own, so there is some kind of problem in the filtfilt_butter. trying to calculate x_uv gives the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<timed exec> in <module>

<ipython-input-86-c470080a0817> in filtfilt_butter(x, filt_year, fs, order_butter, dim)
     20     # x ...... xr data array
     21     # dims .... dimension aong which to apply function
---> 22     filt= xr.apply_ufunc(butter_filt, x,filt_year,fs,order_butter,
     23                          input_core_dims=[[dim], [], [], []],
     24                            dask='parallelized')

~/miniconda3/envs/pyt3_11102018/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, *args)
   1056         )
   1057     elif any(isinstance(a, DataArray) for a in args):
-> 1058         return apply_dataarray_vfunc(
   1059             variables_vfunc,
   1060             *args,

~/miniconda3/envs/pyt3_11102018/lib/python3.8/site-packages/xarray/core/computation.py in 
apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    231 
    232     data_vars = [getattr(a, "variable", a) for a in args]
--> 233     result_var = func(*data_vars)
    234 
    235     if signature.num_outputs > 1:

~/miniconda3/envs/pyt3_11102018/lib/python3.8/site-packages/xarray/core/computation.py in 
 apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, output_sizes, keep_attrs, 
meta, *args)
    621         data = as_compatible_data(data)
    622         if data.ndim != len(dims):
--> 623             raise ValueError(
    624                 "applied function returned data with unexpected "
    625                 "number of dimensions: {} vs {}, for dimensions {}".format(

ValueError: applied function returned data with unexpected number of dimensions: 3 vs 2, for 
dimensions ('deptht', 'km')

How can I fix this?

Jason Aller
  • 3,541
  • 28
  • 38
  • 38
phand
  • 1
  • 2

2 Answers2

1

apply_ufunc expects the output shape to be the initial shape minus the input_core_dims ones plus the output_core_dims ones. In your case you are right to pass time as input core dim as you want to make sure it is moved to be the last dimension so using axis=-1 works correctly.

You therefore need to use output_core_dims to get xarray to expect a 3d output array. You can use time too.

For a more detailed explanation on apply_ufunc arguments, see this other answer

OriolAbril
  • 7,315
  • 4
  • 29
  • 40
0

I found a solution to this problem by defining the input and output dimensions:

def butter_filt(x,filt_year,fs,order_butter):

#filt_year = 1 #1 year
#fs = 12 #monthly data
#fn = fs/2; # Nyquist Frequency
fc = (1/filt_year)/2 # cut off frequency 1sample/ 1year = (1/1)/2 equals 1 year filter (two half cycles/sample)
#fc = (1/2)/2 # cut off frequency 1sample/ 2year = (1/1)/2 equals 2 year filter (two half cycles/sample)
#fc = (1/4)/2 # cut off frequency 1sample/ 4year = (1/1)/2 equals 4 year filter (two half cycles/sample)
b, a = signal.butter(order_butter, fc, 'low', fs=fs, output='ba')

return signal.filtfilt(b, a, x)




def filtfilt_butter(x,filt_year,fs,order_butter,dim='time'):
# x ...... xr data array
# dims .... dimension aong which to apply function    
filt= xr.apply_ufunc(
    butter_filt,  # first the function
    x,# now arguments in the order expected by 'butter_filt'
    filt_year,  # as above
    fs,  # as above
    order_butter,  # as above
    input_core_dims=[["deptht","km","time"], [], [],[]],  # list with one entry per arg
    output_core_dims=[["deptht","km","time"]],  # returned data has 3 dimension
    exclude_dims=set(("time",)),  # dimensions allowed to change size. Must be a set!
    vectorize=True,  # loop over non-core dims
)

return filt
phand
  • 1
  • 2