FBPINNs: domain decomposition for physics-informed neural networks

The problem plain PINNs have with high frequencies

The why section of the README is the clearest statement of the problem, and it is worth restating because it defines the whole project. Physics-informed neural networks are a popular approach for solving forward and inverse problems related to PDEs, but they often struggle with problems that have high frequencies or multi-scale solutions. The cause is given as two things, the spectral bias of neural networks and the heavily increasing complexity of the PINN optimisation problem. Spectral bias is the usual complaint about neural networks, that they learn the smooth part of a function first and the oscillatory part last, and when the answer you need is the oscillatory part, a single network over a whole domain spends its capacity on the wrong scales. The fix FBPINNs propose is a divide and conquer approach, and the README states the empirical claim plainly, that FBPINNs significantly outperform PINNs in terms of accuracy and computational efficiency when solving problems with high frequencies and multi-scale solutions. The comparison figures in the readme are a high frequency one dimensional harmonic oscillator solved both ways, and a two plus one dimensional wave equation with multiscale sources. The underlying paper is the July 2023 Advances in Computational Mathematics article on FBPINNs by Moseley, Nissen-Meyer and Markham, and the slides from the authors' 2023 Maths4DL talk are linked.

Divide the domain, put a network in each, and sum them

The difference section answers the question of what the F in FBPINN means, and the answer has three parts. First, the problem domain is divided into many small, overlapping subdomains. Second, a neural network is placed within each subdomain, and the solution to the PDE is defined as the summation over all the subdomain networks. Third, each subdomain network is locally confined to its subdomain by multiplying it by a smooth, differentiable window function, and the inputs of each network are individually normalised over their subdomain. That is the mechanism in full. The summation is the part that makes the approach different from training one big network, because the global solution is assembled from local pieces rather than learned globally, and the window functions are what stop the local pieces from interfering badly at the boundaries, since the overlap regions are covered twice and the window decides how the two contributions blend. The hypothesis, stated in the readme, is that this divide and conquer approach significantly reduces the complexity of the PINN optimisation problem, and that the individual normalisation ensures the effective frequency each subdomain network sees is low, which reduces the effect of spectral bias. The two claims are separable, and the second is the more surprising one.

Normalisation and windowing are where the accuracy comes from

It is tempting to read the domain decomposition as the whole idea, but the readme's own explanation puts the emphasis elsewhere. The optimisation complexity argument says decomposition makes the problem easier to fit. The spectral bias argument says normalisation is what makes each local network effective, because if a subdomain covers a small patch of a multi-scale solution, then after normalising the inputs of that patch to its own range, the frequencies the network has to represent are the high ones relative to a small domain, and the relative burden is much lower than for the whole domain at once. That is why the readme lists individual subdomain normalisation as a separate ingredient from domain decomposition. The window function is the third ingredient and it exists for a different reason, which is that overlapping subdomains both contribute in the overlap region and the window decides the blend. For a user of the library this matters practically, because the decomposition class and the window behaviour are things you configure, and the readme explicitly says the high-level interface is now flexible enough to define irregular and multilevel decompositions and custom subdomain networks. A researcher who wants a non-uniform mesh, which is what a real domain usually needs, can therefore build one rather than being forced onto a uniform grid.

Subdomain scheduling, and how boundary conditions travel

The scheduling section is the feature that makes FBPINNs usable for time-dependent problems, and it is separate from the decomposition. Another advantage of using domain decomposition is that you can control which parts of the domain are solved at each training step, which is useful if you want to control how boundary conditions are communicated across the domain. The example is a time-stepping scheduler for time-dependent PDEs, where the solution is learned forwards in time from a set of initial conditions, and the readme illustrates it with the time-dependent Burgers equation. The mechanism is a subdomain scheduler taken from fbpinns.schedulers, which defines which subdomains are actively training and which subdomains have fixed parameters at each training step. So instead of every subdomain network updating at every step, you can freeze the ones representing later time and train the ones representing earlier time, and then hand the state forward. This is the difference between a method that solves a time-dependent problem and a method that treats it as one big time-independent fit, and it also explains why the fixed parameters matter, since the frozen subdomain output is the boundary data the next one is trained against.

