xarray.apply_ufunc() with GroupBy: unexpected number of dimensions

user_85248 0
Viewed 1443

I am applying a function to a xarray.DataArray using xarray.apply_ufunc(). It works well with some NetCDFs and fails with others that appear to be comparable in terms of dimensions, coordinates, etc. However there must be something different between the NetCDFs that the code works for and the ones where the code fails, and hopefully someone can comment as to what the problem is after seeing the code and some metadata about the files listed below.

The code I'm running to perform the computation is this:

# open the precipitation NetCDF as an xarray DataSet object
dataset = xr.open_dataset(kwrgs['netcdf_precip'])

# get the precipitation array, over which we'll compute the SPI
da_precip = dataset[kwrgs['var_name_precip']]

# stack the lat and lon dimensions into a new dimension named point, so at each lat/lon
# we'll have a time series for the geospatial point, and group by these points
da_precip_groupby = da_precip.stack(point=('lat', 'lon')).groupby('point')

# apply the SPI function to the data array
da_spi = xr.apply_ufunc(indices.spi,
                        da_precip_groupby)

# unstack the array back into original dimensions
da_spi = da_spi.unstack('point')

The NetCDF that works looks like this:

>>> import xarray as xr
>>> ds_good = xr.open_dataset("good.nc")
>>> ds_good
<xarray.Dataset>
Dimensions:  (lat: 38, lon: 87, time: 1466)
Coordinates:
  * lat      (lat) float32 24.5625 25.229166 25.895834 ... 48.5625 49.229168
  * lon      (lon) float32 -124.6875 -124.020836 ... -68.020836 -67.354164
  * time     (time) datetime64[ns] 1895-01-01 1895-02-01 ... 2017-02-01
Data variables:
    prcp     (lat, lon, time) float32 ...
Attributes:
    Conventions:               CF-1.6, ACDD-1.3
    ncei_template_version:     NCEI_NetCDF_Grid_Template_v2.0
    title:                     nClimGrid
    naming_authority:          gov.noaa.ncei
    standard_name_vocabulary:  Standard Name Table v35
    institution:               National Centers for Environmental Information...
    geospatial_lat_min:        24.5625
    geospatial_lat_max:        49.354168
    geospatial_lon_min:        -124.6875
    geospatial_lon_max:        -67.020836
    geospatial_lat_units:      degrees_north
    geospatial_lon_units:      degrees_east
    NCO:                       4.7.1
    nco_openmp_thread_number:  1
>>> ds_good.prcp
<xarray.DataArray 'prcp' (lat: 38, lon: 87, time: 1466)>
[4846596 values with dtype=float32]
Coordinates:
  * lat      (lat) float32 24.5625 25.229166 25.895834 ... 48.5625 49.229168
  * lon      (lon) float32 -124.6875 -124.020836 ... -68.020836 -67.354164
  * time     (time) datetime64[ns] 1895-01-01 1895-02-01 ... 2017-02-01
Attributes:
    valid_min:      0.0
    units:          millimeter
    valid_max:      2000.0
    standard_name:  precipitation_amount
    long_name:      Precipitation, monthly total

The NetCDF that fails looks like this:

>>> ds_bad = xr.open_dataset("bad.nc")   >>> ds_bad
<xarray.Dataset>
Dimensions:  (lat: 38, lon: 87, time: 1483)
Coordinates:
  * lat      (lat) float32 49.3542 48.687534 48.020866 ... 25.3542 24.687532
  * lon      (lon) float32 -124.6875 -124.020836 ... -68.020836 -67.354164
  * time     (time) datetime64[ns] 1895-01-01 1895-02-01 ... 2018-07-01
Data variables:
    prcp     (lat, lon, time) float32 ...
Attributes:
    date_created:              2018-02-15 10:29:25.485927
    date_modified:             2018-02-15 10:29:25.486042
    Conventions:               CF-1.6, ACDD-1.3
    ncei_template_version:     NCEI_NetCDF_Grid_Template_v2.0
    title:                     nClimGrid
    naming_authority:          gov.noaa.ncei
    standard_name_vocabulary:  Standard Name Table v35
    institution:               National Centers for Environmental Information...
    geospatial_lat_min:        24.562532
    geospatial_lat_max:        49.3542
    geospatial_lon_min:        -124.6875
    geospatial_lon_max:        -67.020836
    geospatial_lat_units:      degrees_north
    geospatial_lon_units:      degrees_east
>>> ds_bad.prcp
<xarray.DataArray 'prcp' (lat: 38, lon: 87, time: 1483)>
[4902798 values with dtype=float32]
Coordinates:
  * lat      (lat) float32 49.3542 48.687534 48.020866 ... 25.3542 24.687532
  * lon      (lon) float32 -124.6875 -124.020836 ... -68.020836 -67.354164
  * time     (time) datetime64[ns] 1895-01-01 1895-02-01 ... 2018-07-01
Attributes:
    valid_min:      0.0
    long_name:      Precipitation, monthly total
    standard_name:  precipitation_amount
    units:          millimeter
    valid_max:      2000.0

When I run the code against the first file above it works without error. When using the second file I get errors like this:

