LARGE LOSS observed in 2-D Navier-Stokes equation

[chrisxbchen]

Hi fellows,

The first issue is about the Adam iteration numbers. We've tried 10k, 20k, and 30k Adam iterations and found that
(1) large number of iterations doesn't help with reducing the loss and
(2) LBFGS doesn't help with further reducing the loss.

The second issue is that the predicted initial velocity looks very different from the exact solution. We believe a good prediction won't have such a big difference. Hereunder are the pictures of the exact initial velocity field and the predicted result from the model. The exact one has several vortices throughout the domain, but the model prediction has no vortices.

The code is attached in which only a small training set and a small number of iterations are used, but the problem would be roughly the same if a large training set and more iterations are utilized.

# This code is to solve 2-D Navier-Stokes equations
# u_t + u*u_x + v*u_y + p_x - nu*(u_xx + u_yy) = f1
# v_t + u*v_x + v*v_y + p_y - nu*(v_xx + v_yy) = f2
# u_x + v_y = 0
# It was written by Edom Belayneh, Xiaobai(Chris) Chen, Ruthwik Nadam, and Xiaoming Zheng
# of Central Michigan University, Mount Pleasant, MI 48859
# July 31, 2024
import deepxde as dde
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import math
import os
os.environ["TF_USE_LEGACY_KERAS"]="1"

#torch.set_default_tensor_type(torch.FloatTensor)
# The space domain is [-L, -D] x [L,D]. The time domain is [0, time]
L = 1
D = 1
time = 0.5
nu = 1

# Define functions f1, f2 with the exact solution in the Navier-Stokes equations
def calculate_functions(xn, yn, t):
    pi = np.pi
    exp_neg_t = tf.exp(-t)
    sin_2pi_xn = tf.sin(2 * pi * xn)
    cos_2pi_yn = tf.cos(2 * pi * yn)
    cos_2pi_xn = tf.cos(2 * pi * xn)
    sin_2pi_yn = tf.sin(2 * pi * yn)
    sin_pi_xn = tf.sin(pi * xn)
    sin_pi_yn = tf.sin(pi * yn)
    cos_pi_xn = tf.cos(pi * xn)
    cos_pi_yn = tf.cos(pi * yn)

    # exact u and v
    u_e = exp_neg_t * sin_2pi_xn * cos_2pi_yn
    v_e = -exp_neg_t * cos_2pi_xn * sin_2pi_yn

    # first order derivatives
    u_x = 2 * pi * exp_neg_t * cos_2pi_xn * cos_2pi_yn
    u_y = -2 * pi * exp_neg_t * sin_2pi_xn * sin_2pi_yn
    u_t = -exp_neg_t * sin_2pi_xn * cos_2pi_yn
    v_x = 2 * pi * exp_neg_t * sin_2pi_xn * sin_2pi_yn
    v_y = -u_x
    v_t = exp_neg_t * cos_2pi_xn * sin_2pi_yn
    p_x = -pi * exp_neg_t * sin_pi_xn * sin_pi_yn
    p_y = pi * exp_neg_t * cos_pi_xn * cos_pi_yn

    # second order derivatives
    u_xx = -4 * pi**2 * exp_neg_t * sin_2pi_xn * cos_2pi_yn
    u_yy = u_xx
    v_xx = 4 * pi**2 * exp_neg_t * cos_2pi_xn * sin_2pi_yn
    v_yy = v_xx

    f1 = u_t + p_x -nu*(u_xx + u_yy) + (u_e * u_x + v_e * u_y)
    f2 = v_t + p_y -nu*(v_xx + v_yy) + (u_e * v_x + v_e * v_y)

    return f1, f2

# x[0] is "x", x[1] is "y", x[2] is time
def uexact(x):
  return np.exp(-x[2])*np.sin(2*np.pi*x[0])*np.cos(2*np.pi*x[1])

def vexact(x):
  return -np.exp(-x[2])*np.cos(2*np.pi*x[0])*np.sin(2*np.pi*x[1])

def pexact(x):
  return np.exp(-x[2])*np.cos(np.pi*x[0])*np.sin(np.pi*x[1])