The JAX rewrite, and what vmap actually buys

The major update note at the top of the readme is the one to read before anything else, because it describes a rewrite rather than an increment. The fbpinns library has been rewritten in JAX, and the readme says it now runs ten to a thousand times faster than the original PyTorch code, by parallelising subdomain computations using jax.vmap, and that it scales to thousands of subdomains. The mechanism is the interesting part and it is stated plainly enough. A PINN over a whole domain is one optimisation problem, and its cost grows with the size of the network and the domain. An FBPINN is a collection of many small problems that are coupled only through the interface conditions, so the subdomain computations are almost embarrassingly parallel, and expressing that with vmap turns a loop over subdomains into a batched operation that the accelerator can execute at once. That is the real reason the speedup can reach up to three orders of magnitude rather than a constant factor. The rewrite also brought inverse problems, arbitrary types of boundary and data constraints, irregular and multilevel decompositions and custom subdomain networks, and the high-level interface is described as much more flexible and easier to use. The v0.2.0 release note is the place the readme sends you for the details of that update.

The six objects you assemble to solve a problem

The getting started section is a six step recipe, and it doubles as the map of the library's structure. You define the problem domain by selecting or defining your own fbpinns.domains.Domain class, then define the PDE to solve and any problem constraints, such as boundary conditions or data constraints, through a fbpinns.problems.Problem class. Third, you define the domain decomposition with a fbpinns.decompositions.Decomposition class, and fourth the neural network placed in each subdomain through a fbpinns.networks.Network class. Fifth, you keep track of the training hyperparameters by passing those classes and their initialisation values to a fbpinns.constants.Constants object, and sixth you start the training by instantiating a fbpinns.trainers.FBPINNTrainer with the Constants object. The worked example is the one dimensional harmonic oscillator from the figures, using RectangularDomainND, HarmonicOscillator1D, RectangularDecompositionND and the FCN network. The point of the decomposition into small classes is that each piece is replaceable, and a research user is expected to write their own rather than configure a monolithic solver, which is also why the examples directory is the more useful documentation for most readers. There are four notebooks, one defining your own problem, one using hard constraints, one solving inverse problems, and one using subdomain scheduling for the one plus one dimensional Burgers equation.

Installing it, and reading the alpha label honestly

Installation is a Python environment, a clone and an editable install, and the readme says fbpinns only requires Python libraries to run, with JAX as the main computational engine.

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

git clone [email protected]:benmoseley/FBPINNs.git

pip install -e .

The editable install is explained in the readme, which notes that it installs the fbpinns package in editable mode so that changes to the source code are immediately present in the package, which is the right choice for a project where you are expected to write your own domains, problems and networks. The packaging metadata is small and modern, with hatchling as the build backend, a MIT licence, a Python requirement written as ~=3.9, and a dependency list of jax at 0.4.8 or newer, optax, numpy, scipy, matplotlib, tensorboardX and ipython. The development status classifier says Alpha, and that is the number to weigh against the last release, v0.2.0 from 2023-07-30, while the last push to the repository was on 2026-05-29. The practical advice is to treat the examples notebooks as the specification, because a library whose API is designed to be subclassed documents itself through worked examples rather than through a reference manual.

Editorial conclusion

Adopt FBPINNs if you are solving a PDE with high frequency or multi-scale behaviour, where a single physics-informed network struggles, and if you are comfortable with JAX and with the loose requirements of a research code release. It does not fit a smooth, low frequency problem, where the plain approach is easier and cheaper, and it does not fit a production service, because the package is classified as Alpha and its only tagged release is v0.2.0 from 2023-07-30. Before you build on it, read the examples notebooks in order, since they are the real documentation, and check the JAX version you have against the jax>=0.4.8 requirement.

Frequently asked questions

Official sources

1. benmoseley/FBPINNs on GitHub
2. Issues
3. License: MIT
4. README
5. Releases