feat(euler): add `.step()` and `.order` to numerical solvers
@@ -7,9 +7,15 @@ class NumericalSolver(ABC):
def solve(self, func, u0, t0, tn, n_steps):
pass
+ @abstractmethod
+ def step(self, func, u, t, h):
+ pass
+
class Euler(NumericalSolver):
- def solve(self, func, u0, t0, tn, n_steps):
+ order = 1
+
+ def solve(self, func, u0, t0, tn, n_steps, hypernet=None):
ts = torch.linspace(t0, tn, n_steps + 1)
if isinstance(u0, (float, int)):
u0 = (u0,)
@@ -19,7 +25,9 @@ class Euler(NumericalSolver):
h = (tn - t0) / n_steps
for t in ts[:-1]:
u_old = us[-1]
- u_new = u_old + h * torch.tensor(func(*u_old, t))
+ u_new = u_old + h * torch.tensor(self.step(func, u_old, t, h))
+ if hypernet is not None:
+ u_new += h ** 2 * hypernet(torch.cat([t.reshape(1, 1), u_old.reshape(1, -1)], dim=1)).flatten()
us.append(u_new)
us = torch.stack(us, dim=0)
@@ -28,3 +36,6 @@ class Euler(NumericalSolver):
ans.append(us[:, j])
return ans
+
+ def step(self, func, u, t, h):
+ return func(*u, t)
