Extracting annual amplitude and phase of xarray dataset for each pixel using xarray.DataArray.curvefit

Question:

I was wondering how can I extract annual amplitude and phase of xarray time series using xarray.DataArray.curvefit. We fit a 1d function in time to return annual and seasonal amplitude and phase with dims (x, y)

We can use formula similar to https://stats.stackexchange.com/questions/77543/how-do-i-get-the-amplitude-and-phase-for-sine-wave-from-lm-summary but I have difficulty to extract a0,a1,a2 pixelwise for the dataset.I have tried also to convert time to julian day to use instead of ds.time in ds.curvefit but it didn’t work. I really appreciate if you can help me. Similar problem with numpy is solved in Fitting a 3D array of data to a 1D function with numpy or scipy

ds = xr.tutorial.open_dataset('air_temperature')
ds['air2'] = ds.air.copy()
timejulian=ds.time.dt.strftime('%y%j')
def timeseries_function_season(x, a0, a1, a2):
    return a0+(a1*np.cos((2*np.pi/365)*x)+a2*np.sin((2*np.pi/365)*x))
dn = ds.curvefit('time', func=timeseries_function_season)

Answers:

I recieved an answer from:

https://github.com/pydata/xarray/issues/6968#issuecomment-1234368670

I don’t think curvefit works well with datetime coordinates at this point, because everything gets coerced to float by apply_ufunc. Probably room for improvement there. At a minimum you would need to specify good guesses (p0) and/or bounds in terms of datetime64[ns] values.

An easy enough workaround is to assign a separate non-datetime coordinate. This works:

ds = xr.tutorial.open_dataset('air_temperature')
ds = ds.assign_coords({'day':(ds.time - ds.time[0]) / np.timedelta64(1, 'D')}).swap_dims({'time':'day'})

def periodic_season(x, a0, a1, a2, a3):
    # periodic function with both phase amplitude and shift parameters
    return a0 + a1 * np.cos(a2 * x - a3)

dn = ds.curvefit(
    'day', 
    func=periodic_season, 
    p0={'a0':275, 'a1':15, 'a2':2*np.pi/365}
)