# Why do we need 3 chains to solve a PDE using NeuralPDE

**URL:** <https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136>\
**Category:** Machine Learning\
**Created:** [December 28, 2023, 4:11pm UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136 "2023-12-28T16:11:54Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![affans](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/affans/32/11911_2.png) [@affans](https://discourse.julialang.org/u/affans)\
**Post date:** [December 28, 2023, 4:11pm UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/1 "2023-12-28T16:11:54Z")

</div>

I am working through examples from `NeuralPDE.jl` package and have a question. I am looking at the [inverse problem example](https://docs.sciml.ai/NeuralPDE/stable/tutorials/param_estim/) (i.e. parameter estimation) and just need to get my thinking straight. This is more of a conceptual/theory question than it is a Julia one.

Let’s consider the Lorenz system as given in the example.

\begin{aligned} x\prime &= \sigma(y - x) \\ y\prime &= x(\rho - z) - y \\ z\prime &= xy - \beta z \end{aligned} 

The parameters to be estimated here are \sigma, \rho, and \beta. From my understanding of the theory of PINNs, the underlying NN would have an input dimension of 3 (corresponding to time t) and output dimension of 3 (corresponding to the outputs x, y, and z. In other words, **there is just one NN that is being trained**.

However, in the example, I see there are **three independent NN** being trained each with an output dimension of 1 (so three neural networks corresponding to x, y, and z).

```julia
 chain1 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                       Dense(n, 1))
 chain2 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                       Dense(n, 1))
 chain3 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                       Dense(n, 1))
 discretization = NeuralPDE.PhysicsInformedNN(
        [chain1, chain2, chain3],
        NeuralPDE.GridTraining(dt), 
        param_estim = true, # whether the parameters of the differential equation should be sent to the additional_loss function
        additional_loss = additional_loss)

```

I don’t really understand why we would need three chains, each with a single output (i.e., dim = 1) instead of just one chain with 3 outputs. The [documentation](https://docs.sciml.ai/NeuralPDE/dev/manual/pinns/#The-PhysicsInformedNN-Discretizer) does say

> `chain` : a vector of Flux.jl or Lux.jl chains with a d-dimensional input and a 1-dimensional output corresponding to each of the dependent variables.

but it does not answer my question. Would anyone recommend suggested readings?

**Edit** : In the dev version of the documentation, there is now an example of parameter estimation of the Volta Lotkerra model (a system with 2 equations, 4 parameters) in a Bayesian framework. In this example, the NN is exactly what I expected it to be; i.e., input of dimension 1 and output of dimension 2.

So now the documentation has two examples of similar models, but one that uses `BNNODE` interface (with a single NN) and the other that uses `PhysicsInformedNN` interface (with multiple NN, each NN corresponding to an equation).

@ChrisRackauckas tagging you for visibility and would appreciate your comment on this.

---

<div class="post-metadata">

**Author:** ![Maucejo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maucejo/32/39090_2.png) [@Maucejo](https://discourse.julialang.org/u/Maucejo)\
**Post date:** [December 28, 2023, 5:29pm UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/2 "2023-12-28T17:29:18Z")

</div>

You can find the answer here [NeuralPDE systems of PDEs](https://docs.sciml.ai/NeuralPDE/stable/tutorials/systems/). In short, it is because you have 3 equations, so 1 NN per equation.

---

<div class="post-metadata">

**Author:** ![affans](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/affans/32/11911_2.png) [@affans](https://discourse.julialang.org/u/affans)\
**Post date:** [December 29, 2023, 6:20pm UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/3 "2023-12-29T18:20:54Z")

</div>

Thanks!

I am wondering why in one of the examples they are able to use a single neural network even though the system has more than 1 equation (see [here](https://docs.sciml.ai/NeuralPDE/dev/tutorials/Lotka_Volterra_BPINNs/)). This example uses the `BNNODE` API but I was wondering if it’s possible to convert this to using the `PhysicsInformedNN` API as well (I am guessing that `PhysicsInformedNN` is more primitive and should be able to handle the Bayesian example as well)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 30, 2023, 10:44am UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/4 "2023-12-30T10:44:58Z")

</div>

> [@affans](#):
>
> I am wondering why in one of the examples they are able to use a single neural network even though the system has more than 1 equation (see [here](https://docs.sciml.ai/NeuralPDE/dev/tutorials/Lotka_Volterra_BPINNs/)).

The ODE formulation specializes on the fact that everything is only differentiated once, and everything is differentiated once.

> [@affans](#):
>
> I was wondering if it’s possible to convert this to using the `PhysicsInformedNN` API as well (I am guessing that `PhysicsInformedNN` is more primitive and should be able to handle the Bayesian example as well)

No. This is addressed in the writeup that explains NeuralPDE.jl.

> **[NeuralPDE: Automating Physics-Informed Neural Networks (PINNs) with Error...](https://arxiv.org/abs/2107.09443)**
>
> Physics-informed neural networks (PINNs) are an increasingly powerful way to solve partial differential equations, generate digital twins, and create neural surrogates of physical models. In this manuscript we detail the inner workings of...

The summary of that part is really that with high order automatic differentiation you get really bad scaling. It’s multiplicative between the number of outputs and the order. So you really only want to differentiate what you need. If you pool everything together in one neural network and only one term needs a second or third derivative, you’re pretty much stuck differentiating everything to third order which has a cubic cost. Splitting the networks gives only a linear cost growth, so it’s much cheaper.

I do plan to make it so it can be a choice in the future, but that’ll take some work.

---

<div class="post-metadata">

**Author:** ![affans](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/affans/32/11911_2.png) [@affans](https://discourse.julialang.org/u/affans)\
**Post date:** [May 7, 2024, 2:43am UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/5 "2024-05-07T02:43:27Z")

</div>

@ChrisRackauckas

I spent some time today reading the write up [NeuralPDE (arxiv.org)](https://arxiv.org/abs/2107.09443) and also [this paper on using ANNs to solve ODEs](https://arxiv.org/abs/physics/9705023) and I have a fairly good understanding now, but I am still not sure about what you mean by “differentiated only once”. I have two questions

1. If I have a system of ODEs with constant parameters, should I stick with the `NNODE` API to solve the system?

2. When is it required to use the `PhysicsInformedNN` API?

And I may create another topic for this, but what if my system of ODEs has a time-dependent parameter? Consider for example, this basic SIR system

```julia
S' = -b(t) S I 
I' = b(t) S I - g I

```

Given data for `I` I’d like to solve the inverse problem of finding `b(t)` but I am not sure how to use the package to do that (it seems to me that I need another NN for `b(t)`)

Thanks in advance for your help, if you have the time!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [May 7, 2024, 6:38am UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/6 "2024-05-07T06:38:37Z")

</div>

> [@affans](#):
>
> - If I have a system of ODEs with constant parameters, should I stick with the `NNODE` API to solve the system?

Yes

> [@affans](#):
>
> 1. When is it required to use the `PhysicsInformedNN` API?

PDEs

> [@affans](#):
>
> And I may create another topic for this, but what if my system of ODEs has a time-dependent parameter?

That’s physics-informed neural operators territory, which should get added soon

> <https://github.com/SciML/NeuralPDE.jl/pull/806>
>
> Here is implementation of the method that combine training data and physics con…straints to learn the solution operator of a given family of parametric Ordinary Differential Equations (ODE).
> 
> https://github.com/SciML/NeuralPDE.jl/issues/575
> 
> Checklist 
> \- \[x\] pino ode 
> \- \[x\] family ode by parameter 
> \- \[x\] physics informed DeepOnet 
> \- \[x\] tests
> \- \[x\] addition loss test
> \- \[\] doc
> 
> 
> https://arxiv.org/abs/2111.03794

---

<div class="post-metadata">

**Author:** ![affans](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/affans/32/11911_2.png) [@affans](https://discourse.julialang.org/u/affans)\
**Post date:** [May 11, 2024, 8:37pm UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/7 "2024-05-11T20:37:56Z")

</div>

Thanks @ChrisRackauckas, so in the documentation there are two examples of inverse problems.

1. Under the ODE section, this [example](https://docs.sciml.ai/NeuralPDE/stable/tutorials/ode_parameter_estimation/) uses the `NNODE` API to optimize the parameters of a system of ordinary differential equations. In this example, a single NN with output dimension equal to the number of equations is used.

2. Also, under the PDE PINN section, this [example](https://docs.sciml.ai/NeuralPDE/stable/tutorials/param_estim/) uses the PINN API instead for the inverse problem (though it’s still a system of ordinary differential equations), with 3 separate chains corresponding to the three solutions `x(t), y(t), z(t)`.

Should we consolidate or get rid of the second example? At the very least, perhaps we can say what’s different between the two. I am happy to create an issue/PR if you let me know what the main differences are here.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [May 11, 2024, 9:04pm UTC](https://discourse.julialang.org/t/why-do-we-need-3-chains-to-solve-a-pde-using-neuralpde/108136/8 "2024-05-11T21:04:21Z")

</div>

> [@affans](#):
>
> Should we consolidate or get rid of the second example? At the very least, perhaps we can say what’s different between the two. I am happy to create an issue/PR if you let me know what the main differences are here.

That one is fine. It should probably be replaced with a PDE one though.