<jupyter_start><jupyter_text>The Cookbook of Neural ODEsTorchdyn implements out-of-the-box a variety of continuous-depth models. In this introductory tutorial, we show how we can flexibly switch to one model to another with minor efforts. We will touch the following neural ODE variants:* **Vanilla** (depth-invariant) (same as the [torchdyn quickstart](./00_quickstart.html) tutorial)* **Vanilla** (depth-variant)* **Galerkin*** **Data-controlled*** **Stacked (Piece-wise Constant Weights)*** **Stacked with Discrete State Transitions*** **Stable** (WiP)--------------------------------------For more advanced models check out * [Hamiltonian Neural Networks](./06_hamiltonian_nn.html)* [Graph Neural ODEs](./07_graph_neural_de.html)* **Lagrangian Neural Networks** (WiP)* **Continuous Normalizing Flows** (WiP)* **ODE-RNN** (WiP)<jupyter_code>from torchdyn.models import *; from torchdyn.data_utils import *
from torchdyn import *<jupyter_output><empty_output><jupyter_text>**Data:** we use again the moons dataset (with some added noise) simply because all the models will be effective to solve the binary classification problem.<jupyter_code>d = ToyDataset()
X, yn = d.generate(n_samples=512, dataset_type='moons', noise=.4)
import matplotlib.pyplot as plt

colors = ['orange', 'blue'] 
fig = plt.figure(figsize=(3,3))
ax = fig.add_subplot(111)
for i in range(len(X)):
    ax.scatter(X[i,0], X[i,1], color=colors[yn[i].int()])
import torch
import torch.utils.data as data
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

X_train = torch.Tensor(X).to(device)
y_train = torch.LongTensor(yn.long()).to(device)
train = data.TensorDataset(X_train, y_train)
trainloader = data.DataLoader(train, batch_size=len(X), shuffle=False)<jupyter_output><empty_output><jupyter_text>**Learner**<jupyter_code>import torch.nn as nn
import pytorch_lightning as pl

class Learner(pl.LightningModule):
    def __init__(self, model:nn.Module, settings:dict={}):
        super().__init__()
        defaults.update(settings)
        self.settings = defaults
        self.model = model
        self.c = 0
    
    def forward(self, x):
        return self.model(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch      
        y_hat = self.model(x)   
        loss = nn.CrossEntropyLoss()(y_hat, y)
        logs = {'train_loss': loss}
        return {'loss': loss, 'log': logs}   
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=0.005)

    def train_dataloader(self):
        return trainloader<jupyter_output><empty_output><jupyter_text>**Note:** In this notebook we will consider the depth domain $[0,1]$, i.e. $s\in[0,1]$. Note that, for most architectures in *static* settings (aka we do not deal with dynamic data) any other depth domain do not actually afftect the absolute expressiveness of neural ODE is it can be seen as a rescaled/shifted version of $[0,1]$. Please note that, however, other choices of the depth domain can indeed affect the training phaseIt can be accessed and modified through the `s_span` setting of `NeuralDE` class  Vanilla Neural ODE (Depth-Invariant) $$ \left\{    \begin{aligned}        \dot{h}(s) &= f(h(s), \theta)\\        h(0) &= x\\        \hat y & = h(1)    \end{aligned}    \right. \quad s\in[0,1]$$ This model is the same used in [torchdyn quickstart](./00_quickstart.html) tutorial. The vector field is parametrized by a neural network $f$ with *static* parameters $\theta$ and taking as input only the state $h(s)$.<jupyter_code># model settings
settings = {'type':'classic', 'controlled':False, 'solver':'dopri5'}

# vector field parametrized by a NN
f = DEFunc(nn.Sequential(
        nn.Linear(2, 64),
        nn.Tanh(), 
        nn.Linear(64, 2)))

# neural ODE
model = NeuralDE(f, settings).to(device)
# train the neural ODE
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=300, max_nb_epochs=600)
trainer.fit(learn)<jupyter_output>INFO:lightning:
  | Name             | Type       | Params
