# NeuralODE

Neural ODE (`astroNN.neuralODE`; Neural Ordinary Differential Equation) module provides numerical integrator implemented in `Tensorflow` for solutions of an ODE system, and can calculate gradient.

## Numerical Integrator

`astroNN` implemented numerical integrator in `Tensorflow`

astroNN.neuralode.odeint.odeint(func=None, x=None, t=None, aux=None, method='dop853', precision=tf.float32, *args, **kwargs)[source]

To computes the numerical solution of a system of first order ordinary differential equations y’=f(x,y). Default precision at float32.

Parameters:
• func (callable) – function of the differential equation, usually take func([position, velocity], time) and return velocity, acceleration

• x (Union([tf.Tensor, numpy.ndarray, list])) – initial x, usually is [position, velocity]

• t (Union([tf.Tensor, numpy.ndarray, list])) – set of times at which one wants the result

• method (str) – numerical integrator to use, available integrators are [‘dop853’, ‘rk4’]

• precision (type) – float precision, tf.float32 or tf.float64

• t – set of times at which one wants the result

Returns:

integrated result

Return type:

tf.Tensor

History:

2020-May-31 - Written - Henry Leung (University of Toronto)

An example integration an ODE for `sin(x)`

``` 1import time
2import pylab as plt
3import numpy as np
4import tensorflow as tf
5from astroNN.shared.nn_tools import cpu_fallback, gpu_memory_manage
6from astroNN.neuralode import odeint
7
8cpu_fallback()
9gpu_memory_manage()
10
11# time array
12t = tf.constant(np.linspace(0, 100, 10000))
13# initial condition
14true_y0 = tf.constant([0., 1.])
15# analytical ODE system for sine wave [x, t] -> [v, a]
16ode_func = lambda y, t: tf.stack([tf.cos(t), tf.sin(t)])
17
18start_t = time.time()
19true_y = odeint(ode_func, true_y0, t, method='dop853')
20print(time.time() - start_t)  # approx. 4.3 seconds on i7-9750H GTX1650
21
22# plot the solution and compare
23plt.figure(dpi=300)
24plt.title("sine(x)")
25plt.plot(t, np.sin(t), label='Analytical')
26plt.plot(t, true_y[:, 0], ls='--', label='astroNN odeint')
27plt.legend(loc='best')
28plt.xlabel("t")
29plt.ylabel("y")
30plt.show()
```

Moreover `odeint` supports numerically integration in parallel, the example below integration the `sin(x)` for 50 initial conditions. You can see the execution time is the same!!

```1start_t = time.time()
2# initial conditions, 50 of them instead of a single initial condition
3true_y0sss = tf.random.normal((50, 2), 0, 1)
4# time array, 50 of them instead of the same time array for every initial condition
5tsss = tf.random.normal((50, 10000), 0, 1)
6true_y = odeint(ode_func, true_y0sss, tsss, method='dop853')
7print(time.time() - start_t)  # also approx. 4.3 seconds on i7-9750H GTX1650
```

## Neural Network model with Numerical Integrator

You can use `odeint` along with neural network model, below is an example

``` 1import numpy as np
2import tensorflow as tf
3from astroNN.shared.nn_tools import gpu_memory_manage, cpu_fallback
4from astroNN.neuralode import odeint
5
6cpu_fallback()
7gpu_memory_manage()
8
9t = tf.constant(np.linspace(0, 1, 20))
10# initial condition
11true_y0 = tf.constant([0., 1.])
12
13class MyModel(tf.keras.Model):
14    def __init__(self):
15        super(MyModel, self).__init__()
16        self.dense1 = tf.keras.layers.Dense(2, activation=tf.nn.relu)
17        self.dense2 = tf.keras.layers.Dense(16, activation=tf.nn.relu)
18        self.dense3 = tf.keras.layers.Dense(2)
19
20    def call(self, inputs, t, *args):
21        inputs = tf.expand_dims(inputs, axis=0)
22        x = self.dense2(self.dense1(inputs))
23        return tf.squeeze(self.dense3(x))
24
25model = MyModel()
26