levimcclenny / sa-pinns

Implementation of the paper "Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism" [AAAI-MLPS 2021]

Self-Adaptive PINN - Official Implementation

Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism [AAAI-MLPS 2021]

Levi McClenny^1,2, Ulisses Braga-Neto¹

Accepted to AAAI-MLPS 2021

Update: The self-adaptive implementations of the Allen-Cahn, Burgers, and Helmholtz PDE systems, shown here, are available in our new package TensorDiffEq

Those examples in TensorDiffEq are available here

Paper: https://arxiv.org/pdf/2009.04544.pdf

Abstract: Physics-Informed Neural Networks (PINNs) have emerged recently as a promising application of deep neural networks to the numerical solution of nonlinear partial differential equations (PDEs). However, the original PINN algorithm is known to suffer from stability and accuracy problems in cases where the solution has sharp spatio-temporal transitions. These stiff PDEs require an unreasonably large number of collocation points to be solved accurately. It has been recognized that adaptive procedures are needed to force the neural network to fit accurately the stubborn spots in the solution of stiff PDEs. To accomplish this, previous approaches have used fixed weights hard-coded over regions of the solution deemed to be important. In this paper, we propose a fundamentally new method to train PINNs adaptively, where the adaptation weights are fully trainable, so the neural network learns by itself which regions of the solution are difficult and is forced to focus on them, which is reminiscent of soft multiplicative-mask attention mechanism used in computer vision. The basic idea behind these Self-Adaptive PINNs is to make the weights increase where the corresponding loss is higher, which is accomplished by training the network to simultaneously minimize the losses and maximize the weights, i.e., to find a saddle point in the cost surface. We show that this is formally equivalent to solving a PDE-constrained optimization problem using a penalty-based method, though in a way where the monotonically-nondecreasing penalty coefficients are trainable. Numerical experiments with an Allen-Cahn stiff PDE, the Self-Adaptive PINN outperformed other state-of-the-art PINN algorithms in L2 error by a wide margin, while using a smaller number of training epochs. An Appendix contains additional results with Burger's and Helmholtz PDEs, which confirmed the trends observed in the Allen-Cahn experiments.

_{¹Texas A&M Dept. of Electrical Engineering, College Station, TX
²US Army CCDC Army Research Lab, Aberdeen Proving Ground/Adelphi, MD}

Requirements

Code was implemented in python 3.7 with the following package versions:

tensorflow version = 2.3
keras version = 2.2.4

and matplotlib 3.1.1 was used for visualization. It is expected that any combination of recent numpy/matplotlib will be sufficient, however issues have been experienced on tensorflow versions <2.3.0

Virtual Environment (Optional)

(Mac) To create a virtual environment to run this code, download the repository either via git clone or by clicking download at the top of github, then navigate to the top-level folder in a terminal window and execute the commands

python3 -m venv --system-site-packages ./venv
source ./venv/bin/activate

This will create a virtual environment named venv in that directory (first line) and drop you into it (second line). At that point you can install/uninstall package versions without effecting your overall environment. You can verify you're in the virtual environment if you see (venv) at the beginning of your terminal line. At this point you can install the exact versions of the packages listed here with the pip into the venv:

pip install tensorflow==2.3 numpy==1.17.2 keras==2.2.4

run

python versiontest.py

And you should see the following output:

Using TensorFlow backend
tensorflow version = 2.3
keras version = 2.2.4
numpy version = 1.17.2

TeX Dependencies

(Debian) Some plots require TeX packages, you can have them installed using the following command:

sudo apt-get -qq install texlive-fonts-recommended texlive-fonts-extra dvipng

Data

The data used in this paper is publicly available in the Raissi implementation of Physics-Informed Neural Networks found here. It has already been copied into the appropriate directories for utilization in the script files.

Usage

You can recreate the results of the paper by simply navigating to the desired system (i.e. opening the Burgers folder) and running the .py script in the folder. After opening the Burgers folder, simply run

python burgers.py

And the training will begin, followed by the plots.

Note

The results in the paper were calculated on GPU. Running for the full 10k/10k training iterations for Adam and L-BFGS will likely take a very long time on CPU.

Citation

Cite using the Bibtex citation below:

