91 lines
2.1 KiB
Python
91 lines
2.1 KiB
Python
# %%
|
|
|
|
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])
|
|
# %%
|