6
$\begingroup$

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). Orange line is the exact solution, blue line is the numerical 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()
$\endgroup$
3
  • $\begingroup$ Have you tried different activation functions and/or deeper networks ? $\endgroup$ Apr 15, 2019 at 13:13
  • $\begingroup$ Try swish activation function with four hidden layers each with 32 nodes. $\endgroup$
    – Karim
    Apr 12, 2020 at 12:52
  • $\begingroup$ You can add in the lost function the two boundary conditions $\endgroup$
    – user149654
    May 7, 2023 at 19:03

1 Answer 1

3
$\begingroup$
  1. Please avoid using piecewise linear activations for second or higher-order ODEs, as the second and higher-order derivatives of these functions are zero. Alternatively, Tanh activations might be a good option.

  2. I recommend using deeper networks. Based on my experience, depth can significantly improve the results when solving differential equations with neural networks.

$\endgroup$

Not the answer you're looking for? Browse other questions tagged or ask your own question.