diff --git a/scratch/johannes/horner_scheme.py b/scratch/johannes/horner_scheme.py new file mode 100644 index 0000000..c454dfe --- /dev/null +++ b/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]) +# %%