### Solve ODE with torchode and JIT compilation

Source: https://github.com/martenlienen/torchode/blob/main/README.md

Example demonstrating how to define an ODE, set up a solver with adaptive step-sizing, compile it with torch.compile, and solve an initial value problem. The output includes solver statistics.

```python
import matplotlib.pyplot as pp
import torch
import torchode as to

def f(t, y):
    return -0.5 * y

y0 = torch.tensor([[1.2], [5.0]])
n_steps = 10
t_eval = torch.stack((torch.linspace(0, 5, n_steps), torch.linspace(3, 4, n_steps)))

term = to.ODETerm(f)
step_method = to.Dopri5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
solver = to.AutoDiffAdjoint(step_method, step_size_controller)
jit_solver = torch.compile(solver)

sol = jit_solver.solve(to.InitialValueProblem(y0=y0, t_eval=t_eval))
print(sol.stats)
# => {'n_f_evals': tensor([26, 26]), 'n_steps': tensor([4, 2]),
# =>  'n_accepted': tensor([4, 2]), 'n_initialized': tensor([10, 10])}

pp.plot(sol.ts[0], sol.ys[0])
pp.plot(sol.ts[1], sol.ys[1])
```

--------------------------------

### Install torchode using pip

Source: https://github.com/martenlienen/torchode/blob/main/README.md

Install the latest released version of torchode from PyPI.

```sh
pip install torchode
```

--------------------------------

### Install torchode development version

Source: https://github.com/martenlienen/torchode/blob/main/README.md

Clone the repository and install torchode in editable mode for development.

```sh
git clone https://github.com/martenlienen/torchode
cd torchode
pip install -e .
```

--------------------------------

### Import necessary libraries

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Imports the required PyTorch and TorchODE modules for JIT compilation examples.

```python
import torch
import torch.nn as nn
import torchode as to

torch.random.manual_seed(180819023);
```

--------------------------------

### Initialize ODE solver components

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

Sets up the ODE model, selects a stepping method (Dopri5), and initializes a step size controller, wrapping it with the custom tracker.

```python
dev = torch.device("cpu")
term = to.ODETerm(model)
step_method = to.Dopri5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
step_size_controller = StepSizeControllerTracker(step_size_controller)
adjoint = to.AutoDiffAdjoint(step_method, step_size_controller).to(dev)
```

--------------------------------

### Construct and configure the ODE solver

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Sets up the ODE solver components, including the term, step method, and controller, ensuring the model is fixed for JIT compilation.

```python
dev = torch.device("cpu")
term = to.ODETerm(model)
step_method = to.Dopri5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
adjoint = to.AutoDiffAdjoint(step_method, step_size_controller).to(dev)
```

--------------------------------

### Initialize data and time points

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Sets up the initial state `y0` for the ODE and the time points `t_eval` for evaluation. `y0` is a tensor of random values, and `t_eval` is a linearly spaced tensor.

```python
batch_size = 16
n_steps = 10
y0 = torch.randn((batch_size, n_features))
t_eval = torch.linspace(0.0, 1.0, n_steps)
```

--------------------------------

### JIT compile the ODE solver

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Compiles the configured ODE solver using `torch.jit.script` for optimized execution.

```python
adjoint_jit = torch.jit.script(adjoint)
```

--------------------------------

### Create a custom step size controller tracker

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

A wrapper class that extends `StepSizeController` to log step times (`t`), step sizes (`dt`), and acceptance status (`accept`). It defers actual adaptation to an underlying controller.

```python
class StepSizeControllerTracker(StepSizeController):
    """A wrapper that collects time step and step acceptance information."""

    def __init__(self, controller: StepSizeController):
        super().__init__()

        self.controller = controller

    def init(self, term, problem, method_order: int, dt0, *, stats, args):
        stats["all_t"] = []
        stats["all_dt"] = []
        stats["all_accept"] = []

        return self.controller.init(
            term, problem, method_order, dt0, stats=stats, args=args
        )

    def adapt_step_size(self, t0, dt, y0, step, state, stats):
        accept, dt_next, state, status = self.controller.adapt_step_size(
            t0, dt, y0, step, state, stats
        )

        stats["all_t"].append(t0)
        stats["all_dt"].append(dt)
        stats["all_accept"].append(accept)

        return accept, dt_next, state, status

    def merge_states(self, running, current, previous):
        return self.controller.merge_states(running, current, previous)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Imports PyTorch, neural network modules, torchode, and torchdiffeq. Sets a manual seed for reproducibility.

```python
import torch
import torch.nn as nn
import torchode as to
import torchdiffeq as tde

