What loss function is used in the tutorial for forward PINN for solving Poisson problem in 1D?

[sgorgiSW]

Hello,
First I would like to quickly mention that the question is in the following paragraph and the big wall of text with code below it is probably not necessary to read for my question.

I am doing PINNs in my bachelor's thesis and to start, I got assigned to try to use them to solve Poisson's problem in 1D with Dirichlet boundary condition. I tried it with keras (high level API for tensorflow), but my neural network wasn't really approximating the real solution. In some paper I saw DeepXDE mentioned so I gave it a try according to the demo https://deepxde.readthedocs.io/en/latest/demos/pinn_forward/poisson.1d.dirichlet.html. And I have to say I am amazed, the same problem with the same set up (same network architecture, optimizer and number of epochs) was solved by your DeepXDE much faster and nearly perfectly. So concerning the precision, could you, please, tell me, what loss functions are optimized and used during the training in the network from the DeepXDE demo? I tried digging in the source code to see if I could find it, but was unsuccessful.

Here below I write how I did the training in my own code using keras. It might not be needed to answer the question above, but I provided it just in case.
I subclassed the base squential model in keras API to write a custom train_step method, the function which processes one batch of data during training, so that I have access to the input, which I wouldn't have with the default keras class for a sequential model. Here is my init and train_step method (I provided a summary and some commentary/explanation below it, so in case you are interested in reading further, you can maybe start there before the code itself):

class Poisson_problem_1D_pointwise(keras.Sequential):
    """
    Model for Poisson problem with fixed right hand side (RHS). Output shape should be the same as the input
    shape, with the shape being (batch_of_points, 1)

    Parameters
    ----------
    RHS_function : Callable function that takes in batch of points as an input
    with shape (batch_size, 1) and returns the right hand side evaluated in those
    points. The output can be of shape (batch_size, 1) or (batch_size).
    Should be made using tensorflow operations (operations on tensors and operations
    from modules tensorflow, tensorflow.math, tensorflow.linalg,and such), so it can be differentiated.
    
    .
    """

# GITHUB DISCUSSION NOTE: in the following I left out the first indentation so that the code fits better
# into the box, but these 2 methods belong under the class above
def __init__(self, RHS_function, left_boundary_point, right_boundary_point, *args, equation_loss_coef=1, boundary_loss_coef=1, **kwargs):
    super().__init__(*args, **kwargs)
    self.RHS_function = RHS_function
    self.equation_loss_coef = equation_loss_coef
    self.boundary_loss_coef = boundary_loss_coef
    self.left_boundary_point = left_boundary_point
    self.right_boundary_point = right_boundary_point

