PINA: a PyTorch Framework for Solving Differential Equations by Deep Learning for Research and Production Environments

Authors: dario  coscia
nicola  demo
gianluigi  rozza
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-01-30 14:47:48

Abstract: We present a versatile software designed for solving differential equations employing neural networks. The package is called PINA, an open-source Python library built upon the robust foundations of PyTorch and Lightning. It allows end-users to formulate their problem and craft their models to effortlessly compute solutions of PDEs by Physics Informed Neural Networks and Neural Operators. The modular structure of PINA permits it to adapt for user specifics, thus offering the freedom to select the most suitable learning techniques for their particular problem domain. Furthermore, by leveraging the capabilities of the Lightning package, PINA adapts to various hardware setups, including GPUs and TPUs. This adaptability positions PINA as an ideal candidate for the transition of these methodologies into production and industrial pipelines, where computational efficiency and scalability are of paramount importance.

Main paper: 
 introduction: 
 In recent years, the world has seen an unprecedented revolution in artificial intelligence (AI) and machine learning (ML), that has permeated numerous sectors, transforming solutions and processes in many different fields of applied sciences. Within the scientific computing community, this revolution has manifested itself as a powerful tool for overcoming the limitations inherent in traditional methods for solving complex differential equations.Among the promising developments in this arena, two standout approaches have emerged as central players for differential equation learning: Physics-Informed Neural Networks (PINNs) Raissi et al. (2019) and Neural Operators (NOs) Li et al. (2020) ; Lu et al. (2021a) ; Bhattacharya et al. (2021) . These methodologies exploit the knowledge of the equations, symmetries, and data to approximate the unknown solution of the differential equation or the differential operator defining the problem. These recent advancements combined with the evolution of open-source frameworks, such as TensorFlow (Abadi et al., 2015) , and PyTorch (Paszke et al., 2019) led to the development of several libraries for solving ODEs and PDEs via PINNs and NOs. PINN TensorFlow-based libraries include DeepXDE (Lu et al., 2021b) (which also supports PyTorch), TensorDiffEq (McClenny et al., 2021) and PyDEns (Koryagin et al., 2019) ; while PyTorch-based libraries include NeuroDiffEq (Chen et al., 2020) , IDRLNet (Peng et al., Figure 1 : PINA package application programming interface. Starting from the problem definition, a specific model is passed to the solver, which defines, together with the trainer, the optimization strategy of the model. 2021). For NO NeuralOperator (Li et al., 2020; Kovachki et al., 2021) is the main library. Finally, hybrid software for PDE learning includes (Bouziani & Ham, 2023; Kidger, 2022) .There are multiple challenges with the packages mentioned above that limit their usage for research and production environments. First, most of the packages lack abstract interfaces which limit the possibility of adding extensions, like new loss functions or training procedures. Additionally, the packages presented are sectorized to only PINNs or NOs, without the possibility of combining the two methodologies, which is a new research direction in the field (Li et al., 2021; Wang et al., 2021b) . Another common problem of the libraries is the absence of common deep learning advancements for training such as multiple device training, modern model compression techniques, gradient accumulation, and so on. Finally, the possibility of inserting common deep-learning loggers into the training for monitoring is missing. For this reason, we present Physics Informed Neural network for Advanced modeling (PINA), an open-source Python library providing an intuitive interface for solving differential equations using PINNs, NOs or both together. The contribution is organized to show features, capabilities, and practical applications of PINA, illustrating how this software tool can be exploited for solving complex differential equations using deep learning. 
