Add Hybrid AutoODE

.gitignore

@@ -13,6 +13,11 @@ lightning_logs/
 # pytest coverage files
 .coverage*
 
+# Sphinx
+docs/build/
+.vscode/
+examples/ode/lorenz63/.ipynb_checkpoints/
+
 # build artifacts
 build/
 dist/

examples/ode/README.md

@@ -0,0 +1,19 @@
+In this folder, we train the ODESolver model on three different ODEs. Our goal is to estimate all the parameters here.
+
+Lorenz63:
+
+![alt text](https://wikimedia.org/api/rest_v1/media/math/render/svg/7928004d58943529a7be774575a62ca436a82a7f)
+
+Parameters to estimate: $\sigma, \rho, \beta$
+
+<!-- Duffing equation:
+
+![alt text](https://wikimedia.org/api/rest_v1/media/math/render/svg/4881d84893e137772068573bb1218fc1e2b295cd)
+
+Parameters to estimate: $\alpha, \beta, \gamma, \delta, \omega$ -->
+
+SEIR:
+
+![alt text](https://miro.medium.com/max/1056/1*dXCHv_pSYiMG90efXiFNPQ.png)
+
+Parameters to estimate: $\alpha, \beta, \gamma$

torchts/nn/models/hybridode.py

@@ -0,0 +1,58 @@
+import torch
+from torch import nn
+
+from torchts.nn.models.ode import ODESolver
+
+
+class HybridODENet(ODESolver):
+    def __init__(
+        self,
+        ode,
+        dnns,
+        init_vars,
+        init_coeffs,
+        dt,
+        solver="euler",
+        outvar=None,
+        **kwargs,
+    ):
+        super().__init__(ode, init_vars, init_coeffs, dt, solver, outvar, **kwargs)
+
+        if ode.keys() != init_vars.keys():
+            raise ValueError("Inconsistent keys in ode and init_vars")
+
+        if solver == "euler":
+            self.solver = self.euler
+        else:
+            raise ValueError(f"Unrecognized solver {solver}")
+
+        for name, value in init_coeffs.items():
+            self.register_parameter(name, nn.Parameter(torch.tensor(value)))
+
+        self.ode = ode
+        self.dnns = dnns
+        self.var_names = ode.keys()
+        self.init_vars = {
+            name: torch.tensor(value, device=self.device)
+            for name, value in init_vars.items()
+        }
+        self.coeffs = {name: param for name, param in self.named_parameters()}
+        self.outvar = self.var_names if outvar is None else outvar
+        self.dt = dt
+
+    def euler(self, nt):
+        pred = {name: value.unsqueeze(0) for name, value in self.init_vars.items()}
+
+        for n in range(nt - 1):
+            # create dictionary containing values from previous time step
+            prev_val = {var: pred[var][[n]] for var in self.var_names}
+
+            for var in self.var_names:
+                new_val = (
+                    prev_val[var]
+                    + self.ode[var](prev_val, self.coeffs, self.dnns) * self.dt
+                )
+                pred[var] = torch.cat([pred[var], new_val])
+
+        # reformat output to contain desired (observed) variables
+        return torch.stack([pred[var] for var in self.outvar], dim=1)

torchts/utils/data.py

@@ -122,3 +122,13 @@ def sliding_window(tensor, lags, horizon=1, dim=0, step=1):
         x, y = data[:, [lag - 1 for lag in lags]], data[:, -1]
 
     return x, y
+
+
+def generate_ode_dataset(data, num_steps):
+    n = data.shape[0]
+    y_2d = data[:num_steps, :]
+    y = y_2d.view(1, y_2d.shape[0], y_2d.shape[1])
+    for i in range(1, n - num_steps):
+        y_2d = data[i : i + num_steps, :]
+        y = torch.cat((y, y_2d.view(1, y_2d.shape[0], y_2d.shape[1])), dim=0)
+    return data[: n - num_steps, :], y