--------------------------------------------
0 | model            | NeuralDE   | 322   
1 | model.defunc     | DEFunc     | 322   
2 | model.defunc.m   | Sequential | 322   
3 | model.defunc.m.0 | Linear     | 192   
4 | model.defunc.m.1 | Tanh       | 0     
5 | model.defunc.m.2 | Linear     | 130   
6 | model.adjoint    | Adjoint    | 0<jupyter_text>**Plots**<jupyter_code># Evaluate the data trajectories
s_span = torch.linspace(0,1,100)
trajectory = model.trajectory(X_train, s_span).detach().cpu()
# Trajectories in the depth domain
plot_2D_depth_trajectory(s_span, trajectory, yn, len(X))
# Trajectories in the state-space
plot_2D_state_space(trajectory, yn, len(X))
# Trajectories in space-depth
plot_2D_space_depth(s_span, trajectory, yn, len(X))<jupyter_output><empty_output><jupyter_text>Vanilla Neural ODE (Depth-Variant) $$ \left\{    \begin{aligned}        \dot{h}(s) &= f(s, h(s), \theta)\\        h(0) &= x\\        \hat y & = h(1)    \end{aligned}    \right. \quad s\in[0,1]$$This model, contemplated by the seminal paper from [[<NAME>. et al, 2018]](https://arxiv.org/abs/1806.07366), is usually obtained by concatenating $s$ to the state $h$ as input of $f$, i.e. $f([h(s),s])$. For a simple and flexible implementation we developed the class `DepthCat`, which takes care automatically of concatenating $s$ to the state at each call of the `DEFunc`. The final user only needs specify to which dimension $s$ should be appended. Specifically, for an MLP, $h\in\mathbb{R}^{batch\times dims}$ and, thus we should use `DepthCat(1)`.<jupyter_code># model settings
settings = {'type':'classic', 'controlled':False, 'solver':'dopri5'}

# vector field parametrized by a NN
f = DEFunc(nn.Sequential(
        DepthCat(1),
        nn.Linear(3, 64),
        nn.Tanh(),
        DepthCat(1),
        nn.Linear(65, 2)))
# neural ODE
model = NeuralDE(f, settings).to(device)
# train the neural ODE
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=300, max_nb_epochs=600)
trainer.fit(learn)<jupyter_output>INFO:lightning:
  | Name             | Type       | Params
--------------------------------------------
0 | model            | NeuralDE   | 388   
1 | model.defunc     | DEFunc     | 388   
2 | model.defunc.m   | Sequential | 388   
3 | model.defunc.m.0 | DepthCat   | 0     
4 | model.defunc.m.1 | Linear     | 256   
5 | model.defunc.m.2 | Tanh       | 0     
6 | model.defunc.m.3 | DepthCat   | 0     
7 | model.defunc.m.4 | Linear     | 132   
8 | model.adjoint    | Adjoint    | 0<jupyter_text>**Plot**<jupyter_code># Evaluate the data trajectories
s_span = torch.linspace(0,1,100)
trajectory = model.trajectory(X_train, s_span).detach().cpu()
# Trajectories in the depth domain
plot_2D_depth_trajectory(s_span, trajectory, yn, len(X))
# Trajectories in the state-space
plot_2D_state_space(trajectory, yn, len(X))
# Trajectories in space-depth
plot_2D_space_depth(s_span, trajectory, yn, len(X))<jupyter_output><empty_output><jupyter_text>Galerkin-Style Neural ODE Galerkin-style neural ODEs proposed in [<NAME>., <NAME>. et al., 2020](https://arxiv.org/abs/2002.08071) make the weights of the neural ODE to be *depth-varying*, i.e. $\theta=\theta(s)$ obtaining a model of type$$    \left\{    \begin{aligned}        \dot{h}(s) &= f(h(s), \theta(s))\\        h(0) &= x\\        \hat y & = h(1)    \end{aligned}    \right. \quad s\in[0,1]$$where the depth evolution of $\theta(s)$ is parametrized by a trucated eigenfunction expasion, usually an orthogonal basis of some functional space, i.e.$$    \forall i \quad \theta_i(s) = \sum_{j=0}^{m}\gamma_j\odot\psi_j(s)$$The model is then trained by optimizing for the eigenvalues $\gamma_j$. Note that if $\theta\in \mathbb{R}^d$ there will be $d\times m$ final model's parameters. In this tutorial, we use a 10th-order polynomial expansion to model $\theta(s)$.**Note:** In `torchdyn 0.1.0` only Fourier `FourierExpansion` and polynomial `PolyExpansion` bases are currently implemented out of the box. Although we planned to add further bases, e.g.  piece-wise constant functions, radial basis functions, etc., users could create custom `Expansions` by their own. A tutorial notebook will be soon released<jupyter_code># model settings (same as previous models)
settings = {'type':'classic', 'controlled':False, 'solver':'dopri5'}

# vector field parametrized by a NN with "GalLinear" layer
f = DEFunc(nn.Sequential(
            DepthCat(1),
            GalLinear(2, 32, expfunc=PolyExpansion, 
                      n_harmonics=10, n_eig=1),
            nn.Tanh(),
            nn.Linear(32, 2)))
# neural ODE
model = NeuralDE(f, settings).to(device)
# train the neural ODE
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=300, max_nb_epochs=600)
trainer.fit(learn)<jupyter_output>INFO:lightning:
  | Name             | Type       | Params
--------------------------------------------
0 | model            | NeuralDE   | 1 K   
1 | model.defunc     | DEFunc     | 1 K   
2 | model.defunc.m   | Sequential | 1 K   
3 | model.defunc.m.0 | DepthCat   | 0     
4 | model.defunc.m.1 | GalLinear  | 960   
5 | model.defunc.m.2 | Tanh       | 0     
6 | model.defunc.m.3 | Linear     | 66    
7 | model.adjoint    | Adjoint    | 0<jupyter_text>**Plot**<jupyter_code># Evaluate the data trajectories
s_span = torch.linspace(0,1,100)
trajectory = model.trajectory(X_train, s_span).detach().cpu()
# Trajectories in the depth domain
plot_2D_depth_trajectory(s_span, trajectory, yn, len(X))
# Trajectories in the state-space
plot_2D_state_space(trajectory, yn, len(X))
# Trajectories in space-depth
plot_2D_space_depth(s_span, trajectory, yn, len(X))<jupyter_output><empty_output><jupyter_text>Data-Controlled Neural ODE Data-controlled neural ODEs, also introduced in [<NAME>., <NAME>. et al., 2020](https://arxiv.org/abs/2002.08071) consist in feeding to the vector field the input data $x$ (the initial condition of the ODE), leading to$$ \left\{    \begin{aligned}        \dot{h}(s) &= f(h(s), x, \theta)\\        h(0) &= x\\        \hat y & = h(1)    \end{aligned}    \right. \quad s\in[0,1]$$In this way, the neural ODE learns a family of vector fields rather than a single one. In practice, we simply concatenate $x$ to $h$ and simply feed $[h, x]$ to $f$, which should indeed be defined to accomodate the extra $dim(x)$ dimensions. Data-controlled neural ODEs type can be implemented by simply setting `controlled=True`<jupyter_code># model settings - here we set 'controlled' to True
settings = {'type':'classic', 'controlled':True, 'solver':'dopri5'}

# vector field parametrized by a NN which takes as input [h, x]
f = DEFunc(nn.Sequential(
        nn.Linear(4, 64),
        nn.Tanh(),
        nn.Linear(64, 2)))

# neural ODE
model = NeuralDE(f, settings).to(device)
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=300, max_nb_epochs=600)
trainer.fit(learn)<jupyter_output>INFO:lightning:
  | Name             | Type       | Params
--------------------------------------------
0 | model            | NeuralDE   | 450   
1 | model.defunc     | DEFunc     | 450   
2 | model.defunc.m   | Sequential | 450   
3 | model.defunc.m.0 | Linear     | 320   
4 | model.defunc.m.1 | Tanh       | 0     
5 | model.defunc.m.2 | Linear     | 130   
6 | model.adjoint    | Adjoint    | 0<jupyter_text>**Plots**<jupyter_code># Evaluate the data trajectories
s_span = torch.linspace(0,1,100)
trajectory = model.trajectory(X_train, s_span).detach().cpu()
# Trajectories in the depth domain
plot_2D_depth_trajectory(s_span, trajectory, yn, len(X))
# Trajectories in the state-space
plot_2D_state_space(trajectory, yn, len(X))
# Trajectories in space-depth
plot_2D_space_depth(s_span, trajectory, yn, len(X))<jupyter_output><empty_output><jupyter_text>Stacked Neural ODE (aka Piece-wise constant parameters) While looking for an "easier" solution than Galererkin Neural ODEs without trading the expressivity of the model, **stacked neural ODEs** may come in handy. Instead of approximating the solution of an infinite-dimensional problem (e.g. Galerkin-style), we can also use piece-wise constant weight functions. Let us subdivide the depth in $N$ intervals    $$        \mathcal{S} = \prod_{i=0}^{N}[s_i,s_{i+1}], \quad s_0 = 0, \quad s_{N+1}=S  \quad(\forall i  \quad s_i\leq s_{i+1})$$We can then define the weights as piece-wise constant functions $$        \forall i  \quad s\in[s_i,s_{i+1}]\Rightarrow\theta(s) = \theta_i, \quad\theta_i\in\{\theta_1,\theta_2,\theta_N\}$$    The solution of the neural ODE becomes     $$    \begin{aligned}        h(s) &= x + \int_0^S f(h(\tau),\theta(\tau))d\tau \\        &= x + \sum_{i=1}^{N-1}\int_{s_i}^{s_{i+1}} f(h(\tau),\theta_i)d\tau    \end{aligned}$$    This is equivalent to stacking $N-1$ neural ODEs with **identical structure** and **disentangled weights**    $$        \begin{aligned}         \dot h =  f( h(s),\theta_i) \quad s\in[s_i,s_{i+1}]        \end{aligned}$$    which are **stacked sequentially** and trained with *classic* adjoint method.In principle, $f$ can be chosen arbitrarily. Hereafter we consider, e.g. the `data-controlled` case.**NOTE:** Let $\Delta s_i = s_{i+1} - s_i$. Since the individual neural ODEs are *depth-invariant*, we can just solve the ODEs in $[0,\Delta s_i]$<jupyter_code># model settings: here we chose Δ𝑠_𝑖 = 1 ∀𝑖.
settings = {'type':'classic', 'controlled':True, 'backprop_style':'autograd', 'solver':'dopri5', 's_span':torch.linspace(0,1,2)}

# We choose to divide the domain [0,1] in num_pieces=5 intervals
num_pieces = 5

# Stacked depth-invariant Neural ODEs aka Neural ODEs with piecewise constant weights 
nde = []
for i in range(num_pieces):
    nde.append(NeuralDE(DEFunc(nn.Sequential(nn.Linear(4,4), 
                                             nn.Tanh(), 
                                             nn.Linear(4,2)), func_type='classic'), settings))
model = nn.Sequential(*nde).to(device)
# train the neural ODE
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=300, max_nb_epochs=600)
trainer.fit(learn)<jupyter_output>INFO:lightning:
   | Name               | Type       | Params