pina: 
 Physics Informed Neural network for Advanced modeling (PINA) is an open-source Python library built-in PyTorch, with PyTorchLightning (Falcon & The PyTorch Lightning team, 2019) as backhand to solve differential equations using artificial intelligence model. Employing PyTorchLightning as the backhand offers professional AI researchers and machine learning engineers the possibility of using advancement training strategies provided by the library. In addition, it provides the possibility to add arbitrary self-contained routines (callbacks) to the training for easy extensions without the need to touch the underlying code. The application programming interface (API) of PINA is schematized in Figure 1 . The pipeline to solve differential equations with PINA follows five steps: problem definition, data generation, model and solver selection, and training. To show the full capabilities of PINA the next sections will follow the prototypical pipeline for solving a problem, highlighting the various features provided by the software. The mathematical notation and a background introduction to PINNs and NOs can be found in Appendix B.The first step is the formalization of the problem. In PINA the problem is formulated by constructing a class inheriting from one or more problem classes (at the moment the available classes are Ab-stractProblem, SpatialProblem, TimeDependentProblem, ParametricProblem, InverseProblem), depending on the nature of the problem treated. For example, a simple ODE that depends only on a spatial variable is defined via a class that inherits only from SpatialProblem. Differently, for a parametric time-independent PDE, the problem class inherits from both SpatialProblem and ParametricProblem. In case the user wants to define its own problem, the AbstractProblem interface must be used as the base class. In the problem formulation class, the user must include information about the domains -e.g. spatial, temporal, parametric -, the output variables, and the conditions that the neural network has to satisfy. Multiple types of geometries are available currently in PINA for defining the domain (see Section 2.2). The output variables are represented by a list of symbols constituting the unknowns of the problem. Indeed, standard PyTorch tensors carry a label (LabelTensor) in PINA, allowing the user an easy way to manipulate the tensors. Finally, for training PINNs and NOs it is essential to give appropriate constraints as a form of loss function. The Condition class encapsulates all the possible ways the loss could be defined, i.e., physical loss, boundary loss, or data loss. The users must use the Condition class to define all the constraints the unknown field needs to satisfy. Moreover, PINA already implements differential operators (e.g. laplacian or grad) and common equations (e.g. Dirichlet boundary conditions, systems of equations) to ease the problem formulation for the users.NO learning procedure uses a finite set of observations and it is trained in a fully supervised manner. These observations, obtained by numerical solver solutions, in PINA can be passed as LabelTensor in the Condition class defined in Section 2.1. Differently, some training strategies, e.g. PINNs, use collocation points sampled inside the domain where the residual of the differential equation (see equation 3) must be evaluated. For these types of solvers, PINA provides a simple sampling strategy for multiple different geometries. In PINA each domain is a Location object, which defines the geometry of the domain. There are already multiple sampling methods in PINA e.g. random uniform, grid sampling, or latin hypercube sampling for the different available multidimensional geometries, e.g. hypercube, hypersphere. In addition to multidimensional geometries, the software also provides set operations (difference, union, intersection, and so on) allowing the user to build a custom domain. Finally, in Condition class the user can also employ available scatter points, and pass them as LabelTensors.Once the model is defined, the user must choose the neural network model to optimize, and the optimization strategy. In PINA the model is represented as a standard torch.nn.Module. The package is equipped with many customizable models and layers (see Table 1) already implemented using PyTorch. The user can then decide to use built-in models (e.g. for benchmarking) or build new models and layers for research purposes.For optimizing the model a specific solver must be used. A solver is a Python object which defines the optimization strategy for the model. In PINA the solver is constructed by inheriting from Solver-Interface, an abstract class wrapping Lightning Modules. Available solvers include a supervised learning solver (SupervisedSolver), particularly crafted for data-driven problems and NO approach, a physics-informed solver (PINN) (Raissi et al., 2019) , and an adversarial solver (GAROM) (Coscia et al., 2023a) . We plan to continuously add solvers as the state-of-the-art evolves. Notice that all solvers are customizable by the user. For example, the PINN solver allows changing the loss (e.g. using a variational loss (Kharazmi et al., 2019) ), or extending the solver with regularization strategies (Yu et al., 2022) , or modifying the optimizer (Davi & Braga-Neto, 2022) . All of these, apparently different solvers, can be changed by a keyword argument in the PINN class.  
