# Solving an ODE using neural networks (via Tensorflow)

I'm very new to deep learning (coming from a math PDE background), but I'm trying to solve some ODEs using a neural network (via tensorflow). I've solved some simple ones like $$u'(x)+u(x) = f(x)$$ with no problem, but I'm trying something a bit harder now: $$u''(x) - xu(x) = 0$$ with the initial conditions $$u(0)=A$$ and $$u'(0)=B$$.

I'm mostly following this paper, and my solution is written as $$u_{N}(x) = A + Bx + x^2N(x,w)$$, where $$N(x,w)$$ is the output of the neural net. The loss function I'm using is just the residual of the ODE in a mean square sense, so it's pretty crude: $$\ell(x,w) = \sum_{j=1}^{N} (u_{N}''(x) - xu_{N}(x))^{2}$$. I'm having a lot of trouble getting a good numerical solution to this particular equation. You can see a typical result below (orange is the exact solution, blue is my solution).

My current setup is just 2 hidden layers of 400 nodes each (one leaky ReLU and one ReLU) followed by a linear activation layer. My input data is an evenly spaced discretization of the domain. I'm using the Adam optimizer with a batch size of 32 and run for 400 epochs. I can't seem to capture the oscillating behavior in the left of the domain properly no matter what parameters I tweak.

Does anyone have any suggestions for how to improve the result? I'm very very new to deep learning and neural networks. If it helps, my code is included below. I should also probably give credit to Emm and vijay m, whose code for 1D approximation was the base for my code

# Load modules
import tensorflow as tf
import numpy as np
import math, random
import matplotlib.pyplot as plt
from scipy import special

######################################################################
# Routine to solve u''(x) - x*u(x) = f(x), u(0)=A, u'(0)=B in the form
#     u(x) = A + B*x + x^2*N(x,w)
# where N(x,w) is the output of the neural network.
######################################################################

# Create the arrays x and y, where x is a discretization of the domain (a,b) and y is the source term f(x)
N = 200
a = -6.0
b = 2.0
x = np.arange(a, b, (b-a)/N).reshape((N,1))
y = np.zeros(N)

# Boundary conditions
A = 1.0
B = 0.0

# Define the number of neurons in each layer
n_nodes_hl1 = 400
n_nodes_hl2 = 400

# Define the number of outputs and the learning rate
n_classes = 1
learn_rate = 0.00003

# Define input / output placeholders
x_ph = tf.placeholder('float', [None, 1],name='input')
y_ph = tf.placeholder('float')

# Define standard deviation for the weights and biases
hl_sigma = 0.02

# Routine to compute the neural network (5 hidden layers)
def neural_network_model(data):
hidden_1_layer = {'weights': tf.Variable(name='w_h1',initial_value=tf.random_normal([1, n_nodes_hl1], stddev=hl_sigma)),
'biases': tf.Variable(name='b_h1',initial_value=tf.random_normal([n_nodes_hl1], stddev=hl_sigma))}

hidden_2_layer = {'weights': tf.Variable(name='w_h2',initial_value=tf.random_normal([n_nodes_hl1, n_nodes_hl2], stddev=hl_sigma)),
'biases': tf.Variable(name='b_h2',initial_value=tf.random_normal([n_nodes_hl2], stddev=hl_sigma))}

output_layer = {'weights': tf.Variable(name='w_o',initial_value=tf.random_normal([n_nodes_hl2, n_classes], stddev=hl_sigma)),
'biases': tf.Variable(name='b_o',initial_value=tf.random_normal([n_classes], stddev=hl_sigma))}

# (input_data * weights) + biases
l1 = tf.nn.leaky_relu(l1)

l2 = tf.nn.relu(l2)

output = tf.add(tf.matmul(l2, output_layer['weights']), output_layer['biases'], name='output')

return output

batch_size = 32

# Feed batch data
def get_batch(inputX, inputY, batch_size):
duration = len(inputX)
for i in range(0,duration//batch_size):
idx = i*batch_size + np.random.randint(0,10,(1))[0]

yield inputX[idx:idx+batch_size], inputY[idx:idx+batch_size]

# Routine to train the neural network
def train_neural_network_batch(x_ph, predict=False):
prediction = neural_network_model(x_ph)

# Compute u and its second derivative
u = A + B*x_ph + (x_ph*x_ph)*prediction
dudx2 = (x_ph*x_ph)*pred_dx2 + 2.0*x_ph*pred_dx + 2.0*x_ph*pred_dx + 2.0*prediction

# The cost function is just the residual of u''(x) - x*u(x) = 0, i.e. residual = u''(x)-x*u(x)
cost = tf.reduce_mean(tf.square(dudx2-x_ph*u - y_ph))

# cycles feed forward + backprop
hm_epochs = 400

with tf.Session() as sess:
sess.run(tf.global_variables_initializer())

# Train in each epoch with the whole data
for epoch in range(hm_epochs):

epoch_loss = 0
for step in range(N//batch_size):
for inputX, inputY in get_batch(x, y, batch_size):
_, l = sess.run([optimizer,cost], feed_dict={x_ph:inputX, y_ph:inputY})
epoch_loss += l
if epoch %10 == 0:
print('Epoch', epoch, 'completed out of', hm_epochs, 'loss:', epoch_loss)

# Predict a new input by adding a random number, to check whether the network has actually learned
x_valid = x + 0.0*np.random.normal(scale=0.1,size=(1))
return sess.run(tf.squeeze(prediction),{x_ph:x_valid}), x_valid

# Train network
tf.set_random_seed(42)
pred, time = train_neural_network_batch(x_ph)

mypred = pred.reshape(N,1)

# Compute Airy functions for exact solution
ai, aip, bi, bip = special.airy(time)

# Numerical solution vs. exact solution
fig = plt.figure()
plt.plot(time, A + B*time + (time*time)*mypred)
plt.plot(time, 0.5*(3.0**(1/6))*special.gamma(2/3)*(3**(1/2)*ai + bi))
plt.show()
`
• Have you tried different activation functions and/or deeper networks ? – Shamit Verma Apr 15 at 13:13