PINN fails to learn from PDEs with clamping functions

[akruel01]

Description:

My specific equation is failing to learn, so I found the smallest working example I could based on your basic diffusion example neuromancer/examples/PDEs/Part_1_PINN_DiffusionEquation.ipynb. I think it stems from a portion of my equation incorporating a clamp max(0,x), as when I add this to the basic diffusion example the results never correspond to the training data. I tried various sharp functions from python max to activation functions like pytorch ReLU. Even the differentiable alternative Softplus shows no improvement. I also tried the following program tweaks suggested by Jan for each function: longer runtime, normalization, decreased batch size, decrease/increase resolutions (temporal or spatial). Note that while the analytical does not reach zero after many timesteps, the solution steeply drops after a few steps, and the aforementioned larger equation performs even worse. If there is more to try, or clamped equations are fundamentally incompatible (even if differentiable), I would like to know.

Illustration of the Results:

Complete Code to Reproduce (ran in the provided colab env):

!pip install "neuromancer[examples] @ git+https://github.com/pnnl/neuromancer.git@master"
!pip install pyDOE

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

torch.set_default_dtype(torch.float)
torch.manual_seed(1234)
np.random.seed(1234)
device = 'cpu'

This is where the analytical function is modified

## instantiate both clamp functions
def SOFTP(x,beta=5): # torch.nn.Softplus() throws error with neuro class 'Variable' when building the PiNN symbols - different issue
    return torch.log(1+torch.exp(beta*x))/beta
RELU = torch.nn.ReLU()

def diffusion_eq(x,t):
    return torch.exp(-t)*RELU(torch.sin(np.pi*x)) # <- change the analytical equation
    # return torch.exp(-t)*SOFTP(torch.sin(np.pi*x)) # differentiable

x_min, x_max = -1, 1
total_points_x = 200
t_min, t_max = 0, 3
total_points_t = 100
x=torch.linspace(x_min,x_max,total_points_x).view(-1,1)
t=torch.linspace(t_min,t_max,total_points_t).view(-1,1)
X,T=torch.meshgrid(x.squeeze(1),t.squeeze(1))

# plot3D(X, T, diffusion_eq(X,T)) # to visualize

from neuromancer.modules import blocks
from neuromancer.system import Node
from neuromancer.constraint import variable
from neuromancer.loss import PenaltyLoss
from neuromancer.problem import Problem
from neuromancer.dataset import DictDataset

# Transform the mesh into a 2-column vector
X_test = X.transpose(1,0).flatten()[:,None].float() # the input dataset of x
T_test = T.transpose(1,0).flatten()[:,None].float() # the input dataset of t
Y_test = diffusion_eq(X,T).transpose(1,0).flatten()[:,None].float() # the real solution over (x,t)

#   Left Edge: y(x,0) = sin(pi*x)
left_X = X[:, [0]]
left_T = T[:, [0]]

Here is where I changed the left_y at t=0

left_Y = RELU(torch.sin(np.pi * left_X[:, 0])).unsqueeze(1)

#   Bottom Edge: x=min; tmin=<t=<max
bottom_X = X[[0], :].T
bottom_T = T[[0], :].T
bottom_Y = torch.zeros(bottom_X.shape[0], 1)
#   Top Edge: x=max; 0=<t=<1
top_X = X[[-1], :].T
top_T = T[[-1], :].T
top_Y = torch.zeros(top_X.shape[0], 1)

# Choose (Nu) Number of training points for initial and boundary conditions
Nu = 200
# Get all the initial and boundary condition data
X_train = torch.vstack([left_X, bottom_X, top_X])
T_train = torch.vstack([left_T, bottom_T, top_T])
Y_train = torch.vstack([left_Y, bottom_Y, top_Y])
# Randomly sample Nu points of our available initial and boundary condition data:
idx = np.sort(np.random.choice(X_train.shape[0], Nu, replace=False))
X_train_Nu = X_train[idx, :].float()  # Training Points  of x at (IC+BC)
T_train_Nu = T_train[idx, :].float()  # Training Points  of t at (IC+BC)
Y_train_Nu = Y_train[idx, :].float()  # Training Points  of y at (IC+BC)