torch.random.manual_seed(180819023);
```

--------------------------------

### Set up torchode components for ODE solving

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Configures the ODE solver in torchode by defining the `ODETerm`, the `Dopri5` step method, and an `IntegralController` for step size control. An `AutoDiffAdjoint` is created for backpropagation.

```python
term = to.ODETerm(model)
step_method = to.Dopri5(term=term)
step_size_controller = to.IntegralController(atol=1e-9, rtol=1e-7, term=term)
adjoint = to.AutoDiffAdjoint(step_method, step_size_controller)
```

--------------------------------

### Warm-up run for JIT-compiled solver

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Performs a warm-up run for the JIT-compiled solver. This is necessary because the second call can trigger further compilation, which should not be included in performance measurements.

```python
# A second warm up run. For some reason the second call to the compiled solver triggers more compilation
# which we don't want to measure.
adjoint_jit.solve(problem);
```

--------------------------------

### Import necessary modules

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

Imports PyTorch and TorchODE modules. Sets a manual seed for reproducibility.

```python
import torch
import torch.nn as nn
import torchode as to
from torchode.step_size_controllers import StepSizeController

torch.random.manual_seed(180819023);
```

--------------------------------

### Solve ODE using torchode

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Creates an `InitialValueProblem` instance and solves it using the configured `adjoint` solver in torchode. Note that `t_eval` is repeated for each sample in the batch.

```python
problem = to.InitialValueProblem(y0=y0, t_eval=t_eval.repeat((batch_size, 1)))
sol = adjoint.solve(problem)
```

--------------------------------

### Define and solve the initial value problem

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

Creates an initial value problem with specified time points and solves it using the configured adjoint method.

```python
t_eval = torch.tile(torch.linspace(0.0, 3.0, 10), (batch_size, 1))
problem = to.InitialValueProblem(y0=torch.zeros((batch_size, n_features)).to(dev), t_eval=t_eval.to(dev))

sol = adjoint.solve(problem)
```

--------------------------------

### Inspect logged step sizes

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

Prints the recorded step sizes from the `sol.stats` dictionary, showing the values collected by the `StepSizeControllerTracker`.

```python
print(torch.stack(sol.stats["all_dt"]))
```

--------------------------------

### PIDController for Stiff Dynamics

Source: https://github.com/martenlienen/torchode/blob/main/docs/step-size-controllers.md

Employ PIDController for stiff dynamics to potentially save solver steps. It offers more general control than IntegralController.

```python
step_size_controller = to.PIDController(atol=1e-6, rtol=1e-3, pcoeff=0.2, icoeff=0.5, dcoeff=0.0, term=term)
```

--------------------------------

### Display torchode solution statistics

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Prints the statistics gathered by the torchode solver, including the number of function evaluations, steps taken, accepted steps, and initialized steps for each sample in the batch.

```python
sol.stats
```

--------------------------------

### IntegralController for Adaptive Step Size

Source: https://github.com/martenlienen/torchode/blob/main/docs/step-size-controllers.md

Use IntegralController for automatic step size control to maintain error within specified tolerances. It's suitable for most ODEs.

```python
import torchode as to

def f(t, y):
    return -0.5 * y

term = to.ODETerm(f)
step_method = to.Tsit5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
solver = to.AutoDiffAdjoint(step_method, step_size_controller)
problem = to.InitialValueProblem(...)
sol = solver.solve(problem)
```

--------------------------------

### Compare torchdiffeq and torchode solutions

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Calculates and prints the mean and maximum absolute error between the solutions obtained from torchdiffeq and torchode to verify their closeness.

```python
abs_err = (sol.ys - sol_tde.transpose(0, 1)).abs()

abs_err.mean().item(), abs_err.max().item()
```

--------------------------------

### Inspect logged step acceptances

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

Prints the recorded step acceptance statuses from the `sol.stats` dictionary, showing the boolean values collected by the `StepSizeControllerTracker`.

```python
print(torch.stack(sol.stats["all_accept"]))
```

--------------------------------

### FixedStepController for Constant Progress

Source: https://github.com/martenlienen/torchode/blob/main/docs/step-size-controllers.md

Utilize FixedStepController when a known step size is desired or constant solver progress is prioritized over error control. An initial step size must be provided to the solver.

```python
step_size_controller = to.FixedStepController()
solver = to.AutoDiffAdjoint(step_method, step_size_controller)
dt0 = torch.full((batch_size,), 0.01)
sol = solver.solve(problem, dt0=dt0)
```

--------------------------------

### Solve ODE with normal and JIT-compiled solvers

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Solves the defined ODE problem using both the standard and JIT-compiled solvers and prints their statistics and the maximum absolute difference between results.

```python
sol = adjoint.solve(problem)
sol_jit = adjoint_jit.solve(problem)

