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.
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 , and this theoretical backing has extended into the PINN framework . 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
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 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¶
CollocationSolverND() solver can be initialized in the following way:
boolthat describes whether the
CollocationSolverNDwill 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(layer_sizes, f_model, domain, bcs,
intsdescribing the size of the input, hidden, and output layers of the FC MLP network
funcdescribing the physics of the problem. More info is provided in this section
domainobject containing the collocation points, defined further here
listof BCs describing the problem
booldescribing whether the problem is solved adaptively using the SA-PINN
tf.Variableobject containing the vector of collocation weights used in self-adaptive training, if enabled via
tf.Variableobject containing the vector of initial boundary weights used in self-adaptive training, if enabled via
funcdescribing the lambda function described in the SA-PINN framework. This defaults to squaring the collocation weights if not explicitly defined. Only applicable if
boolenabling 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.
batch_sz = None
newton_eager = True)
intdictating the number of iterated for the selected tensorflow optimizer
intdictating the number of L-BFGS iterations to be completed following the
intindicating the size of batches of collocation points fed into the solver
boolindicating whether the L-BFGS iterations will be executed eagerly
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.
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.
Allan Pinkus. Approximation theory of the mlp model. Acta Numerica 1999: Volume 8, 8:143–195, 1999.
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.
Yeonjong Shin, Jerome Darbon, and George Em Karniadakis. On the convergence and generalization of physics informed neural networks. arXiv preprint arXiv:2004.01806, 2020.