Gradient calculation in PINN

Hi, everyone,

I am trying to build a Physics Informed Neural Network (PINN) for solving the Navier-Stokes equations like NVidia SimNet using Flux. The input of the neural network is the position X_i, Y_i, Z_i of each points in 3D space and the output is the velocity Ux_i, Uy_i, Uz_i and pressure P_i at each points. I need to check if the output from the neural network satisfies the Navier-Stokes equations. How can I calculate the first derivative of velocity ∂U_i/∂X_i and pressure ∂P/∂X_i, and second derivative of velocity ∂^2U_i/∂X_i^2 included in the Navier-Stokes equations by automatic differentiation using Zygote?

In other PINN programs using TensorFlow, the gradients are calculated in tf.gradients as follows

u_x = tf.gradients(u, x)[0]
u_y = tf.gradients(u, y)[0]
v_x = tf.gradients(v, x)[0]
v_y = tf.gradients(v, y)[0]

However, since tf.gradients calculates the gradient of the sum of the functions grad(sum(f(x))), I assume that it does not correctly calculate the gradient of the physical quantity at each point.

I may have found the information I was looking for on how to find the first derivative in the discussion shown in the link below.

Gradient of Flux model wrt to weights

Based on this information, the part of the PINN programme I would like to create would look like this.

using Flux
model = Flux.Dense(24,32)
baseline = Flux.params(model)
input = [-1 1 -1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1 1 -1 -1 -1 -1 1 1 1 1]
output = model(input)
grad = Flux.jacobian(x -> model(x), input)

In the program, input is location of 8 points in 3-dimensional domain (X[-1, 1], Y[-1, 1], Z[-1, 1]), and output is velocity in X, Y, Z direction and pressure at the each points.

I am looking for a way to calculate the second order derivative of the output to the input of this neural network, and when I use the Hessian function I get the following error message

grad2 = Flux.hessian(x -> model(x), some_input)
ERROR: output an array, so the gradient is not defined. Perhaps you wanted jacobian.

Could you tell me how to solve this error?

The right answer is Diffractor, or to use the modified getindex gradients of Flux. NeuralPDE.jl uses some tricks in the meantime.

@ChrisRackauckas Could you please provide slightly more detail or reference about Diffractor?

Hi Chris,

Thank you for your reply.

Diffractor.jl doesn’t seem to be available to the public yet, when will it be?

I still don’t understand the details of your second solution, “to use the modified getindex gradients of Flux”, but I will try to read the source code of NeuralPDE.jl.