Merge branch 'main' into idm_upgrade

scratch/johannes/horner_scheme.py

@@ -0,0 +1,90 @@
+# %%
+
+import numpy as np
+import torch
+from timeit import default_timer as timer
+
+# %%
+
+def powerseries(x, deg):
+    return torch.stack([x**i for i in range(deg+1)],dim=-1)
+
+def improved_powerseries(x, deg):
+    r = torch.ones(*x.shape, deg+1, dtype=torch.float64)
+    for i in range(1,deg+1):
+        r[:, :, i] = r[:, :, i-1] * x
+    return r
+
+def horner_scheme(x, poly):
+    deg = poly.shape[-1]
+    nsteps = x.shape[-1]
+    r = poly[:, -1:].repeat(1, nsteps)
+    for i in range(2, deg+1):
+        r *= x
+        r += poly[:, -i:1-i]
+    return r
+
+# %%
+
+nv = 151
+delta = 10
+n = 20
+
+state_s = torch.rand((nv, 1))
+nan_idx = np.random.choice([True, False], 151)
+state_s[nan_idx] = np.nan
+
+# %%
+
+n_coef = 21
+xpoly = torch.rand((nv, n_coef),dtype=torch.float64)
+ypoly = torch.rand((nv, n_coef),dtype=torch.float64)
+
+ds = delta * torch.arange(1,n+1).repeat(nv,1)
+
+s = ds + state_s
+s = s.type(torch.float64)
+
+smax = s[:, 0]
+smax = smax.unsqueeze(-1)
+
+
+start = timer()
+for _ in range(100):
+    deg = xpoly.shape[-1] - 1
+    expand_sims = powerseries(s, deg) # (nv, n, deg+1)
+    # print(expand_sims.shape)
+    y = (ypoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+    x = (xpoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+end = timer()
+print("Powerseries: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    deg = xpoly.shape[-1] - 1
+    expand_sims = improved_powerseries(s, deg) # (nv, n, deg+1)
+    # print(expand_sims.shape)
+    yp = (ypoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+    xp = (xpoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+end = timer()
+print("Improved Powerseries: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    x_horner = horner_scheme(s, xpoly)
+    y_horner = horner_scheme(s, ypoly)
+end = timer()
+print("Horner: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    x_max = horner_scheme(smax, xpoly)
+    y_max = horner_scheme(smax, ypoly)
+end = timer()
+# print("Horner smax: {}".format((end-start)*1))
+
+assert np.all(np.isclose(xp,x)[~nan_idx])
+assert np.all(np.isclose(yp,y)[~nan_idx])
+assert np.all(np.isclose(x_horner,x)[~nan_idx])
+assert np.all(np.isclose(y_horner,y)[~nan_idx])
+# %%