def train_step(self, data):
    x, y_true = data # extracts the input and true output from the batch data
    # I don't use the y_true next, as I am trying to train the network without giving it the correct solution
    
    input_type = x.dtype
    left_boundary_point = tf.constant([[self.left_boundary_point]], dtype=input_type) # needs shape (1, 1) so it can be fed into a dense layer
    # left_boundary_point = tf.Variable(initial_value=[[self.left_boundary_point]], trainable=True, dtype=input_type)
    
    right_boundary_point = tf.constant([[self.right_boundary_point]], dtype=input_type) # needs shape (1, 1) so it can be fed into a dense layer
    # right_boundary_point = tf.Variable(initial_value=[[self.right_boundary_point]], trainable=True, dtype=input_type)
    
    with tf.GradientTape() as weights_update_tape:
        with tf.GradientTape(persistent=True) as second_der_tape:
            second_der_tape.watch(x)    # shape of x: (batch_size, 1)
            second_der_tape.watch(left_boundary_point)
            second_der_tape.watch(right_boundary_point)
    
            with tf.GradientTape(persistent=True) as first_der_tape:
                first_der_tape.watch(x) # shape of x: (batch_size, 1)
                first_der_tape.watch(left_boundary_point)
                first_der_tape.watch(right_boundary_point)
                prediction = self(x, training=True)   # shape: (batch_size, 1)
                prediction_left_boundary = self(left_boundary_point, training=True) # shape: (1, 1)
                prediction_right_boundary = self(right_boundary_point, training=True)   # shape: (1, 1)
    
            u_x = first_der_tape.gradient(target=prediction, sources=x)   # shape: (batch_of_points, 1)
            u_x_left_boundary = first_der_tape.gradient(target=prediction_left_boundary, sources=left_boundary_point)   # shape: (1, 1)
            u_x_right_boundary = first_der_tape.gradient(target=prediction_right_boundary, sources=right_boundary_point)    # shape: (1, 1)
            del first_der_tape
    
        u_xx = second_der_tape.batch_jacobian(target=u_x, source=x) # shape: (batch_of_points, 1, 1)
        u_xx_left_boundary = second_der_tape.batch_jacobian(target=u_x_left_boundary, source=left_boundary_point) # shape: (1, 1, 1)
        u_xx_right_boundary = second_der_tape.batch_jacobian(target=u_x_right_boundary, source=right_boundary_point) # shape: (1, 1, 1)
        del second_der_tape
    
        u_xx = tf.squeeze(u_xx) # shape: (batch_size)
        RHS_evaluated = tf.squeeze( self.RHS_function(x) )  # shape: (batch_size)
        residual = u_xx - RHS_evaluated
        residual_loss = self.equation_loss_coef * tf.math.reduce_mean( tf.math.square(residual) )
    
        u_xx_left_boundary = tf.squeeze(u_xx_left_boundary)
        RHS_at_left_boundary = tf.squeeze( self.RHS_function(left_boundary_point) )
        loss_left_boundary = tf.math.squared_difference(u_xx_left_boundary, RHS_at_left_boundary)
    
        u_xx_right_boundary = tf.squeeze(u_xx_right_boundary)
        RHS_at_right_boundary = tf.squeeze( self.RHS_function(right_boundary_point) )
        loss_right_boundary = tf.math.squared_difference(u_xx_right_boundary, RHS_at_right_boundary)
    
        boundary_loss = self.boundary_loss_coef * (loss_left_boundary + loss_right_boundary) / 2
    
        loss = boundary_loss + residual_loss
    
    self.optimizer.minimize(loss, self.trainable_variables, tape=weights_update_tape)
    
    return {
        "loss": loss
    }

In short what I do is I compute the second derivative of the output with respect to the input for each point inside the interval and for both boundary points. For the points inside the interval I then compute the residual, square it and then take the mean over all points inside the interval. Note: I took the equation as du/dx^2 = F instead of -du/dx^2 = F, therefore the residual is computed as +du/dx^2 - F, instead of
-du/dx^2 - F.
For the points on the boundary, I also take the residual, square it and take the mean (hence the division by 2 in boundary_loss).
I also included weight coeficient for the loss inside the interval and on the boundary, so that it can be customized how much the loss for boundary points and the loss for points inside the interval should play role.

In my deepXDE code I used the same code as in the demo, except I didn't include solution=func in data = dde.data.PDE(geom, pde, bc, 16, 2, solution=func, num_test=100) as I am trying to train the network without knowing the solution. I also removed the L2 metric in model.compile("adam", lr=0.001, metrics=["l2 relative error"]) as that threw an error, because it requires the correct solution, which I excluded, as mentioned in the previous sentence.
In my personal keras code I tried a fully connected network with input of shape (1,), then 3 dense layers with 50 neurons and tanh activation, then output with shape (1,) and no activation, which I think should be the same as in the DeepXDE demo (not sure if the activation is also applied to the output). I used Adam optimizer with learning rate of 0.001 and 10 000 epochs, just like in the demo. One possible difference in the setup that I can notice might be that I used numpy.linspace to generate 18 points (same number as in the demo) from -1 to 1 as my training input. But I don't think that should really affect the accuracy that much.

If you wouldn't mind tellling me, please, I am interested in which loss functions uses DeepXDE for the 1D Poisson problem in the demo, as your DeepXDE gives a very accurate solution. Or if you would possible have any other idea, why my simple implementation doesn't give accurate result (below I included a picture with the accuracy of my network).
Thank you very much in advance.
My NN accuracy:

[sgorgiSW]

Ok I found out that my accuracy was off because I wasn't actually putting the correct equation into my code (I was using a different equation before I switched to the one from deepxde tutorial, but I accidentally forgot to change it in my NN setup). And the slower speed was because I was using more points with smaller batch size (therefore much more training iterations took place) + I was using keras, which has a nice progress bar displayed each epoch in terminal, but writing it to the terminal also took quite a bit of time. After I turned off the progress bar and set my number of points also to 18, I was getting similar training speeds as deepxde.