----------------------------------------------
0  | model              | Sequential | 150   
1  | model.0            | NeuralDE   | 30    
2  | model.0.defunc     | DEFunc     | 30    
3  | model.0.defunc.m   | Sequential | 30    
4  | model.0.defunc.m.0 | Linear     | 20    
5  | model.0.defunc.m.1 | Tanh       | 0     
6  | model.0.defunc.m.2 | Linear     | 10    
7  | model.0.adjoint    | Adjoint    | 0     
8  | model.1            | NeuralDE   | 30    
9  | model.1.defunc     | DEFunc     | 30    
10 | model.1.defunc.m   | Sequential | 30    
11 | model.1.defunc.m.0 | Linear     | 20    
12 | model.1.defunc.m.1 | Tanh       | 0     
13 | model.1.defunc.m.2 | Linear     | 10    
14 | model.1.adjoint    | Adjoint    | 0     
15 | model.2            | NeuralDE   | 30    
16 | model.2.defunc     | DEFunc     | 30    
17 | model.2.defunc.m   | Sequential | 30    
18 | model.2.defunc.m.0 | Linear     | 20    
19 | model.2.defu[...]<jupyter_text>**Plots**<jupyter_code># Evaluate the data trajectories
s_span = torch.linspace(0,1,20)
trajectory = [model[0].trajectory(X_train, s_span)]
for i in range(1, num_pieces):
    trajectory.append(
        model[i].trajectory(trajectory[i-1][-1,:,:], s_span))
                      