multiprocessing.pool.RemoteTraceback:
"""
Traceback (most recent call last):
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/multiprocessing/pool.py", line 119, in worker
    result = (True, func(*args, **kwds))
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/multiprocessing/pool.py", line 44, in mapstar
    return list(map(*args))
  File "/home/paperspace/git/climate_indices/scripts/process_grid_ufunc.py", line 278, in compute_write_spi
    kwargs=args_dict)
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/computation.py", line 974, in apply_ufunc
    return apply_groupby_ufunc(this_apply, *args)
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/computation.py", line 432, in apply_groupby_ufunc
    applied_example, applied = peek_at(applied)
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/utils.py", line 133, in peek_at
    peek = next(gen)
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/computation.py", line 431, in <genexpr>
    applied = (func(*zipped_args) for zipped_args in zip(*iterators))
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/computation.py", line 987, in apply_ufunc
    exclude_dims=exclude_dims)
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/computation.py", line 211, in apply_dataarray_ufunc
    result_var = func(*data_vars)
  File "/home/paperspace/anaconda3/envs/climate/lib/python3.6/site-packages/xarray/core/computation.py", line 579, in apply_variable_ufunc
    .format(data.ndim, len(dims), dims))
ValueError: applied function returned data with unexpected number of dimensions: 1 vs 2, for dimensions ('time', 'point')

Can anyone comment as to what may be the issue?

python-xarray
2 Answers

It turned out that the NetCDF files that were problematic as inputs the latitude coordinate values were in descending order. xarray.apply_ufunc() appears to require that coordinate values be in ascending order, at least in order to avoid this particular issue. This is easily remedied by reversing the offending dimension's coordinate values using NCO's ncpdq command before using the NetCDF file as input to xarray.

user_852480

thank you for your reply.

Sometimes, it seems that by sorting the dimensions in ascending mode may properly solve the problem regarding the xr.apply_ufunc. Nevertheless, there are times in which that maneuver is not sufficient.

Another alternative solution would be to stack the coordinates that will be broadcasted by the external user's function into a new dimension (i.e.: stack 'Longitude' and 'Latitude' dimensions into a new one called 'Grid_Point'). After this stacking, one can do a groupby operation over this new dimension "Grid_Point" and apply the xr.apply_ufunc.

Here is an example of how one can derive the respective statistical moments from a gaussian distribution ('mean' and 'standard deviation') from a temperature dataset in netcdf per pixel based.

import xarray as xr

# http://xarray.pydata.org/en/stable/dask.html
from scipy import stats
from dask.diagnostics import ProgressBar
import numpy as np
import warnings

def get_params_from_distribution(data, distribution='exponweib'):

    distribution = getattr(stats, distribution)

    if np.all(np.isnan(data)):

        with warnings.catch_warnings():
            warnings.filterwarnings(action="ignore")
            try:
                temp_data = distribution.rvs(1, size=10)
            except:
                try:
                    temp_data = distribution.rvs(1, 1, size=10)    

                except:
                    temp_data = distribution.rvs(1, 1, 1, size=10)  


            n_params = len(distribution.fit(temp_data))


        return data[:n_params]

    else:
        return list(distribution.fit(data))

def get_params_vectorized_from_stacked(stacked_data, distribution='exponweib', 
                                       dask='allowed',
                                       input_core_dims='time',
                                       output_core_dims = 'stat_moments',
                                       output_dtypes=[xr.core.dataset.Dataset]):

    kwargs = {'distribution': distribution}


    with ProgressBar():

        da_spi = xr.apply_ufunc(get_params_from_distribution,
                                stacked_data, 
                                exclude_dims={input_core_dims},
                                kwargs=kwargs,
                                input_core_dims=[[input_core_dims]],
                                output_core_dims=[[output_core_dims]],
                                dask=dask,
                                output_dtypes=[output_dtypes]).compute()


    return da_spi

def stack_ds(ds, dims=['lon', 'lat'], stacked_dim_name='point'):

    return ds.stack({stacked_dim_name:dims})

def main_pdf_u_function_getter(ds, 
                               dims_to_stack=['lon', 'lat'], 
                               stacked_dim_name='point', 
                               distribution_name='exponweib',
                               dask='allowed',
                               input_core_dims = 'time',
                               output_core_dims = 'stat_moments',
                               output_dtypes=[float]):

    ds_stacked = stack_ds(ds, dims_to_stack, stacked_dim_name) # observation 1

    ds_groupby = ds_stacked.groupby(stacked_dim_name) # observation 1

    results = get_params_vectorized_from_stacked(ds_groupby, 
                                                 distribution=distribution_name, 
                                                 dask=dask,
                                                 output_core_dims=output_core_dims,
                                                 input_core_dims=input_core_dims,
                                                 output_dtypes=output_dtypes)

    return results.unstack(stacked_dim_name)



if '__main__' == __name__:

    ds = xr.tutorial.open_dataset('air_temperature').sortby(['lat', 'lon', 'time'])


    R = main_pdf_u_function_getter(ds, 
                                   dask='parallelized', 
                                   dims_to_stack=['lon', 'lat'],
                                   stacked_dim_name='point',
                                   distribution_name='norm')

    print(R)
    import matplotlib.pyplot as plt


    fig, ax= plt.subplots(1,2)
    ax = ax.ravel()
    for moment in range(R.dims['stat_moments']):
        R['air'].isel({'stat_moments':moment}).plot(ax=ax[moment], cmap='viridis')

Notice that in the code above, there is a commentary line written "observation 1". Those are the main lines that ensure that the whole algorithms works. It does the stacking over the broadcasting dimensions prior to the ufunction operation.

Despite the given solution (which works), I still don't know why one should do the stacking prior to the xr.apply_ufunc. This is a question that remains unanswered.

Sincerely,

user_86780150
Related