Polynomial fits and roots¶
%matplotlib inline
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
Make a curve to be modeled by a polynomial.
t = np.arange(-1, 5, 0.1)
y = np.sin(t)
fig, ax = plt.subplots()
ax.plot(t, y)
Old polynomial API¶
https://numpy.org/doc/stable/reference/routines.polynomials.poly1d.html
The coefficients are ordered from highest power to lowest.
order = 4
p = np.polyfit(t, y, order)
print("Coefficients: ", p)
yp = np.polyval(p, t)
r = np.roots(p)
print("All roots: ", r)
r = [root for root in r if t[0] < root < t[-1]]
print("Roots in original domain: ", r)
fig, ax = plt.subplots()
ax.plot(t, y, label="sine")
ax.plot(t, yp, label=f"poly ({order})")
ax.legend()
for root in r:
ax.axvline(root, color="0.5", lw=0.7)
New polynomial API¶
https://numpy.org/doc/stable/reference/routines.polynomials.html
Coefficients are ordered from lowest power to highest.
This API is more powerful but more complicated, and the documentation is somewhat scattered.
from numpy.polynomial import Polynomial as P
order = 4
pnew = P.fit(t, y, order)
print(pnew)
print("Coefficients: ", pnew.coef)
print("Converted back to original (unscaled) x: ", pnew.convert())
print("Converted coefficients: ", pnew.convert().coef)
print("Domain: ", pnew.domain)
print("All roots: ", pnew.roots())
rnew = [root for root in pnew.roots() if t[0] < root < t[-1]]
print("Roots in original domain: ", rnew)
ypnew = pnew(t)
fig, ax = plt.subplots()
ax.plot(t, y, label="sine")
ax.plot(t, ypnew, label=f"poly ({order})")
ax.legend()
for root in rnew:
ax.axvline(root, color="0.5", lw=0.7)
Starting from xarray¶
At this writing, xarray has polyval and polyfit functions from the old API, but is missing the roots function, so we need to go straight to numpy for that.
For illustration, we will convert t as days from the start of 2022 to datetime64 and then combine it with y in a Dataset.
dt64 = np.datetime64("2022-01-01") + (t * 86400).astype("timedelta64[s]")
ds = xr.Dataset({"y": (["time"], y)}, coords={"time": dt64})
print(dt64[:3])
ds
Notice that xarray is converting our dt64 in seconds to use nanoseconds, which means it cannot be used for applications involving years before 1678.
xpoly = ds.polyfit("time", 4)
I think that warning is because the old API is working directly with time in nanoseconds since the unix epoch, which is very bad from a numeric standpoint.
xpoly
We can get the fit from an xarray function, and then go to numpy for the roots.
xfit = xr.polyval(ds.time, xpoly)
rx = np.roots(xpoly.y_polyfit_coefficients.values)
print("Raw roots: ", rx)
rx_dt64 = np.datetime64("1970-01-01") + rx.astype("timedelta64[ns]")
print("As dt64: ", rx_dt64)
rx_dt64_trimmed = rx_dt64[(dt64[0] < rx_dt64) & (rx_dt64 < dt64[-1])]
print("In the domain: ", rx_dt64_trimmed)
We can convert those datetime64 roots back to days for comparison to the earlier calculations, done with more reasonable numerics.
ns_per_day = 86400 * 1e9
rx_days = (rx_dt64_trimmed - np.datetime64("2022-01-01")).astype(float) / ns_per_day
print("Roots in the domain")
print("Bad numerics with nanoseconds: ", sorted(rx_days))
print("Good numerics (Polynomial): ", sorted(rnew))
print("Good enough (old API): ", sorted(r))
Conclusion: if you are doing a polynomial (or other) fit against a time variable in datetime64, don't use xarray for that calculation. Pull your variables out as numpy arrays, shift and scale as needed, and do your calculations with numpy functions.
Continuing with this example:
dt64_1 = ds.time.values
t_1 = (dt64_1 - np.datetime64("2022-01-01")).astype(float) / ns_per_day
y_1 = ds.y.values
Now you can proceed with either the old or the new API.
# old
p_1 = np.polyfit(t_1, y_1, order)
roots_1 = np.roots(p_1)
# new
p_2 = P.fit(t_1, y_1, order)
roots_2 = p_2.roots()
print("Old API: ", sorted(roots_1)[1:-1])
print("New API: ", sorted(roots_2)[1:-1])
If you want to look at the fit to the original times, you can still do it easily by calculating the fit in numpy (so it's accurate) and then putting it into your original (or another) xarray.
# old API
fit_1 = np.polyval(p_1, t_1)
ds2 = ds.copy()
ds2["y_fit"] = (["time"], fit_1)
ds2["residual"] = ds2["y"] - ds2["y_fit"]
ds2