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()
../_images/odeint_sine.png

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
27with tf.GradientTape() as g:
28    g.watch(true_y0)
29    y = odeint(model, true_y0, t)
30# gradient of the result w.r.t. model's weights
31g.gradient(y, model.trainable_variables)  # well define, no None, no inf or no NaN