# Return a 3D object alongside 1D object in apply\_ufunc

**URL:** <https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008>\
**Category:** Science\
**Created:** [February 1, 2024, 11:11pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008 "2024-02-01T23:11:15Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![ahuang11](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/ahuang11/32/2290_2.png) [@ahuang11](https://discourse.pangeo.io/u/ahuang11)\
**Post date:** [February 1, 2024, 11:11pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008/1 "2024-02-01T23:11:15Z")

</div>

Currently, I have something similar to this, where the `input_lat` is transformed to `new_lat` (here, +0.25, but in real use case, it’s indeterministic).

Since `xarray_ufunc` doesn’t return a dataset with actual coordinates values, I had to return a second output to retain new\_lat to re-assign the values, but this second output is shaped `time, lat, lon` so I have to `ds["lat"] = new_lat.isel(lon=0, time=0).values`, which I think is inefficient; I simply need it to be shaped `lat`.

Any ideas on how I can modify this to make it more efficient?

```python
import xarray as xr
import numpy as np

air = xr.tutorial.open_dataset("air_temperature")["air"]
input_lat = np.arange(20, 45)

def interp1d_np(data, base_lat, input_lat):
    new_lat = input_lat + 0.25
    return np.interp(input_lat, base_lat, data), new_lat

ds, new_lat = xr.apply_ufunc(
    interp1d_np, # first the function
    air,
    air.lat, # as above
    input_lat, # as above
    input_core_dims=[["lat"], ["lat"], ["lat"]], # list with one entry per arg
    output_core_dims=[["lat"], ["lat"]], # returned data has one dimension
    exclude_dims=set(("lat",)), # dimensions allowed to change size. Must be a set!
    vectorize=True, # loop over non-core dims
)
new_lat = new_lat.isel(lon=0, time=0).values
ds["lat"] = new_lat

```

---

<div class="post-metadata">

**Author:** ![ahuang11](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/ahuang11/32/2290_2.png) [@ahuang11](https://discourse.pangeo.io/u/ahuang11)\
**Post date:** [February 1, 2024, 11:18pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008/2 "2024-02-01T23:18:04Z")

</div>

Edit: updating xarray version fixes this dask issue below.

~Also, I’m wondering why once I convert this to dask, it errors out with the following:~

```auto
File ~/miniconda3/envs/czi/lib/python3.10/site-packages/dask/array/gufunc.py:492, in <genexpr>(.0)
    490 for i, (ocd, oax, meta) in enumerate(zip(output_coredimss, output_axes, metas)):
    491 print(core_shapes)
--> 492 core_output_shape = tuple(core_shapes[d] for d in ocd)
    493 core_chunkinds = len(ocd) * (0,)
    494 output_shape = loop_output_shape + core_output_shape

KeyError: 'dim0'

```

Seems like a bug, but it might just be me using it wrong; I printed out the `core_shapes`:  
`{'dim0_0': 25, 'dim0_1': 25, 'dim0_2': 25}`

```python
import xarray as xr
import numpy as np

air = xr.tutorial.open_dataset("air_temperature")["air"].chunk(time=-1)
input_lat = np.arange(20, 45)

def interp1d_np(data, base_lat, input_lat):
    new_lat = input_lat + 0.25
    return np.interp(input_lat, base_lat, data), new_lat

ds, new_lat = xr.apply_ufunc(
    interp1d_np, # first the function
    air,
    air.lat, # as above
    input_lat, # as above
    input_core_dims=[["lat"], ["lat"], ["lat"]], # list with one entry per arg
    output_core_dims=[["lat"], ["lat"]], # returned data has one dimension
    exclude_dims=set(("lat",)), # dimensions allowed to change size. Must be a set!
    vectorize=True, # loop over non-core dims
    dask="parallelized"
)
new_lat = new_lat.isel(lon=0, time=0).values
ds["lat"] = new_lat

```

---

<div class="post-metadata">

**Author:** ![TomNicholas](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/tomnicholas/32/3189_2.png) [@TomNicholas](https://discourse.pangeo.io/u/TomNicholas)\
**Post date:** [February 2, 2024, 4:45pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008/3 "2024-02-02T16:45:24Z")

</div>

That error seems suspicious, or at least unclear - do you want to raise this on the xarray issue tracker?

---

<div class="post-metadata">

**Author:** ![ahuang11](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/ahuang11/32/2290_2.png) [@ahuang11](https://discourse.pangeo.io/u/ahuang11)\
**Post date:** [February 2, 2024, 5:12pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008/4 "2024-02-02T17:12:37Z")

</div>

I was about to raise an issue, but as I was completing it, it asked me to post versions, and I realized I was using an outdated version (mismatch likely between dask + xarray; I tracked it down to \_enumerate). So, I updated, and the error went away.

I’m still wondering about returning a 1D object though (or even better, getting back the new\_lat assigned automatically)

---

<div class="post-metadata">

**Author:** ![dcherian](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/dcherian/32/2235_2.png) [@dcherian](https://discourse.pangeo.io/u/dcherian)\
**Post date:** [February 2, 2024, 6:28pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008/5 "2024-02-02T18:28:44Z")

</div>

Fairly sure this isn’t possible yet without explicitly specifying the broadcast dimensions as core dimensions on all variables. Can you open an issue please?

---

<div class="post-metadata">

**Author:** ![ahuang11](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.pangeo.io/ahuang11/32/2290_2.png) [@ahuang11](https://discourse.pangeo.io/u/ahuang11)\
**Post date:** [February 2, 2024, 6:48pm UTC](https://discourse.pangeo.io/t/return-a-3d-object-alongside-1d-object-in-apply-ufunc/4008/6 "2024-02-02T18:48:45Z")

</div>

Thanks for the responsiveness; here’s the issue: [Return a 3D object alongside 1D object in apply\_ufunc · Issue #8695 · pydata/xarray · GitHub](https://github.com/pydata/xarray/issues/8695)