trajectory = torch.cat(trajectory, 0).detach().cpu() 
tot_s_span = torch.linspace(0, 5, 100)
# Trajectories in the depth domain
plot_2D_depth_trajectory(tot_s_span, trajectory, yn, len(X))
# Trajectories in the state-space
plot_2D_state_space(trajectory, yn, len(X))
# Trajectories in space-depth
plot_2D_space_depth(tot_s_span, trajectory, yn, len(X))<jupyter_output><empty_output><jupyter_text>Stacked Neural ODEs with Discrete State TransitionsThey are basically the `stacked` where, at the end of an interval $[s_i, s_{i+1}]$, the state is also updated by learned transition maps $g(h, \omega_i)$ parametrized by a NN, i.e.$$  \left\{        \begin{aligned}             \dot h &=  f(h(s),\theta_i) \quad s\in[s_i,s_{i+1}]\\             h^+ &= g(h(s),\omega_i) \quad s = s_{i+1}        \end{aligned}  \right.$$**NOTE:** This class of models is still experimental and, altough<jupyter_code># model settings
settings = {'type':'classic', 'controlled':True, 'backprop_style':'autograd', 'solver':'dopri5', 's_span':torch.linspace(0,1,2)}

# We choose to divide the domain [0,1] in num_pieces=5 intervals
num_pieces = 5

# stacked depth-invariant Neural ODEs
nde = []
for i in range(num_pieces):
    nde.append(NeuralDE(DEFunc(nn.Sequential(nn.Linear(4,4), 
                                             nn.Tanh(), 
                                             nn.Linear(4,2)), func_type='classic'), settings)
              )
    # In this case the state "jump" is parametrized by a simple linear layer
    nde.append(
        nn.Linear(2,2)
    )
    
model = nn.Sequential(*nde).to(device)
# train the neural ODE
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=300, max_nb_epochs=600)
trainer.fit(learn)<jupyter_output>INFO:lightning:
   | Name               | Type       | Params