experiments: 
 In this section, we show possible benchmark results obtainable with PINA. We want to highlight that the purpose of this section is not to provide accurate measurements of model performance, but rather to show how easily is to benchmark on PINA. As model cases, we use four different models all implemented in PINA: a standard multilayer perceptron (MLP); the skip connection MLP (Wang et al., 2021a ) (m-MLP); a hard constraint MLP (Lu et al., 2021c ) (hard-MLP); the Deep Operator Network (Lu et al., 2021a ) (DeepONet). The models are benchmarked on four different problems using different PINN's learning methodologies: the Burgers and Parametric Poisson equations, with classical PINN learning; the Poisson's equation using extra features (Demo et al., 2023) ; and the Wave equation, using R3 adaptive refinement (Daw et al., 2023) . For a complete description of training details and differential problems see Appendix D and C. In Table 2 the mean square residual for all the simulations done employing PINA is reported. It is worth mentioning that all simulations have been done by changing just a few lines of code (the problem class, and model definition), which shows the great versatility of the software. Finally, in Figure 2 we show how solutions can be visualized in PINA via the software plotting API with the Poisson problem example.  
conclusions: 
 We present in this contribution PINA, a software framework for learning differential equations leveraging deep learning. With a focus on centralizing research efforts in this domain, PINA aims to expedite the integration of these methodologies into production environments while providing a highly customizable entry point for active research. We introduced the most important features, highlighting the modular structure, the PyTorch and PyTorchLighting inheritance, the extensibility for defining problems and domains, the capability to use several built-in models or crafting from scratch a new one. We showed how PINA can be used to solve different problems, using different benchmarking cases. 
 
 Appendix: The PINA software is available at: https://github.com/mathLab/PINA.ODEs and PDEs are used to describe different physical phenomena in a mathematical form. Local updates expressed by partial or total derivatives are used to represent the evolution of a function characterizing a system. Following the notation presented in (Cuomo et al., 2022) , the general form of a differential equation, which we aim to solve, can be written as:EQUATIONwhere the solution field is u living in a suitable space U, the variables z = [x 1 , . . . , x dz-1 , t] indicate all the spatiotemporal coordinates in a domain Ω ⊂ R dz with ∂Ω its boundaries and d z the space dimension, α ∈ A the physical parameters in the suitable space A ⊂ R dα with d α the space dimension. Finally, F is a differential operator describing the dynamics with forcing term f , and B is the operator which indicates arbitrary initial or boundary conditions, with g the function on the boundaries.Solving ODEs and PDEs of the form in Equation equation 1 is one of the main computational challenges in mathematics and engineering. Numerical solvers, such as finite element methods (FEM), finite difference methods (FDM), or finite volume method (FVM), rely on discretizing the domain Ω (Morton & Mayers, 2005; Quarteroni & Quarteroni, 2009) . For many complex domains, the discretization is not straightforward, and a specific study is needed to ensure the final accuracy of the solver. Moreover, these solvers are often computationally expensive, resulting in high energy consumption, and slow computational time.Over the past decades, multiple deep learning methods have risen for solving the problem formalized in equation 1, aiming to overcome the classical numerical solver issues. Eventually, a dichotomy of methodologies can be made: Neural Operator (NO) methods, which assume knowledge of the system in the form of data; and Physics Informed Neural Networks (PINNs), which use the underlying equation to learn the solution.Neural Operator (NO) methods (Li et al., 2020; Lu et al., 2021a; Bhattacharya et al., 2021; Kovachki et al., 2021; Brandstetter et al., 2022) build a mapping from infinite-dimensional function spaces by using a supervised learning strategy. Given a specific ODE or PDE as in the form of equation 1, a neural operator G :EQUATIONFor example, a NO could map the field at the initial temporal condition of a PDE, to the evolution at a specific time step; or the parameter of a differential equation to its solution for the specific parameter. NO possesses important characteristics: they are discretization invariant, i.e. the model is not defined on a fixed grid; the input and output is a function; the universal approximation theorem for operator holds Chen & Chen (1995).In many situations training data are not available, and NO can not be trained using a supervised loss. As an alternative, PINNs (Raissi et al., 2019) have been proposed. PINNs are trained by approximating the true solution of equation 1 with a neural network u θ ≈ u with parameters θ.In PINNs the network is trained directly with the ODE or PDE itself, ensuring that equation 1 is satisfied by the network:L(θ) = L F + L B .(3) The first term is the physics-informed loss inside the domain Ω, while the second one is a supervised loss for boundary or initial conditions. Different types of losses can be implemented, for example using the MSE loss equation 3 becomes:EQUATIONwhere N f is the number of collocation points sampled inside Ω, and N b the number of collocation points sampled in ∂Ω.Since PINN's inception, many follow-up improvements have been made to improve training stability and convergence. Examples of further research include studying different losses (Kharazmi et al., 2019; Wang et al., 2022; McClenny & Braga-Neto, 2020) , sample strategies for collocation points (Wu et al., 2023; Nabian et al., 2021; Daw et al., 2023) to speed up convergence, or specific network architecture (Wang et al., 2021a; b) and input augmentation (Demo et al., 2023; Lu et al., 2021c) to ease the neural network training.In this section, we provide the mathematical formulations of the problem presented in the experiment section 3.Burger's equation is a convection-diffusion equation widely used in many fields of mathematics.The problem is crafted as the benchmark presented in (Raissi et al., 2019) . Let x = (x, t) be the spatio-temporal variables, and u be the unknown field. The Burger equation is:EQUATIONFor solving the equation we sample 10000 points uniformly random in the domain [-1, 1] × [0, 1].Poisson's equation is an elliptic partial differential equation widely used in physics. The problem is crafted as the benchmark presented in (Demo et al., 2023) . Let x = (x, y) be the spatial variables, u be the unknown field, and Ω = [-1, 1] 2 the domain. The Poisson equation is:EQUATIONwhere ∂Ω indicates the boundary of the domain, and the Laplacian operator ∇ 2 acts on the spatial variables. For solving the equation we sample 10000 points uniformly random in the domain Ω.During problem learning we employ extra features, by augmenting the input with the forcing term, i.e. the model input is given by (x, y, sin(πx) sin(πy)).The Wave's Equation is a linear differential equation vastly used in fluid dynamics. Let x = (x, y, t) be the spatio-temporal variables, u be the unknown field, Ω = [0, 1] 2 the domain, and T = [0, 1] 2 the parameter domain. The Wave equation is:EQUATIONwhere ∂Ω indicates the boundary of the domain, and the Laplacian operator ∇ 2 acts on the spatial variables. For solving the equation we sample 10000 points uniformly random in the domain Ω. We use R3 adaptive refinement for moving the collocation points during training every 100 epochs.C.4 PARAMETRIC POISSON'S EQUATION Parametric Poisson's equation is an example of a Poisson equation where the forcing term depends on external parameters. The problem is crafted as the benchmark presented in (Demo et al., 2023) , where the objective is to learn a function for different parameters. The problem can be considered as a NO problem since we map different initial functions (for different parameters) to the field functions. In the experiment section, we use PINN learning to solve the problem. Let x = (x, y) be the spatial variables, u be the unknown field, Ω = [0, 1] 2 the domain, and Ξ = [-1, 1] 2 the parameter domain. The Poisson equation is:EQUATIONwhere ∂Ω indicates the boundary of the domain, and the Laplacian operator ∇ 2 acts on the spatial variables. For solving the equation we sample 10000 points uniformly random in the domain Ω.In this section, we provide the network specifics for the experiments performed in Section 3. All the models were trained using the Adam optimizer (Kingma & Ba, 2014) , with a learning rate of 0.001 for 10000 epochs minimizing the mean square error loss. The training was done on an Intel CPU.The networks' composition:• MLP: Three linear layers of size [20, 10, 5] with hyperbolic tangent activation on all layers except the last• m-MLP: Three linear layers of size [20, 20, 20] with hyperbolic tangent activation on all layers except the last. The transformer networks were two linear layers mapping the input to the inner size of 20• hard-MLP: Same as MLP. Hard constraints on boundary conditions are imposed by multiplying the network output with the term (1 + x)(1 -x)• DeepONet: The branch and trunk net are the same architecture of two linear layers of size [20, 20] with hyperbolic tangent activation on all layers except the last. The reduction is done by aggregating with a linear layer with input dimension 20 and output dimension 1.The trunk net takes t as input. The branch net takes x as input.The input dimension of the problem is 2 (one spatial + one temporal variables), and the output dimension is 1.The networks' composition:• MLP: Same architecture as Burgers's problem.• m-MLP: Same architecture as Burgers's problem.• hard-MLP: Same architecture as Burgers's problem.• DeepONet: Same architecture as Burgers's problem, but the trunk net takes x, t as input.The branch net takes sin(πx) sin(πy) as input.The input dimension of the problem is 3 (two spatial + augmentation variables), and the output dimension is 1. 