#Define the PDE
def pde(X, Y):
  # Y: first column is u, second col is v and third column is p
  # X: first column is x, second col is y and third column is t
  u_vel = Y[:,0]
  v_vel = Y[:,1]
  p     = Y[:,2]

  # time terms
  du_dt = dde.grad.jacobian(Y,X, i=0, j=2)
  dv_dt = dde.grad.jacobian(Y,X, i=1, j=2)

  # space terms
  du_dx = dde.grad.jacobian(Y,X, i=0, j=0)
  du_dy = dde.grad.jacobian(Y,X, i=0, j=1)
  dv_dx = dde.grad.jacobian(Y,X, i=1, j=0)
  dv_dy = dde.grad.jacobian(Y,X, i=1, j=1)
  dp_dx = dde.grad.jacobian(Y,X, i=2, j=0)
  dp_dy = dde.grad.jacobian(Y,X, i=2, j=1)

  # diffusion terms
  ddu_dxx = dde.grad.hessian(Y,X, i=0, j=0)
  ddu_dyy = dde.grad.hessian(Y,X, i=0, j=1)
  ddv_dxx = dde.grad.hessian(Y,X, i=1, j=0)
  ddv_dyy = dde.grad.hessian(Y,X, i=1, j=1)

  # Calculate the functions
  f1, f2 = calculate_functions(X[:,0], X[:,1], X[:,2])

  #PDE
  pde_u = du_dt + (u_vel*du_dx) + (v_vel*du_dy) - (nu * (ddu_dxx + ddu_dyy)) + dp_dx -f1
  pde_v = dv_dt + (u_vel*dv_dx) + (v_vel*dv_dy) - (nu * (ddv_dxx + ddv_dyy)) + dp_dy -f2
  pde_cont = du_dx + dv_dy

  return [pde_u, pde_v, pde_cont]

#Define the Sapce and Time Domain Geometry
geom = dde.geometry.Rectangle(xmin=[-L, -D], xmax=[L, D])
timedomain = dde.geometry.TimeDomain(0,time)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)

# Dirichlet boundary conditions
bc_wall_u = dde.icbc.DirichletBC(geomtime, lambda x: uexact(x), lambda _, on_boundary: on_boundary, component=0)
bc_wall_v = dde.icbc.DirichletBC(geomtime, lambda x: vexact(x), lambda _, on_boundary: on_boundary, component=1)

#Initial conditions at t=0
ic_u = dde.icbc.IC(geomtime, lambda x: uexact(x), lambda _, on_initial: on_initial, component=0)
ic_v = dde.icbc.IC(geomtime, lambda x: vexact(x), lambda _, on_initial: on_initial, component=1)

train_distribution = "LHS"

data = dde.data.TimePDE(geomtime, pde, [bc_wall_u, bc_wall_v, ic_u, ic_v], train_distribution=train_distribution,
                       num_domain= 1000, num_boundary=200, num_initial=200)

print("data.train_x_all.shape (no duplications)= ", data.train_x_all.shape)

#Plotting 3D domain
fig = plt.figure(figsize=(8, 8))
ax1 = fig.add_subplot(121, projection='3d')
ax1.scatter(data.train_x[:,0], data.train_x[:,1], data.train_x[:,2],color='blue', label='Domain Points', s=5.5)
ax1.set_title('Points in Spatial Domain')
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_zlabel('t')
ax1.legend()
#plt.savefig(f"3D_XYT_plot_{train_distribution}.png")
plt.show()

print("data.num_bcs",data.num_bcs)
 print(type(data.bc_points()))
 print("data.bc_points().shape",data.bc_points().shape)

# Creating the neural network with 3 inputs(x,y,t), 2 hidden layers with 30 neurons and 3 outputs (u,v,p).
# Activation function is tanh
net = dde.nn.FNN([3] + [30] * 2 + [3], "tanh", "Glorot normal")

#Creating the model
model = dde.Model(data, net)

#Compile the model with Adam Optimizer
model.compile("adam", lr=1e-3)

#Train the model
losshistory, train_state = model.train(iterations=200)

#Compile the model again with LBFGS Optimizer after ADAM optimiser is done
dde.optimizers.config.set_LBFGS_options(maxiter=30)
model.compile("L-BFGS")
losshistory, train_state = model.train()

#Train vs test loss plot
dde.saveplot(losshistory, train_state, issave=True, isplot=True)

