About PINNx

PINNx is a library for scientific machine learning and physics-informed learning. It is rewritten according to DeepXDE but is enhanced by our Brain Dynamics Programming (BDP) ecosystem. For example, it leverages

• brainstate for just-in-time compilation,

• brainunit for dimensional analysis,

• braintools for checkpointing, loss functions, and other utilities.

Algorithms

PINNx implements the following algorithms, but with the flexibility and efficiency of JAX:

• Solving different problems using PINN

  • solving forward/inverse ordinary/partial differential equations (ODEs/PDEs) [SIAM Rev.]

  • solving forward/inverse integro-differential equations (IDEs) [SIAM Rev.]

  • fPINN: solving forward/inverse fractional PDEs (fPDEs) [SIAM J. Sci. Comput.]

  • NN-arbitrary polynomial chaos (NN-aPC): solving forward/inverse stochastic PDEs (sPDEs) [J. Comput. Phys.]

  • PINN with hard constraints (hPINN): solving inverse design/topology optimization [SIAM J. Sci. Comput.]

• Improving PINN accuracy

  • residual-based adaptive sampling [SIAM Rev., Comput. Methods Appl. Mech. Eng.]

  • gradient-enhanced PINN (gPINN) [Comput. Methods Appl. Mech. Eng.]

  • PINN with multi-scale Fourier features [Comput. Methods Appl. Mech. Eng.]

• (physics-informed) deep operator network (DeepONet)

  • DeepONet: learning operators [Nat. Mach. Intell.]

  • DeepONet extensions, e.g., POD-DeepONet [Comput. Methods Appl. Mech. Eng.]

  • MIONet: learning multiple-input operators [SIAM J. Sci. Comput.]

  • Fourier-DeepONet [Comput. Methods Appl. Mech. Eng.], Fourier-MIONet [arXiv]

  • physics-informed DeepONet [Sci. Adv.]

  • multifidelity DeepONet [Phys. Rev. Research]

  • DeepM&Mnet: solving multiphysics and multiscale problems [J. Comput. Phys., J. Comput. Phys.]

  • Reliable extrapolation [Comput. Methods Appl. Mech. Eng.]

• multifidelity neural network (MFNN)

  • learning from multifidelity data [J. Comput. Phys., PNAS]

Features

PINNx has implemented many algorithms as shown above and supports many features:

• enables the user code to be compact, resembling closely the mathematical formulation.

• complex domain geometries without tyranny mesh generation. The primitive geometries are interval, triangle, rectangle, polygon, disk, ellipse, star-shaped, cuboid, sphere, hypercube, and hypersphere. Other geometries can be constructed as constructive solid geometry (CSG) using three boolean operations: union, difference, and intersection. PINNx also supports a geometry represented by a point cloud.

• 5 types of boundary conditions (BCs): Dirichlet, Neumann, Robin, periodic, and a general BC, which can be defined on an arbitrary domain or on a point set; and approximate distance functions for hard constraints.

• 3 automatic differentiation (AD) methods to compute derivatives: reverse mode (i.e., backpropagation), forward mode, and zero coordinate shift (ZCS).

• different neural networks: fully connected neural network (FNN), stacked FNN, residual neural network, (spatio-temporal) multi-scale Fourier feature networks, etc.

• many sampling methods: uniform, pseudorandom, Latin hypercube sampling, Halton sequence, Hammersley sequence, and Sobol sequence. The training points can keep the same during training or be resampled (adaptively) every certain iterations.

• 4 function spaces: power series, Chebyshev polynomial, Gaussian random field (1D/2D).

• data-parallel training on multiple GPUs.

• different optimizers: Adam, L-BFGS, etc.

• conveniently save the model during training, and load a trained model.

• callbacks to monitor the internal states and statistics of the model during training: early stopping, etc.

• uncertainty quantification using dropout.

• float16, float32, and float64.

• many other useful features: different (weighted) losses, learning rate schedules, metrics, etc.