# x Domain bounds
x_lb = X_test[0]  # first value
x_ub = X_test[-1]  # last value
# t Domain bounds
t_lb = T_test[0]  # first value
t_ub = T_test[-1]  # last value
#  Choose (Nf) Collocation Points to Evaluate the PDE on
Nf = 1000  # Nf: Number of collocation points (Evaluate PDE)
# generate collocation points (CP)
X_train_CP = torch.FloatTensor(Nf, 1).uniform_(float(x_lb), float(x_ub))
T_train_CP = torch.FloatTensor(Nf, 1).uniform_(float(t_lb), float(t_ub))
# add IC+BC to the collocation points
X_train_Nf = torch.vstack((X_train_CP, X_train_Nu)).float()  # Collocation Points of x (CP)
T_train_Nf = torch.vstack((T_train_CP, T_train_Nu)).float()  # Collocation Points of t (CP)

# turn on gradients for PINN
X_train_Nf.requires_grad=True
T_train_Nf.requires_grad=True
# Training dataset
train_data = DictDataset({'x': X_train_Nf, 't':T_train_Nf}, name='train')
# test dataset
test_data = DictDataset({'x': X_test, 't':T_test, 'y':Y_test}, name='test')
# torch dataloaders
batch_size = X_train_Nf.shape[0]  # full batch training
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,
                                           collate_fn=train_data.collate_fn,
                                           shuffle=False)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size,
                                         collate_fn=test_data.collate_fn,
                                         shuffle=False)

# neural net to solve the PDE problem bounded in the PDE domain
net = blocks.MLP(insize=2, outsize=1, hsizes=[32, 32], nonlin=nn.Tanh)
# symbolic wrapper of the neural net
pde_net = Node(net, ['x', 't'], ['y_hat'], name='net')
# evaluate forward pass on the train data
net_out = pde_net(train_data.datadict)

# symbolic Neuromancer variables
y_hat = variable('y_hat')
t = variable('t')  # temporal domain
x = variable('x')  # spatial domain
# get the symbolic derivatives
dy_dt = y_hat.grad(t)
dy_dx = y_hat.grad(x)
d2y_d2x = dy_dx.grad(x)

This is where I changed the PINN form

# get the PINN form
f_pinn = dy_dt - d2y_d2x + torch.exp(-t)* RELU(torch.sin(torch.pi * x) - torch.pi ** 2 * torch.sin(torch.pi * x))

scaling = 100.
ell_f = scaling*(f_pinn == 0.)^2
ell_u = scaling*(y_hat[-Nu:] == Y_train_Nu)^2  # remember we stacked CP with IC and BC
ymin, ymax = Y_test.min(), Y_test.max()
con_1 = scaling*(y_hat <= ymax)^2
con_2 = scaling*(y_hat >= ymin)^2
pinn_loss = PenaltyLoss(objectives=[ell_f, ell_u], constraints=[con_1, con_2])
problem = Problem(nodes=[pde_net],      # list of nodes (neural nets) to be optimized
                  loss=pinn_loss,       # physics-informed loss function
                  grad_inference=True   # argument for allowing computation of gradients at the inference time)
                 )

from neuromancer.trainer import Trainer

epochs = 5000
trainer = Trainer(
    problem.to('cpu'),
    train_loader,
    optimizer=torch.optim.Adam(problem.parameters(), lr=0.001),
    epochs=epochs,
    epoch_verbose=200,
    train_metric='train_loss',
    dev_metric='train_loss',
    eval_metric="train_loss",
    warmup=epochs,
)
best_model = trainer.train()

# load best trained model
problem.load_state_dict(best_model)

# evaluate trained PINN on test data
PINN = problem.nodes[0]
y1   = PINN(test_data.datadict)['y_hat']

## Plot solution
# y_pinn = y1.reshape(shape=[100,200]).transpose(1,0).detach().cpu()
# plot3D(X, T, y_pinn)

[drgona]

Dear @akruel01, thanks for the inquiry.

Some PDE are hard to solve with vanilla PINNs. There are some well-known challenges associated with numerical ill-conditioning of the loss function and associated gradients. For more details, please see this lecture by Paris Perdikaris

I can also suggest a couple of papers addressing this issue and proposing some remedies.
https://www.sciencedirect.com/science/article/pii/S002199912100663X
https://arxiv.org/abs/2308.08468
https://arxiv.org/abs/2402.01868