# Generate convergence test data
def gen_testdata(L, D, time):
    # Generate evenly distributed midpoints
    h = 0.1
    x = np.linspace(-L + h/2, L-h/2, int(2*L/h))
    y = np.linspace(-D + h/2, D-h/2, int(2*D/h))
    t = np.linspace(0+h/2, time-h/2, int(time/h))
    X, Y, T = np.meshgrid(x, y, t, indexing='ij')
    Mesh = np.stack([X.flatten(), Y.flatten(), T.flatten()], axis=1)

    # Extract xn, yn, t from the dataset
    xn = Mesh[:, 0]
    yn = Mesh[:, 1]
    t  = Mesh[:, 2]

    # Compute the exact values
    exp_neg_t = np.exp(-t)
    sin_2pi_xn = np.sin(2 * np.pi * xn)
    cos_2pi_yn = np.cos(2 * np.pi * yn)
    cos_2pi_xn = np.cos(2 * np.pi * xn)
    sin_2pi_yn = np.sin(2 * np.pi * yn)
    cos_pi_xn = np.cos(np.pi * xn)
    sin_pi_yn = np.sin(np.pi * yn)

    u_e = exp_neg_t * sin_2pi_xn * cos_2pi_yn
    v_e = -exp_neg_t * cos_2pi_xn * sin_2pi_yn
    p_e = exp_neg_t * cos_pi_xn * sin_pi_yn

    # Stack the exact values into a single array
    All_exact = np.stack([u_e, v_e, p_e], axis=1)

    return Mesh, All_exact

# define X, u_exact and u_pred
Mesh, All_exact = gen_testdata(L, D, time)
All_pred = model.predict(Mesh)
# Get the p_tilde for predicted value
p_tilde_pred = All_pred[:, 2] - np.mean(All_pred[:, 2])
# Update u_pred with p_tilde_pred
All_pred[:, 2] = p_tilde_pred

# compute the max errors
diff_u = np.abs(All_pred[:, 0] - All_exact[:, 0])
diff_v = np.abs(All_pred[:, 1] - All_exact[:, 1])
diff_p = np.abs(All_pred[:, 2] - All_exact[:, 2])
max_diff_u = np.max(diff_u)
max_diff_v = np.max(diff_v)
max_diff_p = np.max(diff_p)

print("Max error in u:", max_diff_u)
print("Max error in v:", max_diff_v)
print("Max error in p:", max_diff_p)

# compute L2 error
# Calculate the square of the differences
squared_diff = (All_pred - All_exact) ** 2

# Reshape squared_diff to match the original meshgrid dimensions
h=0.1
num_x_points = int(2*L/h)
num_y_points = int(2*D/h)
num_t_points = int(time/h)
squared_diff = squared_diff.reshape(num_x_points, num_y_points, num_t_points, 3)

# Integration over the domain and time for each of the four dimensions
dx = dy = dt = h

# Integrate over the spatial-temporal domain (x, y, t)
integral_u = np.sum(squared_diff[:, :, :, 0]) * dx * dy * dt
integral_v = np.sum(squared_diff[:, :, :, 1]) * dx * dy * dt
integral_p = np.sum(squared_diff[:, :, :, 2]) * dx * dy * dt

# take square root
L2_Error_u = np.sqrt(integral_u)
L2_Error_v = np.sqrt(integral_v)
L2_Error_p = np.sqrt(integral_p)

# Print the result
print("L2 error for u:", L2_Error_u)
print("L2 error for v:", L2_Error_v)
print("L2 error for p:", L2_Error_p)

# Compute H1-error
# Function to compute the numerical partial derivative with respect to x or y
# using central difference scheme
def compute_partial_derivative(y, h, axis):
    return np.gradient(y, h, axis=axis)

# Reshape u_pred and u_exact to match the original grid dimensions
num_x_points = int(2*L/h)
num_y_points = int(2*D/h)
num_t_points = int(time/h)
All_pred = All_pred.reshape(num_x_points, num_y_points, num_t_points, 3)
All_exact = All_exact.reshape(num_x_points, num_y_points, num_t_points, 3)

# Compute the partial derivatives with respect to x and y
All_pred_x = compute_partial_derivative(All_pred, h, axis=0)
All_exact_x = compute_partial_derivative(All_exact, h, axis=0)
All_pred_y = compute_partial_derivative(All_pred, h, axis=1)
All_exact_y = compute_partial_derivative(All_exact, h, axis=1)

# Calculate the squared differences
squared_diff = (All_pred_x - All_exact_x)**2 + (All_pred_y - All_exact_y)**2

