# What is the best way to implement PINN in Julia
**URL:** https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713
**Category:** New to Julia
**Tags:** pde, neural-network, autodiff
**Created:** [April 4, 2025, 3:34pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713 "2025-04-04T15:34:34Z")
**Posts on this page:** 19
**Page:** 1
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 4, 2025, 3:34pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/1 "2025-04-04T15:34:34Z")
Dear all,
I want to use PINNs to solve a PDE. There are two packages in Julia that can do this [NeuralPDE.jl](https://github.com/SciML/NeuralPDE.jl) and [Sophon.jl](https://github.com/YichengDWu/Sophon.jl)
Both of these packages use the finite difference method to realize derivatives. I want to know if that is possible to use autodiff in PINNs. Since we may have a third-order derivative, it is not an easy task.
I basically have the same problem in this post [https://www.reddit.com/r/MachineLearning/comments/1etkmm6/r\_nested\_ad\_for\_pinns/](https://www.reddit.com/r/MachineLearning/comments/1etkmm6/r_nested_ad_for_pinns/)
I also would like to know if [Autograd.jl](https://github.com/denizyuret/AutoGrad.jl) is a possible choice to do the higher-order auto derivative. I know the Enzyme+Reactant may be the best choice in the future, but it seems not ready to use in the current time.
---
### Author:  [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)
#### Post date: [April 4, 2025, 4:13pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/2 "2025-04-04T16:13:35Z")
Hi,
1. What makes you think NeuralPDE only use Finite Diff ?
2. Zygote and Enzyme supports nested autodiff there is an exemple in Lux doc, did you try ?
3. For high order derivative, I think TaylorDiff.jl and GPTSA.jl are the best but I don’t think it supports nested autodiff
4. Yes Enzyme+Reactant is great, you will just get 1K errors before finally making it work (I did but lost the code 😢 )
5. Why do you need that ? Finite Diff is almost always the right choise actually, the error it comes with are just as big as those from nested autodiff and leads to as performant code and anyway, a Neural approx of a PDE has a big error coming with it compared to numerical solvers and often runs slower, so Finite Diff is far enough in most cases.
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [April 4, 2025, 4:16pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/3 "2025-04-04T16:16:19Z")
As the Reactant Readme warns “This package is under active development at the moment and may change its API and supported end systems at any time. End-users are advised to wait until a corresponding release with broader availability is made. Package developers are suggested to try out Reactant for integration, but should be advised of the same risks.”
We just figure its better to start building it in the open then wait til its feature-complete before open sourcing 🙂 .
That said, do you have any of the errors, we’re pretty…Reactive… so a fix/help shouldn’t be hard to get help with?
@avik-pal also probably has some Lux PINNS lying around
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [April 4, 2025, 4:21pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/4 "2025-04-04T16:21:45Z")
Actually [Training a PINN on 2D PDE | Lux.jl Docs](https://lux.csail.mit.edu/dev/tutorials/intermediate/4_PINN2DPDE#Training-a-PINN-on-2D-PDE) has a PINN using Enzyme+Reactant
---
### Author:  [@Oscar\_Smith](https://discourse.julialang.org/u/Oscar_Smith)
#### Post date: [April 4, 2025, 4:36pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/5 "2025-04-04T16:36:33Z")
> I want to use PINNs to solve a PDE
Why? PINNs tend to be slow inaccurate methods for PDEs.
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 4, 2025, 4:46pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/6 "2025-04-04T16:46:23Z")
Hi, Yolhan
Thanks for your reply.
(1). I find the `numeric_derivative` function in NeuralPDE
[https://github.com/SciML/NeuralPDE.jl/blob/a322d7be3cfcf52304a11967f18063f1b87baf9a/src/pinn\_types.jl#L356](https://github.com/SciML/NeuralPDE.jl/blob/a322d7be3cfcf52304a11967f18063f1b87baf9a/src/pinn_types.jl#L356)
I also take a look at the generated symbolic loss function.
(2). I tried that, it works for me. However, it only supports 2nd AD. But I need 3rd AD(for a second order PDE). I also find this pull request [https://github.com/LuxDL/Lux.jl/pull/1097](https://github.com/LuxDL/Lux.jl/pull/1097). But it does not work well on GPU, I waited about a half hour and stuck in this line [https://github.com/LuxDL/Lux.jl/blob/2c40b8f11d4cdde37d11e341428115e95c5fe47b/docs/src/manual/nested\_autodiff\_reactant.md?plain=1#L54](https://github.com/LuxDL/Lux.jl/blob/2c40b8f11d4cdde37d11e341428115e95c5fe47b/docs/src/manual/nested_autodiff_reactant.md?plain=1#L54)
(5). Thanks for your important suggestion. I think I will use the Finite Diff. I use PINN because I have a very stiff PDE. But you are right, PINN is much slower than the BDF solver+ Finite Diff in DifferentialEquations.jl. Another important reason is I have a high-dimensional PDE.(no one has solved a 3-dimensional case in our field).
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 4, 2025, 4:56pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/7 "2025-04-04T16:56:14Z")
Hi William,
This example works well. If there is a possible for 3rd AD. I find this pull request by avik [https://github.com/LuxDL/Lux.jl/pull/1097](https://github.com/LuxDL/Lux.jl/pull/1097). But it does not work well on GPU, I waited about a half hour. I can try more time tomorrow.
---
### Author:  [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)
#### Post date: [April 4, 2025, 4:56pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/8 "2025-04-04T16:56:25Z")
Have a look at HighDimPDE.jl it’s exactly here for that. Not sûre it supports high order diff though but I think not (doc wise)
Also did you try with MethodOfLines.jl or maybe you can’t?(High Advection ) Nevermind if it’s too stiff it’s a bad idea
Thinking about it I would still go for NeuralPDE 😅
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 4, 2025, 5:10pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/9 "2025-04-04T17:10:37Z")
Hi Oscar,
You are right, PINNs tend to be slow. I think the accuracy may be improved by using the (L-)BFGS optimizers. I want to solve a high-dimensional stiff PDE.(well, actually only 3 dimensional in space and one-dimensional in time is an unsolved problem.) I also want to extend to solve a functional differential equation in the future.
---
### Author:  [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)
#### Post date: [April 4, 2025, 5:18pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/10 "2025-04-04T17:18:17Z")
Something of the same difficulty (maybe with non-linearity) as
\partial\_t u - \sum\_{i=1}^3 \partial\_{x\_i}^2 u + \sum\_{i=1}^3 \partial\_{x\_i}^3 u = 0
With the BC, IC ?
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 4, 2025, 5:27pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/11 "2025-04-04T17:27:33Z")
Thanks again for your nice suggestion.
I have noticed the HighDimPDE.jl before, but I didn’t figure out the math at that time, I will try this package in the future.
I haven’t tried the MethodOfLines.jl but I basically use the same method. I don’t use MethodOfLines.jl for some subtle reason. I don’t have boundary conditions. The common choice is to use the forward and backward difference in the boundary separately.
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 4, 2025, 5:40pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/12 "2025-04-04T17:40:51Z")
This is what I need to solve
k\partial\_k V\_k(\rho)=T k^3\left[\frac{k^2}{k^2+V\_k'(\rho)+2\rho V\_k''(\rho)}+\frac{(N-1)k^2}{k^2+V\_k'(\rho)}\right]
where \rho=\phi^2/2. The IC is simple V(\rho)=\lambda(\rho-\rho\_0)^2.
The solution will be something like the following:
You will see the V'(\rho) becomes more and more flat as k\to0. And in the k=0, it will becomes completly flat and there will be a non smooth point around \rho=5000 in this fig.
The denominator k^2+V\_k'(\rho)+2\rho V\_k''(\rho) and k^2+V\_k'(\rho) will close to 0 as k\to 0.
---
### Author:  [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)
#### Post date: [April 4, 2025, 6:54pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/14 "2025-04-04T18:54:11Z")
Giving a try, and they are tons of mistake because I’m not 100% sure about the eq and bc,
```julia
julia> using ModelingToolkit,MethodOfLines,DomainSets
julia> @parameters k,ρ
2-element Vector{Num}:
k
ρ
julia> @variables V(..)
1-element Vector{Symbolics.CallWithMetadata{SymbolicUtils.FnType{Tuple, Real}, Base.ImmutableDict{DataType, Any}}}:
V⋆
julia> Dρ = Differential(ρ)
Differential(ρ)
julia> D2ρ = Dρ^2
Differential(ρ) ∘ Differential(ρ)
julia> Dk = Differential(k)
Differential(k)
julia> kmin = -4e5
-400000.0
julia> kmax= 3e5
300000.0
julia> rhomin = 0
0
julia> rhomax = 2e4
20000.0
julia> T= 1
1
julia> N = 2
2
julia> eq = [k*Dk(V(k,ρ)) ~ T*k^3*(k^2/(k^2+Dρ(V(k,ρ)) + 2*ρ*D2ρ(V(k,ρ))) + (N-1)*k^2/(k^2 + Dρ(V(k,ρ))))]
1-element Vector{Equation}:
k*Differential(k)(V(k, ρ)) ~ (k^3)*((k^2) / (Differential(ρ)(V(k, ρ)) + k^2 + 2Differential(ρ)(Differential(ρ)(V(k, ρ)))*ρ) + (k^2) / (Differential(ρ)(V(k, ρ)) + k^2))
julia> bcs = [V(kmin,ρ)~λ*ρ, V(k,0)~0,Dρ(V(k,0)) ~ λ]
3-element Vector{Equation}:
V(-400000.0, ρ) ~ 2ρ
V(k, 0) ~ 0
Differential(ρ)(V(k, 0)) ~ 2
julia> @named pdesys = PDESystem(eq,bcs,domains,[k,ρ],[V(k,ρ)])
PDESystem
Equations: Equation[k*Differential(k)(V(k, ρ)) ~ (k^3)*((k^2) / (Differential(ρ)(V(k, ρ)) + k^2 + 2Differential(ρ)(Differential(ρ)(V(k, ρ)))*ρ) + (k^2) / (Differential(ρ)(V(k, ρ)) + k^2))]
Boundary Conditions: Equation[V(-400000.0, ρ) ~ 2ρ, V(k, 0) ~ 0, Differential(ρ)(V(k, 0)) ~ 2]
Domain: Symbolics.VarDomainPairing[Symbolics.VarDomainPairing(k, -400000.0 .. 300000.0), Symbolics.VarDomainPairing(ρ, 0.0 .. 20000.0)]
Dependent Variables: Num[V(k, ρ)]
Independent Variables: Num[k, ρ]
Parameters: SciMLBase.NullParameters()
Default Parameter ValuesDict{Any, Any}()
julia> NN = 32
32
julia> disc = MOLFiniteDifference([ρ=>NN],k,approx_order=order)
MOLFiniteDifference{MethodOfLines.CenterAlignedGrid, MethodOfLines.ScalarizedDiscretization}(Dict{Num, Int64}(ρ => 32), k, 4, UpwindScheme(1), MethodOfLines.CenterAlignedGrid(), true, false, MethodOfLines.ScalarizedDiscretization(), true, Any[], Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
julia> prob = discretize(pdesys,disc)
ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
Non-trivial mass matrix: false
timespan: (-400000.0, 300000.0)
u0: 30-element Vector{Float64}:
2580.6451612903224
3870.967741935484
5161.290322580645
6451.612903225807
7741.935483870968
9032.258064516129
10322.58064516129
11612.90322580645
12903.225806451614
14193.548387096775
15483.870967741936
16774.1935483871
18064.516129032258
19354.83870967742
20645.16129032258
21935.483870967742
23225.8064516129
24516.129032258064
25806.451612903227
27096.774193548386
28387.09677419355
29677.41935483871
30967.74193548387
32258.064516129034
33548.3870967742
34838.709677419356
36129.032258064515
37419.354838709674
38709.67741935484
40000.0
julia> sol = solve(prob,Tsit5(),saveat=1e2)
PDESolution:
Return Code:
Success
Dependent variables:
V(k, ρ): (7001, 32) sized solution
Domain:
k ∈ (-400000.0, 300000.0) with 7001 points, step size 100.0
ρ ∈ (0.0, 20000.0) with 32 points, step size 645.1612903225806
From system:
Equations:
L"\begin{align}
k \frac{\mathrm{d}}{\mathrm{d}k} V\left( k, \rho \right) - k^{3} \left( \frac{k^{2}}{\frac{\mathrm{d}}{\mathrm{d}\rho} V\left( k, \rho \right) + k^{2} + 2 \frac{\mathrm{d}}{\mathrm{d}\rho} \frac{\mathrm{d}}{\mathrm{d}\rho} V\left( k, \rho \right) \rho} + \frac{k^{2}}{\frac{\mathrm{d}}{\mathrm{d}\rho} V\left( k, \rho \right) + k^{2}} \right) &= 0
\end{align}
"
Boundary/Initial Conditions:
L"\begin{align}
V\left( -4 \cdot 10^{5}, \rho \right) &= 2 \rho \\
V\left( k, 0 \right) &= 0 \\
\frac{\mathrm{d}}{\mathrm{d}\rho} V\left( k, 0 \right) &= 2
\end{align}
"
julia> plot(sol.disc_data.discretespace.grid[ρ],sol.u[V(k, ρ)][1,:])
julia> plot!(sol.disc_data.discretespace.grid[ρ],sol.u[V(k, ρ)][2,:])
julia> plot!(sol.disc_data.discretespace.grid[ρ],sol.u[V(k, ρ)][3,:])
julia> plot!(sol.disc_data.discretespace.grid[ρ],sol.u[V(k, ρ)][4,:])
```
one thing that may be odd is that I can’t make V(0, ρ) ~ something
plain code
```julia
using ModelingToolkit,MethodOfLines,DomainSets
@parameters k,ρ
@variables V(..)
Dρ = Differential(ρ)
D2ρ = Dρ^2
Dk = Differential(k)
kmin = -4e5
kmax= 3e5
rhomin = 0
rhomax = 2e4
T= 1
N = 2
eq = [k*Dk(V(k,ρ)) ~ T*k^3*(k^2/(k^2+Dρ(V(k,ρ)) + 2*ρ*D2ρ(V(k,ρ))) + (N-1)*k^2/(k^2 + Dρ(V(k,ρ))))]
bcs = [V(kmin,ρ)~λ*ρ, V(k,0)~0,Dρ(V(k,0)) ~ λ]
@named pdesys = PDESystem(eq,bcs,domains,[k,ρ],[V(k,ρ)])
NN = 32
disc = MOLFiniteDifference([ρ=>NN],k,approx_order=4)
prob = discretize(pdesys,disc)
sol = solve(prob,Tsit5(),saveat=1e2)
plot(sol.disc_data.discretespace.grid[ρ],sol.u[V(k, ρ)][1,:])
plot!(sol.disc_data.discretespace.grid[ρ],sol.u[V(k, ρ)][2,:])
```
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [April 4, 2025, 8:44pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/15 "2025-04-04T20:44:40Z")
Reactant + Lux is by far the best way to do it, the benchmarks aren’t even close. NeuralPDE.jl has not updated to the infrastructure yet. The plan is to do that over the summer.
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 5, 2025, 4:52am UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/16 "2025-04-05T04:52:43Z")
Hi Yolhan,
Thanks a lot for taking the time to look at this equation.
I omitted a constant \nu\_d=1 / \left(2^{d + 1} \pi^{d / 2} \Gamma(2 + d / 2)\right) in the equation.
To lower the stiffness, I take the derivative with respect to ρ from both sides. I solve the equation about the V'(\rho).
The code to solve this equation is
```julia
using DifferentialEquations, Plots, Dierckx, LinearAlgebra, SpecialFunctions
using ProgressLogging, Sundials
function flowV!(dVdt, Vp, p, k)
(Λ, νd, T, d, n, rho_grid) = p
#Vpfun is V'(ρ)
Vpfun = Spline1D(rho_grid, Vp; bc = "nearest")
#Vpp is V''(ρ)
Vpp = Dierckx.derivative(Vpfun, rho_grid)
#Vppp is V'''(ρ)
Vppp = Dierckx.derivative(Vpfun, rho_grid, 2)
@. dVdt = -(k^(1 + d) * T * (2 + d) * νd *
(((-1 + n) * Vpp) / (k^2 + Vp)^2 +
(3 * Vpp + 2 * Vppp * rho_grid) / (k^2 + Vp + 2 * Vpp * rho_grid)^2))
return nothing
end
#spatial dimension d=3
d = 3.0
#Λ is the UV cutoff scale
Λ = 1000.0
#νd is a constant
νd = 1 / (2^(d + 1) * pi^(d / 2) * gamma(2 + d / 2))
#dimensionless ρ grid.
rho_dimensionless = collect(0.0:0.001:0.5)
#critical temperature, above this temperature this equation doesn't have stiffness
Tc = 150.0
#rho_grid is the ρ grid
rho_grid = rho_dimensionless * (Λ^(d - 2) * Tc)
rho0_dimensionless = 0.07147268020712277
# initial condition, V'(ρ)=2λ(ρ-ρ0)
u0 = 35.0 .* (rho_grid .- rho0_dimensionless * (1000^(d - 2) * Tc))
#Temperature. The lower the temperature, the more stiff the equation.
#Temperature =1.0, roughly corresponds to Zero temperature or the Vacuum.
Temperature = 100.0
#n can be 1,2,3,4,5..., but non-integer n is also possible.
n = 4.0
p = (Λ, νd, Temperature, d, n, rho_grid)
#solving range, from UV to IR scale. In ideal case, the IR scale show be 0, but it's very hard, here we chose 5.0.
tspan = (Λ, 5.0)
prob = ODEProblem(flowV!, u0, tspan, p)
# A stiff solver is needed. CVODE_BDF is a good choice in practice.
sol = solve(prob, CVODE_BDF(), ;
reltol = 1e-10,
dtmax = 0.5,
progress = true,
abstol = 1e-10, maxiters = 5 * 10^3,)
plot(sol.u[end][1:30])
plot(sol.u[end])
```
This equation is highly stiff. Recently some researchers have used the Discontinuous Galerkin or Fininte Volum method to solve this equation. But these methods are not easily extended to high dimension cases which means V could be something like V(\rho\_1,\rho\_2,\rho\_3). That is why I try the PINNs.
---
### Author:  [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)
#### Post date: [April 5, 2025, 5:56am UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/17 "2025-04-05T05:56:30Z")
Did you try with SplineGrids.jl which supports multi-dim SPline ?
DG is fine too, but the difficulty is higher and you may need to play a lot with penalty parameters and all (especially if you have a non-smooth solution at a point), only consider it if you need a mesh like handling.
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 5, 2025, 11:51am UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/18 "2025-04-05T11:51:34Z")
I haven’t noticed SplineGrids.jl before. It looks very promising. I will try it.
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [April 5, 2025, 5:03pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/19 "2025-04-05T17:03:14Z")
> [@Yang-yang](#):
>
> . I tried that, it works for me. However, it only supports 2nd AD. But I need 3rd AD(for a second order PDE). I also find this pull request [docs: add 3rd order AD example using Reactant by avik-pal · Pull Request #1097 · LuxDL/Lux.jl · GitHub](https://github.com/LuxDL/Lux.jl/pull/1097). But it does not work well on GPU, I waited about a half hour and stuck in this line [Lux.jl/docs/src/manual/nested\_autodiff\_reactant.md at 2c40b8f11d4cdde37d11e341428115e95c5fe47b · LuxDL/Lux.jl · GitHub](https://github.com/LuxDL/Lux.jl/blob/2c40b8f11d4cdde37d11e341428115e95c5fe47b/docs/src/manual/nested_autodiff_reactant.md?plain=1#L54)
The example there is different from what you need. In that example, we are trying to do a gradient of the sum of hessian. I was exploring that more as an exercise in stressing the Reactant compiler pipeline (as you saw with the high compile time) and hence it hasn’t been merged.
You care about the gradient of a function computing the N-th order partial derivatives which already work quite well in Reactant.
---
### Author:  [@Yang-yang](https://discourse.julialang.org/u/Yang-yang)
#### Post date: [April 6, 2025, 2:37pm UTC](https://discourse.julialang.org/t/what-is-the-best-way-to-implement-pinn-in-julia/127713/20 "2025-04-06T14:37:28Z")
Thanks for the reply, I will try it.