----------------------------------------------
0  | model              | Sequential | 180   
1  | model.0            | NeuralDE   | 30    
2  | model.0.defunc     | DEFunc     | 30    
3  | model.0.defunc.m   | Sequential | 30    
4  | model.0.defunc.m.0 | Linear     | 20    
5  | model.0.defunc.m.1 | Tanh       | 0     
6  | model.0.defunc.m.2 | Linear     | 10    
7  | model.0.adjoint    | Adjoint    | 0     
8  | model.1            | Linear     | 6     
9  | model.2            | NeuralDE   | 30    
10 | model.2.defunc     | DEFunc     | 30    
11 | model.2.defunc.m   | Sequential | 30    
12 | model.2.defunc.m.0 | Linear     | 20    
13 | model.2.defunc.m.1 | Tanh       | 0     
14 | model.2.defunc.m.2 | Linear     | 10    
15 | model.2.adjoint    | Adjoint    | 0     
16 | model.3            | Linear     | 6     
17 | model.4            | NeuralDE   | 30    
18 | model.4.defunc     | DEFunc     | 30    
19 | model.4.defu[...]<jupyter_text>**Plots**<jupyter_code># Evaluate the data trajectories
s_span = torch.linspace(0,1,20)
trajectory = [model[0].trajectory(X_train, s_span)]
i = 2
c = 0
while c < num_pieces-1:
    x0 = model[i-1](trajectory[c][-1,:,:])
    trajectory.append(
        model[i].trajectory(x0, s_span))
    i += 2
    c += 1
                      
trajectory = torch.cat(trajectory, 0).detach().cpu() 
tot_s_span = torch.linspace(0, 5, 100)
# Trajectories in the depth domain
plot_2D_depth_trajectory(tot_s_span, trajectory, yn, len(X))
# Trajectories in the state-space
plot_2D_state_space(trajectory, yn, len(X))
# Trajectories in space-depth
plot_2D_space_depth(tot_s_span, trajectory, yn, len(X))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14790568949917504, 'fraction_numerical': 0.0303901218798233, 'mean_word_length': 2.9784035570611898, 'pattern_counts': {'":': 0, '<': 61, '<?xml version=': 0, '>': 60, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '402392', 'license': 'Apache-2.0', 'n_tokens_mistral': 6887, 'n_tokens_neox': 6424, 'n_words': 1729, 'path': 'tutorials/01_neural_ode_cookbook.ipynb', 'repo_name': 'cfinlay/torchdyn', 'size': '18802'}