# Sum up the squared differences over all points in the matrix
result0 = np.sqrt(np.sum(squared_diff[:,:,:,0])*(h**3))
result1 = np.sqrt(np.sum(squared_diff[:,:,:,1])*(h**3))
result2 = np.sqrt(np.sum(squared_diff[:,:,:,2])*(h**3))
result = np.sqrt(np.sum(squared_diff)*(h**3))

# Print the result
print(squared_diff.shape)
print("H1 error for u:", result0)
print("H1 error for v:", result1)
print("H1 error for p:", result2)
print("H1 error for all:", result)

# Extract initial points (where t=0)
initial_points = data.train_x[data.train_x[:, 2] == 0]

# Evaluate the initial conditions at these points
x = initial_points[:, 0]
y = initial_points[:, 1]
u_initial = np.sin(2 * np.pi * x) * np.cos(2 * np.pi * y)  # Initial condition for u
v_initial = -np.cos(2 * np.pi * x) * np.sin(2 * np.pi * y)  # Initial condition for v

# Make predictions using the trained model on initial points
predictions = model.predict(initial_points)

# Extract x, y coordinates and the predicted u, v values
u = predictions[:, 0]  # Assuming this is the u component of the velocity
v = predictions[:, 1]  # Assuming this is the v component of the velocity

# Create subplots
fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))

# Plot initial conditions
axes[0].quiver(x, y, u_initial, v_initial, color='r', scale=20, scale_units='xy', angles='xy', label='Vector Field')
axes[0].set_title('Vector Plot (Exact) at time 0')
axes[0].set_xlabel('x')
axes[0].set_ylabel('y')
axes[0].legend()
axes[0].grid()

# Plot predicted values
axes[1].quiver(x, y, u, v, color='r', scale=20, scale_units='xy', angles='xy', label='Vector Field')
axes[1].set_title('Vector Plot (Predicted) at time 0')
axes[1].set_xlabel('x')
axes[1].set_ylabel('y')
axes[1].grid()

# Show the plots
plt.tight_layout()
plt.show()

[lululxvi]

It seems that the IC from PINN is not correct. Try a very small time internal such as [0, 0.1]; if IC is still not correct, then the code of IC could be wrong. You need to use x[:, 0:1], not x[0], in IC.

[chrisxbchen]

Dr. Lu, thanks you so much for the reply. I updated the code with these changes:

1. Time interval was changed to [0, 0.1], and change the time step to 0.01 so that there were still be 10 steps
2. used 0:1 format instead of 0 in as many places as possible; and
3. increased the number of iteration from 200 to 10,000.

The loss drops roughly from 1e+03 to 1e+02, but the plot still looks wild and completely different from the exact initial value.

# This code is to solve 2-D Navier-Stokes equations
# u_t + u*u_x + v*u_y + p_x - nu*(u_xx + u_yy) = f1
# v_t + u*v_x + v*v_y + p_y - nu*(v_xx + v_yy) = f2
# u_x + v_y = 0
# It was written by Edom Belayneh, Xiaobai(Chris) Chen, Ruthwik Nadam, and Xiaoming Zheng
# of Central Michigan University, Mount Pleasant, MI 48859
# July 31, 2024
!pip install deepxde
import deepxde as dde
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import math
import os
os.environ["TF_USE_LEGACY_KERAS"]="1"

#torch.set_default_tensor_type(torch.FloatTensor)
# The space domain is [-L, -D] x [L,D]. The time domain is [0, time]
L = 1
D = 1
time = 0.1
nu = 1