Decision: This paper introduces a new PyTorch interface, PINA, for developing ML models based on the PINN and/or NO methodologies. PINA is built upon Lightning, and therefore benefits from its features. This is achieved mainly by providing a Trainer class wrapping Lightning’s trainer. The contribution is positive for the community but the performance of this package is not compared to existing open-source packages and there is limited scientific advances made in this work so with this in mind it will be accepted but only for a poster given some of the limitations of this paper.

Reviews: review: title: title: A PyTorch Lightning's wrapper for PINNs and/or NOs
main_review: main_review: This paper introduces a new PyTorch interface, PINA, for developing ML models based on the PINN and/or NO methodologies. PINA is built upon Lightning, and therefore benefits from its features. This is achieved mainly by providing a Trainer class wrapping Lightning’s trainer. Details are provided about the API, and a few results are shown.

The proposed framework seems useful for implementing various existing ML methods for solving PDE-based problems. The engineering effort behind PINA provides a productive way to develop new methods and can facilitate the deployment of those methods on different hardware setups.

While the engineering effort deployed to compose the interface with Lightning can be useful, the originality/novelty of the proposed framework is questionable and there are a couple of flaws in the mathematical formulation of the problem (see below). More specifically, the main contribution is to provide a wrapper around Lightning’s Trainer class and most of the claimed advantages of PINA over related methods in the literature (e.g., multiple device training, modern model compression techniques, gradient accumulation, logger insertion for monitoring, etc.) are consequences of having a Lightning backend and do not result from new features introduced in this work. Also, the paper mentions that existing packages can’t combine the PINN and NO methodologies. However, the DeepXDE package [1] provides support for PINNs as well as DeepONet, a well-known type of NO. In addition, further experiments comparing PINA against competitors from the literature could improve the dissemination of the package.

