xr.apply_ufunc raise error: Buffer has wrong number of dimensions (expected 2, got 3)

[QINQINKONG]

Hi everybody! I'm trying to combine xarray and dask to calculate cosine zenith angle lazily without loading any big data into memory. My cosine zenith angle function are as follows. It was written by cython and called some other function within it. I didn't include those functions because I think it's irrelevant to the issue.

@cython.wraparound(False)
@cython.boundscheck(False)
def cosz(double[:, ::1] lat,double[:,::1] lon, date):
    doy=((date-date.astype('datetime64[Y]')).astype('timedelta64[D]'))/np.timedelta64(1, 'D')
    hour=((date-date.astype('datetime64[D]')).astype('timedelta64[h]'))/np.timedelta64(1, 'h')
    coz=np.zeros((date.shape[0], lon.shape[0], lon.shape[1]), dtype=np.float64)
    cdef double tod, h, g, Decl
    tod=(np.datetime64(str(date.astype('datetime64[Y]')[0])+'-12-31')-
 np.datetime64(str(date.astype('datetime64[Y]')[0])+'-01-01')+1)/np.timedelta64(1, 'D')
    cdef double[:,:,::1] coz_view=coz
    cdef double[::1] doy_view=doy
    cdef double[::1] hour_view=hour
    cdef Py_ssize_t i, j, k
    cdef Py_ssize_t x_max = date.shape[0]
    cdef Py_ssize_t y_max = lon.shape[0]
    cdef Py_ssize_t z_max = lon.shape[1]
    for i in prange(x_max,nogil=True):
        for j in range(y_max):
            for k in range(z_max):
                h=hourangel(hour_view[i],lon[j,k])
                g=(360.0*((tod)**(-1)))*(doy_view[i]+hour_view[i]/24.0)*(PI/180.0)
                Decl=DECL1-DECL2*math.cos(g)+DECL3*math.sin(g)-DECL4*math.cos(2*g)+DECL5*math.sin(2*g)-DECL6*math.cos(3*g)+DECL7*math.sin(3*g)
                coz_view[i,j,k]=math.sin(Decl)*math.sin(lat[j,k])+math.cos(Decl)*math.cos(lat[j,k])*math.cos(h)
    return coz

Back to python, I want to pass a 2-D lat array, a 2-D lon array, and a 1-D date array (dtype=datetime64). Both the 2-D lat and lon array have the same shape and coordinates which are named as lat and lon as well. The 1-D date array have a coordinate of time, and will be chunked. I hope I can get a 3-D chunked array of cosine zenith angle by cos3d=xr.apply_ufunc(cosz,lat,lon,date,dask="parallelized",output_dtypes=[float]). But, it will raise the error of ValueError('Buffer has wrong number of dimensions (expected 2, got 3)') which really confuse me since my input of 'lat' and 'lon' are indeed 2-D. By the way, the function works well if I directly call it as cosz(lat,lon,date).
I also attached a screen shot of my 'lat', 'lon', and date input here. Appreciate any help! Thanks!

[dcherian]

I think you need to specify input_core_dims=[["lat", "lon"], ["lat","lon"], ["time"]] to prevent "broadcasting" which is converting everything to a 3D variable.

[QINQINKONG]

Thanks @dcherian ! I tried your method. If I only add input_core_dims=[["time"],["lat", "lon"], ["lat","lon"]], it will raise ValueError: applied function returned data with unexpected number of dimensions. Received 3 dimension(s) but expected 0 dimensions with names: (). It seems that the output can only keep "axes corresponding to each broadcast dimension", so after make time,lat,lon input_core_dims, the output is expected to be 0 dimensions.
Then I also add output_core_dims=[["time","lat","lon"]]. But, it still didn't work since dask=parallelized doesn't allow chunks along the core dimension (I chunked date along time).
Finally, the following command works, but it gives me a dask array with only 1 chunk which may mean that I still cannot avoid loading the whole array into memory if I want to use this array for further analysis. I can surely rechunk it, but I wonder whether the rechunk operation itself will need to load the whole data into memory?