# Define functions f1, f2 with the exact solution in the Navier-Stokes equations
def calculate_functions(xn, yn, t):
    pi = np.pi
    exp_neg_t = tf.exp(-t)
    sin_2pi_xn = tf.sin(2 * pi * xn)
    cos_2pi_yn = tf.cos(2 * pi * yn)
    cos_2pi_xn = tf.cos(2 * pi * xn)
    sin_2pi_yn = tf.sin(2 * pi * yn)
    sin_pi_xn = tf.sin(pi * xn)
    sin_pi_yn = tf.sin(pi * yn)
    cos_pi_xn = tf.cos(pi * xn)
    cos_pi_yn = tf.cos(pi * yn)

    # exact u and v
    u_e = exp_neg_t * sin_2pi_xn * cos_2pi_yn
    v_e = -exp_neg_t * cos_2pi_xn * sin_2pi_yn

    # first order derivatives
    u_x = 2 * pi * exp_neg_t * cos_2pi_xn * cos_2pi_yn
    u_y = -2 * pi * exp_neg_t * sin_2pi_xn * sin_2pi_yn
    u_t = -exp_neg_t * sin_2pi_xn * cos_2pi_yn
    v_x = 2 * pi * exp_neg_t * sin_2pi_xn * sin_2pi_yn
    v_y = -u_x
    v_t = exp_neg_t * cos_2pi_xn * sin_2pi_yn
    p_x = -pi * exp_neg_t * sin_pi_xn * sin_pi_yn
    p_y = pi * exp_neg_t * cos_pi_xn * cos_pi_yn

    # second order derivatives
    u_xx = -4 * pi**2 * exp_neg_t * sin_2pi_xn * cos_2pi_yn
    u_yy = u_xx
    v_xx = 4 * pi**2 * exp_neg_t * cos_2pi_xn * sin_2pi_yn
    v_yy = v_xx

    f1 = u_t + p_x -nu*(u_xx + u_yy) + (u_e * u_x + v_e * u_y)
    f2 = v_t + p_y -nu*(v_xx + v_yy) + (u_e * v_x + v_e * v_y)

    return f1, f2

# x[0] is "x", x[1] is "y", x[2] is time
def uexact(x):
  return np.exp(-x[:,2:3])*np.sin(2*np.pi*x[:, 0:1])*np.cos(2*np.pi*x[:,1:2])

def vexact(x):
  return -np.exp(-x[:,2:3])*np.cos(2*np.pi*x[:, 0:1])*np.sin(2*np.pi*x[:,1:2])

def pexact(x):
  return np.exp(-x[:,2:3])*np.cos(np.pi*x[:, 0:1])*np.sin(np.pi*x[:,1:2])

#Define the PDE
def pde(X, Y):
  # Y: first column is u, second col is v and third column is p
  # X: first column is x, second col is y and third column is t
  u_vel = Y[:,0:1]
  v_vel = Y[:,1:2]
  p     = Y[:,2:3]

  # time terms
  du_dt = dde.grad.jacobian(Y,X, i=0, j=2)
  dv_dt = dde.grad.jacobian(Y,X, i=1, j=2)

  # space terms
  du_dx = dde.grad.jacobian(Y,X, i=0, j=0)
  du_dy = dde.grad.jacobian(Y,X, i=0, j=1)
  dv_dx = dde.grad.jacobian(Y,X, i=1, j=0)
  dv_dy = dde.grad.jacobian(Y,X, i=1, j=1)
  dp_dx = dde.grad.jacobian(Y,X, i=2, j=0)
  dp_dy = dde.grad.jacobian(Y,X, i=2, j=1)

  # diffusion terms
  ddu_dxx = dde.grad.hessian(Y,X, i=0, j=0)
  ddu_dyy = dde.grad.hessian(Y,X, i=0, j=1)
  ddv_dxx = dde.grad.hessian(Y,X, i=1, j=0)
  ddv_dyy = dde.grad.hessian(Y,X, i=1, j=1)

  # Calculate the functions
  f1, f2 = calculate_functions(X[:,0:1], X[:,1:2], X[:,2:3])

  #PDE
  pde_u = du_dt + (u_vel*du_dx) + (v_vel*du_dy) - (nu * (ddu_dxx + ddu_dyy)) + dp_dx -f1
  pde_v = dv_dt + (u_vel*dv_dx) + (v_vel*dv_dy) - (nu * (ddv_dxx + ddv_dyy)) + dp_dy -f2
  pde_cont = du_dx + dv_dy

  return [pde_u, pde_v, pde_cont]

#Define the Sapce and Time Domain Geometry
geom = dde.geometry.Rectangle(xmin=[-L, -D], xmax=[L, D])
timedomain = dde.geometry.TimeDomain(0,time)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)

# Dirichlet boundary conditions
bc_wall_u = dde.icbc.DirichletBC(geomtime, lambda x: uexact(x), lambda _, on_boundary: on_boundary, component=0)
bc_wall_v = dde.icbc.DirichletBC(geomtime, lambda x: vexact(x), lambda _, on_boundary: on_boundary, component=1)

