adityalab / pinnsformer

PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks

Publication

Implementation of the paper "PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks."

Authors: Leo Zhiyuan Zhao, Xueying Ding, B.Aditya Prakash

Placement: ICLR 2024 Poster

Paper + Appendix: https://arxiv.org/abs/2307.11833

Training

We also provide demo notebooks for convection, 1d_reaction, 1d_wave, and Navier-Stokes PDEs. The demos include all code for training, testing, and ground truth acquirement.

To visualize the loss landscape, run the above command to train and save the model first, then run the script:

python3 vis_landscape.py

Please adapt the model path accordingly.

Contact

If you have any questions about the code, please contact Leo Zhiyuan Zhao at leozhao1997[at]gatech[dot]edu.

Citation

If you find our work useful, please cite our work:

@article{zhao2023pinnsformer,
  title={PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks},
  author={Zhao, Leo Zhiyuan and Ding, Xueying and Prakash, B Aditya},
  journal={arXiv preprint arXiv:2307.11833},
  year={2023}
}

pinnsformer's Issues

The derivative in the code

I have been closely examining the implementation of the Navier-Stokes equation in your code and I appreciate the effort put into this work. I have a query regarding the computation of the second-order derivative, specifically u_xx.

In the code, the second-order derivative u_xx is calculated directly from u with respect to x_train, as shown below:

u_xx = torch.autograd.grad(u, x_train, ...)

However, I believe that to obtain the second-order derivative u_xx, we should first compute the first-order derivative u_x with respect to x_train, and then differentiate u_x with respect to x_train again, like this:

u_x = torch.autograd.grad(u, x_train, ..., create_graph=True)[0]
u_xx = torch.autograd.grad(u_x, x_train, ...)[0]

Could you please confirm if my understanding is correct? If so, would it be possible to update the code accordingly to reflect this change?

Thank you for your time and consideration. I look forward to your response.

How to add dimension to my input

代码中默认的输入是(x,t)，但是我希望输入是(x,y,t)。我应该如何做

Some confusions in `1d_wave_pinn_ntk`

Hello, great work!

According to the NTK algorithm for PINNs in When and why PINNs fail to train: A neural tangent kernel perspective, I am quite puzzled by the process of solving for J1-J3 below.

Based on my understanding, J1 should be the Jacobian matrix of $u_{tt} - 4u_{xx}$ with respect to the parameters, and J2 should be the Jacobian matrix of $u_t$-initial-conditions with respect to the parameters. Also, the initial-conditions for $u$ should be combined with the boundary-conditions to form J3 ?

...
...
for i in tqdm(range(1000)):
    if i % 50 == 0:
        J1 = torch.zeros((D1, n_params))
        J2 = torch.zeros((D2, n_params))
        J3 = torch.zeros((D3, n_params))

        batch_ind = np.random.choice(len(x_res), kernel_size, replace=False)
        x_train, t_train = x_res[batch_ind], t_res[batch_ind]

        pred_res = model(x_train, t_train)
        pred_left = model(x_left, t_left)
        pred_upper = model(x_upper, t_upper)
        pred_lower = model(x_lower, t_lower)

        for j in range(len(x_train)):
            model.zero_grad()
            pred_res[j].backward(retain_graph=True)
            J1[j, :] = torch.cat([p.grad.view(-1) for p in model.parameters()])

        for j in range(len(x_left)):
            model.zero_grad()
            pred_left[j].backward(retain_graph=True)
            J2[j, :] = torch.cat([p.grad.view(-1) for p in model.parameters()])

        for j in range(len(x_lower)):
            model.zero_grad()
            pred_lower[j].backward(retain_graph=True)
            pred_upper[j].backward(retain_graph=True)
            J3[j, :] = torch.cat([p.grad.view(-1) for p in model.parameters()])
        ...
        ...

---

Here is the code I have roughly modified, I am not sure if it is correct.

        J1 = torch.zeros((D1, n_params))
        J2 = torch.zeros((D2, n_params))
        J3 = torch.zeros((D3, n_params))

        batch_ind = np.random.choice(len(x_res), kernel_size, replace=False)
        x_train, t_train = x_res[batch_ind], t_res[batch_ind]

        pred_res = model(x_train, t_train)
        pred_left = model(x_left, t_left)
        pred_upper = model(x_upper, t_upper)
        pred_lower = model(x_lower, t_lower)

        u_x = torch.autograd.grad(pred_res, x_train, grad_outputs=torch.ones_like(pred_res), retain_graph=True, create_graph=True)[0]
        u_xx = torch.autograd.grad(u_x, x_train, grad_outputs=torch.ones_like(pred_res), retain_graph=True, create_graph=True)[0]
        u_t = torch.autograd.grad(pred_res, t_train, grad_outputs=torch.ones_like(pred_res), retain_graph=True, create_graph=True)[0]
        u_tt = torch.autograd.grad(u_t, t_train, grad_outputs=torch.ones_like(pred_res), retain_graph=True, create_graph=True)[0]
        wave_opt = u_tt - 4 * u_xx  # wave operator
        del u_x, u_xx, u_t, u_tt

        pred_t = torch.autograd.grad(pred_left, t_left, grad_outputs=torch.ones_like(pred_left), retain_graph=True, create_graph=True)[0]

        for j in range(len(x_train)):
            model.zero_grad()
            wave_opt[j].backward(retain_graph=True)
            J1[j, :] = torch.cat([p.grad.view(-1) if p.grad is not None else torch.tensor([0.]).view(-1) for p in model.parameters()])

        for j in range(len(x_left)):
            model.zero_grad()
            pred_t[j].backward(retain_graph=True)
            J2[j, :] = torch.cat([p.grad.view(-1) if p.grad is not None else torch.tensor([0.]).view(-1) for p in model.parameters()])

        for j in range(len(x_lower)):
            model.zero_grad()
            pred_left[j].backward(retain_graph=True)
            pred_lower[j].backward(retain_graph=True)
            pred_upper[j].backward(retain_graph=True)
            J3[j, :] = torch.cat([p.grad.view(-1) for p in model.parameters()])