Accuracy of PINN on 1D Shallow Water Equations

[ryabs07]

Hi all,
I am trying use PINN on 1D Shallow Water Equations. More specifically, I am testing it on 1D Wet Domain Dam Break scenario. But the code is unable to achieve precise solutions for water depth and velocity, specially near the wavefront. I have tried to train the model with very high number of training points (~300,000.) The results of the code can be seen in the video links below:

water_depth_animation.mp4

water_velocity_animation.mp4

The code for the problem as follows:

import deepxde as dde
import numpy as np

dde.config.set_default_float("float64")
 

#Configuration for the simulation
x_start = 0 #Starting point of domain
x_end = 100 #End point of domain
time_start = 0 #Starting time of simulation
time_end = 30 #Ending time of simulation
g = 9.81  # Specify the value for gravity constant g
dam_location = 35 #Location of dam
dam_height = 1  #water depth upstream of dam
downstream_depth = 0.4  #water depth downstream of dam

#Definition of Scaling factors
L_scale = x_end - x_start
T_scale = time_end - time_start
U_scale = 3 #put tentative maximum value you think is appropriate in m/s
H_scale = dam_height
G_scale = (U_scale**2) / H_scale

#Scaling the variables
x_start = x_start/L_scale
x_end = x_end / L_scale
time_start = time_start / T_scale
time_end = time_end / T_scale
g = g / G_scale  # Specify the value for gravity constant g
dam_location = dam_location / L_scale
dam_height = dam_height / H_scale  #water depth upstream of dam
downstream_depth = downstream_depth / H_scale  #water depth downstream of dam

iterations = 100000 #iterations for optimal neural net

# Preparation of Input for Prediction using PINN
total_points = 501 #Number of gridpoints in domain
num_time_steps = 501 #Number of timesteps
x = np.linspace(x_start, x_end, total_points)
t = np.linspace(time_start, time_end, num_time_steps)
X, T = np.meshgrid(x,t) 
input = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))

#Define the simplified 1D Shallow Water Equations
def pde(x, y):
    #x is a 2D array where x[:,0:1] represents spatial domain in x direction and x[:,1:2] represents the temporal domain
    #x is the input of the neural net
    #y is a 2D array where y[:,0:1] represents water depth and y[:,1:2] represents the water velocity
    #y is the output of the neural net

    h, u = y[:, 0:1], y[:, 1:]
    
    dh_dx = dde.grad.jacobian(y, x, i=0, j=0)
    dh_dt = dde.grad.jacobian(y, x, i=0, j=1)
    du_dx = dde.grad.jacobian(y, x, i=1, j=0)
    du_dt = dde.grad.jacobian(y, x, i=1, j=1)

    mass_balance = ((L_scale * dh_dt) / (T_scale * U_scale))  + ( h * du_dx + u * dh_dx )
    momentum_x = ((L_scale * du_dt) / (T_scale * U_scale)) + u * du_dx + g * dh_dx #+ g * S0 #+ (g * n_st ** 2 * u**2 * h**(-4/3))

    return [mass_balance, momentum_x]


# function for  Initial condition of water depth
def h_initial(x):
    h = np.where(x < dam_location, dam_height, downstream_depth)
    return h


geom = dde.geometry.Interval(x_start, x_end)
timedomain = dde.geometry.TimeDomain(time_start, time_end)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)

#Define initial condition of water depth and velocity
ic_h = dde.icbc.IC(
    geomtime, lambda x: h_initial(x[:, 0:1]), lambda _, on_initial: on_initial, component=0
)
ic_u = dde.icbc.IC(
    geomtime, lambda x: 0, lambda _, on_initial: on_initial, component=1
)

data = dde.data.TimePDE(
    geomtime, pde, [ic_h, ic_u], num_domain=350000, num_initial=30000
)
net = dde.nn.FNN([2] + [50] * 8 + [2], "tanh", "Glorot normal")
model = dde.Model(data, net)

model.compile("adam", lr=3e-4, loss_weights = [1, 10, 1000, 1000])
losshistory, train_state = model.train(iterations= iterations)
dde.saveplot(losshistory, train_state, issave=True, isplot=True)
# Save the trained model
model.save("saved_model.h5")

### Post Processing 
output = model.predict(input) #passing the input gridpoints and timestamps through optimized neural network

# Extracting water depth 'h' and velocity 'u' from the result
h = (output[:, 0].reshape((-1, total_points)).T) * H_scale
u = (output[:, 1].reshape((-1, total_points)).T) * U_scale

np.save('h.npy', h)
np.save('u.npy', u)

I would be very grateful if anyone could give insight to increase the accuracy of PINN.
I have also tried the L-BFGS optimizer, but it ends up performing very poorly in minimizing loss for initial conditions.

[DevSaugat]

I would also appreciate more details regarding this issue.

[lululxvi]

Try more dense points around the discontinuity, and less dense points in other places. Also try other activations such as swish.

[ryabs07]

Hi Dr. Lu Lu,
Thank you for you response and amazing library.
How is possible to increase the density of sampling at particular range? Also, since the discontinuity is ever moving, would be it be possible to also move the dense sampling near discontinuity in DeepXDE?

I will definitely try the swish activation function.
Regards!

[lululxvi]

Use anchor points.

See https://www.sciencedirect.com/science/article/abs/pii/S0045782519306814

[ryabs07]

Anchoring the points definitely helped get better accuracy. It is still struggling a little bit in the discontinuity region. Thank you very much for the suggestion :)

[ryabs07]

2 stage optimization (Adam + L-BFGS) and anchoring of the sampling points vastly improved the accuracy of the model. Thank you very much for your kind suggestions :)