Neural Adapter Documentation Example Not Working

[geriatricvibes]

Describe the bug 🐞

The Neural Adapter Example is not working

Expected behavior

First it showed a 'variable not defined error' with the 'chain' variable in 'init_params2'.

init_params2 = Float64.(ComponentArray(Lux.setup(Random.default_rng(), chain)[1]))

I then changed it into chain1 and tried again.

Code 👇

@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2

# 2D PDE
eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)

# Initial and boundary conditions
bcs = [u(0, y) ~ 0.0, u(1, y) ~ -sin(pi * 1) * sin(pi * y),
    u(x, 0) ~ 0.0, u(x, 1) ~ -sin(pi * x) * sin(pi * 1)]
# Space and time domains
domains = [x ∈ Interval(0.0, 1.0),
    y ∈ Interval(0.0, 1.0)]
quadrature_strategy = NeuralPDE.QuadratureTraining(reltol = 1e-2, abstol = 1e-2,
                                                   maxiters = 50, batch = 100)
inner = 8
af = Lux.tanh
chain1 = Chain(Dense(2, inner, af),
               Dense(inner, inner, af),
               Dense(inner, 1))

discretization = NeuralPDE.PhysicsInformedNN(chain1,
                                             quadrature_strategy)

@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
prob = NeuralPDE.discretize(pde_system, discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization)

res = Optimization.solve(prob, BFGS(); maxiters = 3)
phi = discretization.phi

inner_ = 12
af = Lux.tanh
chain2 = Lux.Chain(Dense(2, inner_, af),
                   Dense(inner_, inner_, af),
                   Dense(inner_, inner_, af),
                   Dense(inner_, 1))

init_params2 = Float64.(ComponentArray(Lux.setup(Random.default_rng(), chain1)[1]))

# the rule by which the training will take place is described here in loss function
function loss(cord, θ)
    chain2(cord, θ) .- phi(cord, res.u)
end

strategy = NeuralPDE.QuadratureTraining()

prob_ = NeuralPDE.neural_adapter(loss, init_params2, pde_system, strategy)
callback = function (p, l)
    println("Current loss is: $l")
    return false
end
res_ = Optimization.solve(prob_, BFGS(); callback = callback, maxiters = 3)

phi_ = NeuralPDE.get_phi(chain2)

xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)

u_predict = reshape([first(phi([x, y], res.u)) for x in xs for y in ys],
                    (length(xs), length(ys)))
u_predict_ = reshape([first(phi_([x, y], res_.minimizer)) for x in xs for y in ys],
                     (length(xs), length(ys)))
u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys],
                 (length(xs), length(ys)))
diff_u = u_predict .- u_real
diff_u_ = u_predict_ .- u_real

using Plots
p1 = plot(xs, ys, u_predict, linetype = :contourf, title = "first predict");
p2 = plot(xs, ys, u_predict_, linetype = :contourf, title = "second predict");
p3 = plot(xs, ys, u_real, linetype = :contourf, title = "analytic");
p4 = plot(xs, ys, diff_u, linetype = :contourf, title = "error 1");
p5 = plot(xs, ys, diff_u_, linetype = :contourf, title = "error 2");
plot(p1, p2, p3, p4, p5)

Error & Stacktrace ⚠️

MethodError: no method matching (::Chain{@NamedTuple{layer_1::Dense{true, typeof(tanh_fast), typeof(glorot_uniform), typeof(zeros32)}, layer_2::Dense{true, typeof(tanh_fast), typeof(glorot_uniform), typeof(zeros32)}, layer_3::Dense{true, typeof(tanh_fast), typeof(glorot_uniform), typeof(zeros32)}, layer_4::Dense{true, typeof(identity), typeof(glorot_uniform), typeof(zeros32)}}, Nothing})(::Matrix{Float64}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}})

