Solve Burgers' Equation with PINNยค
In this example, we would like to show you another example of how to use ConFIG method to train a physics informed neural network (PINN) for solving a PDE.
In this example, we will solve the 1D Burgers' equation:
where \(\nu\) is the viscosity and set as \(\nu=0.01/\pi\) in the current study. The spatial-temporal domain is \(t \in [0,1]\) and \(x \in [-1.0,1.0]\) with the corresponding initial and boundary conditions of
Let's start by defining a dataset where we can sample points for the initial/boundary conditions and the internal domain:
import torch
from typing import Callable,Union
import numpy as np
from scipy.stats.qmc import LatinHypercube
class BurgersDataset():
def __init__(self,
x_start:float=-1.0,x_end:float=1.0,ini_boundary:Callable=lambda x: -1*torch.sin(np.pi*x),simulation_time:float=1.0,
n_pde:int=2000,n_initial:int=50,n_boundary:int=50,device:Union[str,torch.device]="cpu",
nu=0.01/np.pi,
seed:int=21339,) -> None:
self.ini_boundary=ini_boundary
self.x_start=x_start
self.x_end=x_end
self.n_pde=n_pde
self.n_boundary=n_boundary
self.n_initial=n_initial
self.nu=nu
self.simulation_time=simulation_time
self.device=device
self.random_engine_initial_boundary=LatinHypercube(d=1,seed=seed)
self.random_engine_pde=LatinHypercube(d=2,seed=seed)
def sample_initial(self):
with torch.no_grad():
x_initial = torch.tensor(self.random_engine_initial_boundary.random(n=self.n_initial)[:,0],device=self.device,dtype=torch.float32)*(self.x_end-self.x_start)+self.x_start
t_initial = torch.zeros_like(x_initial)
value_initial = self.ini_boundary(x_initial)
return (x_initial,t_initial,value_initial)
def sample_boundary(self):
with torch.no_grad():
x_boundary = torch.cat([self.x_start*torch.ones(self.n_boundary//2,device=self.device),self.x_end*torch.ones(self.n_boundary//2,device=self.device)])
t_boundary = torch.tensor(self.random_engine_initial_boundary.random(n=self.n_boundary)[:,0],device=self.device,dtype=torch.float32)
value_boundary = torch.zeros_like(t_boundary,device=self.device)
return (x_boundary,t_boundary,value_boundary)
def sample_pde(self):
sample = self.random_engine_pde.random(n=self.n_pde)
x=torch.tensor(sample[:,0],device=self.device,dtype=torch.float32)*(self.x_end-self.x_start)+self.x_start
x.requires_grad=True
t=torch.tensor(sample[:,1],device=self.device,dtype=torch.float32)*self.simulation_time
t.requires_grad=True
return (x,t,self.nu)
Then, we can define a neural network with a simple MLP architecture:
import torch.nn as nn
class BurgersNet(nn.Module):
def __init__(self,channel_basics=50,n_layers=4, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.ini_net=nn.Sequential(nn.Linear(2, channel_basics), nn.Tanh())
self.net=[]
for i in range(n_layers):
self.net.append(nn.Sequential(nn.Linear(channel_basics, channel_basics), nn.Tanh()))
self.net=nn.Sequential(*self.net)
self.out_net=nn.Sequential(nn.Linear(channel_basics, 1))
def forward(self, x, t):
ini_shape=x.shape
y = torch.stack([x.view(-1), t.view(-1)], dim=-1)
y = torch.stack([x, t], dim=-1)
y = self.ini_net(y)
y = self.net(y)
y = self.out_net(y)
return y.view(ini_shape)
We can also a initialization function for the weights and biases. Here we use the Xavier initialization, which is a common initialization method for training PINNs:
def xavier_init_weights(m):
if type(m) == nn.Linear:
torch.nn.init.xavier_normal_(m.weight, 1)
m.bias.data.fill_(0.001)
Then, lets define the loss functions where we can calculate the loss for the initial/boundary conditions and the PDE separately:
class BurgersLoss(nn.Module):
def _derivative(self, y: torch.Tensor, x: torch.Tensor, order: int = 1) -> torch.Tensor:
for i in range(order):
y = torch.autograd.grad(
y, x, grad_outputs = torch.ones_like(y), create_graph=True, retain_graph=True
)[0]
return y
def pde_loss(self,network,inputs):
x,t,nu = inputs
u=network(x,t)
x.grad=None
t.grad=None
""" Physics-based loss function with Burgers equation """
u_t = self._derivative(u, t, order=1)
u_x = self._derivative(u, x, order=1)
u_xx = self._derivative(u_x, x, order=1)
return (u_t + u*u_x - nu * u_xx).pow(2).mean()
def boundary_loss(self,network,inputs):
x_b,t_b,value_b=inputs
boundary_loss=torch.mean((network(x_b,t_b)-value_b)**2)
return boundary_loss
def initial_loss(self,network,inputs):
x_i,t_i,value_i= inputs
initial_loss=torch.mean((network(x_i,t_i)-value_i)**2)
return initial_loss
Finally, we can define a tester class to record the best performance of the model during the training process:
import matplotlib.pyplot as plt
from conflictfree.utils import get_para_vector,apply_para_vector
class Tester():
def __init__(self,
data_path="../../experiments/PINN/data/burgers/simulation_data.npy") -> None:
self.simulation_data=torch.from_numpy(np.load(data_path)).to(dtype=torch.float32)
self.x=torch.linspace(-1.0+1/256,1.0-1/256,256).unsqueeze(0).repeat(self.simulation_data.shape[0],1)
self.t=torch.linspace(0,1.0,100).unsqueeze(1).repeat(1,self.simulation_data.shape[1])
self.best_weight=None
self.best_mse=100
def _get_se(self,network):
with torch.no_grad():
device=network.parameters().__next__().device
self.x=self.x.to(device)
self.t=self.t.to(device)
self.simulation_data=self.simulation_data.to(device)
prediction=network(self.x,self.t)
return (prediction-self.simulation_data).pow(2),prediction
def test(self,network):
with torch.no_grad():
mse=self._get_se(network)[0].mean().item()
if mse < self.best_mse:
self.best_mse=mse
self.best_weight=get_para_vector(network)
return mse
def plot_best(self,network):
apply_para_vector(network,self.best_weight)
se,prediction=self._get_se(network)
fig,ax=plt.subplots(1,3,figsize=(10,5))
ax[0].imshow(self.simulation_data.cpu().detach().numpy().T)
ax[0].set_title("simulation data")
ax[1].imshow(prediction.cpu().detach().numpy().T)
ax[1].set_title("prediction")
im=ax[2].imshow(se.cpu().detach().numpy().T)
ax[2].set_title("squared error")
plt.colorbar(im,ax=ax)
plt.show()
print(f"best mse: {self.best_mse}")
Now, let's train the network! we can first try with the classic Adam optimizer:
from torch.optim import Adam,SGD
from tqdm import tqdm
seed=19018
device="cuda:0"
torch.manual_seed(seed)
np.random.seed(seed)
net = BurgersNet()
net.apply(xavier_init_weights)
net.to(device)
optimizer=Adam(net.parameters(),lr=1e-4)
dataset=BurgersDataset(device=device,seed=seed)
loss=BurgersLoss()
tester=Tester()
p_bar=tqdm(range(5000))
for i in p_bar:
optimizer.zero_grad()
loss_i=loss.initial_loss(net,dataset.sample_initial())
loss_b=loss.boundary_loss(net,dataset.sample_boundary())
loss_pde=loss.pde_loss(net,dataset.sample_pde())
loss_total=loss_i+loss_b+loss_pde
loss_total.backward()
optimizer.step()
p_bar.set_postfix({"current_mse":tester.test(net),"best_mse":tester.best_mse})
tester.plot_best(net)
Then, let's try to train the network with our ConFIG method:
from conflictfree.grad_operator import ConFIG_update
from conflictfree.utils import get_gradient_vector,apply_gradient_vector
torch.manual_seed(seed)
np.random.seed(seed)
net = BurgersNet()
net.apply(xavier_init_weights)
net.to(device)
optimizer=Adam(net.parameters(),lr=1e-4)
dataset=BurgersDataset(device=device,seed=seed)
loss=BurgersLoss()
tester=Tester()
p_bar=tqdm(range(5000))
for i in p_bar:
grads=[]
losses=[]
for input_fn,loss_fn in zip([dataset.sample_initial,dataset.sample_boundary,dataset.sample_pde],
[loss.initial_loss,loss.boundary_loss,loss.pde_loss]):
optimizer.zero_grad()
loss_i=loss_fn(net,input_fn())
loss_i.backward()
grads.append(get_gradient_vector(net))
losses.append(loss_i.item())
apply_gradient_vector(net,ConFIG_update(grads))
optimizer.step()
p_bar.set_postfix({"current_mse":tester.test(net),"best_mse":tester.best_mse})
tester.plot_best(net)
As the results show, the ConFIG method can significantly improve the training accuracy of the PINN model. Now, lets try the momentum version of the ConFIG method:
from conflictfree.utils import OrderedSliceSelector
from conflictfree.momentum_operator import PseudoMomentumOperator
torch.manual_seed(seed)
np.random.seed(seed)
net = BurgersNet()
net.apply(xavier_init_weights)
net.to(device)
optimizer=SGD(net.parameters(),lr=5e-4)
dataset=BurgersDataset(device=device,seed=seed)
loss=BurgersLoss()
tester=Tester()
momentum_operator=PseudoMomentumOperator(num_vectors=3)
loss_selector=OrderedSliceSelector()
input_loss_fns=[(input_fn,loss_fn) for input_fn,loss_fn in zip(
[dataset.sample_initial,dataset.sample_boundary,dataset.sample_pde],
[loss.initial_loss,loss.boundary_loss,loss.pde_loss])]
p_bar=tqdm(range(5000))
for i in p_bar:
index,input_loss_fn=loss_selector.select(1,input_loss_fns)
optimizer.zero_grad()
loss=input_loss_fn[1](net,input_loss_fn[0]())
loss.backward()
momentum_operator.update_gradient(net,index,grads=get_gradient_vector(net))
optimizer.step()
p_bar.set_postfix({"current_mse":tester.test(net),"best_mse":tester.best_mse})
tester.plot_best(net)
As the result shows, both the training speed and test accuracy are improved by using the momentum version of the ConFIG method. Please note that the momentum version does not always guarantee a better performance than the non-momentum version. The main feature of the momentum version is the acceleration, as it only requires a single gradient update in each iteration. We usually will just give the momentum version more training epochs to improve the performance further.
Click here to have a check in the PINN experiment in our research paper.