PINN Collocation Solver¶

Once you’ve defined your problem, it must be compiled such that TensorDiffEq can build the loss function described by the boundary conditions, initial conditions, and physics defined in the previous sections.

Layer Sizes¶

Here is where we will define the neural network size and depth. Currently, most PINN approaches use dense fully connected neural networks for function approximation. Fully-connected Neural Networks have some level of theoretical backing that they will converge to a solution of the underlying function [1][2], and this theoretical backing has extended into the PINN framework [3]. With that being said, currently the only type of network supported in TensorDiffEq is the fully-connected MLP network.

TensorDiffEq uses the Keras API for neural network construction. All you need to do is define a list of layer sizes for your neural network. So, for a network with an [x,t] input, 4 layers deep, with 128 nodes, one would define a layer size list of [2,128,128,1].

For our problem we have been building in the previous sections, we can define layer sizes as such:

layer_sizes = [2, 128, 128, 128, 128, 1]

Or, if your problem is a function of [x,y,t], then you could define the exact same network with an input layer with 3 nodes, i.e.

layer_sizes = [3, 128, 128, 128, 128, 1]

Build and Train the Model¶

In order to compile the model, we first initialize the model we are interested in. Currently, forward solutions of PINNs are performed by the CollocationSolverND() method.

Collocation Solver¶

The primary method of solving forward problems in TensorDiffEq is the collocation solver. This methodology identifies points in the domain of the problem and collocates them to the solution via a loss function. Therefore, this is a natural application for a neural network function approximation.

Instantiate the Model¶

The CollocationSolverND() solver can be initialized in the following way:

CollocationSolverND(assimilate=False)

Args:

assimilate - a bool that describes whether the CollocationSolverND will be used for data assimilation

Note that very little in the solver is truly initialized when creating the CollocationSolverND instance, most comes later in the compile call.

Methods¶

compile(layer_sizes, f_model, domain, bcs,
    isAdaptive=False,
    col_weights=None,
    u_weights=None,
    g=None,
    dist=False)

Args:

layer_sizes - a list of ints describing the size of the input, hidden, and output layers of the FC MLP network
f_model - a func describing the physics of the problem. More info is provided in this section
domain - a domain object containing the collocation points, defined further here
bcs - a list of BCs describing the problem
isAdaptive - a bool describing whether the problem is solved adaptively using the SA-PINN
col_weights - a tf.Variable object containing the vector of collocation weights used in self-adaptive training, if enabled via isAdaptive
u_weights - a tf.Variable object containing the vector of initial boundary weights used in self-adaptive training, if enabled via isAdaptive
g - a func describing the lambda function described in the SA-PINN framework. This defaults to squaring the collocation weights if not explicitly defined. Only applicable if isAdaptive is enabled.
dist - a bool enabling distributed training across multiple GPUs

Model compilation is truly where the rubber meets the road in defining an inference model in TensorDiffEq. We compile the model using the compile method on the CollocationSolverND method. This will build out the loss function in the solver by iterating through the BCs and the IC that define your problem. The compile function will also pull in the collocation points and optimize your f_model function for running in graph-mode in Tensorflow.

fit(tf_iter, newton_iter,
    batch_sz = None
    newton_eager = True)

Args:

tf_iter - an int dictating the number of iterated for the selected tensorflow optimizer
newton_iter - and int dictating the number of L-BFGS iterations to be completed following the tf_iter iterations
batch_sz - an int indicating the size of batches of collocation points fed into the solver
newton_eager - a bool indicating whether the L-BFGS iterations will be executed eagerly

Note

Currently, newton_eager is demonstrating some stability issues as of release v0.1.1. Neither form of newton optimization (eager or graph) is currently supported in a distributed environment.

Additionally, batch_sz is not a tunable parameter in a distributed environment when using the SA-PINN methodology.

Specific notes on optimizing for multi-GPU environments will be addressed later in this document.

References¶

1: Allan Pinkus. Approximation theory of the mlp model. Acta Numerica 1999: Volume 8, 8:143–195, 1999.
2: Tianping Chen and Hong Chen. Approximations of continuous functionals by neural networks with application to dynamic systems. IEEE Transactions on Neural Networks, 4(6):910–918, 1993.
3: Yeonjong Shin, Jerome Darbon, and George Em Karniadakis. On the convergence and generalization of physics informed neural networks. arXiv preprint arXiv:2004.01806, 2020.