Closest candidates are:
  (::Chain)(::Any, ::Any, ::NamedTuple)
   @ Lux [C:\Users\rnath\.julia\packages\Lux\Bs5su\src\layers\](file:///C:/Users/rnath/.julia/packages/Lux/Bs5su/src/layers/)containers.jl:479


Stacktrace:
  [1] loss(cord::Matrix{Float64}, θ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}})
    @ Main [.\In](https://file+.vscode-resource.vscode-cdn.net/c%3A/Users/rnath/roboticai/git/RoboticAI/roboticai/jl_roboticai/jl_roboticai/In)[8]:47
  [2] (::NeuralPDE.var"#integrand#138"{typeof(loss)})(x::Matrix{Float64}, θ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}})
    @ NeuralPDE [C:\Users\rnath\.julia\packages\NeuralPDE\z18Qg\src\training_strategies.jl:316](file:///C:/Users/rnath/.julia/packages/NeuralPDE/z18Qg/src/training_strategies.jl:316)
  [3] BatchIntegralFunction
    @ [C:\Users\rnath\.julia\packages\SciMLBase\m3AcC\src\scimlfunctions.jl:2207](file:///C:/Users/rnath/.julia/packages/SciMLBase/m3AcC/src/scimlfunctions.jl:2207) [inlined]
  [4] __solvebp_call(prob::IntegralProblem{false, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, BatchIntegralFunction{false, SciMLBase.FullSpecialize, NeuralPDE.var"#integrand#138"{typeof(loss)}, Nothing}, Tuple{Vector{Float64}, Vector{Float64}}, @Kwargs{}}, alg::Integrals.CubatureJLh, sensealg::Integrals.ReCallVJP{Integrals.ZygoteVJP}, domain::Tuple{Vector{Float64}, Vector{Float64}}, p::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}; reltol::Float64, abstol::Float64, maxiters::Int64)
    @ IntegralsCubatureExt [C:\Users\rnath\.julia\packages\Integrals\uahDt\ext\IntegralsCubatureExt.jl:17](file:///C:/Users/rnath/.julia/packages/Integrals/uahDt/ext/IntegralsCubatureExt.jl:17)
  [5] __solvebp_call
    @ [C:\Users\rnath\.julia\packages\Integrals\uahDt\ext\IntegralsCubatureExt.jl:7](file:///C:/Users/rnath/.julia/packages/Integrals/uahDt/ext/IntegralsCubatureExt.jl:7) [inlined]
  [6] #__solvebp_call#4
    @ [C:\Users\rnath\.julia\packages\Integrals\uahDt\src\common.jl:115](file:///C:/Users/rnath/.julia/packages/Integrals/uahDt/src/common.jl:115) [inlined]
  [7] __solvebp_call
    @ [C:\Users\rnath\.julia\packages\Integrals\uahDt\src\common.jl:114](file:///C:/Users/rnath/.julia/packages/Integrals/uahDt/src/common.jl:114) [inlined]
  [8] #rrule#15
    @ [C:\Users\rnath\.julia\packages\Integrals\uahDt\ext\IntegralsZygoteExt.jl:21](file:///C:/Users/rnath/.julia/packages/Integrals/uahDt/ext/IntegralsZygoteExt.jl:21) [inlined]
  [9] rrule
...
    @ OptimizationOptimJL [C:\Users\rnath\.julia\packages\OptimizationOptimJL\yMF3E\src\OptimizationOptimJL.jl:212](file:///C:/Users/rnath/.julia/packages/OptimizationOptimJL/yMF3E/src/OptimizationOptimJL.jl:212)
 [50] solve!(cache::OptimizationCache{OptimizationFunction{true, AutoZygote, NeuralPDE.var"#loss_function_#387"{NeuralPDE.var"#135#139"{typeof(loss), Vector{Float64}, Vector{Float64}, NeuralPDE.var"#134#137"{QuadratureTraining{Integrals.CubatureJLh, Float64}}, Float64}}, OptimizationZygoteExt.var"#38#56"{OptimizationZygoteExt.var"#37#55"{OptimizationFunction{true, AutoZygote, NeuralPDE.var"#loss_function_#387"{NeuralPDE.var"#135#139"{typeof(loss), Vector{Float64}, Vector{Float64}, NeuralPDE.var"#134#137"{QuadratureTraining{Integrals.CubatureJLh, Float64}}, Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, SciMLBase.NullParameters}}}, OptimizationZygoteExt.var"#41#59"{OptimizationZygoteExt.var"#37#55"{OptimizationFunction{true, AutoZygote, NeuralPDE.var"#loss_function_#387"{NeuralPDE.var"#135#139"{typeof(loss), Vector{Float64}, Vector{Float64}, NeuralPDE.var"#134#137"{QuadratureTraining{Integrals.CubatureJLh, Float64}}, Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, SciMLBase.NullParameters}}}, OptimizationZygoteExt.var"#45#63", Nothing, OptimizationZygoteExt.var"#49#67"{OptimizationFunction{true, AutoZygote, NeuralPDE.var"#loss_function_#387"{NeuralPDE.var"#135#139"{typeof(loss), Vector{Float64}, Vector{Float64}, NeuralPDE.var"#134#137"{QuadratureTraining{Integrals.CubatureJLh, Float64}}, Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, SciMLBase.NullParameters}}, OptimizationZygoteExt.var"#53#71"{OptimizationFunction{true, AutoZygote, NeuralPDE.var"#loss_function_#387"{NeuralPDE.var"#135#139"{typeof(loss), Vector{Float64}, Vector{Float64}, NeuralPDE.var"#134#137"{QuadratureTraining{Integrals.CubatureJLh, Float64}}, Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, SciMLBase.NullParameters}}, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, SciMLBase.NullParameters}, Nothing, Nothing, Nothing, Nothing, Nothing, BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}, Bool, var"#56#57"})
    @ SciMLBase [C:\Users\rnath\.julia\packages\SciMLBase\m3AcC\src\solve.jl:180](file:///C:/Users/rnath/.julia/packages/SciMLBase/m3AcC/src/solve.jl:180)
 [51] solve(::OptimizationProblem{true, OptimizationFunction{true, AutoZygote, NeuralPDE.var"#loss_function_#387"{NeuralPDE.var"#135#139"{typeof(loss), Vector{Float64}, Vector{Float64}, NeuralPDE.var"#134#137"{QuadratureTraining{Integrals.CubatureJLh, Float64}}, Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:24, Axis(weight = ViewAxis(1:16, ShapedAxis((8, 2))), bias = ViewAxis(17:24, ShapedAxis((8, 1))))), layer_2 = ViewAxis(25:96, Axis(weight = ViewAxis(1:64, ShapedAxis((8, 8))), bias = ViewAxis(65:72, ShapedAxis((8, 1))))), layer_3 = ViewAxis(97:105, Axis(weight = ViewAxis(1:8, ShapedAxis((1, 8))), bias = ViewAxis(9:9, ShapedAxis((1, 1))))))}}}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, @Kwargs{}}, ::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}; kwargs::@Kwargs{callback::var"#56#57", maxiters::Int64})
    @ SciMLBase [C:\Users\rnath\.julia\packages\SciMLBase\m3AcC\src\solve.jl:96](file:///C:/Users/rnath/.julia/packages/SciMLBase/m3AcC/src/solve.jl:96)

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, DiffEqBase
import ModelingToolkit: Interval, infimum, supremum

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

• Output of versioninfo()

Additional context

Add any other context about the problem here.

[ChrisRackauckas]

@sathvikbhagavan I thought you had just fixed that?

[geriatricvibes]

I’m running two weeks old release don’t know the exact version since I don’t have my laptop rn, if the fix is in updates after that, I’ll update it and try again tomorrow.

[geriatricvibes]

I'm running

[[deps.NeuralPDE]] deps = ["Adapt", "AdvancedHMC", "ArrayInterface", "ChainRulesCore", "ComponentArrays", "Cubature", "DiffEqBase", "DiffEqNoiseProcess", "Distributions", "DocStringExtensions", "DomainSets", "ForwardDiff", "Functors", "Integrals", "LinearAlgebra", "LogDensityProblems", "Lux", "MCMCChains", "ModelingToolkit", "MonteCarloMeasurements", "Optim", "Optimization", "OptimizationOptimisers", "QuasiMonteCarlo", "Random", "Reexport", "RuntimeGeneratedFunctions", "SciMLBase", "Statistics", "SymbolicUtils", "Symbolics", "UnPack", "Zygote"] git-tree-sha1 = "68b85e65255922385c08c9b31aa65711e62a7736" uuid = "315f7962-48a3-4962-8226-d0f33b1235f0" version = "5.12.0

The latest release

[geriatricvibes]

Hi, can you tell me what the problem or suggest some solutions so I could try to to solve it locally? Need neural adapter for a project

[sathvikbhagavan]

Hey, I think I have fixed it in #815 - https://github.com/SciML/NeuralPDE.jl/pull/815/files#diff-9cdb22a5deb20e6c866490a292335fbf406604b23d9a0602a4d230cbf588850e

or you can refer to the tests in https://github.com/SciML/NeuralPDE.jl/blob/master/test/neural_adapter_tests.jl

[geriatricvibes]

I'll take a look, thank you!