From 095027ed40f2e0e5dcc9cdacb52bca8130444e29 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 16 Oct 2024 09:30:55 -0400 Subject: [PATCH] refactor: more cleanup of neural adapter --- docs/Project.toml | 12 ++++----- src/neural_adapter.jl | 22 +++------------ src/training_strategies.jl | 2 +- test/BPINN_Tests.jl | 2 +- test/neural_adapter_tests.jl | 52 +++++++++++++++--------------------- 5 files changed, 33 insertions(+), 57 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index 3e62098b0a..eaf86e1c71 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -35,18 +35,18 @@ DiffEqBase = "6.148" Distributions = "0.25.107" Documenter = "1" DomainSets = "0.6, 0.7" -Flux = "0.14.11" +Flux = "0.14.17" Integrals = "4" LineSearches = "7.2" -Lux = "0.5.22" +Lux = "1" LuxCUDA = "0.3.2" MethodOfLines = "0.11" ModelingToolkit = "9.7" MonteCarloMeasurements = "1" -NeuralPDE = "5.14" -Optimization = "3.24, 4" -OptimizationOptimJL = "0.2.1, 0.3, 0.4" -OptimizationOptimisers = "0.2.1, 0.3" +NeuralPDE = "5" +Optimization = "4" +OptimizationOptimJL = "0.4" +OptimizationOptimisers = "0.3" OptimizationPolyalgorithms = "0.2" OrdinaryDiffEq = "6.74" Plots = "1.36" diff --git a/src/neural_adapter.jl b/src/neural_adapter.jl index 9c9580fee5..7fb811b2e2 100644 --- a/src/neural_adapter.jl +++ b/src/neural_adapter.jl @@ -1,11 +1,11 @@ function generate_training_sets(domains, dx, eqs, eltypeθ) dxs = dx isa Array ? dx : fill(dx, length(domains)) spans = [infimum(d.domain):dx:supremum(d.domain) for (d, dx) in zip(domains, dxs)] - return hcat(vec(map(points -> collect(points), Iterators.product(spans...)))...) |> + return reduce(hcat, vec(map(collect, Iterators.product(spans...)))) |> EltypeAdaptor{eltypeθ}() end -function get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, strategy) +function get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, _) dict_span = Dict([Symbol(d.variables) => [infimum(d.domain), supremum(d.domain)] for d in domains]) args = get_argument(eqs, dict_indvars, dict_depvars) @@ -13,23 +13,7 @@ function get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, strateg bounds = first(map(args) do pd return get.((dict_span,), pd, pd) |> EltypeAdaptor{eltypeθ}() end) - return [getindex.(bounds, 1), getindex.(bounds, 2)] -end - -function get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, - ::QuadratureTraining) - dict_lower_bound = Dict([Symbol(d.variables) => infimum(d.domain) for d in domains]) - dict_upper_bound = Dict([Symbol(d.variables) => supremum(d.domain) for d in domains]) - - args = get_argument(eqs, dict_indvars, dict_depvars) - - lower_bounds = map(args) do pd - return get.((dict_lower_bound,), pd, pd) |> EltypeAdaptor{eltypeθ}() - end - upper_bounds = map(args) do pd - return get.((dict_upper_bound,), pd, pd) |> EltypeAdaptor{eltypeθ}() - end - return lower_bounds, upper_bounds + return first.(bounds), last.(bounds) end function get_loss_function_neural_adapter( diff --git a/src/training_strategies.jl b/src/training_strategies.jl index ae31a350f6..974f2529fa 100644 --- a/src/training_strategies.jl +++ b/src/training_strategies.jl @@ -285,7 +285,7 @@ function get_loss_function(init_params, loss_function, lb, ub, eltypeθ, return solve(prob, strategy.quadrature_alg; strategy.reltol, strategy.abstol, strategy.maxiters)[1] end - return (θ) -> 1 / area * f_(lb, ub, loss_function, θ) + return (θ) -> f_(lb, ub, loss_function, θ) / area end """ diff --git a/test/BPINN_Tests.jl b/test/BPINN_Tests.jl index 73f571353f..c011e8fe9b 100644 --- a/test/BPINN_Tests.jl +++ b/test/BPINN_Tests.jl @@ -278,7 +278,7 @@ end @test_broken abs(param2 - p) < abs(0.25 * p) param1 = mean(i[62] for i in fhsampleslux12[750:length(fhsampleslux12)]) - @test_broken abs(param1 - p) < abs(0.75 * p) + @test abs(param1 - p) < abs(0.8 * p) @test abs(param2 - p) < abs(param1 - p) #-------------------------- solve() call diff --git a/test/neural_adapter_tests.jl b/test/neural_adapter_tests.jl index 108d1a4142..9b39a6260d 100644 --- a/test/neural_adapter_tests.jl +++ b/test/neural_adapter_tests.jl @@ -1,5 +1,5 @@ using Test, NeuralPDE, Optimization, Lux, OptimizationOptimisers, Statistics, - ComponentArrays, Random + ComponentArrays, Random, LinearAlgebra import ModelingToolkit: Interval, infimum, supremum Random.seed!(100) @@ -17,11 +17,11 @@ end Dyy = Differential(y)^2 # 2D PDE - eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y) + eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sinpi(x) * sinpi(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)] + bcs = [u(0, y) ~ 0.0, u(1, y) ~ -sinpi(1) * sinpi(y), + u(x, 0) ~ 0.0, u(x, 1) ~ -sinpi(x) * sinpi(1)] # Space and time domains domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] quadrature_strategy = QuadratureTraining( @@ -29,14 +29,12 @@ end inner = 8 af = tanh chain1 = Chain(Dense(2, inner, af), Dense(inner, inner, af), Dense(inner, 1)) - init_params = Lux.initialparameters(Random.default_rng(), chain1) |> - ComponentArray{Float64} - discretization = PhysicsInformedNN(chain1, quadrature_strategy; init_params) + discretization = PhysicsInformedNN(chain1, quadrature_strategy) @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) - prob = NeuralPDE.discretize(pde_system, discretization) + prob = discretize(pde_system, discretization) println("Poisson equation, strategy: $(nameof(typeof(quadrature_strategy)))") - @time res = solve(prob, Optimisers.Adam(5e-3); callback, maxiters = 2000) + @time res = solve(prob, Optimisers.Adam(5e-3); callback, maxiters = 10000) phi = discretization.phi inner_ = 8 @@ -58,7 +56,7 @@ end quasirandom_strategy]) do strategy_ println("Neural adapter Poisson equation, strategy: $(nameof(typeof(strategy_)))") prob_ = neural_adapter(loss, init_params2, pde_system, strategy_) - @time res_ = solve(prob_, Optimisers.Adam(5e-3); callback, maxiters = 2000) + @time res_ = solve(prob_, Optimisers.Adam(5e-3); callback, maxiters = 10000) end discretizations = map( @@ -67,7 +65,7 @@ end phis = map(discret -> discret.phi, discretizations) 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) + analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2) u_predict = [first(phi([x, y], res.u)) for x in xs for y in ys] @@ -90,10 +88,10 @@ end Dxx = Differential(x)^2 Dyy = Differential(y)^2 - eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y) + eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sinpi(x) * sinpi(y) - 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)] + bcs = [u(0, y) ~ 0.0, u(1, y) ~ -sinpi(1) * sinpi(y), + u(x, 0) ~ 0.0, u(x, 1) ~ -sinpi(x) * sinpi(1)] # Space x_0 = 0.0 @@ -108,9 +106,6 @@ end inner = 12 chains = [Chain(Dense(2, inner, af), Dense(inner, inner, af), Dense(inner, 1)) for _ in 1:count_decomp] - init_params = map( - c -> ComponentArray{Float64}(Lux.initialparameters(Random.default_rng(), c)), - chains) xs_ = infimum(x_domain):(1 / count_decomp):supremum(x_domain) xs_domain = [(xs_[i], xs_[i + 1]) for i in 1:(length(xs_) - 1)] @@ -119,16 +114,16 @@ end domains_ = [x ∈ x_domain_, y ∈ y_domain] end - analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2) + analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2) function create_bcs(x_domain_, phi_bound) x_0, x_e = x_domain_.left, x_domain_.right if x_0 == 0.0 bcs = [u(0, y) ~ 0.0, u(x_e, y) ~ analytic_sol_func(x_e, y), - u(x, 0) ~ 0.0, u(x, 1) ~ -sin(pi * x) * sin(pi * 1)] + u(x, 0) ~ 0.0, u(x, 1) ~ -sinpi(x) * sinpi(1)] return bcs end bcs = [u(x_0, y) ~ phi_bound(x_0, y), u(x_e, y) ~ analytic_sol_func(x_e, y), - u(x, 0) ~ 0.0, u(x, 1) ~ -sin(pi * x) * sin(pi * 1)] + u(x, 0) ~ 0.0, u(x, 1) ~ -sinpi(x) * sinpi(1)] bcs end @@ -138,6 +133,7 @@ end for i in 1:count_decomp println("decomposition $i") + domains_ = domains_map[i] phi_in(cord) = phis[i - 1](cord, reses[i - 1].u) phi_bound(x, y) = phi_in(vcat(x, y)) @@ -147,13 +143,12 @@ end @named pde_system_ = PDESystem(eq, bcs_, domains_, [x, y], [u(x, y)]) push!(pde_system_map, pde_system_) strategy = GridTraining([0.1 / count_decomp, 0.1]) - discretization = PhysicsInformedNN( - chains[i], strategy; init_params = init_params[i]) + discretization = PhysicsInformedNN(chains[i], strategy) prob = discretize(pde_system_, discretization) - @time res_ = solve( - prob, OptimizationOptimisers.Adam(5e-3); callback, maxiters = 2000) + @time res_ = solve(prob, Optimisers.Adam(5e-3); callback, maxiters = 2000) @show res_.objective phi = discretization.phi + push!(reses, res_) push!(phis, phi) end @@ -198,10 +193,7 @@ end @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) losses = map(1:count_decomp) do i - function loss(cord, θ) - ch2, st = chain2(cord, θ, st) - ch2 .- phis[i](cord, reses[i].u) - end + loss(cord, θ) = first(chain2(cord, θ, st)) .- phis[i](cord, reses[i].u) end prob_ = neural_adapter( @@ -220,6 +212,6 @@ end [analytic_sol_func(x, y) for x in xs for y in ys], (length(xs), length(ys))) diff_u_ = u_predict_ .- u_real - @test u_predict≈u_real rtol=1e-1 - @test u_predict_≈u_real rtol=1e-1 + @test u_predict≈u_real atol=5e-2 norm=Base.Fix2(norm, Inf) + @test u_predict_≈u_real atol=5e-2 norm=Base.Fix2(norm, Inf) end