#Initial conditions at t=0
ic_u = dde.icbc.IC(geomtime, lambda x: uexact(x), lambda _, on_initial: on_initial, component=0)
ic_v = dde.icbc.IC(geomtime, lambda x: vexact(x), lambda _, on_initial: on_initial, component=1)

train_distribution = "LHS"

data = dde.data.TimePDE(geomtime, pde, [bc_wall_u, bc_wall_v, ic_u, ic_v], train_distribution=train_distribution,
                       num_domain= 1000, num_boundary=200, num_initial=200)

print("data.train_x_all.shape (no duplications)= ", data.train_x_all.shape)

#Plotting 3D domain
fig = plt.figure(figsize=(8, 8))
ax1 = fig.add_subplot(121, projection='3d')
ax1.scatter(data.train_x[:,0], data.train_x[:,1], data.train_x[:,2],color='blue', label='Domain Points', s=5.5)
ax1.set_title('Points in Spatial Domain')
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_zlabel('t')
ax1.legend()
#plt.savefig(f"3D_XYT_plot_{train_distribution}.png")
plt.show()

print("data.num_bcs",data.num_bcs)
print(type(data.bc_points()))
print("data.bc_points().shape",data.bc_points().shape)

# Creating the neural network with 3 inputs(x,y,t), 2 hidden layers with 30 neurons and 3 outputs (u,v,p).
# Activation function is tanh
net = dde.nn.FNN([3] + [30] * 2 + [3], "tanh", "Glorot normal")

#Creating the model
model = dde.Model(data, net)

#Compile the model with Adam Optimizer
model.compile("adam", lr=1e-3)

#Train the model
losshistory, train_state = model.train(iterations=10000)

#Compile the model again with LBFGS Optimizer after ADAM optimiser is done
dde.optimizers.config.set_LBFGS_options(maxiter=30)
model.compile("L-BFGS")
losshistory, train_state = model.train()

#Train vs test loss plot
dde.saveplot(losshistory, train_state, issave=True, isplot=True)

# Generate convergence test data
def gen_testdata(L, D, time):
    # Generate evenly distributed midpoints
    h = 0.1
    h_t = 0.01
    x = np.linspace(-L + h/2, L-h/2, int(2*L/h))
    y = np.linspace(-D + h/2, D-h/2, int(2*D/h))
    t = np.linspace(0+h_t/2, time-h_t/2, int(time/h_t))
    X, Y, T = np.meshgrid(x, y, t, indexing='ij')
    Mesh = np.stack([X.flatten(), Y.flatten(), T.flatten()], axis=1)

    # Extract xn, yn, t from the dataset
    xn = Mesh[:, 0:1]
    yn = Mesh[:, 1:2]
    t  = Mesh[:, 2:3]

    # Compute the exact values
    exp_neg_t = np.exp(-t)
    sin_2pi_xn = np.sin(2 * np.pi * xn)
    cos_2pi_yn = np.cos(2 * np.pi * yn)
    cos_2pi_xn = np.cos(2 * np.pi * xn)
    sin_2pi_yn = np.sin(2 * np.pi * yn)
    cos_pi_xn = np.cos(np.pi * xn)
    sin_pi_yn = np.sin(np.pi * yn)

    u_e = exp_neg_t * sin_2pi_xn * cos_2pi_yn
    v_e = -exp_neg_t * cos_2pi_xn * sin_2pi_yn
    p_e = exp_neg_t * cos_pi_xn * sin_pi_yn

    # Stack the exact values into a single array and remove the extra dimension
    All_exact = np.stack([u_e.flatten(), v_e.flatten(), p_e.flatten()], axis=1)

    return Mesh, All_exact

# define X, u_exact and u_pred
Mesh, All_exact = gen_testdata(L, D, time)
All_pred = model.predict(Mesh)
# Get the p_tilde for predicted value
p_tilde_pred = All_pred[:, 2:3] - np.mean(All_pred[:, 2:3])
# Update u_pred with p_tilde_pred
All_pred[:, 2:3] = p_tilde_pred

# compute the max errors
diff_u = np.abs(All_pred[:, 0:1] - All_exact[:, 0:1])
diff_v = np.abs(All_pred[:, 1:2] - All_exact[:, 1:2])
diff_p = np.abs(All_pred[:, 2:3] - All_exact[:, 2:3])
max_diff_u = np.max(diff_u)
max_diff_v = np.max(diff_v)
max_diff_p = np.max(diff_p)