print(sol.stats)
print(sol_jit.stats)
print("Max absolute difference", float((sol.ys - sol_jit.ys).abs().max()))
```

--------------------------------

### Measure runtime of the JIT-compiled solver

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Measures the execution time of the `solve` method for the JIT-compiled ODE solver after a warm-up run.

```python
%%timeit
```

```python
adjoint_jit.solve(problem)
```

--------------------------------

### Solve ODE using torchdiffeq

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Solves the ordinary differential equation using the `odeint` function from the `torchdiffeq` library with the defined model, initial state, and time points.

```python
sol_tde = tde.odeint(model, y0, t_eval)
```

--------------------------------

### Inspect logged step times

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

Prints the recorded time points from the `sol.stats` dictionary, showing the values collected by the `StepSizeControllerTracker`.

```python
print(torch.stack(sol.stats["all_t"]))
```

--------------------------------

### Measure runtime of the normal solver

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Measures the execution time of the `solve` method for the standard (non-JIT compiled) ODE solver.

```python
%%timeit
```

```python
adjoint.solve(problem)
```

--------------------------------

### Define the initial value problem

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Combines the initial condition `y0` and evaluation points `t_eval` into a problem instance for the ODE solver.

```python
batch_size = 3
t_eval = torch.tile(torch.linspace(0.0, 3.0, 10), (batch_size, 1))
problem = to.InitialValueProblem(y0=torch.zeros((batch_size, 5)).to(dev), t_eval=t_eval.to(dev))
```

--------------------------------

### Define a simple neural ODE model

Source: https://github.com/martenlienen/torchode/blob/main/docs/jit.ipynb

Defines a neural network model that represents the dynamics of an ODE. This model will be part of the JIT-compiled solver.

```python
class Model(nn.Module):
    def __init__(self, n_features, n_hidden):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(n_features, n_hidden),
            nn.Softplus(),
            nn.Linear(n_hidden, n_hidden),
            nn.Softplus(),
            nn.Linear(n_hidden, n_features)
        )
    
    def forward(self, t, y):
        return self.layers(y)

n_features = 5
model = Model(n_features=n_features, n_hidden=32)
```

--------------------------------

### Flatten and Reshape Images for torchode

Source: https://github.com/martenlienen/torchode/blob/main/docs/nd-data.md

Use this pattern when your data, like RGB images, has more than two dimensions. Flatten the input before passing it to torchode and reshape it within the ODE function for CNN processing.

```python
b, c, w, h = 8, 3, 72, 72
y0 = torch.randn(b, w, h)

# Setup solver etc.
# ...

def f(t, y, args)
    c, w, h = args

    # Reshape 2D to B x C x W x H tensor and apply convolutional neural network
    y = some_cnn(y.reshape((-1, c, w, h)))

    # Flatten back to 2D when we return control back to torchode
    return y.flatten(start_dim=1)

term = to.ODETerm(f, with_args=True)
problem = to.InitialValueProblem(y0=y0.flatten(start_dim=1), t_eval=..)
sol = solver.solve(problem, args=(c, w, h))
```

--------------------------------

### Passing Extra Arguments to Dynamics Function

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-args.md

Instantiate `to.ODETerm` with `with_args=True` to pass an arbitrary object as a third argument to your dynamics function `f`. Ensure correct broadcasting of arguments yourself.

```python
import matplotlib.pyplot as pp
import torch
import torchode as to

def f(t, y, decay_rate):
    return -decay_rate[:, None] * y

term = to.ODETerm(f, with_args=True)
step_method = to.Tsit5(term=term)
step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
solver = to.AutoDiffAdjoint(step_method, step_size_controller)

y0 = torch.tensor([[1.2], [5.0]])
n_steps = 10
t_eval = torch.stack((torch.linspace(0, 5, n_steps), torch.linspace(3, 4, n_steps)))
decay_rate = torch.tensor([0.1, 5.0])

problem = to.InitialValueProblem(y0=y0, t_eval=t_eval)
sol = solver.solve(problem, args=decay_rate)

pp.plot(sol.ts[0], sol.ys[0])
pp.plot(sol.ts[1], sol.ys[1])
pp.show()
```

--------------------------------

### Define a two-layer MLP model

Source: https://github.com/martenlienen/torchode/blob/main/docs/torchdiffeq.ipynb

Defines a simple Multi-Layer Perceptron (MLP) with two hidden layers using PyTorch's nn.Module. The model takes `n_features` as input and outputs `n_features`, with `n_hidden` neurons in each hidden layer. Softplus activation is used.

```python
class Model(nn.Module):
    def __init__(self, n_features, n_hidden):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(n_features, n_hidden),
            nn.Softplus(),
            nn.Linear(n_hidden, n_hidden),
            nn.Softplus(),
            nn.Linear(n_hidden, n_features)
        )
    
    def forward(self, t, y):
        return self.layers(y)
    
n_features = 5
model = Model(n_features, 16)
```

--------------------------------

### Define a generic ODE model

Source: https://github.com/martenlienen/torchode/blob/main/docs/extra-stats.ipynb

A simple feedforward neural network model for the ODE. It takes `n_features` and `n_hidden` as input.

```python
class Model(nn.Module):
    def __init__(self, n_features, n_hidden):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(n_features, n_hidden),
            nn.Softplus(),
            nn.Linear(n_hidden, n_hidden),
            nn.Softplus(),
            nn.Linear(n_hidden, n_features)
        )
    
    def forward(self, t, y):
        return self.layers(y)
```

--------------------------------

### torchode citation

Source: https://github.com/martenlienen/torchode/blob/main/README.md

BibTeX entry for citing the torchode library in academic work.

```bibtex
@inproceedings{lienen2022torchode,
  title = {torchode: A Parallel {ODE} Solver for PyTorch},
  author = {Marten Lienen and Stephan G{"u}nnemann},
  booktitle = {The Symbiosis of Deep Learning and Differential Equations II, NeurIPS},
  year = {2022},
  url = {https://openreview.net/forum?id=uiKVKTiUYB0}
}
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.