As for the related works, the paper misses out on prior works on the subject of software for PDE learning, see [2, 3], which should be included.

Lastly, the paper contains several flaws in the mathematical formulation of the problem:

- There are a few problems with equation 1:
	- z is said to indicate the spatio-temporal coordinates, however $z \in \Omega$, where $\Omega$ is the spatial domain. In fact, equation 1 is effectively a steady problem, i.e. with no time dependency, which is a pretty restrictive assumption for software tackling PDEs.
	- $\mathscr{B}$ cannot be used for initial conditions given that the problem is time-independent.
	- $d_{z}$ is not defined
	- $\Omega \in \mathbb{R}^{d_{z}}$ → $\Omega \subset \mathbb{R}^{d_{z}}$
	- the spaces $A$ and $U$ are not defined

- In equation 2, $u^{\prime}$ and $u$ are not defined. I assume you meant $G(u\prime_{i}) = u_{i} \quad \forall 1 \le i \le N$ ?


[1] Lu et. al, DeepXDE: A deep learning library for solving differential equations, SIAM Review, arXiv.1907.04502, 2021.

[2] N. Bouziani and D. A. Ham, Physics-driven machine learning models coupling PyTorch and Firedrake, ICLR 2023 Workshop on Physics for Machine Learning, arXiv.2303.06871, 2023.

[3] P. Kidger, Diffrax - On Neural Differential Equations, University of Oxford, arXiv:2202.02435, 2021.
title: title: paper introducing a pytorch-based open-source package for solving PDEs and operator learning
main_review: main_review: The paper introduces PINA (Physics Informed Neural network for Advanced modeling) which is a PyTorch-based open-source package for solving PDEs and operator learning via neural networks. As argued in the paper, the main advantage of PINA compared to existing packages is that it streamlines the training process by adding interfaces for logging or developing new loss functions. 

I would like to congratulate the authors for their efforts in building an open-source package. However, in my professional opinion, the technical contributions of this work are marginal. Specifically, while having a streamlined pipeline is useful, the already available libraries make it very easy to implement any of the models that the authors have exemplified in Sec. 4. The performance of this package (and many of its claimed benefits) are also not compared to existing open-source packages so it is difficult to evaluate the authors’ claims.