print("Max error in u:", max_diff_u)
print("Max error in v:", max_diff_v)
print("Max error in p:", max_diff_p)

# compute L2 error
# Calculate the square of the differences
squared_diff = (All_pred - All_exact) ** 2

# Reshape squared_diff to match the original meshgrid dimensions
h=0.1
h_t=0.01
num_x_points = int(2*L/h)
num_y_points = int(2*D/h)
num_t_points = int(time/h_t)
squared_diff = squared_diff.reshape(num_x_points, num_y_points, num_t_points, 3)

# Integration over the domain and time for each of the four dimensions
dx = dy = h
dt = h_t

# Integrate over the spatial-temporal domain (x, y, t)
integral_u = np.sum(squared_diff[:, :, :, 0]) * dx * dy * dt
integral_v = np.sum(squared_diff[:, :, :, 1]) * dx * dy * dt
integral_p = np.sum(squared_diff[:, :, :, 2]) * dx * dy * dt

# take square root
L2_Error_u = np.sqrt(integral_u)
L2_Error_v = np.sqrt(integral_v)
L2_Error_p = np.sqrt(integral_p)

# Print the result
print("L2 error for u:", L2_Error_u)
print("L2 error for v:", L2_Error_v)
print("L2 error for p:", L2_Error_p)

# Compute H1-error
# Function to compute the numerical partial derivative with respect to x or y
# using central difference scheme
def compute_partial_derivative(y, h, axis):
    return np.gradient(y, h, axis=axis)

# Reshape u_pred and u_exact to match the original grid dimensions
num_x_points = int(2*L/h)
num_y_points = int(2*D/h)
num_t_points = int(time/h_t)
All_pred = All_pred.reshape(num_x_points, num_y_points, num_t_points, 3)
All_exact = All_exact.reshape(num_x_points, num_y_points, num_t_points, 3)

# Compute the partial derivatives with respect to x and y
All_pred_x = compute_partial_derivative(All_pred, h, axis=0)
All_exact_x = compute_partial_derivative(All_exact, h, axis=0)
All_pred_y = compute_partial_derivative(All_pred, h, axis=1)
All_exact_y = compute_partial_derivative(All_exact, h, axis=1)

# Calculate the squared differences
squared_diff = (All_pred_x - All_exact_x)**2 + (All_pred_y - All_exact_y)**2

# Sum up the squared differences over all points in the matrix
result0 = np.sqrt(np.sum(squared_diff[:,:,:,0])*(h**3))
result1 = np.sqrt(np.sum(squared_diff[:,:,:,1])*(h**3))
result2 = np.sqrt(np.sum(squared_diff[:,:,:,2])*(h**3))
result = np.sqrt(np.sum(squared_diff)*(h**3))

# Print the result
print(squared_diff.shape)
print("H1 error for u:", result0)
print("H1 error for v:", result1)
print("H1 error for p:", result2)
print("H1 error for all:", result)

# Extract initial points (where t=0)
initial_points = data.train_x[data.train_x[:, 2] == 0]

# Evaluate the initial conditions at these points
x = initial_points[:, 0:1]
y = initial_points[:, 1:2]
u_initial = np.sin(2 * np.pi * x) * np.cos(2 * np.pi * y)  # Initial condition for u
v_initial = -np.cos(2 * np.pi * x) * np.sin(2 * np.pi * y)  # Initial condition for v

# Make predictions using the trained model on initial points
predictions = model.predict(initial_points)

# Extract x, y coordinates and the predicted u, v values
u = predictions[:, 0:1]  # Assuming this is the u component of the velocity
v = predictions[:, 1:2]  # Assuming this is the v component of the velocity

# Create subplots
fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))

# Plot initial conditions
axes[0].quiver(x, y, u_initial, v_initial, color='r', scale=20, scale_units='xy', angles='xy', label='Vector Field')
axes[0].set_title('Vector Plot (Exact) at time 0')
axes[0].set_xlabel('x')
axes[0].set_ylabel('y')
axes[0].legend()
axes[0].grid()

# Plot predicted values
axes[1].quiver(x, y, u, v, color='r', scale=20, scale_units='xy', angles='xy', label='Vector Field')
axes[1].set_title('Vector Plot (Predicted) at time 0')
axes[1].set_xlabel('x')
axes[1].set_ylabel('y')
axes[1].grid()

# Show the plots
plt.tight_layout()
plt.show()