hybrid model level to pressure level: AxisError for time dim expanded

[johannespletzer]

Describe the bug

I did calculate the pressure levels for a DataArray with hybrid model levels with interp_hybrid_to_pressure(data,pres,hyai,hybi,p0=1.). I have to mention that I did a time selection with .isel(time=0) to keep it simple and computation time at a minimum for testing. This worked well and the next step was to calculate it without timestep selection. Here the error message did pop up.

To Reproduce

Working:

import xarray as xr
from geocat.comp import interp_hybrid_to_pressure

ds = xr.open_dataset(<path_to_data_file>)

data = ds[<your_variable>].isel(time=0)
pres = ds[<press_var>].isel(time=0)
hyai = ds.hyai.isel(time=0)
hybi = ds.hybi.isel(time=0)

interp_hybrid_to_pressure(data,pres,hyai,hybi,p0=1.)

Not working:

import xarray as xr
from geocat.comp import interp_hybrid_to_pressure

ds = xr.open_dataset(<path_to_data_file>)

data = ds[<your_variable>].isel(time=[0,1])
pres = ds[<press_var>].isel(time=[0,1])
hyai = ds.hyai.isel(time=0)
hybi = ds.hybi.isel(time=0)

interp_hybrid_to_pressure(data,pres,hyai,hybi,p0=1.)

