apply_ufunc for time dependent function

[zxdawn]

I have one DataFrame which has three columns: time, lon, and lat.

The goal is to predict the location of each point at the end_time, using a look up Dataset (ds, three dimension [time, lon, lat]) which has two DataArrays: wdir, and wspd.

Here's the example code:

import pandas as pd
import xarray as xr
import numpy as np

# !!!! edit len_time to test speed !!!!
len_time = 10

# -------------------------------------- Sample Data --------------------------------------

# Two functions for creating sample data
def random_dates(start, end, n=10):
    # random dates based on length (n)
    start_u = start.value//10**9
    end_u = end.value//10**9

    return pd.to_datetime(np.random.randint(start_u, end_u, n), unit='s')

def predict_loc(lon, lat, wdir, wspd, delta):
    # some calculation depends on inputs, i just simplify them here
    lon2 = lon/2; lat2 = lat/2

    return lon2, lat2

# create 3d DataArray as the look up array
da = xr.DataArray(np.abs(np.random.randn(600).reshape(6, 10, 10)),
                [("time",  pd.date_range("20130101", periods=6, freq="1H")),
                 ("lon", range(10)),
                 ("lat", range(10))
                 ],
)

# create two DataArrays (wind direction and wind speed) and merge them into one Dataset
ds = xr.merge([da.rename('wdir')/2, da.rename('wspd')])

# the end time of prediction
end_time = pd.Timestamp('2013-01-01 03:10')

# create times
times = random_dates(pd.to_datetime('2013-01-01 00:00'),
                     pd.to_datetime('2013-01-01 02:00'),
                     n=len_time)

# create DataFrame
df = pd.DataFrame(times, columns=['time'])
df['lon'] = np.random.randint(low=0, high=9, size=(len_time))
df['lat'] = np.random.randint(low=0, high=9, size=(len_time))

# -------------------------------------- main func --------------------------------------
# get the initial wind data at each location
ds_interp = ds.interp(time=df['time'].to_xarray(), lon=df['lon'].to_xarray(), lat=df['lat'].to_xarray())

# loop over the time dimension and calculate by each time step
# create emtpy list for saving results
lons, lats = [], []

for index,time in enumerate(ds_interp.time):
    ds_init = ds_interp.sel(index=index)

    # because the ds is hourly data, we need to create the hourly time step
    times = np.concatenate(([pd.Timestamp(time.values).to_pydatetime()],
                        pd.date_range(time.dt.ceil('h').values,
                                        end_time.floor('h'), freq='H').to_pydatetime(),
                        [end_time.to_pydatetime()]
    ))

    # calculate the delta seconds
    delta_wind = [t.total_seconds() for t in np.diff(times)]

    lon, lat = ds_init.lon, ds_init.lat

    # predict the location by each time step
    for t_index, time in enumerate(times[:-1]):
        # interpolate to the location at each time
        lon, lat = predict_loc(lon, lat, ds_init['wdir'], ds_init['wspd'], delta_wind[t_index])

    # save the new location at the end of time
    lons.append(float(lon))
    lats.append(float(lat))

It seems the apply_ufunc can't meet this condition. Any suggestions are welcome!

[zxdawn]

I suppose this is related to pandas. And I have tried different methods to speed it up: https://stackoverflow.com/a/71179755/7347925.