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.add(tf.matmul(data, hidden_1_layer['weights']), hidden_1_layer['biases'])
l1 = tf.nn.leaky_relu(l1)
l2 = tf.add(tf.matmul(l1, hidden_2_layer['weights']), hidden_2_layer['biases'])
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)
pred_dx = tf.gradients(prediction, x_ph)
pred_dx2 = tf.gradients(tf.gradients(prediction, x_ph),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))
optimizer = tf.train.AdamOptimizer(learn_rate).minimize(cost)
# 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()