@article{mcclenny2020self,
  title={Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism},
  author={McClenny, Levi and Braga-Neto, Ulisses},
  journal={arXiv preprint arXiv:2009.04544},
  year={2020}
}

sa-pinns's Issues

How to solve Schrödinger equation with SA-PINN?

Thanks for your important contribution! I want to solve the Schrödinger equation,but i have no idea how to implement it.Could you please give some advice about it.Thank you in advance!

$$ \begin{aligned} &i h_t+0.5 h_{x x}+|h|^2 h=0, \quad x \in[-5,5], \quad t \in[0, \pi / 2], \\ &h(0, x)=2 \operatorname{sech}(x) \\ &h(t,-5)=h(t, 5) \\ &h_x(t,-5)=h_x(t, 5) \end{aligned} $$

the problem in Helmholtz ntk

Dear @levimcclenny ,
Thanks for sharing the code, it is awesome. I have a question about helmholtz ntk, the uploaded code doesn't use ntk, could you please update the code? In addition, do you think I can use it to deal with the ultrasound data e.g. frequency = 5MHz.
Thanks a lot

How to run this on GPU?

Just setting the device is not useful. Is there any special trick to let it run on GPU?

Bug (Error) Report

Hi @levimcclenny

A very large weight factor (100) is used in the simulations:

i.e. u_weights = tf.Variable(100*tf.random.uniform([N0, 1]))

However, If we set the weight factor to 1.0 (default value)

This proposed Self-Adaptive PINN will NOT work any more, see the output figure from AC example with the following modification:
u_weights = tf.Variable(1*tf.random.uniform([N0, 1]))

at t = 0.99, obviouly, the numerical error is (very) large.

Best
Libo

Is there a wrong?

mse_0_u = tf.reduce_mean(tf.square(u_weights*(u0 - u0_pred)))
mse_f_u = tf.reduce_mean(tf.square(col_weights*f_u_pred))

order your paper,

mse_0_u and mse_f_u should be:

mse_0_u = tf.reduce_mean(u_weightstf.square((u0 - u0_pred)))
mse_f_u = tf.reduce_mean(col_weightstf.square(f_u_pred))

I am not sure for this, is it right?

loss increased in 1D wave equation

Dear @levimcclenny ,
Thanks for the innovative contribution. I tried to run the wave equation 1D but the loss increased all the time. I don't know why. Could you please help me check it? Thanks a lot!!

Implementing the 1D wave equation

Thank you for the great contribution. PINNs have been an amazing tool in scientific computing and this enhancement makes it better. I replicated the code of burger's equation and made changes to model the 1D wave equation with a source term:

@tf.function
def f_model(x,t):
    c = tf.constant(1, dtype = tf.float32)
    Amp = tf.constant(1, dtype = tf.float32)
    frequency = tf.constant(1, dtype = tf.float32)
    sigma = tf.constant(0.5, dtype = tf.float32)
    source_x_coord = tf.constant(0, dtype = tf.float32)
    
    Gaussian_impulse =  Amp * tf.exp(-(1/(sigma**2))*(x-source_x_coord)**2)
    S = Gaussian_impulse * tf.sin( 1 * tf.constant(math.pi)  * frequency * t )
    u = u_model(tf.concat([x,t], 1))
    u_x = tf.gradients(u,x)
    u_xx = tf.gradients(u_x, x)
    u_t = tf.gradients(u,t)
    u_tt = tf.gradients(u_t,t)
    f_u = u_tt + (c**2) * u_xx - S
    
    return f_u

However, the output is not quite what I expected.

I tried a couple of network architectures and learning rates, but the result seems to be the same. What would you suggest to get more accurate results?
Also, I saw how you placed the boundary conditions but I'm not sure how to specify the initial condition? Is it extracted from the exact solution?

Thanks again! Awaiting your feedback.

Question about Viscous Burgers equation

 Thank you very much for sharing your code. It has been extremely helpful and inspiring. I appreciate your contributions to the development of PINN. 
 I have downloaded your source code and tried running the code related to the Viscous Burgers equation. However, I'm experiencing gradient explosions and ending up with NaN values for the loss when using the L-BFGS algorithm. Do you have any suggestions on how to address this issue?