The error appears for .isel(time=[0] as well. It seems it is related to the time dimension being expanded or not, which also makes a difference. E.g.

data = ds.flxs.isel(time=0).expand_dims('time')
pres =  ds.aps.isel(time=0).expand_dims('time')
hyai = ds.hyai.isel(time=0)
hybi = ds.hybi.isel(time=0)

Some other pieces of information:

1. I tried the same with a conda geocat python3.6 installation. Same error.
2. Since my hyai and hybi vary along time I initially selected one time step only, but then also tried with the whole time line, which did not work (Example of data declaration below).

data = ds.flxs
pres =  ds.aps
hyai = ds.hyai
hybi = ds.hybi

Error message:

---

AxisError Traceback (most recent call last)

in ----> 1 interp_hybrid_to_pressure(data,pres,hyai,hybi,p0=1.) /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/geocat/comp/interp_hybrid_to_pressure.py in interp_hybrid_to_pressure(data, ps, hyam, hybm, p0, new_levels, lev_dim, method) 89 # Apply vertical interpolation 90 # Apply Dask parallelization with xarray.apply_ufunc ---> 91 output = xr.apply_ufunc( 92 _vertical_remap, 93 data, /some_directory/miniconda3/envs/geocat/lib/python3.9/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) 1136 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc 1137 elif any(isinstance(a, DataArray) for a in args): -> 1138 return apply_dataarray_vfunc( 1139 variables_vfunc, 1140 *args, /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/xarray/core/computation.py in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args) 269 270 data_vars = [getattr(a, "variable", a) for a in args] --> 271 result_var = func(*data_vars) 272 273 if signature.num_outputs > 1: /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/xarray/core/computation.py in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args) 722 ) 723 --> 724 result_data = func(*input_data) 725 726 if signature.num_outputs == 1: /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/numpy/lib/function_base.py in __call__(self, *args, **kwargs) 2111 vargs.extend([kwargs[_n] for _n in names]) 2112 -> 2113 return self._vectorize_call(func=func, args=vargs) 2114 2115 def _get_ufunc_and_otypes(self, func, args): /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/numpy/lib/function_base.py in _vectorize_call(self, func, args) 2185 """Vectorized call to `func` over positional `args`.""" 2186 if self.signature is not None: -> 2187 res = self._vectorize_call_with_signature(func, args) 2188 elif not args: 2189 res = func() /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/numpy/lib/function_base.py in _vectorize_call_with_signature(self, func, args) 2226 2227 for index in np.ndindex(*broadcast_shape): -> 2228 results = func(*(arg[index] for arg in args)) 2229 2230 n_results = len(results) if isinstance(results, tuple) else 1 /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/geocat/comp/interp_hybrid_to_pressure.py in _vertical_remap(data, pressure) 85 """Define interpolation function.""" 86 ---> 87 return func_interpolate(new_levels, pressure, data, axis=interp_axis) 88 89 # Apply vertical interpolation /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/metpy/xarray.py in wrapper(*args, **kwargs) 1214 1215 # Evaluate inner calculation -> 1216 result = func(*bound_args.args, **bound_args.kwargs) 1217 1218 # Wrap output based on match and match_unit /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/metpy/interpolate/one_dimension.py in interpolate_1d(x, xp, axis, fill_value, return_list_always, *args) 107 108 # Sort input data --> 109 sort_args = np.argsort(xp, axis=axis) 110 sort_x = np.argsort(x) 111 <__array_function__ internals> in argsort(*args, **kwargs) /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/numpy/core/fromnumeric.py in argsort(a, axis, kind, order) 1110 1111 """ -> 1112 return _wrapfunc(a, 'argsort', axis=axis, kind=kind, order=order) 1113 1114 /some_directory/miniconda3/envs/geocat/lib/python3.9/site-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds) 56 57 try: ---> 58 return bound(*args, **kwds) 59 except TypeError: 60 # A TypeError occurs if the object does have such a method in its AxisError: axis 1 is out of bounds for array of dimension 1

Expected behavior I expected that DaraArrays with time dimension can be calculated as well. This is not stated explicitely in the documentation, but the following sentence from documentation hints at it.

"ps (xarray.DataArray) – A multi-dimensional array of surface pressures (Pa), same time/space shape as data."

Screenshots

OS:

Unix

Environment

Name Version Build Channel

_libgcc_mutex 0.1 conda_forge conda-forge _openmp_mutex 4.5 1_gnu conda-forge aerocalc3 0.10 pypi_0 pypi appdirs 1.4.4 pyh9f0ad1d_0 conda-forge argon2-cffi 20.1.0 py39h3811e60_2 conda-forge async_generator 1.10 py_0 conda-forge attrs 21.2.0 pyhd8ed1ab_0 conda-forge backcall 0.2.0 pyh9f0ad1d_0 conda-forge backports 1.0 py_2 conda-forge backports.functools_lru_cache 1.6.4 pyhd8ed1ab_0 conda-forge binutils_impl_linux-64 2.35.1 h193b22a_2 conda-forge binutils_linux-64 2.35 h67ddf6f_30 conda-forge bleach 3.3.0 pyh44b312d_0 conda-forge bokeh 2.3.2 py39hf3d152e_0 conda-forge brotlipy 0.7.0 py39h3811e60_1001 conda-forge c-ares 1.17.1 h7f98852_1 conda-forge ca-certificates 2020.12.5 ha878542_0 conda-forge cairo 1.14.12 h8948797_3 anaconda cartopy 0.19.0.post1 py39h3b23250_0 conda-forge certifi 2020.12.5 py39hf3d152e_1 conda-forge cf_xarray 0.5.2 pyh6c4a22f_0 conda-forge cffi 1.14.5 py39he32792d_0 conda-forge cftime 1.4.1 pypi_0 pypi chardet 4.0.0 py39hf3d152e_1 conda-forge click 8.0.0 py39hf3d152e_0 conda-forge cloudpickle 1.6.0 py_0 conda-forge cryptography 3.4.7 py39hbca0aa6_0 conda-forge cycler 0.10.0 py_2 conda-forge cytoolz 0.11.0 py39h3811e60_3 conda-forge dask 2021.5.0 pyhd8ed1ab_0 conda-forge dask-core 2021.5.0 pyhd8ed1ab_0 conda-forge dbus 1.13.18 hb2f20db_0 anaconda decorator 5.0.9 pyhd8ed1ab_0 conda-forge defusedxml 0.7.1 pyhd8ed1ab_0 conda-forge distributed 2021.5.0 py39hf3d152e_0 conda-forge entrypoints 0.3 pyhd8ed1ab_1003 conda-forge eofs 1.4.0 py_0 conda-forge expat 2.3.0 h9c3ff4c_0 conda-forge fontconfig 2.13.0 h9420a91_0 anaconda freetype 2.10.4 h0708190_1 conda-forge fribidi 1.0.10 h7b6447c_0 anaconda fsspec 2021.5.0 pyhd8ed1ab_0 conda-forge gcc_impl_linux-64 9.3.0 h70c0ae5_19 conda-forge gcc_linux-64 9.3.0 hf25ea35_30 conda-forge geocat-comp 2021.05.1 py_0 ncar geocat-f2py 2021.05.1 py39h3d0eb6f_0 ncar geos 3.9.1 h9c3ff4c_2 conda-forge gettext 0.19.8.1 h0b5b191_1005 conda-forge gfortran_impl_linux-64 9.3.0 hc4a2995_19 conda-forge gfortran_linux-64 9.3.0 hdc58fab_30 conda-forge glib 2.56.2 hd408876_0 anaconda graphite2 1.3.14 h23475e2_0 anaconda graphviz 2.40.1 h21bd128_2 anaconda gst-plugins-base 1.14.0 hbbd80ab_1 gstreamer 1.14.0 hb453b48_1 anaconda harfbuzz 1.8.8 hffaf4a1_0 anaconda heapdict 1.0.1 py_0 conda-forge icu 58.2 hf484d3e_1000 conda-forge idna 2.10 pyh9f0ad1d_0 conda-forge importlib-metadata 4.0.1 py39hf3d152e_0 conda-forge importlib_metadata 4.0.1 hd8ed1ab_0 conda-forge importlib_resources 5.1.3 py39hf3d152e_0 conda-forge ipykernel 5.5.5 py39hef51801_0 conda-forge ipython 7.23.1 py39hef51801_0 conda-forge ipython_genutils 0.2.0 py_1 conda-forge ipywidgets 7.6.3 pyhd3deb0d_0 conda-forge jedi 0.18.0 py39hf3d152e_2 conda-forge jinja2 3.0.0 pyhd8ed1ab_0 conda-forge jpeg 9d h36c2ea0_0 conda-forge jsonschema 3.2.0 pyhd8ed1ab_3 conda-forge jupyter 1.0.0 py39hf3d152e_6 conda-forge jupyter_client 6.1.12 pyhd8ed1ab_0 conda-forge jupyter_console 6.4.0 pyhd8ed1ab_0 conda-forge jupyter_core 4.7.1 py39hf3d152e_0 conda-forge jupyterlab_pygments 0.1.2 pyh9f0ad1d_0 conda-forge jupyterlab_widgets 1.0.0 pyhd8ed1ab_1 conda-forge kernel-headers_linux-64 2.6.32 h77966d4_13 conda-forge kiwisolver 1.3.1 py39h1a9c180_1 conda-forge krb5 1.19.1 hcc1bbae_0 conda-forge lcms2 2.12 hddcbb42_0 conda-forge ld_impl_linux-64 2.35.1 hea4e1c9_2 conda-forge libblas 3.9.0 9_openblas conda-forge libcblas 3.9.0 9_openblas conda-forge libcurl 7.76.1 h2574ce0_2 conda-forge libedit 3.1.20191231 he28a2e2_2 conda-forge libev 4.33 h516909a_1 conda-forge libffi 3.3 h58526e2_2 conda-forge libgcc-devel_linux-64 9.3.0 h7864c58_19 conda-forge libgcc-ng 9.3.0 h2828fa1_19 conda-forge libgfortran-ng 9.3.0 hff62375_19 conda-forge libgfortran5 9.3.0 hff62375_19 conda-forge libgomp 9.3.0 h2828fa1_19 conda-forge libiconv 1.16 h516909a_0 conda-forge liblapack 3.9.0 9_openblas conda-forge libnghttp2 1.43.0 h812cca2_0 conda-forge libopenblas 0.3.15 pthreads_h8fe5266_0 conda-forge libpng 1.6.37 h21135ba_2 conda-forge libsodium 1.0.18 h36c2ea0_1 conda-forge libssh2 1.9.0 ha56f1ee_6 conda-forge libstdcxx-ng 9.3.0 h6de172a_19 conda-forge libtiff 4.2.0 hdc55705_1 conda-forge libuuid 1.0.3 h1bed415_2 anaconda libwebp-base 1.2.0 h7f98852_2 conda-forge libxcb 1.13 h7f98852_1003 conda-forge libxml2 2.9.10 hb55368b_3 locket 0.2.0 py_2 conda-forge lz4-c 1.9.3 h9c3ff4c_0 conda-forge markupsafe 2.0.0 py39h3811e60_0 conda-forge matplotlib 3.4.2 py39hf3d152e_0 conda-forge matplotlib-base 3.4.2 py39h2fa2bec_0 conda-forge matplotlib-inline 0.1.2 pyhd8ed1ab_2 conda-forge metpy 1.0.1 pyhd8ed1ab_0 conda-forge mistune 0.8.4 py39h3811e60_1003 conda-forge msgpack-python 1.0.2 py39h1a9c180_1 conda-forge nbclient 0.5.3 pyhd8ed1ab_0 conda-forge nbconvert 6.0.7 py39hf3d152e_3 conda-forge nbformat 5.1.3 pyhd8ed1ab_0 conda-forge ncurses 6.2 h58526e2_4 conda-forge nest-asyncio 1.5.1 pyhd8ed1ab_0 conda-forge netcdf4 1.5.6 pypi_0 pypi notebook 6.3.0 pyha770c72_1 conda-forge numpy 1.20.2 py39hdbf815f_0 conda-forge olefile 0.46 pyh9f0ad1d_1 conda-forge openjpeg 2.4.0 hb52868f_1 conda-forge openssl 1.1.1k h7f98852_0 conda-forge packaging 20.9 pyh44b312d_0 conda-forge pandas 1.2.4 py39hde0f152_0 conda-forge pandoc 2.13 h7f98852_0 conda-forge pandocfilters 1.4.2 py_1 conda-forge pango 1.42.4 h049681c_0 anaconda parso 0.8.2 pyhd8ed1ab_0 conda-forge partd 1.2.0 pyhd8ed1ab_0 conda-forge pcre 8.44 he1b5a44_0 conda-forge pexpect 4.8.0 pyh9f0ad1d_2 conda-forge pickleshare 0.7.5 py_1003 conda-forge pillow 8.1.2 py39hf95b381_1 conda-forge pint 0.17 pyhd8ed1ab_0 conda-forge pip 21.1.1 pyhd8ed1ab_0 conda-forge pixman 0.40.0 h7b6447c_0 anaconda pooch 1.3.0 pyhd8ed1ab_0 conda-forge proj 7.2.0 h277dcde_2 conda-forge prometheus_client 0.10.1 pyhd8ed1ab_0 conda-forge prompt-toolkit 3.0.18 pyha770c72_0 conda-forge prompt_toolkit 3.0.18 hd8ed1ab_0 conda-forge psutil 5.8.0 py39h3811e60_1 conda-forge pthread-stubs 0.4 h36c2ea0_1001 conda-forge ptyprocess 0.7.0 pyhd3deb0d_0 conda-forge pycparser 2.20 pyh9f0ad1d_2 conda-forge pygments 2.9.0 pyhd8ed1ab_0 conda-forge pyopenssl 20.0.1 pyhd8ed1ab_0 conda-forge pyparsing 2.4.7 pyh9f0ad1d_0 conda-forge pyproj 3.0.1 py39h50a64a9_0 conda-forge pyqt 5.9.2 py39h2531618_6 pyrsistent 0.17.3 py39h3811e60_2 conda-forge pyshp 2.1.3 pyh44b312d_0 conda-forge pysocks 1.7.1 py39hf3d152e_3 conda-forge python 3.9.4 hffdb5ce_0_cpython conda-forge python-dateutil 2.8.1 py_0 conda-forge python-graphviz 0.16 pyhd3deb0d_1 conda-forge python_abi 3.9 1_cp39 conda-forge pytz 2021.1 pyhd8ed1ab_0 conda-forge pyyaml 5.4.1 py39h3811e60_0 conda-forge pyzmq 22.0.3 py39h37b5a0c_1 conda-forge qt 5.9.7 h5867ecd_1 qtconsole 5.1.0 pyhd8ed1ab_0 conda-forge qtpy 1.9.0 py_0 conda-forge readline 8.1 h46c0cb4_0 conda-forge requests 2.25.1 pyhd3deb0d_0 conda-forge scipy 1.6.3 py39hee8e79c_0 conda-forge send2trash 1.5.0 py_0 conda-forge setuptools 49.6.0 py39hf3d152e_3 conda-forge shapely 1.7.1 py39ha61afbd_4 conda-forge sip 4.19.13 py39h2531618_0 six 1.16.0 pyh6c4a22f_0 conda-forge sortedcontainers 2.4.0 pyhd8ed1ab_0 conda-forge sqlite 3.35.5 h74cdb3f_0 conda-forge sysroot_linux-64 2.12 h77966d4_13 conda-forge tblib 1.7.0 pyhd8ed1ab_0 conda-forge terminado 0.9.4 py39h06a4308_0 testpath 0.4.4 py_0 conda-forge tk 8.6.10 h21135ba_1 conda-forge toolz 0.11.1 py_0 conda-forge tornado 6.1 py39h3811e60_1 conda-forge traitlets 5.0.5 py_0 conda-forge typing_extensions 3.7.4.3 py_0 conda-forge tzdata 2021a he74cb21_0 conda-forge urllib3 1.26.4 pyhd8ed1ab_0 conda-forge wcwidth 0.2.5 pyh9f0ad1d_2 conda-forge webencodings 0.5.1 py_1 conda-forge wheel 0.36.2 pyhd3deb0d_0 conda-forge widgetsnbextension 3.5.1 py39hf3d152e_4 conda-forge xarray 0.18.0 pyhd8ed1ab_0 conda-forge xorg-libxau 1.0.9 h7f98852_0 conda-forge xorg-libxdmcp 1.1.3 h7f98852_0 conda-forge xz 5.2.5 h516909a_1 conda-forge yaml 0.2.5 h516909a_0 conda-forge yapf 0.31.0 pyhd3eb1b0_0 zeromq 4.3.4 h9c3ff4c_0 conda-forge zict 2.0.0 py_0 conda-forge zipp 3.4.1 pyhd8ed1ab_0 conda-forge zlib 1.2.11 h516909a_1010 conda-forge zstd 1.4.9 ha95c52a_0 conda-forge

[erogluorhan]

Hi @JPletzer ,

Thank you very much for sharing all the details!

We were already working on another bug (maybe the same reason) on this function. We will be sure to address your bug when fixing the function and may need further information from you soon.

Best,

Orhan - GeoCAT

[johannespletzer]

Great, thank you for your quick answer!

[erogluorhan]

Hi @JPletzer ,

We think we addressed this issue and will soon incorporate the fix in a pull request (PR). Can I kindly ask you to give it a try when the PR is up (or maybe merged to the main branch) via source code or via conda distribution when we release the newer version (hopefullylate May or early June).

Best

[johannespletzer]

Hi @erogluorhan ,

thank you for the quick fix! I will verify the new version with my data and share my results. Please let me know when it is ready.

Best

[erogluorhan]

@JPletzer Please see the PR #155 for this issue for now, and I will also let you know when we have released the newer version that reflects this PR.

[johannespletzer]

@erogluorhan

Thank you very much for the notification and your work. I applied the codes changes to interp_hybrid_to_pressure.py and tested my previous posted example. It worked!

A sidenote: There was a calculation time difference, when I used my own pressure levels instead of the __pres_lev_mandatory__. This did not come as a surprise because the latter are approx. 20+ levels, while my own pressure levels number 90. However my workaround time_loop_interp_hybrid_to_pressure that I used in the meantime also relies on the same 90 levels and is nearly as fast as with mandatory levels. Here is an example:

Thank you for solving the issue! I will close it now.

[erogluorhan]

Hi @JPletzer thank you very much for your feedback! It is very helpful.

Could you please share your time_loop_interp_hybrid_to_pressure workaround with us? We could reflect it to the code after investigating and acknowledge your contribution in the docs.

[johannespletzer]

Hi @erogluorhan

here is a simplified version of time_loop_interp_hybrid_to_pressure, where pressure levels have to be added to the function and are not taken from the dataset itself. The run time approx. stays the same. Let me know wether you can confirm my initial test and if your investigation brought someting useful eventually.

Unfortunately the function is not presented correctly in the comment section and the first line is not marked correctly.

def time_loop_interp_hybrid_to_pressure(_ds, _var, _new_levels):

# default pressure value
_p0 = 1. # could be different for other data!

# apply interp_hybrid_to_pressure() at every time step and append result to list
list_val = []
for i in range(len(_ds.time)):

    # declare data for function
    data = _var.isel(time=i)
    pres = _ds.aps.isel(time=i)
    hyax = _ds.hyai.isel(time=i)
    hybx = _ds.hybi.isel(time=i)

    # interpolate hybrid to pressure levels
    interp_val_at_time_step = interp_hybrid_to_pressure(data,
                                                        pres,
                                                        hyax,
                                                        hybx,
                                                        p0=_p0,
                                                        new_levels=_new_levels)

    # append to list
    list_val.append(interp_val_at_time_step)

# concatenate list
interp_var = xr.concat(list_val, dim='time')

return interp_var

[erogluorhan]

Thanks @JPletzer ! I will be looking into this while finalizing the code of the original function.