<jupyter_start><jupyter_text>Approximating a Reflection Map In this tutorial we show how to painlessly train a neural ODE for approximating the **reflection map** $$    y = -x$$However, this tutorial also serves as a warning on limitations of *vanilla* ODE models which should always be considered when designing your task-specific architecture.\In fact, vanilla neural ODEs cannot approximate (in 1D) functions which require the flows to cross, e.g. the reflection map $y=-x$ as they would brake the uniqueness of solutions (and thus the determinism).\As we show later, one way to overcome this issue is to employ **data-controlled** models.<jupyter_code># allows imports from parent folders
from prep import prep_nbook
prep_nbook()

from torchdyn.models.defunc import *; 
from torchdyn.models.neuralde import *; from torchdyn import *<jupyter_output><empty_output><jupyter_text>**Data**<jupyter_code>import torch
import torch.utils.data as data

n_points = 100
X = torch.linspace(-1,1, n_points).reshape(-1,1)
y = -X

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
X_train, y_train = torch.Tensor(X).to(device), torch.Tensor(y).to(device)

bs = len(X)
train = data.TensorDataset(X_train, y_train)
trainloader = data.DataLoader(train, batch_size=bs, 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.MSELoss()(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.01)

    def train_dataloader(self):
        return trainloader<jupyter_output><empty_output><jupyter_text>Uncontrolled Neural ODE models We first consider the following neural ODE variants: `depth invariant`, `depth variant` ("cat"), `stable`. As we expect, these models will not be able to approximate the reflection map.<jupyter_code>settings = {'type':'classic', 'controlled':False, 'solver':'dopri5'}
# vanilla depth-invariant
func = DEFunc(nn.Sequential(
        nn.Linear(1, 64),
        nn.Tanh(),
        nn.Linear(64,1)
        )
)

# vanilla depth-variant
func_dv = DEFunc(nn.Sequential(DepthCat(1),
        nn.Linear(2, 64),
        nn.Tanh(),
        nn.Linear(64,1)
        )
)

# stable neural flows
func_s = DEFunc(nn.Sequential(
        nn.Linear(1, 64),
        nn.Tanh(),
        nn.Linear(64,1)
        )
)
funcs = [func, func_dv, func_s]
import matplotlib.pyplot as plt

plt.figure(figsize=(24,4))
plot_settings = {'n_grid':30, 'x_span':[-1,1], 'device':device}

for i, f in enumerate(funcs):
    # define the model
    model = NeuralDE(f, settings).to(device)
    if i==2:
        model.settings['type'] = 'stable'
    # train the neural ODE
    learn = Learner(model)
    trainer = pl.Trainer(min_nb_epochs=500, max_nb_epochs=550)
    trainer.fit(learn)
    
    # plot the learned flows
    plt.subplot(1,3,1+i)
    s_span = torch.linspace(0,1,100)
    traj = model.trajectory(X_train, s_span).cpu().detach()
    plot_traj_vf_1D(model, s_span, traj, n_grid=30, x_span=[-1,1], device=device);<jupyter_output>INFO:lightning:
  | Name             | Type       | Params
--------------------------------------------
0 | model            | NeuralDE   | 193   
1 | model.defunc     | DEFunc     | 193   
2 | model.defunc.m   | Sequential | 193   
3 | model.defunc.m.0 | Linear     | 128   
4 | model.defunc.m.1 | Tanh       | 0     
5 | model.defunc.m.2 | Linear     | 65    
6 | model.adjoint    | Adjoint    | 0<jupyter_text>Controlled Neural ODE modelsFollowing the work in [<NAME>., <NAME>., et al., 2020](https://arxiv.org/abs/2002.08071), we can easily approximate the reflection map by **data-controlled neural ODEs**<jupyter_code># define the data-controlled model
settings = {'type':'classic', 'controlled':True, 'solver':'dopri5'}

f = DEFunc(nn.Sequential(
        nn.Linear(2, 64),
        nn.Tanh(),
        nn.Linear(64,1)
        )
)

model = NeuralDE(f, settings).to(device)
# train the neural ODE
learn = Learner(model)
trainer = pl.Trainer(min_nb_epochs=200, max_nb_epochs=250)
trainer.fit(learn)<jupyter_output>INFO:lightning:
  | Name             | Type       | Params
--------------------------------------------
0 | model            | NeuralDE   | 257   
1 | model.defunc     | DEFunc     | 257   
2 | model.defunc.m   | Sequential | 257   
3 | model.defunc.m.0 | Linear     | 192   
4 | model.defunc.m.1 | Tanh       | 0     
5 | model.defunc.m.2 | Linear     | 65    
6 | model.adjoint    | Adjoint    | 0<jupyter_text>**Plots**<jupyter_code># evaluate the trajectories of each data point
s_span = torch.linspace(0,1,100)
traj = model.trajectory(X_train, s_span).cpu().detach()
# plot the depth evolution of the data
fig = plt.figure(figsize=(6,3))
ax = fig.add_subplot(111)
ax.plot(s_span, traj[:,::5,0], color='blue', alpha=.3);
ax.set_xlabel(r"$s$ [Depth]")
ax.set_ylabel(r"$h(s)$ [State]")
ax.set_title(r"Depth-Trajectories of Controlled Neural ODEs")
# plot the evolution of the data in the s-x-h space
from mpl_toolkits.mplot3d import Axes3D

n_grid=30
x_span=[-1,1]
fig = plt.figure(figsize=(5,5))
ax = fig.add_subplot(111, projection='3d')

ss = torch.linspace(s_span[0], s_span[-1], n_grid)
xx = torch.linspace(x_span[0], x_span[-1], n_grid)
S, X = torch.meshgrid(ss,xx)
u_traj = traj[0,:,0].repeat(traj.shape[1],1)
e = torch.abs(y.T - traj[:,:,0])
color = plt.cm.RdYlBu(e)
for i in range(traj.shape[1]):
    tr = ax.scatter(s_span, u_traj[:,i],traj[:,i,0],
                c=color[:,i],alpha=1, cmap=color[:,i],zdir='z')
norm = mpl.colors.Normalize(e.min(),e.max())
plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap='RdYlBu'),
     label='Approximation Error', orientation='horizontal')
ax.set_xlabel(r"$s$ [depth]");
ax.set_ylabel(r"$x$");
ax.set_zlabel(r"$h(s)$");
ax.xaxis._axinfo["grid"]['color'] =  (1,1,1,0)
ax.yaxis._axinfo["grid"]['color'] =  (1,1,1,0)
ax.zaxis._axinfo["grid"]['color'] =  (1,1,1,0)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15876224537396982, 'fraction_numerical': 0.02907790390296999, 'mean_word_length': 3.423658872077029, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '199913', 'license': 'Apache-2.0', 'n_tokens_mistral': 2441, 'n_tokens_neox': 2326, 'n_words': 524, 'path': 'tutorials/03_crossing_trajectories.ipynb', 'repo_name': 'cfinlay/torchdyn', 'size': '6436'}