FBINNs are a new, scalable approach for solving large problems relating to differential equations called finite basis physics-informed neural networks. FBPINNs are inspired by classical finite element methods, where the solution of the differential equation is expressed as the sum of a finite set of basis functions with compact support. In FBPINNs, neural networks are used to learn these basis functions, which are defined over small, overlapping subdomains. FBINNs are designed to address the spectral bias of neural networks by using separate input normalisation over each subdomain and reduce the complexity of the underlying optimisation problem by using many smaller neural networks in a parallel divide-and-conquer approach.

Why FBPINNs?

• Physics-informed neural networks (PINNs) are a popular approach for solving forward and inverse problems related to PDEs
• However, PINNs often struggle to solve problems with high frequencies and/or multi-scale solutions
• This is due to the spectral bias of neural networks and the heavily increasing complexity of the PINN optimisation problem
• FBPINNs improve the performance of PINNs in this regime by combining them with domain decomposition, individual subdomain normalisation and flexible subdomain training schedules
• Empirically, FBPINNs significantly outperform PINNs (in terms of accuracy and computational efficiency) when solving problems with high frequencies and multi-scale solutions (Fig 1 and 2)

How are FBPINNs different to PINNs?

To improve the scalability of PINNs to high frequency/ multiscale solutions:

• FBPINNs divide the problem domain into many small, overlapping subdomains.

• A neural network is placed within each subdomain, and the solution to the PDE is defined as the summation over all subdomain networks.

• Each subdomain network is locally confined to its subdomain by multiplying it by a smooth, differentiable window function.

• Finally, the inputs of each network are individually normalised over their subdomain.

The hypothesis is that this "divide and conquer" approach significantly reduces the complexity of the PINN optimisation problem. Furthermore, individual subdomain normalisation ensures the "effective" frequency each subdomain network sees is low, reducing the effect of spectral bias.

Another advantage of using domain decomposition is that we can control which parts of the domain are solved at each training step.

This is useful if we want to control how boundary conditions are communicated across the domain.

For example, we can define a time-stepping scheduler to solve time-dependent PDEs, and learn the solution forwards in time from a set of initial conditions.

This is done by specifying a subdomain scheduler (from fbpinns.schedulers), which defines which subdomains are actively training and which subdomains have fixed parameters at each training step.

Installation

fbpinns only requires Python libraries to run.

JAX is used as the main computational engine for fbpinns.

To install fbpinns, we recommend setting up a new Python environment, for example:

conda create -n fbpinns python=3  # Using conda
conda activate fbpinns

then cloning this repository:

git clone git@github.com:benmoseley/FBPINNs.git

and running this command in the base FBPINNs/ directory (will also install all of the dependencies):

pip install -e .

Note this installs the fbpinns package in "editable mode" - you can make changes to the source code and they are immediately present in the package.

Getting started

Forward and inverse PDE problems are defined and solved by carrying out the following steps:

1. Define the problem domain, by selecting or defining your own fbpinns.domains.Domain class
2. Define the PDE to solve, and any problem constraints (such as boundary conditions or data constraints), by selecting or defining your own fbpinns.problems.Problem class
3. Define the domain decomposition used by the FBPINN, by selecting or defining your own fbpinns.decompositions.Decomposition class
4. Define the neural network placed in each subdomain, by selecting or defining your own fbpinns.networks.Network class
5. Keep track of all the training hyperparameters by passing these classes and their initialisation values to a fbpinns.constants.Constants object
6. Start the FBPINN training by instantiating a fbpinns.trainers.FBPINNTrainer using the Constants object.

For example, to solve the 1D harmonic oscillator problem:

import numpy as np

from fbpinns.domains import RectangularDomainND
from fbpinns.problems import HarmonicOscillator1D
from fbpinns.decompositions import RectangularDecompositionND
from fbpinns.networks import FCN
from fbpinns.constants import Constants
from fbpinns.trainers import FBPINNTrainer

c = Constants(
    domain=RectangularDomainND,# use a 1D problem domain [0, 1]
    domain_init_kwargs=dict(
        xmin=np.array([0,]),
        xmax=np.array([1,]),
    ),
    problem=HarmonicOscillator1D,# solve the 1D harmonic oscillator problem
    problem_init_kwargs=dict(
        d=2, w0=80,# define the ODE parameters
    ),
    decomposition=RectangularDecompositionND,# use a rectangular domain decomposition
    decomposition_init_kwargs=dict(
        subdomain_xs=[np.linspace(0,1,15)],# use 15 equally spaced subdomains
        subdomain_ws=[0.15*np.ones((15,))],# with widths of 0.15
        unnorm=(0.,1.),# define unnormalisation of the subdomain networks
    ),
    network=FCN,# place a fully-connected network in each subdomain
    network_init_kwargs=dict(
        layer_sizes=[1,32,1],# with 2 hidden layers
    ),
    ns=((200,),),# use 200 collocation points for training
    n_test=(500,),# use 500 points for testing
    n_steps=20000,# number of training steps
    optimiser_kwargs=dict(learning_rate=1e-3),
    show_figures=True,# display plots during training
)

run = FBPINNTrainer(c)
run.train()# start training the FBPINN

The FBPINNTrainer will automatically start outputting training statistics, plots and tensorboard summaries. The tensorboard summaries can be viewed by installing tensorboard and then running tensorboard --logdir results/summaries/

Comparing to PINNs

You can easily train a PINN using the same hyperparameters above, using:

from fbpinns.trainers import PINNTrainer

c["network_init_kwargs"] = dict(layer_sizes=[1,64,64,1])# use a larger neural network
run = PINNTrainer(c)
run.train()# start training a PINN on the same problem

Going further

See the examples folder for more advanced examples covering:

• how to define your own Problem class
• how to use hard boundary constraints
• how to solve an inverse problem
• how to use subdomain scheduling

Workshop on Scalable PINNs and minimal FBPINN implementation

Check out our scalable PINNs workshop, where we show how to code a minimal FBPINN and ELM-FBPINN in JAX.

FAQs

Installation

I get the error: RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. when using Apple GPUs.

• As of this commit, JAX only has experimental support for Apple GPUs. Either build JAX from source or install a CPU-only version using conda: pip uninstall jax jaxlib and conda install jax -c conda-forge

Using GPUs

How do I train FBPINNs using a GPU?

• Exactly the same code should run on a GPU automatically, without needing any modification. Make sure you have installed the GPU version of JAX, and that JAX can see your GPU devices (e.g. by checking jax.devices())

Understanding the repository

But I don't know JAX!?

• We highly recommend becoming familiar with JAX - it is a fantastic, general-purpose library for accelerated differentiable computing. But even if you don't want to learn JAX, that's ok - all of the front-end classes (Domain, Problem, Decomposition, and Network) can be defined with only basic understanding of jax.numpy (which is essentially the same as numpy anyway).

Methodology

How are FBPINNs different to other PINN + domain decomposition methods?

• In contrast to other PINN + domain decomposition methods (such as XPINNs), FBPINNs by their mathematical construction do not require additional interface terms in their loss function, and their solution is continuous across subdomain interfaces. Essentially, FBPINNs can just be thought of as defining a custom neural network architecture for PINNs - everything else stays the same.