Solving differential equations using neural networks with PyDEns
Partial differential equations (PDEs) and ordinary differential equations (ODEs) bother researchers from all domains of applied sciences, including engineering, biology and economics. More often that not, carrying out analytical solutions to PDE and ODE problems is an impossible task.
As a result, a lot of effort has been put into the development of methods to obtain numerical solutions. Now, you will learn how to solve PDEs and ODEs via neural networks on Python using PyDEns framework.
Differential equations & neural networks
After the initial development of the approach, it turned out that neural networks-based algorithms can be easily modified to deal with new PDE-related setups: (i) parametric families of PDEs and (ii) PDEs with trainable coefficients. This, of course, is an impossible feat for classical numerical algorithms. Conveniently enough, we developed an open-source PyDEns-framework on Python for handling all these PDE-problems.
We organise the rest of the article as follows:
Let’s now dive into the basis of PyDEns.
Deep Galerkin-algorithm: math behind PyDEns
Consider a simple Laplace problem with a constant boundary condition:
The basic premise of the DeepGalerkin-algorithm is to parametrize the trial function (a space of functions in which we look for a solution to the problem) by a neural network with number of inputs equal to the number of variables:
Direct substitution of the network into the equation gives us the measure of how much does a specific set of neural network parameters not fit the problem:
It is only natural to apply the standard loss function (say, L2) to this measure and minimise it on batches of points using SGD:
Let’s now get into applying the algorithm with PyDEns.
PyDEns for simple PDEs
To write down problems in PyDEns, one uses mathematical token D for differential operations and common mathematical operators from torch:
We can now make a Solver-instance and run the optimisation procedure.
Use predict-method to get the solution in a set of points:
You can also obtain the solution on a 2D-grid and check out the resulting surface:
Note that we haven’t supplied any neural network architecture in a Solver-instance. In such cases Solver uses fully-connected architecture with two hidden layers and Sigmoid-activations. Naturally, you might want to change the model — add skip connections and additional layers or change activations. You can do so by supplying layout, units and activation when defining Solver-instance:
The syntax comes from BatchFlow’s ConvBlock. The class is suited for creation of complex neural network architectures in a line of code and worth checking out. For our purposes here, it is a way to experiment with fully connected architectures with skip connections. Layout-argument configures the sequence of layers (f stands for fully connected while a — for activation). Units-argument controls the amount of features in these layers:
Do not forget to supply these arguments into Solver and press .fit:
Yet, there is much more to PyDEns than working with common PDEs.
Dealing with a parametric family of PDEs
Imagine a setup where one doesn’t have a precise measure of a parameter. A simple example could be modelling cooling down of a metal plate. In this case, however, the diffusivity (inverse value of a in the formula) is only known with large error up to an interval of values:
To handle this problem in a classical manner one needs to run a set of simulations, each for different value of diffusivity. Alternatively, one can use a neural-network based approach. As it turns out, extending DeepGalerkin-algorithm to solve a parametric problem comes down to adding a parameter-input to a neural network:
To train one Pydens-model to deal with the parametric family of equations the user needs to add parameter a in the equation:
then indicate to a Solver-instance that there are total three variables x, y, t and one parameter a by supplying additional arguments ndims and nparams:
Remember also that our neural network is trained on batches. To customise the sampling procedure we use NumpySampler from BatchFlow-library. To learn more about Sampler’s powerful capabilities for points’ generation, read about it here. In short, we declare the sampling procedure for making points uniformly (hence, letter ‘u’ in code below) from the spatial domain (square [0, 1] X [0, 1]), time-range [0, 0.5] and parameter-range [.1, 4]:
Let’s take a look at how the model approximates the solution. The lesser the diffusivity value, the slower the plate cools down. That’s how mean temperature drops through time:
The harder the problem, the more complex neural network you need. Here, for instance, you might want to make a neural network with additional attention branch that depends on the parameter a. Yo can easily do this by inheriting base class TorchModel from PyDEns:
Supply the model-class in Solver and run .fit:
Consider how the whole plate cools down with time for different values of diffusivity:
Note that we’ve learned all of that in one run of a training-procedure. Clearly, using classical numerical methods we would have to run thousands of simulations — each for specific value of the parameter.
In the next section we will handle an even more interesting setup.
Handling an inverse problem
Consider the problem of modelling the dynamics of oil&gas fields in porous-media of earth subsurface. To successfully determine the fluid flow one has to deal with a system of PDEs with unknown distributions of media properties (porosity for instance) and unknown initial state of the system. Yet, to make the problem tractable, there’s data present about the pressure of one of the fluids (water component) in an injection well:
In a similar manner, assume you’ve got a first order ODE at hand with unknown initial state:
To compensate for the missing knowledge, we add an additional constraint binding the value of an unknown function in a domain’s point (similar to known pressure in injection well for oil&gas-problem):
For writing down the problem in PyDEns one needs to supply the additional constraint into the Solver and wrap trainable initial condition in letter V (stands for variable). Note also that we need to supply the name of the variable (init) along with its initial value:
The training procedure now consists of two steps: (i) training the model to satisfy the differential equation and (ii) adjusting the initial condition to satisfy the constraint in t=0.5. To simplify training, during the first stage we do not train the variable. For that, call freeze_layers-method.
A hundred iterations is needed to solve the equation and another hundred for adjusting the trainable variable to bind the constraint:
Conclusion
All in all, we’ve covered how to solve PDEs in Python via neural networks using recently developed framework PyDEns. Not only that — we’ve also discussed and handled PDE-problems of new type: parametric families of equations and equations with trainable coefficients. It’s possible for you now to take on interesting and novel setups: Black–Scholes equation from financial math with uncertainty in interest rate — or fetching the density of a vibrating string by solving a corresponding inverse problem.
Hopefully, you can now make your contribution into the emerging area of research of solving differential equations via neural networks. For more information about installing PyDEns visit the GitHub-page of the framework.
