Variational loss of hp-Variational Physics Informed Neural Networks for 2D-Poisson Equation in Tensorflow

Question

I am trying to reproduce the results from the hp-VPINN paper (https://arxiv.org/pdf/2003.05385.pdf) on tensorflow (v1) for Poisson's equation, particularly the two-dimensional Poisson equation.

In one dimension the variational loss is given by $$ L^p = \sum_{e=1}^{N_{el}}\frac{1}{K^{(e)}}\sum_{k=1}^{K^{(e)}}\left|\mathcal{R}_k^{(e)}\right|^2 $$

where $N_{el}$ is the number of elements and $K^{(e)}$ is the number of test functions used (in this case corresponding to different degrees of Legendre polynomials)

where $\mathcal{R}_k^{(e)}$ is $$ \mathcal{R}_k^{(e)} = \int_{x_{e-1}}^{x_e} \frac{du_{NN}(x)}{dx} \frac{dv_k^{(e)}(x)}{dx}dx - \left. \frac{du_{NN}(x)}{dx}v_k^{(e)}(x)\right|_{x_{e-1}}^{x_{e}} - F_k^{(e)} $$

with $u_{NN}$ being the network's output and and $v_k^{(e)}$ the test function. The boundary term vanishes since the test functions are such that they are zero on the element boundaries and $F_k^{(e)} = \int_{x_{e-1}}^{x_e} f(x)v_k^{(e)}dx$.

To calculate it Gauss-Lobatto quadrature is used giving the following code

import tensorflow as tf
import numpy as np
u_nn_element = tf.reshape(
    tf.stack(
        [
            tf.reduce_sum(
                w_quadrature
                * du_nn_quad_element
                * dtest_quad_element_dx[i]
            )
            for i in range(n_test_functions)
        ]
    ),
    (-1, 1),
)
f_element = jacobian * np.asarray(
    [sum(w_quadrature * f(x_quadrature) * t(x_quadrature) for t in test_functions]
)
residual_element = u_nn_element - f_element
loss_element = tf.reduce_mean(tf.square(residual_nn_element))
total_varloss += loss_element

When combined with the boundary loss (mean squared error for the boundary points at $[-1, 1]$)

In two dimensions the variational loss is given by

$$ L= \sum_{e_x=1}^{N_{el_x}}\sum_{e_y=1}^{N_{el_y}}\frac{1}{K_1K_2}\sum_{k=1}^{K_1K_2}\left|\mathcal{R}_k^{(e_xe_y)}\right|^2 $$

with a residual given by

$$ {}^{(2)} \mathcal{R}_{k_1k_2}^{(e_xe_y)} = - \int_{x_{e-1}}^{x_e}\int_{y_{e-1}}^{y_e} \left(\frac{\partial u_{NN}(x)}{\partial x} \frac{d\phi_{k_1}^{(e_x)}(x)}{dx}\phi_{k_2}^{(e_y)}(y) + \frac{\partial u_{NN}}{\partial y}\phi_{k_1}^{(e_x)}(x)\frac{d\phi_{k_2}^{(e_y)}(y)}{dy} \right) dxdy \\ - F_k^{(e)} $$

with the test functions $\phi_{k_1} \phi_{k2}$

My code attempt has been

l = [
    tf.reduce_sum(
        -wy_quadrature
        * wx_quadrature
        * (
            du_nn_dx * dtest_x_quad_element[i] * y_test_quad_element[j]
            + du_nn_dy
            * dtest_y_quad_element[j]
            * x_test_quad_element[i]
        )
    )
    for i in range(n_test_functions_x)
    for j in range(n_test_functions_y)
]
u_nn_element = tf.reshape(
    tf.stack(l),
    (-1, 1),
)

# element boundaries
jacobian_x = (x1 - x0) / 2
jacobian_y = (y1 - y0) / 2

f = [
    sum(xw_quad * yw_quad * f_quad_element * t1 * t2)
    for t1 in self.test_functions
    for t2 in self.test_functions
]
self.f = jacobian_x * jacobian_y * np.asarray(f)
self.f = self.f.flatten()
self.f = self.f[:, None]


residual_nn_element = u_nn_element - f_element
loss_element = tf.reduce_mean(tf.square(residual_nn_element))
self.varloss_total += loss_element

though my variational error doesn't converge and the net only learns the boundary values. How is the double integral supposed to look in code?

Answer 1

Here are the few observations based on your code.

1, Add a jacobian transformation to the 2D integral as you are doing it for the forcing term as shown in below example ( taken from the 2D hp-PINNS Github )

U_NN_element_1 = tf.convert_to_tensor([[jacobian/jacobian_x*tf.reduce_sum(\
                                    self.wquad[:,0:1]*d1testx_quad_element[r]*self.wquad[:,1:2]*testy_quad_element[k]*d1xu_NN_quad_element) \
                                    for r in range(Ntest_elementx)] for k in range(Ntest_elementy)], dtype= tf.float64)
                    U_NN_element_2 = tf.convert_to_tensor([[jacobian/jacobian_y*tf.reduce_sum(\
                                    self.wquad[:,0:1]*testx_quad_element[r]*self.wquad[:,1:2]*d1testy_quad_element[k]*d1yu_NN_quad_element) \
                                    for r in range(Ntest_elementx)] for k in range(Ntest_elementy)], dtype= tf.float64)
                    U_NN_element = - U_NN_element_1 - U_NN_element_2

As far as learning goes, please try with a very low learning rate and large iterations once to check the model performance for a given problem once.