What is the best way to implement PINN in Julia

Dear all,
I want to use PINNs to solve a PDE. There are two packages in Julia that can do this NeuralPDE.jl and 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/

I also would like to know if 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.

1 Like

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 :cry: )
  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.
2 Likes

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 :slight_smile: .

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

4 Likes

Actually Training a PINN on 2D PDE | Lux.jl Docs has a PINN using Enzyme+Reactant

2 Likes

I want to use PINNs to solve a PDE

Why? PINNs tend to be slow inaccurate methods for PDEs.

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
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. 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

(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).

1 Like

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. But it does not work well on GPU, I waited about a half hour. I can try more time tomorrow.

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 :sweat_smile:

1 Like

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.

1 Like

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 ?

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.

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.

2 Likes

Giving a try, and they are tons of mistake because I’m not 100% sure about the eq and bc,

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

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,:])
2 Likes

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.

4 Likes

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

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.

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.

1 Like

I haven’t noticed SplineGrids.jl before. It looks very promising. I will try it.

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.

1 Like

Thanks for the reply, I will try it.