From 9f501c82be8f1c962a9850ac825f0644cac413c1 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 26 Jun 2022 11:53:49 -0400 Subject: [PATCH 1/2] clean up some NNODE and training strategies docs --- docs/pages.jl | 3 ++- docs/src/examples/ode.md | 3 +-- docs/src/solvers/training_strategies.md | 26 +++++++++++++++++++++++++ src/ode_solve.jl | 6 ++++++ 4 files changed, 35 insertions(+), 3 deletions(-) create mode 100644 docs/src/solvers/training_strategies.md diff --git a/docs/pages.jl b/docs/pages.jl index a3c68b9626..109046c123 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -24,10 +24,11 @@ pages = [ "examples/nnrode_example.md", ], "API Documentation" => Any[ + "solvers/ode.md", "solvers/pinns.md", + "solvers/training_strategies.md", "solvers/kolmogorovbackwards_solver.md", "solvers/optimal_stopping.md",#TODO - "solvers/ode.md", "solvers/nnrode.md",#TODO ] ] \ No newline at end of file diff --git a/docs/src/examples/ode.md b/docs/src/examples/ode.md index 19eba83133..a705228abe 100644 --- a/docs/src/examples/ode.md +++ b/docs/src/examples/ode.md @@ -64,8 +64,7 @@ Once these pieces are together, we call `solve` just like with any other `ODEPro Let's turn on `verbose` so we can see the loss over time during the training process: ```@example nnode1 -sol = solve(prob, NeuralPDE.NNODE(chain, opt), dt=1 / 20f0, verbose=true, - abstol=1e-10, maxiters=200) +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose=true, abstol=1f-6, maxiters=200) ``` And that's it: the neural network solution was computed by training the neural network and diff --git a/docs/src/solvers/training_strategies.md b/docs/src/solvers/training_strategies.md new file mode 100644 index 0000000000..ff6bfb1f33 --- /dev/null +++ b/docs/src/solvers/training_strategies.md @@ -0,0 +1,26 @@ +# Training Strategies + +Training strategies are the choices for how the points are sampled for the definition +the physics-informed loss. + +## Recommendations + +`QuasiRandomTraining` with its default `LatinHyperCubeSample()` is a well-rounded training +strategy which can be used for most situations. It scales well for high dimensional +spaces and is GPU-compatible. `QuadratureTraining` can lead to faster or more robust convergence +with one of the H-Cubature or P-Cubature methods, but are not currently GPU compatible. +For very high dimensional cases, `QuadratureTraining` with an adaptive Monte Carlo quadrature +method, such as `CubaVegas`, can be beneficial for difficult or stiff problems. + +`GridTraining` should only be used for testing purposes and should not be relied upon for real +training cases. `StochasticTraining` achieves a lower convergence rate the quasi-Monte Carlo +methods and thus `QuasiRandomTraining` should be preferred in most cases. + +## API + +```@docs +GridTraining +StochasticTraining +QuasiRandomTraining +QuadratureTraining +``` \ No newline at end of file diff --git a/src/ode_solve.jl b/src/ode_solve.jl index ad405bed40..04b8dabfe6 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -7,6 +7,12 @@ NNODE(chain, opt=OptimizationPolyalgorithms.PolyOpt(), init_params = nothing; Algorithm for solving ordinary differential equations using a neural network. This is a specialization of the physics-informed neural network which is used as a solver for a standard `ODEProblem`. +!!! warn + + Note that NNODE only supports ODEs which are written in the out-of-place form, i.e. + `du = f(u,p,t)`, and not `f(du,u,p,t)`. If not declared out-of-place then the NNODE + will exit with an error. + ## Positional Arguments * `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.Chain`. From 7c3c4049fc36d95fe7df43256f38f0ec4ec1e866 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 26 Jun 2022 12:03:07 -0400 Subject: [PATCH 2/2] format SciML Style --- .github/workflows/FormatCheck.yml | 42 + docs/make.jl | 42 +- docs/pages.jl | 58 +- lib/NeuralPDELogging/src/NeuralPDELogging.jl | 11 +- .../test/adaptive_loss_log_tests.jl | 79 +- lib/NeuralPDELogging/test/runtests.jl | 31 +- src/NeuralPDE.jl | 116 +- src/adaptive_losses.jl | 99 +- src/kolmogorov_solve.jl | 74 +- src/neural_adapter.jl | 93 +- src/ode_solve.jl | 577 ++++---- src/param_kolmogorov_solve.jl | 185 +-- src/pinns_pde_solve.jl | 1158 ++++++++++------- src/rode_solve.jl | 93 +- src/stopping_solve.jl | 80 +- src/training_strategies.jl | 29 +- src/transform_inf_integral.jl | 68 +- test/IDE_tests.jl | 159 ++- test/NNKolmogorov_tests.jl | 99 +- test/NNODE_tests.jl | 146 ++- test/NNPDE_tests.jl | 577 ++++---- test/NNPDE_tests_gpu.jl | 189 +-- test/NNParamKolmogorov_tests.jl | 63 +- test/NNRODE_tests.jl | 43 +- test/Stopping_tests.jl | 53 +- test/adaptive_loss_tests.jl | 66 +- test/additional_loss_tests.jl | 158 +-- test/direct_function_tests.jl | 97 +- test/forward_tests.jl | 108 +- test/neural_adapter_tests.jl | 224 ++-- test/runtests.jl | 78 +- 31 files changed, 2644 insertions(+), 2251 deletions(-) create mode 100644 .github/workflows/FormatCheck.yml diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml new file mode 100644 index 0000000000..2a3517a0f3 --- /dev/null +++ b/.github/workflows/FormatCheck.yml @@ -0,0 +1,42 @@ +name: format-check + +on: + push: + branches: + - 'master' + - 'release-' + tags: '*' + pull_request: + +jobs: + build: + runs-on: ${{ matrix.os }} + strategy: + matrix: + julia-version: [1] + julia-arch: [x86] + os: [ubuntu-latest] + steps: + - uses: julia-actions/setup-julia@latest + with: + version: ${{ matrix.julia-version }} + + - uses: actions/checkout@v1 + - name: Install JuliaFormatter and format + # This will use the latest version by default but you can set the version like so: + # + # julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))' + run: | + julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))' + julia -e 'using JuliaFormatter; format(".", verbose=true)' + - name: Format check + run: | + julia -e ' + out = Cmd(`git diff --name-only`) |> read |> String + if out == "" + exit(0) + else + @error "Some files have not been formatted !!!" + write(stdout, out) + exit(1) + end' diff --git a/docs/make.jl b/docs/make.jl index 4005174cbb..01c56501cd 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,27 +2,23 @@ using Documenter, NeuralPDE include("pages.jl") -makedocs( - sitename="NeuralPDE.jl", - authors="#", - clean=true, - doctest=false, - modules=[NeuralPDE], - strict=[ - :doctest, - :linkcheck, - :parse_error, - :example_block, - # Other available options are - # :autodocs_block, :cross_references, :docs_block, :eval_block, :example_block, :footnote, :meta_block, :missing_docs, :setup_block - ], - format=Documenter.HTML( analytics = "UA-90474609-3", - assets=["assets/favicon.ico"], - canonical="https://neuralpde.sciml.ai/stable/"), - pages=pages -) +makedocs(sitename = "NeuralPDE.jl", + authors = "#", + clean = true, + doctest = false, + modules = [NeuralPDE], + strict = [ + :doctest, + :linkcheck, + :parse_error, + :example_block, + # Other available options are + # :autodocs_block, :cross_references, :docs_block, :eval_block, :example_block, :footnote, :meta_block, :missing_docs, :setup_block + ], + format = Documenter.HTML(analytics = "UA-90474609-3", + assets = ["assets/favicon.ico"], + canonical = "https://neuralpde.sciml.ai/stable/"), + pages = pages) -deploydocs( - repo="github.com/SciML/NeuralPDE.jl.git"; - push_preview=true -) +deploydocs(repo = "github.com/SciML/NeuralPDE.jl.git"; + push_preview = true) diff --git a/docs/pages.jl b/docs/pages.jl index 109046c123..d8910f0674 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -1,34 +1,26 @@ pages = [ - "NeuralPDE.jl: Scientific Machine Learning (SciML) for Partial Differential Equations" => "index.md", - "Physics-Informed Neural Network Tutorials" => Any[ - "pinn/poisson.md", - "pinn/wave.md", - "pinn/2D.md", - "pinn/system.md", - "pinn/3rd.md", - "pinn/low_level.md", - "pinn/ks.md", - "pinn/fp.md", - "pinn/parm_estim.md", - "pinn/heterogeneous.md", - "pinn/integro_diff.md", - "pinn/debugging.md", - "pinn/neural_adapter.md", - ], - "Specialized Neural PDE Tutorials" => Any[ - "examples/kolmogorovbackwards.md", - "examples/optimal_stopping_american.md", - ], - "Specialized Neural ODE Tutorials" => Any[ - "examples/ode.md", - "examples/nnrode_example.md", - ], - "API Documentation" => Any[ - "solvers/ode.md", - "solvers/pinns.md", - "solvers/training_strategies.md", - "solvers/kolmogorovbackwards_solver.md", - "solvers/optimal_stopping.md",#TODO - "solvers/nnrode.md",#TODO - ] - ] \ No newline at end of file + "NeuralPDE.jl: Scientific Machine Learning (SciML) for Partial Differential Equations" => "index.md", + "Physics-Informed Neural Network Tutorials" => Any["pinn/poisson.md", + "pinn/wave.md", + "pinn/2D.md", + "pinn/system.md", + "pinn/3rd.md", + "pinn/low_level.md", + "pinn/ks.md", + "pinn/fp.md", + "pinn/parm_estim.md", + "pinn/heterogeneous.md", + "pinn/integro_diff.md", + "pinn/debugging.md", + "pinn/neural_adapter.md"], + "Specialized Neural PDE Tutorials" => Any["examples/kolmogorovbackwards.md", + "examples/optimal_stopping_american.md"], + "Specialized Neural ODE Tutorials" => Any["examples/ode.md", + "examples/nnrode_example.md"], + "API Documentation" => Any["solvers/ode.md", + "solvers/pinns.md", + "solvers/training_strategies.md", + "solvers/kolmogorovbackwards_solver.md", + "solvers/optimal_stopping.md",#TODO + "solvers/nnrode.md"], +] diff --git a/lib/NeuralPDELogging/src/NeuralPDELogging.jl b/lib/NeuralPDELogging/src/NeuralPDELogging.jl index cafba1625f..24a4cda4e6 100644 --- a/lib/NeuralPDELogging/src/NeuralPDELogging.jl +++ b/lib/NeuralPDELogging/src/NeuralPDELogging.jl @@ -5,19 +5,20 @@ using TensorBoardLogger """This function overrides the empty function in NeuralPDE in order to use TensorBoardLogger in that package This is light type piracy but it should be alright since this is a subpackage of NeuralPDE""" -function NeuralPDE.logvector(logger::TBLogger, vector::AbstractVector{R}, name::AbstractString, step::Integer) where R <: Real +function NeuralPDE.logvector(logger::TBLogger, vector::AbstractVector{R}, + name::AbstractString, step::Integer) where {R <: Real} for j in 1:length(vector) - log_value(logger, "$(name)/$(j)", vector[j], step=step) + log_value(logger, "$(name)/$(j)", vector[j], step = step) end nothing end """This function overrides the empty function in NeuralPDE in order to use TensorBoardLogger in that package. This is light type piracy but it should be alright since this is a subpackage of NeuralPDE""" -function NeuralPDE.logscalar(logger::TBLogger, scalar::R, name::AbstractString, step::Integer) where R <: Real - log_value(logger, "$(name)", scalar, step=step) +function NeuralPDE.logscalar(logger::TBLogger, scalar::R, name::AbstractString, + step::Integer) where {R <: Real} + log_value(logger, "$(name)", scalar, step = step) nothing end - end diff --git a/lib/NeuralPDELogging/test/adaptive_loss_log_tests.jl b/lib/NeuralPDELogging/test/adaptive_loss_log_tests.jl index b79c6ff037..91003e37c6 100644 --- a/lib/NeuralPDELogging/test/adaptive_loss_log_tests.jl +++ b/lib/NeuralPDELogging/test/adaptive_loss_log_tests.jl @@ -7,15 +7,18 @@ using Random #using Plots @info "Starting Soon!" -nonadaptive_loss = NeuralPDE.NonAdaptiveLoss(pde_loss_weights=1, bc_loss_weights=1) -gradnormadaptive_loss = NeuralPDE.GradientScaleAdaptiveLoss(100, pde_loss_weights=1e3, bc_loss_weights=1) -adaptive_loss = NeuralPDE.MiniMaxAdaptiveLoss(100; pde_loss_weights=1, bc_loss_weights=1) -adaptive_losses = [nonadaptive_loss, gradnormadaptive_loss,adaptive_loss] -maxiters=800 -seed=60 +nonadaptive_loss = NeuralPDE.NonAdaptiveLoss(pde_loss_weights = 1, bc_loss_weights = 1) +gradnormadaptive_loss = NeuralPDE.GradientScaleAdaptiveLoss(100, pde_loss_weights = 1e3, + bc_loss_weights = 1) +adaptive_loss = NeuralPDE.MiniMaxAdaptiveLoss(100; pde_loss_weights = 1, + bc_loss_weights = 1) +adaptive_losses = [nonadaptive_loss, gradnormadaptive_loss, adaptive_loss] +maxiters = 800 +seed = 60 ## 2D Poisson equation -function test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, outdir, haslogger; seed=60, maxiters=800) +function test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, outdir, haslogger; + seed = 60, maxiters = 800) logdir = joinpath(outdir, string(run)) if haslogger logger = TBLogger(logdir) @@ -24,8 +27,9 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, outdir, hasl end Random.seed!(seed) hid = 40 - chain_ = FastChain(FastDense(2,hid,Flux.σ),FastDense(hid,hid,Flux.σ),FastDense(hid,1)) - strategy_ = NeuralPDE.StochasticTraining(256) + chain_ = FastChain(FastDense(2, hid, Flux.σ), FastDense(hid, hid, Flux.σ), + FastDense(hid, 1)) + strategy_ = NeuralPDE.StochasticTraining(256) @info "adaptive reweighting test logdir: $(logdir), maxiters: $(maxiters), 2D Poisson equation, adaptive_loss: $(nameof(typeof(adaptive_loss))) " @parameters x y @variables u(..) @@ -33,14 +37,14 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, outdir, hasl 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)) ~ -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)] + 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)] + domains = [x ∈ Interval(0.0, 1.0), + y ∈ Interval(0.0, 1.0)] initθ = Float64.(DiffEqFlux.initial_params(chain_)) iteration = [0] @@ -49,43 +53,46 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, outdir, hasl init_params = initθ, adaptive_loss = adaptive_loss, logger = logger, - iteration=iteration) + iteration = iteration) - - @named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)]) - prob = NeuralPDE.discretize(pde_system,discretization) + @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) + prob = NeuralPDE.discretize(pde_system, discretization) phi = discretization.phi - sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) - + sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) - 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_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys))) + 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_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys], + (length(xs), length(ys))) - callback = function (p,l) + callback = function (p, l) iteration[1] += 1 if iteration[1] % 100 == 0 @info "Current loss is: $l, iteration is $(iteration[1])" end if haslogger - log_value(logger, "outer_error/loss", l, step=iteration[1]) + log_value(logger, "outer_error/loss", l, step = iteration[1]) if iteration[1] % 30 == 0 - u_predict = reshape([first(phi([x,y],p)) for x in xs for y in ys],(length(xs),length(ys))) + u_predict = reshape([first(phi([x, y], p)) for x in xs for y in ys], + (length(xs), length(ys))) diff_u = abs.(u_predict .- u_real) total_diff = sum(diff_u) - log_value(logger, "outer_error/total_diff", total_diff, step=iteration[1]) + log_value(logger, "outer_error/total_diff", total_diff, step = iteration[1]) total_u = sum(abs.(u_real)) total_diff_rel = total_diff / total_u - log_value(logger, "outer_error/total_diff_rel", total_diff_rel, step=iteration[1]) + log_value(logger, "outer_error/total_diff_rel", total_diff_rel, + step = iteration[1]) total_diff_sq = sum(diff_u .^ 2) - log_value(logger, "outer_error/total_diff_sq", total_diff_sq, step=iteration[1]) + log_value(logger, "outer_error/total_diff_sq", total_diff_sq, + step = iteration[1]) end end return false end - res = Optimization.solve(prob, ADAM(0.03); maxiters=maxiters, callback=callback) + res = Optimization.solve(prob, ADAM(0.03); maxiters = maxiters, callback = callback) - u_predict = reshape([first(phi([x,y],res.minimizer)) 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))) diff_u = abs.(u_predict .- u_real) total_diff = sum(diff_u) total_u = sum(abs.(u_real)) @@ -95,7 +102,7 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, outdir, hasl #p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict"); #p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error"); #(plot=plot(p1,p2,p3), error=total_diff, total_diff_rel=total_diff_rel) - (error=total_diff, total_diff_rel=total_diff_rel) + (error = total_diff, total_diff_rel = total_diff_rel) end possible_logger_dir = mktempdir() @@ -115,8 +122,12 @@ end @info "has logger: $(haslogger), expected log folders: $(expected_log_folders)" -test_2d_poisson_equation_adaptive_loss_run_seediters(adaptive_loss, run) = test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, possible_logger_dir, haslogger; seed=seed, maxiters=maxiters) -error_results = map(test_2d_poisson_equation_adaptive_loss_run_seediters, adaptive_losses, 1:length(adaptive_losses)) +function test_2d_poisson_equation_adaptive_loss_run_seediters(adaptive_loss, run) + test_2d_poisson_equation_adaptive_loss(adaptive_loss, run, possible_logger_dir, + haslogger; seed = seed, maxiters = maxiters) +end +error_results = map(test_2d_poisson_equation_adaptive_loss_run_seediters, adaptive_losses, + 1:length(adaptive_losses)) @test length(readdir(possible_logger_dir)) == expected_log_folders if expected_log_folders > 0 diff --git a/lib/NeuralPDELogging/test/runtests.jl b/lib/NeuralPDELogging/test/runtests.jl index 7e8fdb6ddc..f00b4eb825 100644 --- a/lib/NeuralPDELogging/test/runtests.jl +++ b/lib/NeuralPDELogging/test/runtests.jl @@ -3,16 +3,14 @@ using SafeTestsets const GROUP = get(ENV, "GROUP", "All") -const is_APPVEYOR = Sys.iswindows() && haskey(ENV,"APPVEYOR") +const is_APPVEYOR = Sys.iswindows() && haskey(ENV, "APPVEYOR") -const is_TRAVIS = haskey(ENV,"TRAVIS") +const is_TRAVIS = haskey(ENV, "TRAVIS") -is_CI = haskey(ENV,"CI") +is_CI = haskey(ENV, "CI") - -@time begin - if GROUP == "All" || GROUP == "Logging" - @time @safetestset "AdaptiveLossLogNoImport" begin +@time begin if GROUP == "All" || GROUP == "Logging" + @time @safetestset "AdaptiveLossLogNoImport" begin using Pkg neuralpde_dir = dirname(abspath(joinpath(@__DIR__, "..", "..", ".."))) @info "loading neuralpde package at : $(neuralpde_dir)" @@ -20,9 +18,9 @@ is_CI = haskey(ENV,"CI") Pkg.develop(neuralpde) @info "making sure that there are no logs without having imported NeuralPDELogging" ENV["LOG_SETTING"] = "NoImport" - include("adaptive_loss_log_tests.jl") - end - @time @safetestset "AdaptiveLossLogImportNoUse" begin + include("adaptive_loss_log_tests.jl") + end + @time @safetestset "AdaptiveLossLogImportNoUse" begin using Pkg neuralpde_dir = dirname(abspath(joinpath(@__DIR__, "..", "..", ".."))) @info "loading neuralpde package at : $(neuralpde_dir)" @@ -30,9 +28,9 @@ is_CI = haskey(ENV,"CI") Pkg.develop(neuralpde) @info "making sure that there are still no logs now that we have imported NeuralPDELogging" ENV["LOG_SETTING"] = "ImportNoUse" - include("adaptive_loss_log_tests.jl") - end - @time @safetestset "AdaptiveLossLogImportUse" begin + include("adaptive_loss_log_tests.jl") + end + @time @safetestset "AdaptiveLossLogImportUse" begin using Pkg neuralpde_dir = dirname(abspath(joinpath(@__DIR__, "..", "..", ".."))) @info "loading neuralpde package at : $(neuralpde_dir)" @@ -40,7 +38,6 @@ is_CI = haskey(ENV,"CI") Pkg.develop(neuralpde) ENV["LOG_SETTING"] = "ImportUse" @info "making sure that logs are generated now if we use a logger" - include("adaptive_loss_log_tests.jl") - end - end -end + include("adaptive_loss_log_tests.jl") + end +end end diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index eb7837bb54..f7cc3126ce 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -40,30 +40,33 @@ Consider `du/dt = l(u) + f(u)`; where l is the nonlinear Lipschitz function * `x0`: The initial X for the problem. * `tspan`: The timespan of the problem. """ -struct TerminalPDEProblem{G,F,Mu,Sigma,X,T,P,A,UD,K} <: SciMLBase.SciMLProblem +struct TerminalPDEProblem{G, F, Mu, Sigma, X, T, P, A, UD, K} <: SciMLBase.SciMLProblem g::G f::F μ::Mu σ::Sigma X0::X - tspan::Tuple{T,T} + tspan::Tuple{T, T} p::P A::A u_domain::UD kwargs::K - TerminalPDEProblem(g,f,μ,σ,X0,tspan,p=nothing;A=nothing,u_domain=nothing,kwargs...) = new{typeof(g),typeof(f), - typeof(μ),typeof(σ), - typeof(X0),eltype(tspan), - typeof(p),typeof(A),typeof(u_domain),typeof(kwargs)}( - g,f,μ,σ,X0,tspan,p,A,u_domain,kwargs) + function TerminalPDEProblem(g, f, μ, σ, X0, tspan, p = nothing; A = nothing, + u_domain = nothing, kwargs...) + new{typeof(g), typeof(f), + typeof(μ), typeof(σ), + typeof(X0), eltype(tspan), + typeof(p), typeof(A), typeof(u_domain), typeof(kwargs)}(g, f, μ, σ, X0, tspan, + p, A, u_domain, kwargs) + end end Base.summary(prob::TerminalPDEProblem) = string(nameof(typeof(prob))) function Base.show(io::IO, A::TerminalPDEProblem) - println(io,summary(A)) - print(io,"timespan: ") - show(io,A.tspan) + println(io, summary(A)) + print(io, "timespan: ") + show(io, A.tspan) end """ @@ -78,45 +81,52 @@ A standard Kolmogorov PDE Problem. * `d`: The dimensions of the input x. * `noise_rate_prototype`: A prototype type instance for the noise rates, that is the output g. """ -struct KolmogorovPDEProblem{ F, G, Phi, X , T , D ,P,U0, ND} <: DiffEqBase.DEProblem +struct KolmogorovPDEProblem{F, G, Phi, X, T, D, P, U0, ND} <: DiffEqBase.DEProblem f::F g::G phi::Phi - xspan::Tuple{X,X} - tspan::Tuple{T,T} + xspan::Tuple{X, X} + tspan::Tuple{T, T} d::D p::P u0::U0 noise_rate_prototype::ND - KolmogorovPDEProblem( f, g, phi , xspan , tspan , d, p=nothing, u0=0 , noise_rate_prototype= nothing) = new{typeof(f),typeof(g),typeof(phi),eltype(tspan),eltype(xspan),typeof(d),typeof(p),typeof(u0),typeof(noise_rate_prototype)}(f,g,phi,xspan,tspan,d,p,u0,noise_rate_prototype) + function KolmogorovPDEProblem(f, g, phi, xspan, tspan, d, p = nothing, u0 = 0, + noise_rate_prototype = nothing) + new{typeof(f), typeof(g), typeof(phi), eltype(tspan), eltype(xspan), typeof(d), + typeof(p), typeof(u0), typeof(noise_rate_prototype)}(f, g, phi, xspan, tspan, d, + p, u0, + noise_rate_prototype) + end end Base.summary(prob::KolmogorovPDEProblem) = string(nameof(typeof(prob))) function Base.show(io::IO, A::KolmogorovPDEProblem) - println(io,summary(A)) - print(io,"timespan: ") - show(io,A.tspan) - print(io,"xspan: ") - show(io,A.xspan) - println(io , "μ") - show(io , A.f) - println(io,"Sigma") - show(io , A.g) + println(io, summary(A)) + print(io, "timespan: ") + show(io, A.tspan) + print(io, "xspan: ") + show(io, A.xspan) + println(io, "μ") + show(io, A.f) + println(io, "Sigma") + show(io, A.g) end abstract type ParametersDomain end -struct KolmogorovParamDomain{T} <:ParametersDomain - sigma::Tuple{T,T} - mu::Tuple{T,T} - phi::Tuple{T,T} +struct KolmogorovParamDomain{T} <: ParametersDomain + sigma::Tuple{T, T} + mu::Tuple{T, T} + phi::Tuple{T, T} end -struct ParamKolmogorovPDEProblem{ F, G, Phi, X , T , D, YD , P,U0 , YSP , YMP ,YPH, NP} <: DiffEqBase.DEProblem +struct ParamKolmogorovPDEProblem{F, G, Phi, X, T, D, YD, P, U0, YSP, YMP, YPH, NP} <: + DiffEqBase.DEProblem f::F g::G phi::Phi - xspan::Tuple{X,X} - tspan::Tuple{T,T} + xspan::Tuple{X, X} + tspan::Tuple{T, T} d::D Y_domain::YD p::P @@ -125,20 +135,39 @@ struct ParamKolmogorovPDEProblem{ F, G, Phi, X , T , D, YD , P,U0 , YSP , YMP ,Y Y_mu_prototype::YMP Y_phi_prototype::YPH noise_rate_prototype::NP - ParamKolmogorovPDEProblem( f, g, phi , xspan , tspan , d, Y_domain, p=nothing, u0=0 ; Y_sigma_prototype=nothing , Y_mu_prototype=nothing , Y_phi_prototype=nothing , noise_rate_prototype= nothing) = new{typeof(f),typeof(g),typeof(phi),eltype(tspan),eltype(xspan),typeof(d),typeof(Y_domain),typeof(p),typeof(u0),typeof(Y_sigma_prototype),typeof(Y_mu_prototype),typeof(Y_phi_prototype),typeof(noise_rate_prototype)}(f,g,phi,xspan,tspan,d,Y_domain,p,u0 , Y_sigma_prototype,Y_mu_prototype,Y_phi_prototype,noise_rate_prototype) + function ParamKolmogorovPDEProblem(f, g, phi, xspan, tspan, d, Y_domain, p = nothing, + u0 = 0; Y_sigma_prototype = nothing, + Y_mu_prototype = nothing, Y_phi_prototype = nothing, + noise_rate_prototype = nothing) + new{typeof(f), typeof(g), typeof(phi), eltype(tspan), eltype(xspan), typeof(d), + typeof(Y_domain), typeof(p), typeof(u0), typeof(Y_sigma_prototype), + typeof(Y_mu_prototype), typeof(Y_phi_prototype), typeof(noise_rate_prototype)}(f, + g, + phi, + xspan, + tspan, + d, + Y_domain, + p, + u0, + Y_sigma_prototype, + Y_mu_prototype, + Y_phi_prototype, + noise_rate_prototype) + end end Base.summary(prob::ParamKolmogorovPDEProblem) = string(nameof(typeof(prob))) function Base.show(io::IO, A::ParamKolmogorovPDEProblem) - println(io,summary(A)) - print(io,"timespan: ") - show(io,A.tspan) - print(io,"xspan: ") - show(io,A.xspan) - println(io , "μ") - show(io , A.f) - println(io,"Sigma") - show(io , A.g) + println(io, summary(A)) + print(io, "timespan: ") + show(io, A.tspan) + print(io, "xspan: ") + show(io, A.xspan) + println(io, "μ") + show(io, A.f) + println(io, "Sigma") + show(io, A.g) end include("training_strategies.jl") @@ -153,15 +182,16 @@ include("neural_adapter.jl") include("param_kolmogorov_solve.jl") export NNODE, TerminalPDEProblem, NNPDEHan, NNPDENS, NNRODE, - KolmogorovPDEProblem, NNKolmogorov, NNStopping,ParamKolmogorovPDEProblem,KolmogorovParamDomain, NNParamKolmogorov, + KolmogorovPDEProblem, NNKolmogorov, NNStopping, ParamKolmogorovPDEProblem, + KolmogorovParamDomain, NNParamKolmogorov, PhysicsInformedNN, discretize, GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, build_loss_function, get_loss_function, generate_training_sets, get_variables, get_argument, get_bounds, get_phi, get_numeric_derivative, get_numeric_integral, build_symbolic_equation, build_symbolic_loss_function, symbolic_discretize, - AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, MiniMaxAdaptiveLoss, + AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, + MiniMaxAdaptiveLoss, LogOptions - end # module diff --git a/src/adaptive_losses.jl b/src/adaptive_losses.jl index 7ae9056044..54a0e74b6f 100644 --- a/src/adaptive_losses.jl +++ b/src/adaptive_losses.jl @@ -10,15 +10,24 @@ A way of weighting the components of the loss function in the total sum that doe mutable struct NonAdaptiveLoss{T <: Real} <: AbstractAdaptiveLoss pde_loss_weights::Vector{T} bc_loss_weights::Vector{T} - additional_loss_weights::Vector{T} - SciMLBase.@add_kwonly function NonAdaptiveLoss{T}(;pde_loss_weights=1, bc_loss_weights=1, additional_loss_weights=1) where T <: Real - new(vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T), vectorify(additional_loss_weights, T)) + additional_loss_weights::Vector{T} + SciMLBase.@add_kwonly function NonAdaptiveLoss{T}(; pde_loss_weights = 1, + bc_loss_weights = 1, + additional_loss_weights = 1) where { + T <: + Real + } + new(vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T), + vectorify(additional_loss_weights, T)) end end # default to Float64 -SciMLBase.@add_kwonly function NonAdaptiveLoss(;pde_loss_weights=1, bc_loss_weights=1, additional_loss_weights=1) - NonAdaptiveLoss{Float64}(;pde_loss_weights=pde_loss_weights, bc_loss_weights=bc_loss_weights, additional_loss_weights=additional_loss_weights) +SciMLBase.@add_kwonly function NonAdaptiveLoss(; pde_loss_weights = 1, bc_loss_weights = 1, + additional_loss_weights = 1) + NonAdaptiveLoss{Float64}(; pde_loss_weights = pde_loss_weights, + bc_loss_weights = bc_loss_weights, + additional_loss_weights = additional_loss_weights) end """ @@ -40,20 +49,35 @@ https://github.com/PredictiveIntelligenceLab/GradientPathologiesPINNs mutable struct GradientScaleAdaptiveLoss{T <: Real} <: AbstractAdaptiveLoss reweight_every::Int64 weight_change_inertia::T - pde_loss_weights::Vector{T} - bc_loss_weights::Vector{T} - additional_loss_weights::Vector{T} - SciMLBase.@add_kwonly function GradientScaleAdaptiveLoss{T}(reweight_every; weight_change_inertia=0.9, pde_loss_weights=1, bc_loss_weights=1, additional_loss_weights=1) where T <: Real - new(convert(Int64, reweight_every), convert(T, weight_change_inertia), vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T), vectorify(additional_loss_weights, T)) + pde_loss_weights::Vector{T} + bc_loss_weights::Vector{T} + additional_loss_weights::Vector{T} + SciMLBase.@add_kwonly function GradientScaleAdaptiveLoss{T}(reweight_every; + weight_change_inertia = 0.9, + pde_loss_weights = 1, + bc_loss_weights = 1, + additional_loss_weights = 1) where { + T <: + Real + } + new(convert(Int64, reweight_every), convert(T, weight_change_inertia), + vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T), + vectorify(additional_loss_weights, T)) end end # default to Float64 -SciMLBase.@add_kwonly function GradientScaleAdaptiveLoss(reweight_every; weight_change_inertia=0.9, pde_loss_weights=1, bc_loss_weights=1, additional_loss_weights=1) - GradientScaleAdaptiveLoss{Float64}(reweight_every; weight_change_inertia=weight_change_inertia, - pde_loss_weights=pde_loss_weights, bc_loss_weights=bc_loss_weights, additional_loss_weights=additional_loss_weights) +SciMLBase.@add_kwonly function GradientScaleAdaptiveLoss(reweight_every; + weight_change_inertia = 0.9, + pde_loss_weights = 1, + bc_loss_weights = 1, + additional_loss_weights = 1) + GradientScaleAdaptiveLoss{Float64}(reweight_every; + weight_change_inertia = weight_change_inertia, + pde_loss_weights = pde_loss_weights, + bc_loss_weights = bc_loss_weights, + additional_loss_weights = additional_loss_weights) end - """ A way of adaptively reweighting the components of the loss function in the total sum such that the loss weights are maximized by an internal optimiser, which leads to a behavior where loss functions that have not been satisfied get a greater weight, @@ -69,23 +93,46 @@ Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism Levi McClenny, Ulisses Braga-Neto https://arxiv.org/abs/2009.04544 """ -mutable struct MiniMaxAdaptiveLoss{T <: Real, PDE_OPT <: Flux.Optimise.AbstractOptimiser, BC_OPT <: Flux.Optimise.AbstractOptimiser} <: AbstractAdaptiveLoss +mutable struct MiniMaxAdaptiveLoss{T <: Real, PDE_OPT <: Flux.Optimise.AbstractOptimiser, + BC_OPT <: Flux.Optimise.AbstractOptimiser} <: + AbstractAdaptiveLoss reweight_every::Int64 pde_max_optimiser::PDE_OPT bc_max_optimiser::BC_OPT - pde_loss_weights::Vector{T} - bc_loss_weights::Vector{T} - additional_loss_weights::Vector{T} - SciMLBase.@add_kwonly function MiniMaxAdaptiveLoss{T, PDE_OPT, BC_OPT}(reweight_every; pde_max_optimiser=Flux.ADAM(1e-4), bc_max_optimiser=Flux.ADAM(0.5), - pde_loss_weights=1, bc_loss_weights=1, additional_loss_weights=1) where {T <: Real, PDE_OPT <: Flux.Optimise.AbstractOptimiser, BC_OPT <: Flux.Optimise.AbstractOptimiser} - new(convert(Int64, reweight_every), convert(PDE_OPT, pde_max_optimiser), convert(BC_OPT, bc_max_optimiser), - vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T), vectorify(additional_loss_weights, T)) + pde_loss_weights::Vector{T} + bc_loss_weights::Vector{T} + additional_loss_weights::Vector{T} + SciMLBase.@add_kwonly function MiniMaxAdaptiveLoss{T, PDE_OPT, BC_OPT}(reweight_every; + pde_max_optimiser = Flux.ADAM(1e-4), + bc_max_optimiser = Flux.ADAM(0.5), + pde_loss_weights = 1, + bc_loss_weights = 1, + additional_loss_weights = 1) where { + T <: + Real, + PDE_OPT <: + Flux.Optimise.AbstractOptimiser, + BC_OPT <: + Flux.Optimise.AbstractOptimiser + } + new(convert(Int64, reweight_every), convert(PDE_OPT, pde_max_optimiser), + convert(BC_OPT, bc_max_optimiser), + vectorify(pde_loss_weights, T), vectorify(bc_loss_weights, T), + vectorify(additional_loss_weights, T)) end end # default to Float64, ADAM, ADAM -SciMLBase.@add_kwonly function MiniMaxAdaptiveLoss(reweight_every; pde_max_optimiser=Flux.ADAM(1e-4), bc_max_optimiser=Flux.ADAM(0.5), pde_loss_weights=1, bc_loss_weights=1, additional_loss_weights=1) - MiniMaxAdaptiveLoss{Float64, typeof(pde_max_optimiser), typeof(bc_max_optimiser)}( - reweight_every; pde_max_optimiser=pde_max_optimiser, bc_max_optimiser=bc_max_optimiser, - pde_loss_weights=pde_loss_weights, bc_loss_weights=bc_loss_weights, additional_loss_weights=additional_loss_weights) +SciMLBase.@add_kwonly function MiniMaxAdaptiveLoss(reweight_every; + pde_max_optimiser = Flux.ADAM(1e-4), + bc_max_optimiser = Flux.ADAM(0.5), + pde_loss_weights = 1, + bc_loss_weights = 1, + additional_loss_weights = 1) + MiniMaxAdaptiveLoss{Float64, typeof(pde_max_optimiser), typeof(bc_max_optimiser)}(reweight_every; + pde_max_optimiser = pde_max_optimiser, + bc_max_optimiser = bc_max_optimiser, + pde_loss_weights = pde_loss_weights, + bc_loss_weights = bc_loss_weights, + additional_loss_weights = additional_loss_weights) end diff --git a/src/kolmogorov_solve.jl b/src/kolmogorov_solve.jl index c18edd1e8a..7bdabeab9d 100644 --- a/src/kolmogorov_solve.jl +++ b/src/kolmogorov_solve.jl @@ -16,28 +16,28 @@ Arguments: documentation for more details. [1]Beck, Christian, et al. "Solving stochastic differential equations and Kolmogorov equations by means of deep learning." arXiv preprint arXiv:1806.00421 (2018). """ -struct NNKolmogorov{C,O,S,E} <: NeuralPDEAlgorithm +struct NNKolmogorov{C, O, S, E} <: NeuralPDEAlgorithm chain::C opt::O sdealg::S ensemblealg::E end -NNKolmogorov(chain ; opt=Flux.ADAM(0.1) , sdealg = EM() , ensemblealg = EnsembleThreads()) = NNKolmogorov(chain , opt , sdealg , ensemblealg) - -function DiffEqBase.solve( - prob::Union{KolmogorovPDEProblem,SDEProblem}, - alg::NNKolmogorov; - abstol = 1f-6, - verbose = false, - maxiters = 300, - trajectories = 1000, - save_everystep = false, - use_gpu = false, - dt, - dx, - kwargs... - ) +function NNKolmogorov(chain; opt = Flux.ADAM(0.1), sdealg = EM(), + ensemblealg = EnsembleThreads()) + NNKolmogorov(chain, opt, sdealg, ensemblealg) +end +function DiffEqBase.solve(prob::Union{KolmogorovPDEProblem, SDEProblem}, + alg::NNKolmogorov; + abstol = 1.0f-6, + verbose = false, + maxiters = 300, + trajectories = 1000, + save_everystep = false, + use_gpu = false, + dt, + dx, + kwargs...) tspan = prob.tspan sigma = prob.g μ = prob.f @@ -46,7 +46,7 @@ function DiffEqBase.solve( xspan = prob.kwargs.data.xspan d = prob.kwargs.data.d u0 = prob.u0 - phi(xi) = pdf(u0 , xi) + phi(xi) = pdf(u0, xi) else xspan = prob.xspan d = prob.d @@ -58,38 +58,42 @@ function DiffEqBase.solve( T = tspan[2] #hidden layer - chain = alg.chain - opt = alg.opt + chain = alg.chain + opt = alg.opt sdealg = alg.sdealg ensemblealg = alg.ensemblealg - ps = Flux.params(chain) - xi = rand(xs ,d ,trajectories) + ps = Flux.params(chain) + xi = rand(xs, d, trajectories) #Finding Solution to the SDE having initial condition xi. Y = Phi(S(X , T)) - sdeproblem = SDEProblem(μ,sigma,xi,tspan,noise_rate_prototype = noise_rate_prototype) - function prob_func(prob,i,repeat) - SDEProblem(prob.f , prob.g , xi[: , i] , prob.tspan ,noise_rate_prototype = prob.noise_rate_prototype) + sdeproblem = SDEProblem(μ, sigma, xi, tspan, + noise_rate_prototype = noise_rate_prototype) + function prob_func(prob, i, repeat) + SDEProblem(prob.f, prob.g, xi[:, i], prob.tspan, + noise_rate_prototype = prob.noise_rate_prototype) end - output_func(sol,i) = (sol[end],false) - ensembleprob = EnsembleProblem(sdeproblem , prob_func = prob_func , output_func = output_func) - sim = solve(ensembleprob, sdealg, ensemblealg , dt=dt, trajectories=trajectories,adaptive=false) - x_sde = reshape([],d,0) + output_func(sol, i) = (sol[end], false) + ensembleprob = EnsembleProblem(sdeproblem, prob_func = prob_func, + output_func = output_func) + sim = solve(ensembleprob, sdealg, ensemblealg, dt = dt, trajectories = trajectories, + adaptive = false) + x_sde = reshape([], d, 0) # sol = solve(sdeproblem, sdealg ,dt=0.01 , save_everystep=false , kwargs...) # x_sde = sol[end] for u in sim.u - x_sde = hcat(x_sde , u) + x_sde = hcat(x_sde, u) end y = phi(x_sde) if use_gpu == true - y = y |>gpu + y = y |> gpu xi = xi |> gpu end - data = Iterators.repeated((xi , y), maxiters) + data = Iterators.repeated((xi, y), maxiters) if use_gpu == true - data = data |>gpu + data = data |> gpu end #MSE Loss Function - loss(x , y) =Flux.mse(chain(x), y) + loss(x, y) = Flux.mse(chain(x), y) callback = function () l = loss(xi, y) @@ -99,5 +103,5 @@ function DiffEqBase.solve( Flux.train!(loss, ps, data, opt; callback = callback) chainout = chain(xi) - xi , chainout - end #solve + xi, chainout +end #solve diff --git a/src/neural_adapter.jl b/src/neural_adapter.jl index cb6793e001..61827ee19e 100644 --- a/src/neural_adapter.jl +++ b/src/neural_adapter.jl @@ -1,101 +1,104 @@ -function generate_training_sets(domains,dx,eqs,eltypeθ) +function generate_training_sets(domains, dx, eqs, eltypeθ) if dx isa Array dxs = dx else - dxs = fill(dx,length(domains)) + dxs = fill(dx, length(domains)) end - spans = [infimum(d.domain):dx:supremum(d.domain) for (d,dx) in zip(domains,dxs)] - train_set = adapt(eltypeθ, hcat(vec(map(points -> collect(points), Iterators.product(spans...)))...)) + spans = [infimum(d.domain):dx:supremum(d.domain) for (d, dx) in zip(domains, dxs)] + train_set = adapt(eltypeθ, + hcat(vec(map(points -> collect(points), Iterators.product(spans...)))...)) end -function get_loss_function_(loss,initθ,pde_system,strategy::GridTraining) +function get_loss_function_(loss, initθ, pde_system, strategy::GridTraining) eqs = pde_system.eqs if !(eqs isa Array) eqs = [eqs] end domains = pde_system.domain - depvars,indvars,dict_indvars,dict_depvars = get_vars(pde_system.indvars, pde_system.depvars) + depvars, indvars, dict_indvars, dict_depvars = get_vars(pde_system.indvars, + pde_system.depvars) eltypeθ = eltype(initθ) dx = strategy.dx - train_set = generate_training_sets(domains,dx,eqs,eltypeθ) - get_loss_function(loss,train_set,eltypeθ,strategy) + train_set = generate_training_sets(domains, dx, eqs, eltypeθ) + get_loss_function(loss, train_set, eltypeθ, strategy) end +function get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, strategy) + dict_span = Dict([Symbol(d.variables) => [infimum(d.domain), supremum(d.domain)] + for d in domains]) + args = get_argument(eqs, dict_indvars, dict_depvars) -function get_bounds_(domains,eqs,eltypeθ,dict_indvars,dict_depvars,strategy) - dict_span = Dict([Symbol(d.variables) => [infimum(d.domain), supremum(d.domain)] for d in domains]) - args = get_argument(eqs,dict_indvars,dict_depvars) - - bounds= map(args) do pd + bounds = map(args) do pd span = map(p -> get(dict_span, p, p), pd) - map(s -> adapt(eltypeθ,s), span) + map(s -> adapt(eltypeθ, s), span) end bounds end -function get_loss_function_(loss,initθ,pde_system,strategy::StochasticTraining) +function get_loss_function_(loss, initθ, pde_system, strategy::StochasticTraining) eqs = pde_system.eqs if !(eqs isa Array) eqs = [eqs] end domains = pde_system.domain - depvars,indvars,dict_indvars,dict_depvars = get_vars(pde_system.indvars, pde_system.depvars) + depvars, indvars, dict_indvars, dict_depvars = get_vars(pde_system.indvars, + pde_system.depvars) eltypeθ = eltype(initθ) - bound = get_bounds_(domains,eqs,eltypeθ,dict_indvars,dict_depvars,strategy)[1] + bound = get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, strategy)[1] - get_loss_function(loss,bound,eltypeθ,strategy) + get_loss_function(loss, bound, eltypeθ, strategy) end - -function get_loss_function_(loss,initθ,pde_system,strategy::QuasiRandomTraining) +function get_loss_function_(loss, initθ, pde_system, strategy::QuasiRandomTraining) eqs = pde_system.eqs if !(eqs isa Array) eqs = [eqs] end domains = pde_system.domain - depvars,indvars,dict_indvars,dict_depvars = get_vars(pde_system.indvars, pde_system.depvars) + depvars, indvars, dict_indvars, dict_depvars = get_vars(pde_system.indvars, + pde_system.depvars) eltypeθ = eltype(initθ) - bound = get_bounds_(domains,eqs,eltypeθ,dict_indvars,dict_depvars,strategy)[1] + bound = get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, strategy)[1] - get_loss_function(loss,bound,eltypeθ,strategy) + get_loss_function(loss, bound, eltypeθ, strategy) end -function get_bounds_(domains,eqs,eltypeθ,dict_indvars,dict_depvars,strategy::QuadratureTraining) +function get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, + strategy::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) + args = get_argument(eqs, dict_indvars, dict_depvars) - lower_bounds= map(args) do pd + lower_bounds = map(args) do pd span = map(p -> get(dict_lower_bound, p, p), pd) - map(s -> adapt(eltypeθ,s), span) + map(s -> adapt(eltypeθ, s), span) end - upper_bounds= map(args) do pd + upper_bounds = map(args) do pd span = map(p -> get(dict_upper_bound, p, p), pd) - map(s -> adapt(eltypeθ,s), span) + map(s -> adapt(eltypeθ, s), span) end - bound = lower_bounds,upper_bounds + bound = lower_bounds, upper_bounds end - - -function get_loss_function_(loss,initθ,pde_system,strategy::QuadratureTraining) +function get_loss_function_(loss, initθ, pde_system, strategy::QuadratureTraining) eqs = pde_system.eqs if !(eqs isa Array) eqs = [eqs] end domains = pde_system.domain - depvars,indvars,dict_indvars,dict_depvars = get_vars(pde_system.indvars, pde_system.depvars) + depvars, indvars, dict_indvars, dict_depvars = get_vars(pde_system.indvars, + pde_system.depvars) eltypeθ = eltype(initθ) - bound = get_bounds_(domains,eqs,eltypeθ,dict_indvars,dict_depvars,strategy) - lb,ub = bound - get_loss_function(loss,lb[1],ub[1],eltypeθ,strategy) + bound = get_bounds_(domains, eqs, eltypeθ, dict_indvars, dict_depvars, strategy) + lb, ub = bound + get_loss_function(loss, lb[1], ub[1], eltypeθ, strategy) end """ @@ -108,10 +111,10 @@ Arguments: * `strategy`: determines which training strategy will be used. """ -function neural_adapter(loss,initθ,pde_system,strategy) - loss_function__ = get_loss_function_(loss,initθ,pde_system,strategy) +function neural_adapter(loss, initθ, pde_system, strategy) + loss_function__ = get_loss_function_(loss, initθ, pde_system, strategy) - function loss_function_(θ,p) + function loss_function_(θ, p) loss_function__(θ) end f_ = OptimizationFunction(loss_function_, Optimization.AutoZygote()) @@ -128,12 +131,12 @@ Arguments: * `strategy`: determines which training strategy will be used. """ -function neural_adapter(losses::Array,initθ,pde_systems::Array,strategy) - loss_functions_ = map(zip(losses,pde_systems)) do (l,p) - get_loss_function_(l,initθ,p,strategy) +function neural_adapter(losses::Array, initθ, pde_systems::Array, strategy) + loss_functions_ = map(zip(losses, pde_systems)) do (l, p) + get_loss_function_(l, initθ, p, strategy) end - loss_function__ = θ -> sum(map(l->l(θ) ,loss_functions_)) - function loss_function_(θ,p) + loss_function__ = θ -> sum(map(l -> l(θ), loss_functions_)) + function loss_function_(θ, p) loss_function__(θ) end diff --git a/src/ode_solve.jl b/src/ode_solve.jl index 04b8dabfe6..9e94803804 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -1,285 +1,292 @@ -""" -```julia -NNODE(chain, opt=OptimizationPolyalgorithms.PolyOpt(), init_params = nothing; - autodiff=false, batch=0, kwargs...) -``` - -Algorithm for solving ordinary differential equations using a neural network. This is a specialization -of the physics-informed neural network which is used as a solver for a standard `ODEProblem`. - -!!! warn - - Note that NNODE only supports ODEs which are written in the out-of-place form, i.e. - `du = f(u,p,t)`, and not `f(du,u,p,t)`. If not declared out-of-place then the NNODE - will exit with an error. - -## Positional Arguments - -* `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.Chain`. -* `opt`: The optimizer to train the neural network. Defaults to `OptimizationPolyalgorithms.PolyOpt()` -* `initθ`: The initial parameter of the neural network. By default this is `nothing` - which thus uses the random initialization provided by the neural network library. - -## Keyword Arguments - -* `autodiff`: The switch between automatic and numerical differentiation for - the PDE operators. The reverse mode of the loss function is always - automatic differentation (via Zygote), this is only for the derivative - in the loss function (the derivative with respect to time). -* `batch`: The batch size to use for the internal quadrature. Defaults to `0`, which - means the application of the neural network is done at individual time points one - at a time. `batch>0` means the neural network is applied at a row vector of values - `t` simultaniously, i.e. it's the batch size for the neural network evaluations. - This requires a neural network compatible with batched data. -* `strategy`: The training strategy used to choose the points for the evaluations. - Default of `nothing` means that `QuadratureTraining` with QuadGK is used if no - `dt` is given, and `GridTraining` is used with `dt` if given. -* `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. - -## Example - -```julia -f(u,p,t) = cos(2pi*t) -tspan = (0.0f0, 1.0f0) -u0 = 0.0f0 -prob = ODEProblem(linear, u0 ,tspan) -chain = Flux.Chain(Dense(1,5,σ),Dense(5,1)) -opt = Flux.ADAM(0.1) -sol = solve(prob, NeuralPDE.NNODE(chain,opt), dt=1/20f0, verbose = true, - abstol=1e-10, maxiters = 200) -``` - -## Solution Notes - -Note that the solution is evaluated at fixed time points according to standard output handlers -such as `saveat` and `dt`. However, the neural network is a fully continuous solution so `sol(t)` -is an accuate interpolation (up to the neural network training result). In addition, the -`OptimizationSolution` is returned as `sol.k` for further analysis. - -## References - -Lagaris, Isaac E., Aristidis Likas, and Dimitrios I. Fotiadis. "Artificial neural networks for solving -ordinary and partial differential equations." IEEE Transactions on Neural Networks 9, no. 5 (1998): 987-1000. -""" -struct NNODE{C,O,P,K,S <: Union{Nothing,AbstractTrainingStrategy}} <: NeuralPDEAlgorithm - chain::C - opt::O - initθ::P - autodiff::Bool - batch::Int - strategy::S - kwargs::K -end -function NNODE(chain, opt, init_params = nothing; - strategy = nothing, - autodiff=false, batch = 0, kwargs...) - NNODE(chain,opt,init_params,autodiff,batch,strategy,kwargs) -end - -""" -```julia -ODEPhi(chain::FastChain, t, u0) -``` - -Internal, used as a constructor used for representing the ODE solution as a -neural network in a form that respects boundary conditions, i.e. -`phi(t) = u0 + t*NN(t)`. -""" -struct ODEPhi{C, T, U} - chain::C - t0::T - u0::U - - ODEPhi(chain::FastChain, t::Number, u0) = new{typeof(chain), typeof(t), typeof(u0)}(chain, t, u0) - function ODEPhi(chain::Flux.Chain, t, u0) - p,re = Flux.destructure(chain) - new{typeof(re), typeof(t), typeof(u0)}(re, t, u0) - end -end - -function (f::ODEPhi{C, T, U})(t,θ) where {C<:FastChain, T, U<:Number} - f.u0 + (t-f.t0) * first(f.chain(adapt(parameterless_type(θ),[t]),θ)) -end - -function (f::ODEPhi{C, T, U})(t,θ) where {C<:FastChain, T, U} - f.u0 + (t-f.t0) * f.chain(adapt(parameterless_type(θ),[t]),θ) -end - -function (f::ODEPhi{C, T, U})(t,θ) where {C<:Optimisers.Restructure, T, U<:Number} - f.u0 + (t-f.t0)*first(f.chain(θ)(adapt(parameterless_type(θ),[t]))) -end - -function (f::ODEPhi{C, T, U})(t,θ) where {C<:Optimisers.Restructure, T, U} - f.u0 + (t-f.t0) * f.chain(θ)(adapt(parameterless_type(θ),[t])) -end - -""" -Computes u' using either forward-mode automatic differentiation or -numerical differentiation. -""" -function ode_dfdx(phi::ODEPhi,t::Number,θ,autodiff::Bool) - if autodiff - ForwardDiff.derivative(t->phi(t,θ),t) - else - (phi(t+sqrt(eps(typeof(t))),θ) - phi(t,θ))/sqrt(eps(typeof(t))) - end -end - -""" -Simple L2 inner loss at a time `t` with parameters θ -""" -function inner_loss(phi,f,autodiff::Bool,t,θ,p) - sum(abs2,ode_dfdx(phi,t,θ,autodiff) - f(phi(t,θ),p,t)) -end - -""" -Representation of the loss function, paramtric on the training strategy `strategy` -""" -function generate_loss(strategy::QuadratureTraining,phi,f,autodiff::Bool,tspan,p) - integrand(t::Number,θ) = abs2(inner_loss(phi,f,autodiff,t,θ,p)) - integrand(ts,θ) = [abs2(inner_loss(phi,f,autodiff,t,θ,p)) for t in ts] - - function loss(θ,p) - intprob = IntegralProblem(integrand, tspan[1], tspan[2], θ) - sol = solve(intprob,QuadGKJL();abstol=strategy.abstol,reltol=strategy.reltol) - sol.u - end - - # Default this to ForwardDiff until Integrals.jl autodiff is sorted out - OptimizationFunction(loss,Optimization.AutoForwardDiff()) -end - -function generate_loss(strategy::GridTraining,phi,f,autodiff::Bool,tspan,p) - ts = tspan[1]:strategy.dx:tspan[2] - - # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken - function loss(θ,p) - sum(abs2,[inner_loss(phi,f,autodiff,t,θ,p) for t in ts]) - end - optf = OptimizationFunction(loss,Optimization.AutoZygote()) -end - -function generate_loss(strategy::StochasticTraining,phi,f,autodiff::Bool,tspan,p) - # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken - function loss(θ,p) - # (tspan[2]-tspan[1])*rand() + tspan[1] gives Uniform(tspan[1],tspan[2]) - sum(abs2,[inner_loss(phi,f,autodiff,(tspan[2]-tspan[1])*rand() + tspan[1],θ,p) for i in 1:strategy.points]) - end - optf = OptimizationFunction(loss,Optimization.AutoZygote()) -end - -function generate_loss(strategy::QuasiRandomTraining,phi,f,autodiff::Bool,tspan) - error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.") -end - -struct NNODEInterpolation{T<:ODEPhi,T2} - phi::T - θ::T2 -end -(f::NNODEInterpolation)(t, idxs::Nothing, ::Type{Val{0}}, p, continuity) = f.phi(t,f.θ) -(f::NNODEInterpolation)(t, idxs, ::Type{Val{0}}, p, continuity) = f.phi(t,f.θ)[idxs] - -SciMLBase.interp_summary(::NNODEInterpolation) = "Trained neural network interpolation" - -function DiffEqBase.__solve( - prob::DiffEqBase.AbstractODEProblem, - alg::NNODE, - args...; - dt = nothing, - timeseries_errors = true, - save_everystep=true, - adaptive=false, - abstol = 1f-6, - reltol = 1f-3, - verbose = false, - saveat = nothing, - maxiters = 100) - - u0 = prob.u0 - tspan = prob.tspan - f = prob.f - p = prob.p - t0 = tspan[1] - - #hidden layer - chain = alg.chain - opt = alg.opt - autodiff = alg.autodiff - - #train points generation - initθ = alg.initθ - - if initθ === nothing - if chain isa FastChain - initθ = DiffEqFlux.initial_params(chain) - else - initθ,re = Flux.destructure(chain) - end - else - initθ = initθ - end - - if isinplace(prob) - throw(error("The NNODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) - end - - phi = ODEPhi(chain,t0,u0) - - try - phi(t0 , initθ) - catch err - if isa(err , DimensionMismatch) - throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) - else - throw(err) - end - end - - strategy = if alg.strategy === nothing - if dt !== nothing - GridTraining(dt) - else - QuadratureTraining(;quadrature_alg=QuadGKJL(), reltol = convert(eltype(u0),reltol), - abstol = convert(eltype(u0),abstol), maxiters = maxiters, batch=alg.batch) - end - else - alg.strategy - end - - optf = generate_loss(strategy,phi,f,autodiff::Bool,tspan,p) - - callback = function (p,l) - verbose && println("Current loss is: $l") - l < abstol - end - - optprob = OptimizationProblem(optf, initθ) - res = solve(optprob, opt; callback, maxiters, alg.kwargs...) - - #solutions at timepoints - if saveat isa Number - ts = tspan[1]:saveat:tspan[2] - elseif saveat isa AbstractArray - ts = saveat - elseif dt !== nothing - ts = tspan[1]:dt:tspan[2] - elseif save_everystep - ts = range(tspan[1],tspan[2],length = 100) - else - ts = [tspan[1],tspan[2]] - end - - if u0 isa Number - u = [first(phi(t,res.u)) for t in ts] - else - u = [phi(t,res.u) for t in ts] - end - - - sol = DiffEqBase.build_solution(prob,alg,ts,u; - k = res, dense = true, - interp = NNODEInterpolation(phi,res.u), calculate_error = false, - retcode = :Success) - DiffEqBase.has_analytic(prob.f) && DiffEqBase.calculate_solution_errors!(sol;timeseries_errors=true,dense_errors=false) - sol -end #solve +""" +```julia +NNODE(chain, opt=OptimizationPolyalgorithms.PolyOpt(), init_params = nothing; + autodiff=false, batch=0, kwargs...) +``` + +Algorithm for solving ordinary differential equations using a neural network. This is a specialization +of the physics-informed neural network which is used as a solver for a standard `ODEProblem`. + +!!! warn + + Note that NNODE only supports ODEs which are written in the out-of-place form, i.e. + `du = f(u,p,t)`, and not `f(du,u,p,t)`. If not declared out-of-place then the NNODE + will exit with an error. + +## Positional Arguments + +* `chain`: A neural network architecture, defined as either a `Flux.Chain` or a `Lux.Chain`. +* `opt`: The optimizer to train the neural network. Defaults to `OptimizationPolyalgorithms.PolyOpt()` +* `initθ`: The initial parameter of the neural network. By default this is `nothing` + which thus uses the random initialization provided by the neural network library. + +## Keyword Arguments + +* `autodiff`: The switch between automatic and numerical differentiation for + the PDE operators. The reverse mode of the loss function is always + automatic differentation (via Zygote), this is only for the derivative + in the loss function (the derivative with respect to time). +* `batch`: The batch size to use for the internal quadrature. Defaults to `0`, which + means the application of the neural network is done at individual time points one + at a time. `batch>0` means the neural network is applied at a row vector of values + `t` simultaniously, i.e. it's the batch size for the neural network evaluations. + This requires a neural network compatible with batched data. +* `strategy`: The training strategy used to choose the points for the evaluations. + Default of `nothing` means that `QuadratureTraining` with QuadGK is used if no + `dt` is given, and `GridTraining` is used with `dt` if given. +* `kwargs`: Extra keyword arguments are splatted to the Optimization.jl `solve` call. + +## Example + +```julia +f(u,p,t) = cos(2pi*t) +tspan = (0.0f0, 1.0f0) +u0 = 0.0f0 +prob = ODEProblem(linear, u0 ,tspan) +chain = Flux.Chain(Dense(1,5,σ),Dense(5,1)) +opt = Flux.ADAM(0.1) +sol = solve(prob, NeuralPDE.NNODE(chain,opt), dt=1/20f0, verbose = true, + abstol=1e-10, maxiters = 200) +``` + +## Solution Notes + +Note that the solution is evaluated at fixed time points according to standard output handlers +such as `saveat` and `dt`. However, the neural network is a fully continuous solution so `sol(t)` +is an accuate interpolation (up to the neural network training result). In addition, the +`OptimizationSolution` is returned as `sol.k` for further analysis. + +## References + +Lagaris, Isaac E., Aristidis Likas, and Dimitrios I. Fotiadis. "Artificial neural networks for solving +ordinary and partial differential equations." IEEE Transactions on Neural Networks 9, no. 5 (1998): 987-1000. +""" +struct NNODE{C, O, P, K, S <: Union{Nothing, AbstractTrainingStrategy}} <: + NeuralPDEAlgorithm + chain::C + opt::O + initθ::P + autodiff::Bool + batch::Int + strategy::S + kwargs::K +end +function NNODE(chain, opt, init_params = nothing; + strategy = nothing, + autodiff = false, batch = 0, kwargs...) + NNODE(chain, opt, init_params, autodiff, batch, strategy, kwargs) +end + +""" +```julia +ODEPhi(chain::FastChain, t, u0) +``` + +Internal, used as a constructor used for representing the ODE solution as a +neural network in a form that respects boundary conditions, i.e. +`phi(t) = u0 + t*NN(t)`. +""" +struct ODEPhi{C, T, U} + chain::C + t0::T + u0::U + + function ODEPhi(chain::FastChain, t::Number, u0) + new{typeof(chain), typeof(t), typeof(u0)}(chain, t, u0) + end + function ODEPhi(chain::Flux.Chain, t, u0) + p, re = Flux.destructure(chain) + new{typeof(re), typeof(t), typeof(u0)}(re, t, u0) + end +end + +function (f::ODEPhi{C, T, U})(t, θ) where {C <: FastChain, T, U <: Number} + f.u0 + (t - f.t0) * first(f.chain(adapt(parameterless_type(θ), [t]), θ)) +end + +function (f::ODEPhi{C, T, U})(t, θ) where {C <: FastChain, T, U} + f.u0 + (t - f.t0) * f.chain(adapt(parameterless_type(θ), [t]), θ) +end + +function (f::ODEPhi{C, T, U})(t, θ) where {C <: Optimisers.Restructure, T, U <: Number} + f.u0 + (t - f.t0) * first(f.chain(θ)(adapt(parameterless_type(θ), [t]))) +end + +function (f::ODEPhi{C, T, U})(t, θ) where {C <: Optimisers.Restructure, T, U} + f.u0 + (t - f.t0) * f.chain(θ)(adapt(parameterless_type(θ), [t])) +end + +""" +Computes u' using either forward-mode automatic differentiation or +numerical differentiation. +""" +function ode_dfdx(phi::ODEPhi, t::Number, θ, autodiff::Bool) + if autodiff + ForwardDiff.derivative(t -> phi(t, θ), t) + else + (phi(t + sqrt(eps(typeof(t))), θ) - phi(t, θ)) / sqrt(eps(typeof(t))) + end +end + +""" +Simple L2 inner loss at a time `t` with parameters θ +""" +function inner_loss(phi, f, autodiff::Bool, t, θ, p) + sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(phi(t, θ), p, t)) +end + +""" +Representation of the loss function, paramtric on the training strategy `strategy` +""" +function generate_loss(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p) + integrand(t::Number, θ) = abs2(inner_loss(phi, f, autodiff, t, θ, p)) + integrand(ts, θ) = [abs2(inner_loss(phi, f, autodiff, t, θ, p)) for t in ts] + + function loss(θ, p) + intprob = IntegralProblem(integrand, tspan[1], tspan[2], θ) + sol = solve(intprob, QuadGKJL(); abstol = strategy.abstol, reltol = strategy.reltol) + sol.u + end + + # Default this to ForwardDiff until Integrals.jl autodiff is sorted out + OptimizationFunction(loss, Optimization.AutoForwardDiff()) +end + +function generate_loss(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p) + ts = tspan[1]:(strategy.dx):tspan[2] + + # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken + function loss(θ, p) + sum(abs2, [inner_loss(phi, f, autodiff, t, θ, p) for t in ts]) + end + optf = OptimizationFunction(loss, Optimization.AutoZygote()) +end + +function generate_loss(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p) + # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken + function loss(θ, p) + # (tspan[2]-tspan[1])*rand() + tspan[1] gives Uniform(tspan[1],tspan[2]) + sum(abs2, + [inner_loss(phi, f, autodiff, (tspan[2] - tspan[1]) * rand() + tspan[1], θ, p) + for i in 1:(strategy.points)]) + end + optf = OptimizationFunction(loss, Optimization.AutoZygote()) +end + +function generate_loss(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan) + error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.") +end + +struct NNODEInterpolation{T <: ODEPhi, T2} + phi::T + θ::T2 +end +(f::NNODEInterpolation)(t, idxs::Nothing, ::Type{Val{0}}, p, continuity) = f.phi(t, f.θ) +(f::NNODEInterpolation)(t, idxs, ::Type{Val{0}}, p, continuity) = f.phi(t, f.θ)[idxs] + +SciMLBase.interp_summary(::NNODEInterpolation) = "Trained neural network interpolation" + +function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, + alg::NNODE, + args...; + dt = nothing, + timeseries_errors = true, + save_everystep = true, + adaptive = false, + abstol = 1.0f-6, + reltol = 1.0f-3, + verbose = false, + saveat = nothing, + maxiters = 100) + u0 = prob.u0 + tspan = prob.tspan + f = prob.f + p = prob.p + t0 = tspan[1] + + #hidden layer + chain = alg.chain + opt = alg.opt + autodiff = alg.autodiff + + #train points generation + initθ = alg.initθ + + if initθ === nothing + if chain isa FastChain + initθ = DiffEqFlux.initial_params(chain) + else + initθ, re = Flux.destructure(chain) + end + else + initθ = initθ + end + + if isinplace(prob) + throw(error("The NNODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) + end + + phi = ODEPhi(chain, t0, u0) + + try + phi(t0, initθ) + catch err + if isa(err, DimensionMismatch) + throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) + else + throw(err) + end + end + + strategy = if alg.strategy === nothing + if dt !== nothing + GridTraining(dt) + else + QuadratureTraining(; quadrature_alg = QuadGKJL(), + reltol = convert(eltype(u0), reltol), + abstol = convert(eltype(u0), abstol), maxiters = maxiters, + batch = alg.batch) + end + else + alg.strategy + end + + optf = generate_loss(strategy, phi, f, autodiff::Bool, tspan, p) + + callback = function (p, l) + verbose && println("Current loss is: $l") + l < abstol + end + + optprob = OptimizationProblem(optf, initθ) + res = solve(optprob, opt; callback, maxiters, alg.kwargs...) + + #solutions at timepoints + if saveat isa Number + ts = tspan[1]:saveat:tspan[2] + elseif saveat isa AbstractArray + ts = saveat + elseif dt !== nothing + ts = tspan[1]:dt:tspan[2] + elseif save_everystep + ts = range(tspan[1], tspan[2], length = 100) + else + ts = [tspan[1], tspan[2]] + end + + if u0 isa Number + u = [first(phi(t, res.u)) for t in ts] + else + u = [phi(t, res.u) for t in ts] + end + + sol = DiffEqBase.build_solution(prob, alg, ts, u; + k = res, dense = true, + interp = NNODEInterpolation(phi, res.u), + calculate_error = false, + retcode = :Success) + DiffEqBase.has_analytic(prob.f) && + DiffEqBase.calculate_solution_errors!(sol; timeseries_errors = true, + dense_errors = false) + sol +end #solve diff --git a/src/param_kolmogorov_solve.jl b/src/param_kolmogorov_solve.jl index 4dc813b4ec..415e5b6c44 100644 --- a/src/param_kolmogorov_solve.jl +++ b/src/param_kolmogorov_solve.jl @@ -1,26 +1,26 @@ -struct NNParamKolmogorov{C,O,S,E} <: NeuralPDEAlgorithm +struct NNParamKolmogorov{C, O, S, E} <: NeuralPDEAlgorithm chain::C opt::O sdealg::S ensemblealg::E end -NNParamKolmogorov(chain ; opt=Flux.ADAM(0.1) , sdealg = EM() , ensemblealg = EnsembleThreads()) = NNParamKolmogorov(chain , opt , sdealg , ensemblealg) - -function DiffEqBase.solve( - prob::ParamKolmogorovPDEProblem, - alg::NNParamKolmogorov; - abstol = 1f-6, - verbose = false, - maxiters = 300, - trajectories = 1000, - save_everystep = false, - use_gpu = false, - dγ = 0.01, - dt, - dx, - kwargs... - ) +function NNParamKolmogorov(chain; opt = Flux.ADAM(0.1), sdealg = EM(), + ensemblealg = EnsembleThreads()) + NNParamKolmogorov(chain, opt, sdealg, ensemblealg) +end +function DiffEqBase.solve(prob::ParamKolmogorovPDEProblem, + alg::NNParamKolmogorov; + abstol = 1.0f-6, + verbose = false, + maxiters = 300, + trajectories = 1000, + save_everystep = false, + use_gpu = false, + dγ = 0.01, + dt, + dx, + kwargs...) tspan = prob.tspan g = prob.g f = prob.f @@ -39,127 +39,135 @@ function DiffEqBase.solve( γ_phi_prototype = prob.Y_phi_prototype #hidden layer - chain = alg.chain - opt = alg.opt + chain = alg.chain + opt = alg.opt sdealg = alg.sdealg ensemblealg = alg.ensemblealg - ps = Flux.params(chain) - t = rand(ts , 1 , trajectories) - x = rand(xs , d , 1 , trajectories) + ps = Flux.params(chain) + t = rand(ts, 1, trajectories) + x = rand(xs, d, 1, trajectories) #Sample all the parameters from uniform distributions total_dims = d + 1 if isnothing(γ_sigma_prototype) - γ_sigma = nothing + γ_sigma = nothing else - γ_sigma = rand(γs_sigma , size(γ_sigma_prototype)[1] , size(γ_sigma_prototype)[2] , size(γ_sigma_prototype)[3] , trajectories) - total_dims += d^2*(size(γ_sigma_prototype))[3] + γ_sigma = rand(γs_sigma, size(γ_sigma_prototype)[1], size(γ_sigma_prototype)[2], + size(γ_sigma_prototype)[3], trajectories) + total_dims += d^2 * (size(γ_sigma_prototype))[3] end if isnothing(γ_mu_prototype) - γ_mu_1 = nothing - γ_mu_2 = nothing + γ_mu_1 = nothing + γ_mu_2 = nothing else - if !isnothing(γ_mu_prototype[1]) - γ_mu_1 = rand(γs_mu , size(γ_mu_prototype[1])[1], size(γ_mu_prototype[1])[2], 1 , trajectories) - total_dims += d^2 - elseif !isnothing(γ_mu_prototype[1]) - γ_mu_2 = rand(γs_mu , size(γ_mu_prototype[2])[1], size(γ_mu_prototype[2])[2] , 1 ,trajectories) - total_dims += d - end + if !isnothing(γ_mu_prototype[1]) + γ_mu_1 = rand(γs_mu, size(γ_mu_prototype[1])[1], size(γ_mu_prototype[1])[2], 1, + trajectories) + total_dims += d^2 + elseif !isnothing(γ_mu_prototype[1]) + γ_mu_2 = rand(γs_mu, size(γ_mu_prototype[2])[1], size(γ_mu_prototype[2])[2], 1, + trajectories) + total_dims += d + end end if isnothing(γ_phi_prototype) - γ_phi = nothing + γ_phi = nothing else - γ_phi = rand(γs_phi , size(γ_phi_prototype)[1] , 1 , trajectories) - total_dims += size(γ_phi_prototype)[1] + γ_phi = rand(γs_phi, size(γ_phi_prototype)[1], 1, trajectories) + total_dims += size(γ_phi_prototype)[1] end #Declare wrapper functions to mu and sigma - function sigma_(dx , x , γ_sigma , i) - if isnothing(γ_sigma) - return g(x , γ_sigma) - else - return g(x , γ_sigma[: , : , : , i]) - end + function sigma_(dx, x, γ_sigma, i) + if isnothing(γ_sigma) + return g(x, γ_sigma) + else + return g(x, γ_sigma[:, :, :, i]) + end end - function mu_(dx, x , γ_mu_1, γ_mu_2 , i) - if isnothing(γ_mu_prototype) - return f(x , γ_mu_1 , γ_mu_2) - else - if isnothing(γ_mu_1) && isnothing(γ_mu_2) - return f(x , γ_mu_1 , γ_mu_2) - elseif isnothing(γ_mu_1) && !isnothing(γ_mu_2) - return f(x , γ_mu_1 , γ_mu_2[: , : , : , i]) - elseif !isnothing(γ_mu_1) && isnothing(γ_mu_2) - return f(x , γ_mu_1[: , : , : , i] , γ_mu_2) + function mu_(dx, x, γ_mu_1, γ_mu_2, i) + if isnothing(γ_mu_prototype) + return f(x, γ_mu_1, γ_mu_2) else - return f(x , γ_mu_1[: , : , : , i] , γ_mu_2[: , : , : , i]) + if isnothing(γ_mu_1) && isnothing(γ_mu_2) + return f(x, γ_mu_1, γ_mu_2) + elseif isnothing(γ_mu_1) && !isnothing(γ_mu_2) + return f(x, γ_mu_1, γ_mu_2[:, :, :, i]) + elseif !isnothing(γ_mu_1) && isnothing(γ_mu_2) + return f(x, γ_mu_1[:, :, :, i], γ_mu_2) + else + return f(x, γ_mu_1[:, :, :, i], γ_mu_2[:, :, :, i]) + end end - end end #Preparing training data - X = zeros(total_dims ,trajectories) + X = zeros(total_dims, trajectories) for i in 1:length(t) - K = vcat(t[i] , x[: , : , i]) - if !isnothing(γ_sigma_prototype) - K = vcat(K , reshape(γ_sigma[: , : , : , i] , d^2*(size(γ_sigma_prototype))[3] , 1)) - end - if !isnothing(γ_mu_prototype) - K = vcat(K , reshape(γ_mu_1[: , : , : , i] , d^2 , 1) , reshape(γ_mu_2[: , : , : , i] , d , 1)) - end - if !isnothing(γ_phi_prototype) - K = vcat(K , reshape(γ_phi[: , : , i] , size(γ_phi_prototype)[1] , 1)) - end - X[: , i] = K + K = vcat(t[i], x[:, :, i]) + if !isnothing(γ_sigma_prototype) + K = vcat(K, reshape(γ_sigma[:, :, :, i], d^2 * (size(γ_sigma_prototype))[3], 1)) + end + if !isnothing(γ_mu_prototype) + K = vcat(K, reshape(γ_mu_1[:, :, :, i], d^2, 1), + reshape(γ_mu_2[:, :, :, i], d, 1)) + end + if !isnothing(γ_phi_prototype) + K = vcat(K, reshape(γ_phi[:, :, i], size(γ_phi_prototype)[1], 1)) + end + X[:, i] = K end print("Data Prepared to train the model") - function sigma(dx , x , p , t) - sigma_(dx , x , γ_sigma , 1) + function sigma(dx, x, p, t) + sigma_(dx, x, γ_sigma, 1) end - function mu(dx, x , p , t) - mu_(dx, x , γ_mu_1, γ_mu_2 , 1) + function mu(dx, x, p, t) + mu_(dx, x, γ_mu_1, γ_mu_2, 1) end - sdeproblem = SDEProblem(mu,sigma,x[: , : , 1], (0.00 , t[1])) - function prob_func(prob,i,repeat) - sigma(dx , x , p , t) = sigma_(dx , x , γ_sigma , i) - mu(dx , x , p , t) = mu_(dx , x , γ_mu_1, γ_mu_2 , i) - SDEProblem(prob.f , prob.g , x[: , : , i] , (0.00 , t[i]) ,noise_rate_prototype = prob.noise_rate_prototype) + sdeproblem = SDEProblem(mu, sigma, x[:, :, 1], (0.00, t[1])) + function prob_func(prob, i, repeat) + sigma(dx, x, p, t) = sigma_(dx, x, γ_sigma, i) + mu(dx, x, p, t) = mu_(dx, x, γ_mu_1, γ_mu_2, i) + SDEProblem(prob.f, prob.g, x[:, :, i], (0.00, t[i]), + noise_rate_prototype = prob.noise_rate_prototype) end - output_func(sol,i) = (sol[end],false) - ensembleprob = EnsembleProblem(sdeproblem , prob_func = prob_func , output_func = output_func) - sim = solve(ensembleprob, sdealg, ensemblealg , dt=dt, trajectories=trajectories,adaptive=false) - x_sde = reshape([],d,0) + output_func(sol, i) = (sol[end], false) + ensembleprob = EnsembleProblem(sdeproblem, prob_func = prob_func, + output_func = output_func) + sim = solve(ensembleprob, sdealg, ensemblealg, dt = dt, trajectories = trajectories, + adaptive = false) + x_sde = reshape([], d, 0) # sol = solve(sdeproblem, sdealg ,dt=0.01 , save_everystep=false , kwargs...) # x_sde = sol[end] - x_sde = zeros(d,trajectories) + x_sde = zeros(d, trajectories) for i in 1:length(sim.u) - x_sde[: , i] = sim.u[i] + x_sde[:, i] = sim.u[i] end - y = phi(x_sde , γ_phi) + y = phi(x_sde, γ_phi) #Check the dimensions for the target if size(y)[1] != 1 && size(y)[1] != trajectories - errMsg = "The output from phi function has dimensions "*string(size(y))*" while it should be (1, trajectories)" - throw(DimensionMismatch(errMsg)) + errMsg = "The output from phi function has dimensions " * string(size(y)) * + " while it should be (1, trajectories)" + throw(DimensionMismatch(errMsg)) end if use_gpu == true y = y |> gpu X = X |> gpu end - data = Iterators.repeated((X , y), maxiters) + data = Iterators.repeated((X, y), maxiters) if use_gpu == true data = data |> gpu end #MSE Loss Function - loss(x , y) =Flux.mse(chain(x) , y) + loss(x, y) = Flux.mse(chain(x), y) callback = function () l = loss(X, y) @@ -168,5 +176,4 @@ function DiffEqBase.solve( end Flux.train!(loss, ps, data, opt; cb = callback) - - end #solve +end #solve diff --git a/src/pinns_pde_solve.jl b/src/pinns_pde_solve.jl index 12a9ab28e8..0bd319c71f 100644 --- a/src/pinns_pde_solve.jl +++ b/src/pinns_pde_solve.jl @@ -25,8 +25,9 @@ function _dot_(x::Expr) if x.head === :call && dottable_(x.args[1]) Expr(:., dotargs[1], Expr(:tuple, dotargs[2:end]...)) elseif x.head === :comparison - Expr(:comparison, (iseven(i) && dottable_(arg) && arg isa Symbol && isoperator(arg) ? - Symbol('.', arg) : arg for (i, arg) in pairs(dotargs))...) + Expr(:comparison, + (iseven(i) && dottable_(arg) && arg isa Symbol && isoperator(arg) ? + Symbol('.', arg) : arg for (i, arg) in pairs(dotargs))...) elseif x.head === :$ x.args[1] elseif x.head === :let # don't add dots to `let x=...` assignments @@ -53,7 +54,7 @@ RuntimeGeneratedFunctions.init(@__MODULE__) """ """ -struct LogOptions +struct LogOptions log_frequency::Int64 # TODO: add in an option for saving plots in the log. this is currently not done because the type of plot is dependent on the PDESystem # possible solution: pass in a plot function? @@ -62,22 +63,22 @@ struct LogOptions # we'd want the plot & log to happen internally as well # plots of the learned function can happen in the outer callback, but we might want to offer that here too - SciMLBase.@add_kwonly function LogOptions(;log_frequency=50) + SciMLBase.@add_kwonly function LogOptions(; log_frequency = 50) new(convert(Int64, log_frequency)) end end """This function is defined here as stubs to be overriden by the subpackage NeuralPDELogging if imported""" -function logvector(logger, v::AbstractVector{R}, name::AbstractString, step::Integer) where R <: Real +function logvector(logger, v::AbstractVector{R}, name::AbstractString, + step::Integer) where {R <: Real} nothing end """This function is defined here as stubs to be overriden by the subpackage NeuralPDELogging if imported""" -function logscalar(logger, s::R, name::AbstractString, step::Integer) where R <: Real +function logscalar(logger, s::R, name::AbstractString, step::Integer) where {R <: Real} nothing end - abstract type AbstractPINN end """ @@ -90,33 +91,33 @@ Arguments: * `phi`: a trial solution, * `derivative`: method that calculates the derivative. """ -struct PhysicsInformedNN{T,P,PH,DER,PE,AL,ADA,LOG,K} <: AbstractPINN - strategy::T - init_params::P - phi::PH - derivative::DER - param_estim::PE - additional_loss::AL - adaptive_loss::ADA - logger::LOG - log_options::LogOptions - iteration::Vector{Int64} - self_increment::Bool - multioutput::Bool - kwargs::K - - @add_kwonly function PhysicsInformedNN(chain, - strategy; - init_params = nothing, - phi = nothing, - derivative = nothing, - param_estim=false, - additional_loss=nothing, - adaptive_loss=nothing, - logger=nothing, - log_options=LogOptions(), - iteration=nothing, - kwargs...) where iip +struct PhysicsInformedNN{T, P, PH, DER, PE, AL, ADA, LOG, K} <: AbstractPINN + strategy::T + init_params::P + phi::PH + derivative::DER + param_estim::PE + additional_loss::AL + adaptive_loss::ADA + logger::LOG + log_options::LogOptions + iteration::Vector{Int64} + self_increment::Bool + multioutput::Bool + kwargs::K + + @add_kwonly function PhysicsInformedNN(chain, + strategy; + init_params = nothing, + phi = nothing, + derivative = nothing, + param_estim = false, + additional_loss = nothing, + adaptive_loss = nothing, + logger = nothing, + log_options = LogOptions(), + iteration = nothing, + kwargs...) where {iip} if init_params === nothing if chain isa AbstractArray initθ = DiffEqFlux.initial_params.(chain) @@ -129,7 +130,8 @@ struct PhysicsInformedNN{T,P,PH,DER,PE,AL,ADA,LOG,K} <: AbstractPINN multioutput = typeof(chain) <: AbstractArray - type_initθ = multioutput ? Base.promote_typeof.(initθ)[1] : Base.promote_typeof(initθ) + type_initθ = multioutput ? Base.promote_typeof.(initθ)[1] : + Base.promote_typeof(initθ) if phi === nothing if multioutput @@ -162,17 +164,27 @@ struct PhysicsInformedNN{T,P,PH,DER,PE,AL,ADA,LOG,K} <: AbstractPINN self_increment = true end - new{typeof(strategy),typeof(initθ),typeof(_phi),typeof(_derivative),typeof(param_estim), - typeof(additional_loss),typeof(adaptive_loss),typeof(logger),typeof(kwargs)}( - strategy,initθ,_phi,_derivative,param_estim,additional_loss,adaptive_loss,logger, - log_options,iteration,self_increment,multioutput,kwargs) + new{typeof(strategy), typeof(initθ), typeof(_phi), typeof(_derivative), + typeof(param_estim), + typeof(additional_loss), typeof(adaptive_loss), typeof(logger), typeof(kwargs)}(strategy, + initθ, + _phi, + _derivative, + param_estim, + additional_loss, + adaptive_loss, + logger, + log_options, + iteration, + self_increment, + multioutput, + kwargs) end end -function vectorify(x, t::Type{T}) where T <: Real +function vectorify(x, t::Type{T}) where {T <: Real} convertfunc(y) = convert(t, y) - returnval = - if x isa Vector + returnval = if x isa Vector convertfunc.(x) else t[convertfunc(x)] @@ -195,17 +207,25 @@ Dict{Symbol,Int64} with 3 entries: :u1 => 1 :u2 => 2 """ -get_dict_vars(vars) = Dict( [Symbol(v) .=> i for (i,v) in enumerate(vars)]) +get_dict_vars(vars) = Dict([Symbol(v) .=> i for (i, v) in enumerate(vars)]) # Wrapper for _transform_expression -function transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput::Bool,eltypeθ,strategy,phi,derivative,integral,initθ;is_integral=false, dict_transformation_vars = nothing, transformation_vars = nothing) +function transform_expression(ex, indvars, depvars, dict_indvars, dict_depvars, + dict_depvar_input, multioutput::Bool, eltypeθ, strategy, phi, + derivative, integral, initθ; is_integral = false, + dict_transformation_vars = nothing, + transformation_vars = nothing) if ex isa Expr - ex = _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput,eltypeθ,strategy,phi,derivative,integral,initθ;is_integral = is_integral, dict_transformation_vars = dict_transformation_vars, transformation_vars = transformation_vars) + ex = _transform_expression(ex, indvars, depvars, dict_indvars, dict_depvars, + dict_depvar_input, multioutput, eltypeθ, strategy, phi, + derivative, integral, initθ; is_integral = is_integral, + dict_transformation_vars = dict_transformation_vars, + transformation_vars = transformation_vars) end return ex end -function get_ε(dim, der_num,eltypeθ) +function get_ε(dim, der_num, eltypeθ) epsilon = cbrt(eps(eltypeθ)) ε = zeros(eltypeθ, dim) ε[der_num] = epsilon @@ -216,7 +236,8 @@ function get_limits(domain) if domain isa AbstractInterval return [leftendpoint(domain)], [rightendpoint(domain)] elseif domain isa ProductDomain - return collect(map(leftendpoint , DomainSets.components(domain))), collect(map(rightendpoint , DomainSets.components(domain))) + return collect(map(leftendpoint, DomainSets.components(domain))), + collect(map(rightendpoint, DomainSets.components(domain))) end end @@ -238,9 +259,13 @@ where order - order of derivative θ - weight in neural network """ -function _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput::Bool,eltypeθ,strategy,phi,derivative_,integral,initθ;is_integral=false, dict_transformation_vars = nothing, transformation_vars = nothing) +function _transform_expression(ex, indvars, depvars, dict_indvars, dict_depvars, + dict_depvar_input, multioutput::Bool, eltypeθ, strategy, phi, + derivative_, integral, initθ; is_integral = false, + dict_transformation_vars = nothing, + transformation_vars = nothing) _args = ex.args - for (i,e) in enumerate(_args) + for (i, e) in enumerate(_args) if !(e isa Expr) if e in keys(dict_depvars) depvar = _args[1] @@ -250,7 +275,12 @@ function _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict ex.args = if !multioutput [var_, Symbol(:cord, num_depvar), :($θ), :phi] else - [var_, Symbol(:cord, num_depvar), Symbol(:($θ), num_depvar), Symbol(:phi, num_depvar)] + [ + var_, + Symbol(:cord, num_depvar), + Symbol(:($θ), num_depvar), + Symbol(:phi, num_depvar), + ] end break elseif e isa ModelingToolkit.Differential @@ -264,18 +294,27 @@ function _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict depvar = _args[1] num_depvar = dict_depvars[depvar] indvars = _args[2:end] - dict_interior_indvars = Dict([indvar .=> j for (j, indvar) in enumerate(dict_depvar_input[depvar])]) + dict_interior_indvars = Dict([indvar .=> j + for (j, indvar) in enumerate(dict_depvar_input[depvar])]) dim_l = length(dict_interior_indvars) var_ = is_integral ? :(derivative) : :($(Expr(:$, :derivative))) εs = [get_ε(dim_l, d, eltypeθ) for d in 1:dim_l] - undv = [dict_interior_indvars[d_p] for d_p in derivative_variables] + undv = [dict_interior_indvars[d_p] for d_p in derivative_variables] εs_dnv = [εs[d] for d in undv] ex.args = if !multioutput [var_, :phi, :u, Symbol(:cord, num_depvar), εs_dnv, order, :($θ)] else - [var_, Symbol(:phi, num_depvar), :u, Symbol(:cord, num_depvar), εs_dnv, order, Symbol(:($θ), num_depvar)] + [ + var_, + Symbol(:phi, num_depvar), + :u, + Symbol(:cord, num_depvar), + εs_dnv, + order, + Symbol(:($θ), num_depvar), + ] end break elseif e isa Symbolics.Integral @@ -289,33 +328,48 @@ function _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict end integrating_depvars = [] - integrand_expr =_args[2] + integrand_expr = _args[2] for d in depvars - d_ex = find_thing_in_expr(integrand_expr,d) + d_ex = find_thing_in_expr(integrand_expr, d) if !isempty(d_ex) push!(integrating_depvars, d_ex[1].args[1]) end end lb, ub = get_limits(_args[1].domain.domain) - lb, ub, _args[2], dict_transformation_vars, transformation_vars = transform_inf_integral(lb, ub, _args[2],integrating_depvars, dict_depvar_input, dict_depvars, integrating_variable, eltypeθ) - - num_depvar = map(int_depvar -> dict_depvars[int_depvar], integrating_depvars) - integrand_ = transform_expression(_args[2],indvars,depvars,dict_indvars,dict_depvars, - dict_depvar_input, multioutput,eltypeθ,strategy, - phi,derivative_,integral,initθ; is_integral = false, - dict_transformation_vars = dict_transformation_vars, - transformation_vars = transformation_vars) + lb, ub, _args[2], dict_transformation_vars, transformation_vars = transform_inf_integral(lb, + ub, + _args[2], + integrating_depvars, + dict_depvar_input, + dict_depvars, + integrating_variable, + eltypeθ) + + num_depvar = map(int_depvar -> dict_depvars[int_depvar], + integrating_depvars) + integrand_ = transform_expression(_args[2], indvars, depvars, dict_indvars, + dict_depvars, + dict_depvar_input, multioutput, eltypeθ, + strategy, + phi, derivative_, integral, initθ; + is_integral = false, + dict_transformation_vars = dict_transformation_vars, + transformation_vars = transformation_vars) integrand__ = _dot_(integrand_) - integrand = build_symbolic_loss_function(nothing, indvars,depvars,dict_indvars,dict_depvars, - dict_depvar_input, phi, derivative_, nothing, multioutput, - initθ, strategy, integrand = integrand__, - integrating_depvars=integrating_depvars, - eq_params=SciMLBase.NullParameters(), - dict_transformation_vars = dict_transformation_vars, + integrand = build_symbolic_loss_function(nothing, indvars, depvars, + dict_indvars, dict_depvars, + dict_depvar_input, phi, + derivative_, nothing, multioutput, + initθ, strategy, + integrand = integrand__, + integrating_depvars = integrating_depvars, + eq_params = SciMLBase.NullParameters(), + dict_transformation_vars = dict_transformation_vars, transformation_vars = transformation_vars, - param_estim =false, default_p = nothing) + param_estim = false, + default_p = nothing) # integrand = repr(integrand) lb = toexpr.(lb) ub = toexpr.(ub) @@ -325,12 +379,19 @@ function _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict if l isa Number push!(lb_, l) else - l_expr = NeuralPDE.build_symbolic_loss_function(nothing, indvars,depvars, - dict_indvars,dict_depvars, - dict_depvar_input, phi, derivative_, - nothing, multioutput, initθ, strategy, - integrand = _dot_(l), integrating_depvars=integrating_depvars, - param_estim =false, default_p = nothing) + l_expr = NeuralPDE.build_symbolic_loss_function(nothing, indvars, + depvars, + dict_indvars, + dict_depvars, + dict_depvar_input, + phi, derivative_, + nothing, + multioutput, initθ, + strategy, + integrand = _dot_(l), + integrating_depvars = integrating_depvars, + param_estim = false, + default_p = nothing) l_f = @RuntimeGeneratedFunction(l_expr) push!(lb_, l_f) end @@ -339,23 +400,45 @@ function _transform_expression(ex,indvars,depvars,dict_indvars,dict_depvars,dict if u_ isa Number push!(ub_, u_) else - u_expr = NeuralPDE.build_symbolic_loss_function(nothing, indvars,depvars, - dict_indvars,dict_depvars, - dict_depvar_input, phi, derivative_, - nothing, multioutput, initθ, strategy, - integrand = _dot_(u_), integrating_depvars=integrating_depvars, - param_estim =false, default_p = nothing) + u_expr = NeuralPDE.build_symbolic_loss_function(nothing, indvars, + depvars, + dict_indvars, + dict_depvars, + dict_depvar_input, + phi, derivative_, + nothing, + multioutput, initθ, + strategy, + integrand = _dot_(u_), + integrating_depvars = integrating_depvars, + param_estim = false, + default_p = nothing) u_f = @RuntimeGeneratedFunction(u_expr) push!(ub_, u_f) end end integrand_func = @RuntimeGeneratedFunction(integrand) - ex.args = [:($(Expr(:$, :integral))), :u, Symbol(:cord, num_depvar[1]), :phi, integrating_var_id, integrand_func, lb_, ub_, :($θ)] + ex.args = [ + :($(Expr(:$, :integral))), + :u, + Symbol(:cord, num_depvar[1]), + :phi, + integrating_var_id, + integrand_func, + lb_, + ub_, + :($θ), + ] break end else - ex.args[i] = _transform_expression(ex.args[i],indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput,eltypeθ,strategy,phi,derivative_,integral,initθ; is_integral = is_integral, dict_transformation_vars = dict_transformation_vars, transformation_vars = transformation_vars) + ex.args[i] = _transform_expression(ex.args[i], indvars, depvars, dict_indvars, + dict_depvars, dict_depvar_input, multioutput, + eltypeθ, strategy, phi, derivative_, + integral, initθ; is_integral = is_integral, + dict_transformation_vars = dict_transformation_vars, + transformation_vars = transformation_vars) end end return ex @@ -389,19 +472,25 @@ Example: (derivative(phi2, u2, [x, y], [[ε,0]], 1, θ2) + 9 * derivative(phi1, u, [x, y], [[0,ε]], 1, θ1)) - 0] """ -function build_symbolic_equation(eq,_indvars,_depvars,multioutput::Bool,eltypeθ,strategy,phi,derivative,initθ) - depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) - parse_equation(eq,indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput,eltypeθ,strategy,phi,derivative,integral,initθ) +function build_symbolic_equation(eq, _indvars, _depvars, multioutput::Bool, eltypeθ, + strategy, phi, derivative, initθ) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) + parse_equation(eq, indvars, depvars, dict_indvars, dict_depvars, dict_depvar_input, + multioutput, eltypeθ, strategy, phi, derivative, integral, initθ) end -function parse_equation(eq,indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput::Bool,eltypeθ, - strategy,phi,derivative,integral,initθ) +function parse_equation(eq, indvars, depvars, dict_indvars, dict_depvars, dict_depvar_input, + multioutput::Bool, eltypeθ, + strategy, phi, derivative, integral, initθ) eq_lhs = isequal(expand_derivatives(eq.lhs), 0) ? eq.lhs : expand_derivatives(eq.lhs) eq_rhs = isequal(expand_derivatives(eq.rhs), 0) ? eq.rhs : expand_derivatives(eq.rhs) - left_expr = transform_expression(toexpr(eq_lhs),indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput, - eltypeθ,strategy,phi,derivative,integral,initθ) - right_expr = transform_expression(toexpr(eq_rhs),indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput, - eltypeθ,strategy,phi,derivative,integral,initθ) + left_expr = transform_expression(toexpr(eq_lhs), indvars, depvars, dict_indvars, + dict_depvars, dict_depvar_input, multioutput, + eltypeθ, strategy, phi, derivative, integral, initθ) + right_expr = transform_expression(toexpr(eq_rhs), indvars, depvars, dict_indvars, + dict_depvars, dict_depvar_input, multioutput, + eltypeθ, strategy, phi, derivative, integral, initθ) left_expr = _dot_(left_expr) right_expr = _dot_(right_expr) loss_func = :($left_expr .- $right_expr) @@ -432,44 +521,47 @@ to end end) """ -function build_symbolic_loss_function(eqs,_indvars,_depvars,dict_depvar_input, - phi, derivative,integral,multioutput::Bool,initθ,strategy; - bc_indvars=nothing, +function build_symbolic_loss_function(eqs, _indvars, _depvars, dict_depvar_input, + phi, derivative, integral, multioutput::Bool, initθ, + strategy; + bc_indvars = nothing, eq_params = SciMLBase.NullParameters(), param_estim = false, - default_p=nothing, - integrand=nothing, - dict_transformation_vars = nothing, + default_p = nothing, + integrand = nothing, + dict_transformation_vars = nothing, transformation_vars = nothing, - integrating_depvars=nothing) + integrating_depvars = nothing) # dictionaries: variable -> unique number - depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) bc_indvars = bc_indvars == nothing ? indvars : bc_indvars integrating_depvars = integrating_depvars == nothing ? depvars : integrating_depvars - return build_symbolic_loss_function(eqs,indvars,depvars, - dict_indvars,dict_depvars,dict_depvar_input, - phi,derivative,integral,multioutput,initθ,strategy; + return build_symbolic_loss_function(eqs, indvars, depvars, + dict_indvars, dict_depvars, dict_depvar_input, + phi, derivative, integral, multioutput, initθ, + strategy; bc_indvars = bc_indvars, eq_params = eq_params, param_estim = param_estim, - default_p=default_p, - integrand=integrand, - dict_transformation_vars = dict_transformation_vars, + default_p = default_p, + integrand = integrand, + dict_transformation_vars = dict_transformation_vars, transformation_vars = transformation_vars, - integrating_depvars=integrating_depvars) + integrating_depvars = integrating_depvars) end function get_indvars_ex(bc_indvars) # , dict_this_eq_indvars) - i_=1 + i_ = 1 indvars_ex = map(bc_indvars) do u if u isa Symbol - # i = dict_this_eq_indvars[u] - # ex = :($:cord[[$i],:]) - ex = :($:cord[[$i_],:]) - i_+=1 - ex + # i = dict_this_eq_indvars[u] + # ex = :($:cord[[$i],:]) + ex = :($:cord[[$i_], :]) + i_ += 1 + ex else - :(fill($u,size($:cord[[1],:]))) + :(fill($u, size($:cord[[1], :]))) end end indvars_ex @@ -481,25 +573,25 @@ Finds which dependent variables are being used in an equation. function pair(eq, depvars, dict_depvars, dict_depvar_input) expr = toexpr(eq) pair_ = map(depvars) do depvar - if !isempty(find_thing_in_expr(expr, depvar)) + if !isempty(find_thing_in_expr(expr, depvar)) dict_depvars[depvar] => dict_depvar_input[depvar] end end Dict(filter(p -> p !== nothing, pair_)) end -function build_symbolic_loss_function(eqs,indvars,depvars, - dict_indvars,dict_depvars,dict_depvar_input, - phi,derivative,integral,multioutput::Bool,initθ,strategy; +function build_symbolic_loss_function(eqs, indvars, depvars, + dict_indvars, dict_depvars, dict_depvar_input, + phi, derivative, integral, multioutput::Bool, initθ, + strategy; eq_params = SciMLBase.NullParameters(), param_estim = param_estim, - default_p=default_p, - bc_indvars=indvars, - integrand=nothing, - dict_transformation_vars = nothing, + default_p = default_p, + bc_indvars = indvars, + integrand = nothing, + dict_transformation_vars = nothing, transformation_vars = nothing, - integrating_depvars=depvars, - ) + integrating_depvars = depvars) if multioutput eltypeθ = eltype(initθ[1]) else @@ -507,30 +599,34 @@ function build_symbolic_loss_function(eqs,indvars,depvars, end if integrand isa Nothing - loss_function = parse_equation(eqs,indvars,depvars,dict_indvars,dict_depvars,dict_depvar_input,multioutput,eltypeθ,strategy,phi,derivative,integral,initθ) + loss_function = parse_equation(eqs, indvars, depvars, dict_indvars, dict_depvars, + dict_depvar_input, multioutput, eltypeθ, strategy, + phi, derivative, integral, initθ) this_eq_pair = pair(eqs, depvars, dict_depvars, dict_depvar_input) this_eq_indvars = unique(vcat(values(this_eq_pair)...)) - else - this_eq_pair = Dict(map(intvars -> dict_depvars[intvars] => dict_depvar_input[intvars], integrating_depvars)) - this_eq_indvars = transformation_vars isa Nothing ? unique(vcat(values(this_eq_pair)...)) : transformation_vars + else + this_eq_pair = Dict(map(intvars -> dict_depvars[intvars] => dict_depvar_input[intvars], + integrating_depvars)) + this_eq_indvars = transformation_vars isa Nothing ? + unique(vcat(values(this_eq_pair)...)) : transformation_vars loss_function = integrand end - vars = :(cord, $θ, phi, derivative, integral,u,p) + vars = :(cord, $θ, phi, derivative, integral, u, p) ex = Expr(:block) if multioutput θ_nums = Symbol[] phi_nums = Symbol[] for v in depvars num = dict_depvars[v] - push!(θ_nums,:($(Symbol(:($θ),num)))) - push!(phi_nums,:($(Symbol(:phi,num)))) + push!(θ_nums, :($(Symbol(:($θ), num)))) + push!(phi_nums, :($(Symbol(:phi, num)))) end expr_θ = Expr[] expr_phi = Expr[] - acum = [0;accumulate(+, length.(initθ))] - sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1] + acum = [0; accumulate(+, length.(initθ))] + sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)] for i in eachindex(depvars) push!(expr_θ, :($θ[$(sep[i])])) @@ -538,181 +634,195 @@ function build_symbolic_loss_function(eqs,indvars,depvars, end vars_θ = Expr(:(=), build_expr(:tuple, θ_nums), build_expr(:tuple, expr_θ)) - push!(ex.args, vars_θ) + push!(ex.args, vars_θ) vars_phi = Expr(:(=), build_expr(:tuple, phi_nums), build_expr(:tuple, expr_phi)) - push!(ex.args, vars_phi) + push!(ex.args, vars_phi) end #Add an expression for parameter symbols if param_estim == true && eq_params != SciMLBase.NullParameters() param_len = length(eq_params) - last_indx = [0;accumulate(+, length.(initθ))][end] + last_indx = [0; accumulate(+, length.(initθ))][end] params_symbols = Symbol[] expr_params = Expr[] - for (i,eq_param) in enumerate(eq_params) - push!(expr_params, :($θ[$(i+last_indx:i+last_indx)])) + for (i, eq_param) in enumerate(eq_params) + push!(expr_params, :($θ[$((i + last_indx):(i + last_indx))])) push!(params_symbols, Symbol(:($eq_param))) end - params_eq = Expr(:(=), build_expr(:tuple, params_symbols), build_expr(:tuple, expr_params)) - push!(ex.args, params_eq) + params_eq = Expr(:(=), build_expr(:tuple, params_symbols), + build_expr(:tuple, expr_params)) + push!(ex.args, params_eq) end if eq_params != SciMLBase.NullParameters() && param_estim == false params_symbols = Symbol[] expr_params = Expr[] - for (i , eq_param) in enumerate(eq_params) - push!(expr_params, :(ArrayInterfaceCore.allowed_getindex(p,$i:$i))) + for (i, eq_param) in enumerate(eq_params) + push!(expr_params, :(ArrayInterfaceCore.allowed_getindex(p, ($i):($i)))) push!(params_symbols, Symbol(:($eq_param))) end - params_eq = Expr(:(=), build_expr(:tuple, params_symbols), build_expr(:tuple, expr_params)) - push!(ex.args, params_eq) + params_eq = Expr(:(=), build_expr(:tuple, params_symbols), + build_expr(:tuple, expr_params)) + push!(ex.args, params_eq) end eq_pair_expr = Expr[] for i in keys(this_eq_pair) - push!(eq_pair_expr, :( $(Symbol(:cord, :($i))) = vcat($(this_eq_pair[i]...)))) + push!(eq_pair_expr, :($(Symbol(:cord, :($i))) = vcat($(this_eq_pair[i]...)))) end vcat_expr = Expr(:block, :($(eq_pair_expr...))) vcat_expr_loss_functions = Expr(:block, vcat_expr, loss_function) # TODO rename if strategy isa QuadratureTraining indvars_ex = get_indvars_ex(bc_indvars) - left_arg_pairs, right_arg_pairs = this_eq_indvars, indvars_ex - vars_eq = Expr(:(=), build_expr(:tuple, left_arg_pairs), build_expr(:tuple, right_arg_pairs)) + left_arg_pairs, right_arg_pairs = this_eq_indvars, indvars_ex + vars_eq = Expr(:(=), build_expr(:tuple, left_arg_pairs), + build_expr(:tuple, right_arg_pairs)) else - indvars_ex = [:($:cord[[$i],:]) for (i, u) ∈ enumerate(this_eq_indvars)] + indvars_ex = [:($:cord[[$i], :]) for (i, u) in enumerate(this_eq_indvars)] left_arg_pairs, right_arg_pairs = this_eq_indvars, indvars_ex - vars_eq = Expr(:(=), build_expr(:tuple, left_arg_pairs), build_expr(:tuple, right_arg_pairs)) + vars_eq = Expr(:(=), build_expr(:tuple, left_arg_pairs), + build_expr(:tuple, right_arg_pairs)) end if !(dict_transformation_vars isa Nothing) transformation_expr_ = Expr[] - for (i,u) in dict_transformation_vars + for (i, u) in dict_transformation_vars push!(transformation_expr_, :($i = $u)) end transformation_expr = Expr(:block, :($(transformation_expr_...))) - vcat_expr_loss_functions = Expr(:block, transformation_expr, vcat_expr, loss_function) + vcat_expr_loss_functions = Expr(:block, transformation_expr, vcat_expr, + loss_function) end let_ex = Expr(:let, vars_eq, vcat_expr_loss_functions) - push!(ex.args, let_ex) + push!(ex.args, let_ex) expr_loss_function = :(($vars) -> begin $ex end) end -function build_loss_function(eqs,_indvars,_depvars,phi,derivative,integral, - multioutput::Bool,initθ,strategy; - bc_indvars=nothing, - eq_params=SciMLBase.NullParameters(), - param_estim=false, - default_p=nothing) +function build_loss_function(eqs, _indvars, _depvars, phi, derivative, integral, + multioutput::Bool, initθ, strategy; + bc_indvars = nothing, + eq_params = SciMLBase.NullParameters(), + param_estim = false, + default_p = nothing) # dictionaries: variable -> unique number - depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) - bc_indvars = bc_indvars==nothing ? indvars : bc_indvars - return build_loss_function(eqs,indvars,depvars, - dict_indvars,dict_depvars,dict_depvar_input, - phi,derivative,integral,multioutput,initθ,strategy; - bc_indvars=bc_indvars, - integration_indvars=indvars, - eq_params=eq_params, - param_estim=param_estim, - default_p=default_p) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) + bc_indvars = bc_indvars == nothing ? indvars : bc_indvars + return build_loss_function(eqs, indvars, depvars, + dict_indvars, dict_depvars, dict_depvar_input, + phi, derivative, integral, multioutput, initθ, strategy; + bc_indvars = bc_indvars, + integration_indvars = indvars, + eq_params = eq_params, + param_estim = param_estim, + default_p = default_p) end -function build_loss_function(eqs,indvars,depvars, - dict_indvars,dict_depvars,dict_depvar_input, - phi,derivative,integral,multioutput::Bool,initθ,strategy; +function build_loss_function(eqs, indvars, depvars, + dict_indvars, dict_depvars, dict_depvar_input, + phi, derivative, integral, multioutput::Bool, initθ, strategy; bc_indvars = indvars, - integration_indvars=indvars, - eq_params=SciMLBase.NullParameters(), - param_estim=false, - default_p=nothing) - expr_loss_function = build_symbolic_loss_function(eqs,indvars,depvars, - dict_indvars,dict_depvars, dict_depvar_input, - phi,derivative,integral,multioutput,initθ,strategy; - bc_indvars = bc_indvars, - eq_params = eq_params, - param_estim=param_estim,default_p=default_p) + integration_indvars = indvars, + eq_params = SciMLBase.NullParameters(), + param_estim = false, + default_p = nothing) + expr_loss_function = build_symbolic_loss_function(eqs, indvars, depvars, + dict_indvars, dict_depvars, + dict_depvar_input, + phi, derivative, integral, + multioutput, initθ, strategy; + bc_indvars = bc_indvars, + eq_params = eq_params, + param_estim = param_estim, + default_p = default_p) u = get_u() _loss_function = @RuntimeGeneratedFunction(expr_loss_function) - loss_function = (cord, θ) -> begin - _loss_function(cord, θ, phi, derivative, integral, u, default_p) - end + loss_function = (cord, θ) -> begin _loss_function(cord, θ, phi, derivative, integral, u, + default_p) end return loss_function end function get_vars(indvars_, depvars_) indvars = ModelingToolkit.getname.(indvars_) depvars = Symbol[] - dict_depvar_input = Dict{Symbol,Vector{Symbol}}() + dict_depvar_input = Dict{Symbol, Vector{Symbol}}() for d in depvars_ if ModelingToolkit.value(d) isa Term dname = ModelingToolkit.getname(d) push!(depvars, dname) - push!(dict_depvar_input, dname => [nameof(ModelingToolkit.value(argument)) for argument in ModelingToolkit.value(d).arguments]) + push!(dict_depvar_input, + dname => [nameof(ModelingToolkit.value(argument)) + for argument in ModelingToolkit.value(d).arguments]) else dname = ModelingToolkit.getname(d) push!(depvars, dname) push!(dict_depvar_input, dname => indvars) # default to all inputs if not given end - end + end dict_indvars = get_dict_vars(indvars) dict_depvars = get_dict_vars(depvars) return depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input end function get_integration_variables(eqs, _indvars::Array, _depvars::Array) - depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) get_integration_variables(eqs, dict_indvars, dict_depvars) end function get_integration_variables(eqs, dict_indvars, dict_depvars) exprs = toexpr.(eqs) vars = map(exprs) do expr - _vars = Symbol.(filter(indvar -> length(find_thing_in_expr(expr, indvar)) > 0, sort(collect(keys(dict_indvars))))) + _vars = Symbol.(filter(indvar -> length(find_thing_in_expr(expr, indvar)) > 0, + sort(collect(keys(dict_indvars))))) end end function get_variables(eqs, _indvars::Array, _depvars::Array) - depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) return get_variables(eqs, dict_indvars, dict_depvars) end -function get_variables(eqs,dict_indvars,dict_depvars) - bc_args = get_argument(eqs,dict_indvars,dict_depvars) +function get_variables(eqs, dict_indvars, dict_depvars) + bc_args = get_argument(eqs, dict_indvars, dict_depvars) return map(barg -> filter(x -> x isa Symbol, barg), bc_args) end -function get_number(eqs,dict_indvars,dict_depvars) - bc_args = get_argument(eqs,dict_indvars,dict_depvars) +function get_number(eqs, dict_indvars, dict_depvars) + bc_args = get_argument(eqs, dict_indvars, dict_depvars) return map(barg -> filter(x -> x isa Number, barg), bc_args) end function find_thing_in_expr(ex::Expr, thing; ans = []) if thing in ex.args - push!(ans,ex) + push!(ans, ex) end for e in ex.args if e isa Expr if thing in e.args - push!(ans,e) + push!(ans, e) end - find_thing_in_expr(e,thing; ans=ans) + find_thing_in_expr(e, thing; ans = ans) end end return collect(Set(ans)) end # Get arguments from boundary condition functions -function get_argument(eqs,_indvars::Array,_depvars::Array) - depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) - get_argument(eqs,dict_indvars,dict_depvars) +function get_argument(eqs, _indvars::Array, _depvars::Array) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) + get_argument(eqs, dict_indvars, dict_depvars) end -function get_argument(eqs,dict_indvars,dict_depvars) +function get_argument(eqs, dict_indvars, dict_depvars) exprs = toexpr.(eqs) vars = map(exprs) do expr - _vars = map(depvar -> find_thing_in_expr(expr, depvar), collect(keys(dict_depvars))) + _vars = map(depvar -> find_thing_in_expr(expr, depvar), collect(keys(dict_depvars))) f_vars = filter(x -> !isempty(x), _vars) map(x -> first(x), f_vars) end @@ -735,84 +845,96 @@ function get_argument(eqs,dict_indvars,dict_depvars) return args_ # TODO for all arguments end - -function generate_training_sets(domains,dx,eqs,bcs,eltypeθ,_indvars::Array,_depvars::Array) - depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = get_vars(_indvars, _depvars) - return generate_training_sets(domains,dx,eqs,bcs,eltypeθ,dict_indvars,dict_depvars) +function generate_training_sets(domains, dx, eqs, bcs, eltypeθ, _indvars::Array, + _depvars::Array) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) + return generate_training_sets(domains, dx, eqs, bcs, eltypeθ, dict_indvars, + dict_depvars) end # Generate training set in the domain and on the boundary -function generate_training_sets(domains,dx,eqs,bcs,eltypeθ,dict_indvars::Dict,dict_depvars::Dict) +function generate_training_sets(domains, dx, eqs, bcs, eltypeθ, dict_indvars::Dict, + dict_depvars::Dict) if dx isa Array dxs = dx else - dxs = fill(dx,length(domains)) + dxs = fill(dx, length(domains)) end - spans = [infimum(d.domain):dx:supremum(d.domain) for (d,dx) in zip(domains,dxs)] - dict_var_span = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) for (d,dx) in zip(domains,dxs)]) + spans = [infimum(d.domain):dx:supremum(d.domain) for (d, dx) in zip(domains, dxs)] + dict_var_span = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) + for (d, dx) in zip(domains, dxs)]) - bound_args = get_argument(bcs,dict_indvars,dict_depvars) - bound_vars = get_variables(bcs,dict_indvars,dict_depvars) + bound_args = get_argument(bcs, dict_indvars, dict_depvars) + bound_vars = get_variables(bcs, dict_indvars, dict_depvars) - dif = [eltypeθ[] for i=1:size(domains)[1]] + dif = [eltypeθ[] for i in 1:size(domains)[1]] for _args in bound_args - for (i,x) in enumerate(_args) + for (i, x) in enumerate(_args) if x isa Number - push!(dif[i],x) + push!(dif[i], x) end end end cord_train_set = collect.(spans) - bc_data = map(zip(dif,cord_train_set)) do (d,c) + bc_data = map(zip(dif, cord_train_set)) do (d, c) setdiff(c, d) end - dict_var_span_ = Dict([Symbol(d.variables) => bc for (d,bc) in zip(domains,bc_data)]) + dict_var_span_ = Dict([Symbol(d.variables) => bc for (d, bc) in zip(domains, bc_data)]) bcs_train_sets = map(bound_args) do bt span = map(b -> get(dict_var_span, b, b), bt) - _set = adapt(eltypeθ,hcat(vec(map(points -> collect(points), Iterators.product(span...)))...)) + _set = adapt(eltypeθ, + hcat(vec(map(points -> collect(points), Iterators.product(span...)))...)) end - pde_vars = get_variables(eqs,dict_indvars,dict_depvars) - pde_args = get_argument(eqs,dict_indvars,dict_depvars) + pde_vars = get_variables(eqs, dict_indvars, dict_depvars) + pde_args = get_argument(eqs, dict_indvars, dict_depvars) - pde_train_set = adapt(eltypeθ, hcat(vec(map(points -> collect(points), Iterators.product(bc_data...)))...)) + pde_train_set = adapt(eltypeθ, + hcat(vec(map(points -> collect(points), + Iterators.product(bc_data...)))...)) pde_train_sets = map(pde_args) do bt span = map(b -> get(dict_var_span_, b, b), bt) - _set = adapt(eltypeθ,hcat(vec(map(points -> collect(points), Iterators.product(span...)))...)) + _set = adapt(eltypeθ, + hcat(vec(map(points -> collect(points), Iterators.product(span...)))...)) end - [pde_train_sets,bcs_train_sets] + [pde_train_sets, bcs_train_sets] end -function get_bounds(domains,eqs,bcs,eltypeθ,_indvars::Array,_depvars::Array,strategy) - depvars,indvars,dict_indvars,dict_depvars,dict_depvar_input = get_vars(_indvars, _depvars) - return get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy) +function get_bounds(domains, eqs, bcs, eltypeθ, _indvars::Array, _depvars::Array, strategy) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) + return get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, strategy) end -function get_bounds(domains,eqs,bcs,eltypeθ,_indvars::Array,_depvars::Array,strategy::QuadratureTraining) - depvars,indvars,dict_indvars,dict_depvars,dict_depvar_input = get_vars(_indvars, _depvars) - return get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy) +function get_bounds(domains, eqs, bcs, eltypeθ, _indvars::Array, _depvars::Array, + strategy::QuadratureTraining) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(_indvars, + _depvars) + return get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, strategy) end -function get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy::QuadratureTraining) +function get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, + strategy::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]) - pde_args = get_argument(eqs,dict_indvars,dict_depvars) + pde_args = get_argument(eqs, dict_indvars, dict_depvars) - pde_lower_bounds= map(pde_args) do pd + pde_lower_bounds = map(pde_args) do pd span = map(p -> get(dict_lower_bound, p, p), pd) - map(s -> adapt(eltypeθ,s) + cbrt(eps(eltypeθ)), span) + map(s -> adapt(eltypeθ, s) + cbrt(eps(eltypeθ)), span) end - pde_upper_bounds= map(pde_args) do pd + pde_upper_bounds = map(pde_args) do pd span = map(p -> get(dict_upper_bound, p, p), pd) - map(s -> adapt(eltypeθ,s) - cbrt(eps(eltypeθ)), span) + map(s -> adapt(eltypeθ, s) - cbrt(eps(eltypeθ)), span) end - pde_bounds= [pde_lower_bounds,pde_upper_bounds] + pde_bounds = [pde_lower_bounds, pde_upper_bounds] - bound_vars = get_variables(bcs,dict_indvars,dict_depvars) + bound_vars = get_variables(bcs, dict_indvars, dict_depvars) bcs_lower_bounds = map(bound_vars) do bt map(b -> dict_lower_bound[b], bt) @@ -820,30 +942,34 @@ function get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy: bcs_upper_bounds = map(bound_vars) do bt map(b -> dict_upper_bound[b], bt) end - bcs_bounds= [bcs_lower_bounds,bcs_upper_bounds] + bcs_bounds = [bcs_lower_bounds, bcs_upper_bounds] [pde_bounds, bcs_bounds] end -function get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy) +function get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, strategy) dx = 1 / strategy.points - dict_span = Dict([Symbol(d.variables) => [infimum(d.domain)+dx, supremum(d.domain)-dx] for d in domains]) + dict_span = Dict([Symbol(d.variables) => [ + infimum(d.domain) + dx, + supremum(d.domain) - dx, + ] for d in domains]) # pde_bounds = [[infimum(d.domain),supremum(d.domain)] for d in domains] - pde_args = get_argument(eqs,dict_indvars,dict_depvars) + pde_args = get_argument(eqs, dict_indvars, dict_depvars) - pde_bounds= map(pde_args) do pd + pde_bounds = map(pde_args) do pd span = map(p -> get(dict_span, p, p), pd) - map(s -> adapt(eltypeθ,s), span) + map(s -> adapt(eltypeθ, s), span) end - bound_args = get_argument(bcs,dict_indvars,dict_depvars) - dict_span = Dict([Symbol(d.variables) => [infimum(d.domain), supremum(d.domain)] for d in domains]) + bound_args = get_argument(bcs, dict_indvars, dict_depvars) + dict_span = Dict([Symbol(d.variables) => [infimum(d.domain), supremum(d.domain)] + for d in domains]) - bcs_bounds= map(bound_args) do bt + bcs_bounds = map(bound_args) do bt span = map(b -> get(dict_span, b, b), bt) - map(s -> adapt(eltypeθ,s), span) + map(s -> adapt(eltypeθ, s), span) end - [pde_bounds,bcs_bounds] + [pde_bounds, bcs_bounds] end """ @@ -858,127 +984,130 @@ Fields: struct Phi{C} f::C Phi(chain::FastChain) = new{typeof(chain)}(chain) - Phi(chain::Flux.Chain) = (re = Flux.destructure(chain)[2]; new{typeof(re)}(re)) + Phi(chain::Flux.Chain) = (re = Flux.destructure(chain)[2]; new{typeof(re)}(re)) end -(f::Phi{<:FastChain})(x,θ) = f.f(adapt(parameterless_type(θ),x),θ) -(f::Phi{<:Optimisers.Restructure})(x,θ) = f.f(θ)(adapt(parameterless_type(θ),x)) +(f::Phi{<:FastChain})(x, θ) = f.f(adapt(parameterless_type(θ), x), θ) +(f::Phi{<:Optimisers.Restructure})(x, θ) = f.f(θ)(adapt(parameterless_type(θ), x)) function get_u() - u = (cord, θ, phi)-> phi(cord, θ) + u = (cord, θ, phi) -> phi(cord, θ) end # the method to calculate the derivative -function numeric_derivative(phi,u,x,εs,order,θ) - _epsilon = one(eltype(θ)) / (2*cbrt(eps(eltype(θ)))) +function numeric_derivative(phi, u, x, εs, order, θ) + _epsilon = one(eltype(θ)) / (2 * cbrt(eps(eltype(θ)))) ε = εs[order] - ε = adapt(parameterless_type(θ),ε) - x = adapt(parameterless_type(θ),x) + ε = adapt(parameterless_type(θ), ε) + x = adapt(parameterless_type(θ), x) if order > 1 - return (numeric_derivative(phi,u,x .+ ε,εs,order-1,θ) - .- numeric_derivative(phi,u,x .- ε,εs,order-1,θ)) .* _epsilon + return (numeric_derivative(phi, u, x .+ ε, εs, order - 1, θ) + .- + numeric_derivative(phi, u, x .- ε, εs, order - 1, θ)) .* _epsilon else - return (u(x .+ ε,θ,phi) .- u(x .- ε,θ,phi)) .* _epsilon + return (u(x .+ ε, θ, phi) .- u(x .- ε, θ, phi)) .* _epsilon end end function get_numeric_integral(strategy, _indvars, _depvars, multioutput::Bool, derivative) - depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars) - integral = - (u, cord, phi, integrating_var_id, integrand_func, lb, ub, θ ;strategy=strategy, indvars=indvars, depvars=depvars, dict_indvars=dict_indvars, dict_depvars=dict_depvars)-> - begin - - function integration_(cord, lb, ub, θ) - cord_ = cord - function integrand_(x , p) - @Zygote.ignore @views(cord_[integrating_var_id]) .= x - return integrand_func(cord_, p, phi, derivative, nothing, u, nothing) - end - prob_ = IntegralProblem(integrand_,lb, ub ,θ) - sol = solve(prob_,CubatureJLh(),reltol=1e-3,abstol=1e-3)[1] + depvars, indvars, dict_indvars, dict_depvars = get_vars(_indvars, _depvars) + integral = (u, cord, phi, integrating_var_id, integrand_func, lb, ub, θ; strategy = strategy, indvars = indvars, depvars = depvars, dict_indvars = dict_indvars, dict_depvars = dict_depvars) -> begin + function integration_(cord, lb, ub, θ) + cord_ = cord + function integrand_(x, p) + Zygote.@ignore @views(cord_[integrating_var_id]) .= x + return integrand_func(cord_, p, phi, derivative, nothing, u, nothing) + end + prob_ = IntegralProblem(integrand_, lb, ub, θ) + sol = solve(prob_, CubatureJLh(), reltol = 1e-3, abstol = 1e-3)[1] - return sol - end + return sol + end - lb_ = zeros(size(lb)[1], size(cord)[2]) - ub_ = zeros(size(ub)[1], size(cord)[2]) - for (i, l) in enumerate(lb) - if l isa Number - @Zygote.ignore lb_[i, :] = fill(l, 1, size(cord)[2]) - else - @Zygote.ignore lb_[i, :] = l(cord , θ, phi, derivative, nothing, u, nothing) - end - end - for (i, u_) in enumerate(ub) - if u_ isa Number - @Zygote.ignore ub_[i, :] = fill(u_, 1, size(cord)[2]) - else - @Zygote.ignore ub_[i, :] = u_(cord , θ, phi, derivative, nothing, u, nothing) - end - end - integration_arr = Matrix{Float64}(undef, 1, 0) - for i in 1:size(cord)[2] - # ub__ = @Zygote.ignore getindex(ub_, :, i) - # lb__ = @Zygote.ignore getindex(lb_, :, i) - integration_arr = hcat(integration_arr ,integration_(cord[:, i], lb_[:, i], ub_[:, i], θ)) - end - return integration_arr + lb_ = zeros(size(lb)[1], size(cord)[2]) + ub_ = zeros(size(ub)[1], size(cord)[2]) + for (i, l) in enumerate(lb) + if l isa Number + Zygote.@ignore lb_[i, :] = fill(l, 1, size(cord)[2]) + else + Zygote.@ignore lb_[i, :] = l(cord, θ, phi, derivative, nothing, u, nothing) end + end + for (i, u_) in enumerate(ub) + if u_ isa Number + Zygote.@ignore ub_[i, :] = fill(u_, 1, size(cord)[2]) + else + Zygote.@ignore ub_[i, :] = u_(cord, θ, phi, derivative, nothing, u, nothing) + end + end + integration_arr = Matrix{Float64}(undef, 1, 0) + for i in 1:size(cord)[2] + # ub__ = @Zygote.ignore getindex(ub_, :, i) + # lb__ = @Zygote.ignore getindex(lb_, :, i) + integration_arr = hcat(integration_arr, + integration_(cord[:, i], lb_[:, i], ub_[:, i], θ)) + end + return integration_arr + end end -function get_loss_function(loss_function, train_set, eltypeθ,strategy::GridTraining;τ=nothing) - loss = (θ) -> mean(abs2,loss_function(train_set, θ)) +function get_loss_function(loss_function, train_set, eltypeθ, strategy::GridTraining; + τ = nothing) + loss = (θ) -> mean(abs2, loss_function(train_set, θ)) end - + @nograd function generate_random_points(points, bound, eltypeθ) function f(b) - if b isa Number - fill(eltypeθ(b),(1,points)) - else - lb, ub = b[1], b[2] - lb .+ (ub .- lb) .* rand(eltypeθ,1,points) - lb .+ (ub .- lb) .* rand(eltypeθ,1,points) - end + if b isa Number + fill(eltypeθ(b), (1, points)) + else + lb, ub = b[1], b[2] + lb .+ (ub .- lb) .* rand(eltypeθ, 1, points) + lb .+ (ub .- lb) .* rand(eltypeθ, 1, points) + end end vcat(f.(bound)...) end -function get_loss_function(loss_function, bound, eltypeθ, strategy::StochasticTraining;τ=nothing) +function get_loss_function(loss_function, bound, eltypeθ, strategy::StochasticTraining; + τ = nothing) points = strategy.points loss = (θ) -> begin - sets = generate_random_points(points, bound,eltypeθ) - sets_ = adapt(parameterless_type(θ),sets) - mean(abs2,loss_function(sets_, θ)) + sets = generate_random_points(points, bound, eltypeθ) + sets_ = adapt(parameterless_type(θ), sets) + mean(abs2, loss_function(sets_, θ)) end return loss end @nograd function generate_quasi_random_points(points, bound, eltypeθ, sampling_alg) function f(b) - if b isa Number - fill(eltypeθ(b),(1,points)) - else - lb, ub = eltypeθ[b[1]], [b[2]] - QuasiMonteCarlo.sample(points,lb,ub,sampling_alg) - end + if b isa Number + fill(eltypeθ(b), (1, points)) + else + lb, ub = eltypeθ[b[1]], [b[2]] + QuasiMonteCarlo.sample(points, lb, ub, sampling_alg) + end end vcat(f.(bound)...) end -function generate_quasi_random_points_batch(points, bound, eltypeθ, sampling_alg,minibatch) +function generate_quasi_random_points_batch(points, bound, eltypeθ, sampling_alg, minibatch) map(bound) do b if !(b isa Number) - lb, ub = [b[1]], [b[2]] - set_ = QuasiMonteCarlo.generate_design_matrices(points,lb,ub,sampling_alg,minibatch) - set = map(s -> adapt(eltypeθ,s), set_) + lb, ub = [b[1]], [b[2]] + set_ = QuasiMonteCarlo.generate_design_matrices(points, lb, ub, sampling_alg, + minibatch) + set = map(s -> adapt(eltypeθ, s), set_) else - set = fill(eltypeθ(b),(1,points)) + set = fill(eltypeθ(b), (1, points)) end end end -function get_loss_function(loss_function, bound, eltypeθ,strategy::QuasiRandomTraining;τ=nothing) +function get_loss_function(loss_function, bound, eltypeθ, strategy::QuasiRandomTraining; + τ = nothing) sampling_alg = strategy.sampling_alg points = strategy.points resampling = strategy.resampling @@ -986,227 +1115,253 @@ function get_loss_function(loss_function, bound, eltypeθ,strategy::QuasiRandomT point_batch = nothing point_batch = if resampling == false - generate_quasi_random_points_batch(points, bound,eltypeθ,sampling_alg,minibatch) + generate_quasi_random_points_batch(points, bound, eltypeθ, sampling_alg, minibatch) end - loss = - if resampling == true - θ -> begin - sets = generate_quasi_random_points(points, bound, eltypeθ, sampling_alg) - sets_ = adapt(parameterless_type(θ),sets) - mean(abs2,loss_function(sets_, θ)) - end - else - θ -> begin - sets = [point_batch[i] isa Array{eltypeθ,2} ? - point_batch[i] : point_batch[i][rand(1:minibatch)] - for i in 1:length(point_batch)] #TODO - sets_ = vcat(sets...) - sets__ = adapt(parameterless_type(θ),sets_) - mean(abs2,loss_function(sets__, θ)) - end + loss = if resampling == true + θ -> begin + sets = generate_quasi_random_points(points, bound, eltypeθ, sampling_alg) + sets_ = adapt(parameterless_type(θ), sets) + mean(abs2, loss_function(sets_, θ)) end + else + θ -> begin + sets = [point_batch[i] isa Array{eltypeθ, 2} ? + point_batch[i] : point_batch[i][rand(1:minibatch)] + for i in 1:length(point_batch)] #TODO + sets_ = vcat(sets...) + sets__ = adapt(parameterless_type(θ), sets_) + mean(abs2, loss_function(sets__, θ)) + end + end return loss end -function get_loss_function(loss_function, lb,ub ,eltypeθ, strategy::QuadratureTraining;τ=nothing) - +function get_loss_function(loss_function, lb, ub, eltypeθ, strategy::QuadratureTraining; + τ = nothing) if length(lb) == 0 - loss = (θ) -> mean(abs2,loss_function(rand(eltypeθ,1,10), θ)) + loss = (θ) -> mean(abs2, loss_function(rand(eltypeθ, 1, 10), θ)) return loss end - area = eltypeθ(prod(abs.(ub .-lb))) - f_ = (lb,ub,loss_,θ) -> begin + area = eltypeθ(prod(abs.(ub .- lb))) + f_ = (lb, ub, loss_, θ) -> begin # last_x = 1 - function _loss(x,θ) + function _loss(x, θ) # last_x = x # mean(abs2,loss_(x,θ), dims=2) # size_x = fill(size(x)[2],(1,1)) - x = adapt(parameterless_type(θ),x) - sum(abs2,loss_(x,θ), dims=2) #./ size_x + x = adapt(parameterless_type(θ), x) + sum(abs2, loss_(x, θ), dims = 2) #./ size_x end - prob = IntegralProblem(_loss,lb,ub,θ,batch = strategy.batch,nout=1) + prob = IntegralProblem(_loss, lb, ub, θ, batch = strategy.batch, nout = 1) solve(prob, strategy.quadrature_alg, reltol = strategy.reltol, abstol = strategy.abstol, maxiters = strategy.maxiters)[1] end - loss = (θ) -> 1/area* f_(lb,ub,loss_function,θ) + loss = (θ) -> 1 / area * f_(lb, ub, loss_function, θ) return loss end -function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::PhysicsInformedNN) +function SciMLBase.symbolic_discretize(pde_system::PDESystem, + discretization::PhysicsInformedNN) eqs = pde_system.eqs bcs = pde_system.bcs - + domains = pde_system.domain eq_params = pde_system.ps defaults = pde_system.defaults - default_p = eq_params == SciMLBase.NullParameters() ? nothing : [defaults[ep] for ep in eq_params] + default_p = eq_params == SciMLBase.NullParameters() ? nothing : + [defaults[ep] for ep in eq_params] param_estim = discretization.param_estim additional_loss = discretization.additional_loss # dimensionality of equation dim = length(domains) - depvars,indvars,dict_indvars,dict_depvars,dict_depvar_input = get_vars(pde_system.indvars, pde_system.depvars) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(pde_system.indvars, + pde_system.depvars) multioutput = discretization.multioutput initθ = discretization.init_params - flat_initθ = multioutput ? reduce(vcat,initθ) : initθ + flat_initθ = multioutput ? reduce(vcat, initθ) : initθ eltypeθ = eltype(flat_initθ) phi = discretization.phi derivative = discretization.derivative strategy = discretization.strategy - integral = get_numeric_integral(strategy, pde_system.indvars, pde_system.depvars, multioutput, derivative) + integral = get_numeric_integral(strategy, pde_system.indvars, pde_system.depvars, + multioutput, derivative) if !(eqs isa Array) eqs = [eqs] end - pde_indvars = if strategy isa QuadratureTraining - get_argument(eqs,dict_indvars,dict_depvars) + pde_indvars = if strategy isa QuadratureTraining + get_argument(eqs, dict_indvars, dict_depvars) else - get_variables(eqs,dict_indvars,dict_depvars) + get_variables(eqs, dict_indvars, dict_depvars) end - pde_integration_vars = get_integration_variables(eqs,dict_indvars,dict_depvars) - - symbolic_pde_loss_functions = [build_symbolic_loss_function(eq,indvars,depvars, - dict_indvars,dict_depvars,dict_depvar_input, - phi, derivative,integral, multioutput,initθ,strategy; - eq_params=eq_params,param_estim=param_estim,default_p=default_p, - bc_indvars=pde_indvar - ) for (eq, pde_indvar) in zip(eqs, pde_indvars, pde_integration_vars)] + pde_integration_vars = get_integration_variables(eqs, dict_indvars, dict_depvars) + + symbolic_pde_loss_functions = [build_symbolic_loss_function(eq, indvars, depvars, + dict_indvars, dict_depvars, + dict_depvar_input, + phi, derivative, integral, + multioutput, initθ, + strategy; + eq_params = eq_params, + param_estim = param_estim, + default_p = default_p, + bc_indvars = pde_indvar) + for (eq, pde_indvar) in zip(eqs, pde_indvars, + pde_integration_vars)] bc_indvars = if strategy isa QuadratureTraining - get_argument(bcs,dict_indvars,dict_depvars) + get_argument(bcs, dict_indvars, dict_depvars) else - get_variables(bcs,dict_indvars,dict_depvars) + get_variables(bcs, dict_indvars, dict_depvars) end bc_integration_vars = get_integration_variables(bcs, dict_indvars, dict_depvars) - symbolic_bc_loss_functions = [build_symbolic_loss_function(bc,indvars,depvars, - dict_indvars,dict_depvars, dict_depvar_input, - phi, derivative,integral,multioutput,initθ,strategy, - eq_params=eq_params, - param_estim=param_estim, - default_p=default_p; - bc_indvars=bc_indvar) - for (bc, bc_indvar) in zip(bcs, bc_indvars, bc_integration_vars)] + symbolic_bc_loss_functions = [build_symbolic_loss_function(bc, indvars, depvars, + dict_indvars, dict_depvars, + dict_depvar_input, + phi, derivative, integral, + multioutput, initθ, strategy, + eq_params = eq_params, + param_estim = param_estim, + default_p = default_p; + bc_indvars = bc_indvar) + for (bc, bc_indvar) in zip(bcs, bc_indvars, + bc_integration_vars)] symbolic_pde_loss_functions, symbolic_bc_loss_functions end - -function discretize_inner_functions(pde_system::PDESystem, discretization::PhysicsInformedNN) +function discretize_inner_functions(pde_system::PDESystem, + discretization::PhysicsInformedNN) eqs = pde_system.eqs bcs = pde_system.bcs domains = pde_system.domain eq_params = pde_system.ps defaults = pde_system.defaults - default_p = eq_params == SciMLBase.NullParameters() ? nothing : [defaults[ep] for ep in eq_params] + default_p = eq_params == SciMLBase.NullParameters() ? nothing : + [defaults[ep] for ep in eq_params] param_estim = discretization.param_estim additional_loss = discretization.additional_loss adaloss = discretization.adaptive_loss - + # dimensionality of equation dim = length(domains) #TODO fix it in MTK 6.0.0+v: ModelingToolkit.get_ivs(pde_system) - depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = get_vars(pde_system.ivs, - pde_system.dvs) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = get_vars(pde_system.ivs, + pde_system.dvs) multioutput = discretization.multioutput initθ = discretization.init_params - flat_initθ = multioutput ? reduce(vcat,initθ) : initθ + flat_initθ = multioutput ? reduce(vcat, initθ) : initθ eltypeθ = eltype(flat_initθ) - flat_initθ = param_estim == false ? flat_initθ : vcat(flat_initθ, adapt(typeof(flat_initθ),default_p)) + flat_initθ = param_estim == false ? flat_initθ : + vcat(flat_initθ, adapt(typeof(flat_initθ), default_p)) phi = discretization.phi derivative = discretization.derivative strategy = discretization.strategy - integral = get_numeric_integral(strategy, pde_system.indvars, pde_system.depvars, multioutput, derivative) + integral = get_numeric_integral(strategy, pde_system.indvars, pde_system.depvars, + multioutput, derivative) if !(eqs isa Array) eqs = [eqs] end pde_indvars = if strategy isa QuadratureTraining - get_argument(eqs,dict_indvars,dict_depvars) + get_argument(eqs, dict_indvars, dict_depvars) else - get_variables(eqs,dict_indvars,dict_depvars) + get_variables(eqs, dict_indvars, dict_depvars) end - pde_integration_vars = get_integration_variables(eqs, dict_indvars, dict_depvars) - _pde_loss_functions = [build_loss_function(eq,indvars,depvars, - dict_indvars,dict_depvars,dict_depvar_input, - phi, derivative,integral, multioutput, initθ,strategy; - eq_params=eq_params,param_estim=param_estim,default_p=default_p, - bc_indvars=pde_indvar, integration_indvars=integration_indvar - ) for (eq, pde_indvar, integration_indvar) in zip(eqs, pde_indvars, pde_integration_vars)] + pde_integration_vars = get_integration_variables(eqs, dict_indvars, dict_depvars) + _pde_loss_functions = [build_loss_function(eq, indvars, depvars, + dict_indvars, dict_depvars, + dict_depvar_input, + phi, derivative, integral, multioutput, + initθ, strategy; + eq_params = eq_params, + param_estim = param_estim, + default_p = default_p, + bc_indvars = pde_indvar, + integration_indvars = integration_indvar) + for (eq, pde_indvar, integration_indvar) in zip(eqs, pde_indvars, + pde_integration_vars)] bc_indvars = if strategy isa QuadratureTraining - get_argument(bcs,dict_indvars,dict_depvars) + get_argument(bcs, dict_indvars, dict_depvars) else - get_variables(bcs,dict_indvars,dict_depvars) + get_variables(bcs, dict_indvars, dict_depvars) end bc_integration_vars = get_integration_variables(bcs, dict_indvars, dict_depvars) - _bc_loss_functions = [build_loss_function(bc,indvars,depvars, - dict_indvars,dict_depvars, dict_depvar_input, - phi,derivative,integral,multioutput,initθ,strategy; - eq_params=eq_params, - param_estim=param_estim, - default_p=default_p, - bc_indvars=bc_indvar, - integration_indvars=integration_indvar) for (bc, bc_indvar, integration_indvar) in zip(bcs, bc_indvars, bc_integration_vars)] - - pde_loss_functions, bc_loss_functions = - if strategy isa GridTraining + _bc_loss_functions = [build_loss_function(bc, indvars, depvars, + dict_indvars, dict_depvars, dict_depvar_input, + phi, derivative, integral, multioutput, initθ, + strategy; + eq_params = eq_params, + param_estim = param_estim, + default_p = default_p, + bc_indvars = bc_indvar, + integration_indvars = integration_indvar) + for (bc, bc_indvar, integration_indvar) in zip(bcs, bc_indvars, + bc_integration_vars)] + + pde_loss_functions, bc_loss_functions = if strategy isa GridTraining dx = strategy.dx - train_sets = generate_training_sets(domains,dx,eqs,bcs,eltypeθ, - dict_indvars,dict_depvars) + train_sets = generate_training_sets(domains, dx, eqs, bcs, eltypeθ, + dict_indvars, dict_depvars) # the points in the domain and on the boundary pde_train_sets, bcs_train_sets = train_sets - pde_train_sets = adapt.(typeof(flat_initθ),pde_train_sets) - bcs_train_sets = adapt.(typeof(flat_initθ),bcs_train_sets) - pde_loss_functions = [get_loss_function(_loss,_set,eltypeθ,strategy) - for (_loss,_set) in zip(_pde_loss_functions,pde_train_sets)] + pde_train_sets = adapt.(typeof(flat_initθ), pde_train_sets) + bcs_train_sets = adapt.(typeof(flat_initθ), bcs_train_sets) + pde_loss_functions = [get_loss_function(_loss, _set, eltypeθ, strategy) + for (_loss, _set) in zip(_pde_loss_functions, pde_train_sets)] - bc_loss_functions = [get_loss_function(_loss,_set,eltypeθ,strategy) - for (_loss,_set) in zip(_bc_loss_functions, bcs_train_sets)] + bc_loss_functions = [get_loss_function(_loss, _set, eltypeθ, strategy) + for (_loss, _set) in zip(_bc_loss_functions, bcs_train_sets)] (pde_loss_functions, bc_loss_functions) elseif strategy isa StochasticTraining - bounds = get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy) - pde_bounds, bcs_bounds = bounds + bounds = get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, + strategy) + pde_bounds, bcs_bounds = bounds - pde_loss_functions = [get_loss_function(_loss,bound,eltypeθ,strategy) - for (_loss,bound) in zip(_pde_loss_functions, pde_bounds)] + pde_loss_functions = [get_loss_function(_loss, bound, eltypeθ, strategy) + for (_loss, bound) in zip(_pde_loss_functions, pde_bounds)] - bc_loss_functions = [get_loss_function(_loss,bound,eltypeθ,strategy) - for (_loss, bound) in zip(_bc_loss_functions, bcs_bounds)] - (pde_loss_functions, bc_loss_functions) + bc_loss_functions = [get_loss_function(_loss, bound, eltypeθ, strategy) + for (_loss, bound) in zip(_bc_loss_functions, bcs_bounds)] + (pde_loss_functions, bc_loss_functions) elseif strategy isa QuasiRandomTraining - bounds = get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy) - pde_bounds, bcs_bounds = bounds - - pde_loss_functions = [get_loss_function(_loss,bound,eltypeθ,strategy) - for (_loss,bound) in zip(_pde_loss_functions, pde_bounds)] - - strategy_ = QuasiRandomTraining(strategy.bcs_points; - sampling_alg = strategy.sampling_alg, - resampling = strategy.resampling, - minibatch = strategy.minibatch) - bc_loss_functions = [get_loss_function(_loss,bound,eltypeθ,strategy_) - for (_loss, bound) in zip(_bc_loss_functions, bcs_bounds)] - (pde_loss_functions, bc_loss_functions) + bounds = get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, + strategy) + pde_bounds, bcs_bounds = bounds + + pde_loss_functions = [get_loss_function(_loss, bound, eltypeθ, strategy) + for (_loss, bound) in zip(_pde_loss_functions, pde_bounds)] + + strategy_ = QuasiRandomTraining(strategy.bcs_points; + sampling_alg = strategy.sampling_alg, + resampling = strategy.resampling, + minibatch = strategy.minibatch) + bc_loss_functions = [get_loss_function(_loss, bound, eltypeθ, strategy_) + for (_loss, bound) in zip(_bc_loss_functions, bcs_bounds)] + (pde_loss_functions, bc_loss_functions) elseif strategy isa QuadratureTraining - bounds = get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy) + bounds = get_bounds(domains, eqs, bcs, eltypeθ, dict_indvars, dict_depvars, + strategy) pde_bounds, bcs_bounds = bounds - lbs,ubs = pde_bounds - pde_loss_functions = [get_loss_function(_loss,lb,ub,eltypeθ,strategy) - for (_loss,lb,ub) in zip(_pde_loss_functions, lbs,ubs )] - lbs,ubs = bcs_bounds - bc_loss_functions = [get_loss_function(_loss,lb,ub,eltypeθ,strategy) - for (_loss,lb,ub) in zip(_bc_loss_functions, lbs,ubs)] + lbs, ubs = pde_bounds + pde_loss_functions = [get_loss_function(_loss, lb, ub, eltypeθ, strategy) + for (_loss, lb, ub) in zip(_pde_loss_functions, lbs, ubs)] + lbs, ubs = bcs_bounds + bc_loss_functions = [get_loss_function(_loss, lb, ub, eltypeθ, strategy) + for (_loss, lb, ub) in zip(_bc_loss_functions, lbs, ubs)] (pde_loss_functions, bc_loss_functions) end @@ -1215,7 +1370,7 @@ function discretize_inner_functions(pde_system::PDESystem, discretization::Physi num_pde_losses = length(pde_loss_functions) num_bc_losses = length(bc_loss_functions) # assume one single additional loss function if there is one. this means that the user needs to lump all their functions into a single one, - num_additional_loss = additional_loss isa Nothing ? 0 : 1 + num_additional_loss = additional_loss isa Nothing ? 0 : 1 adaloss_T = eltype(adaloss.pde_loss_weights) logger = discretization.logger @@ -1226,43 +1381,54 @@ function discretize_inner_functions(pde_system::PDESystem, discretization::Physi # this will error if the user has provided a number of initial weights that is more than 1 and doesn't match the number of loss functions adaloss.pde_loss_weights = ones(adaloss_T, num_pde_losses) .* adaloss.pde_loss_weights adaloss.bc_loss_weights = ones(adaloss_T, num_bc_losses) .* adaloss.bc_loss_weights - adaloss.additional_loss_weights = ones(adaloss_T, num_additional_loss) .* adaloss.additional_loss_weights - + adaloss.additional_loss_weights = ones(adaloss_T, num_additional_loss) .* + adaloss.additional_loss_weights # this is the function that gets called to do the adaptive reweighting, a function specific to the # type of adaptive reweighting being performed. # TODO: I'd love to pull this out and then dispatch on it via the AbstractAdaptiveLoss, so that users can implement their own # currently this is kind of tricky since the different methods need different types of information, and the loss functions # are generated internal to the code - reweight_losses_func = - if adaloss isa GradientScaleAdaptiveLoss + reweight_losses_func = if adaloss isa GradientScaleAdaptiveLoss weight_change_inertia = discretization.adaptive_loss.weight_change_inertia function run_loss_gradients_adaptive_loss(θ) if iteration[1] % adaloss.reweight_every == 0 # the paper assumes a single pde loss function, so here we grab the maximum of the maximums of each pde loss function - pde_grads_maxes = [maximum(abs.(Zygote.gradient(pde_loss_function, θ)[1])) for pde_loss_function in pde_loss_functions] + pde_grads_maxes = [maximum(abs.(Zygote.gradient(pde_loss_function, θ)[1])) + for pde_loss_function in pde_loss_functions] pde_grads_max = maximum(pde_grads_maxes) - bc_grads_mean = [mean(abs.(Zygote.gradient(bc_loss_function, θ)[1])) for bc_loss_function in bc_loss_functions] - - nonzero_divisor_eps = adaloss_T isa Float64 ? Float64(1e-11) : convert(adaloss_T, 1e-7) - bc_loss_weights_proposed = pde_grads_max ./ (bc_grads_mean .+ nonzero_divisor_eps) - adaloss.bc_loss_weights .= weight_change_inertia .* adaloss.bc_loss_weights .+ (1 .- weight_change_inertia) .* bc_loss_weights_proposed + bc_grads_mean = [mean(abs.(Zygote.gradient(bc_loss_function, θ)[1])) + for bc_loss_function in bc_loss_functions] + + nonzero_divisor_eps = adaloss_T isa Float64 ? Float64(1e-11) : + convert(adaloss_T, 1e-7) + bc_loss_weights_proposed = pde_grads_max ./ + (bc_grads_mean .+ nonzero_divisor_eps) + adaloss.bc_loss_weights .= weight_change_inertia .* + adaloss.bc_loss_weights .+ + (1 .- weight_change_inertia) .* + bc_loss_weights_proposed logscalar(logger, pde_grads_max, "adaptive_loss/pde_grad_max", iteration[1]) - logvector(logger, pde_grads_maxes, "adaptive_loss/pde_grad_maxes", iteration[1]) + logvector(logger, pde_grads_maxes, "adaptive_loss/pde_grad_maxes", + iteration[1]) logvector(logger, bc_grads_mean, "adaptive_loss/bc_grad_mean", iteration[1]) - logvector(logger, adaloss.bc_loss_weights, "adaptive_loss/bc_loss_weights", iteration[1]) + logvector(logger, adaloss.bc_loss_weights, "adaptive_loss/bc_loss_weights", + iteration[1]) end nothing end elseif adaloss isa MiniMaxAdaptiveLoss pde_max_optimiser = adaloss.pde_max_optimiser bc_max_optimiser = adaloss.bc_max_optimiser - function run_minimax_adaptive_loss(θ, pde_losses, bc_losses) + function run_minimax_adaptive_loss(θ, pde_losses, bc_losses) if iteration[1] % adaloss.reweight_every == 0 - Flux.Optimise.update!(pde_max_optimiser, adaloss.pde_loss_weights, -pde_losses) + Flux.Optimise.update!(pde_max_optimiser, adaloss.pde_loss_weights, + -pde_losses) Flux.Optimise.update!(bc_max_optimiser, adaloss.bc_loss_weights, -bc_losses) - logvector(logger, adaloss.pde_loss_weights, "adaptive_loss/pde_loss_weights", iteration[1]) - logvector(logger, adaloss.bc_loss_weights, "adaptive_loss/bc_loss_weights", iteration[1]) + logvector(logger, adaloss.pde_loss_weights, + "adaptive_loss/pde_loss_weights", iteration[1]) + logvector(logger, adaloss.bc_loss_weights, "adaptive_loss/bc_loss_weights", + iteration[1]) end nothing end @@ -1272,7 +1438,7 @@ function discretize_inner_functions(pde_system::PDESystem, discretization::Physi end end - function loss_function_(θ,p) + function loss_function_(θ, p) # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them pde_losses = [pde_loss_function(θ) for pde_loss_function in pde_loss_functions] @@ -1280,17 +1446,15 @@ function discretize_inner_functions(pde_system::PDESystem, discretization::Physi # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized # that's why we prefer the user to maintain the increment in the outer loop callback during optimization - Zygote.@ignore if self_increment + Zygote.@ignore if self_increment iteration[1] += 1 end - Zygote.@ignore begin - if adaloss isa MiniMaxAdaptiveLoss - reweight_losses_func(θ, pde_losses, bc_losses) - else - reweight_losses_func(θ) - end - end + Zygote.@ignore begin if adaloss isa MiniMaxAdaptiveLoss + reweight_losses_func(θ, pde_losses, bc_losses) + else + reweight_losses_func(θ) + end end weighted_pde_losses = adaloss.pde_loss_weights .* pde_losses weighted_bc_losses = adaloss.bc_loss_weights .* bc_losses @@ -1298,47 +1462,55 @@ function discretize_inner_functions(pde_system::PDESystem, discretization::Physi sum_weighted_bc_losses = sum(weighted_bc_losses) weighted_loss_before_additional = sum_weighted_pde_losses + sum_weighted_bc_losses - full_weighted_loss = - if additional_loss isa Nothing + full_weighted_loss = if additional_loss isa Nothing weighted_loss_before_additional else - function _additional_loss(phi,θ) - (θ_,p_) = if (param_estim == true) - θ[1:end - length(default_p)], θ[(end - length(default_p) + 1):end] + function _additional_loss(phi, θ) + (θ_, p_) = if (param_estim == true) + θ[1:(end - length(default_p))], θ[(end - length(default_p) + 1):end] else θ, nothing end return additional_loss(phi, θ, p_) end - weighted_additional_loss_val = adaloss.additional_loss_weights[1] * _additional_loss(phi, θ) + weighted_additional_loss_val = adaloss.additional_loss_weights[1] * + _additional_loss(phi, θ) weighted_loss_before_additional + weighted_additional_loss_val end - Zygote.@ignore begin - if iteration[1] % log_frequency == 0 - logvector(logger, pde_losses, "unweighted_loss/pde_losses", iteration[1]) - logvector(logger, bc_losses, "unweighted_loss/bc_losses", iteration[1]) - logvector(logger, weighted_pde_losses, "weighted_loss/weighted_pde_losses", iteration[1]) - logvector(logger, weighted_bc_losses, "weighted_loss/weighted_bc_losses", iteration[1]) - logscalar(logger, sum_weighted_pde_losses, "weighted_loss/sum_weighted_pde_losses", iteration[1]) - logscalar(logger, sum_weighted_bc_losses, "weighted_loss/sum_weighted_bc_losses", iteration[1]) - logscalar(logger, full_weighted_loss, "weighted_loss/full_weighted_loss", iteration[1]) - logvector(logger, adaloss.pde_loss_weights, "adaptive_loss/pde_loss_weights", iteration[1]) - logvector(logger, adaloss.bc_loss_weights, "adaptive_loss/bc_loss_weights", iteration[1]) - end - end + Zygote.@ignore begin if iteration[1] % log_frequency == 0 + logvector(logger, pde_losses, "unweighted_loss/pde_losses", iteration[1]) + logvector(logger, bc_losses, "unweighted_loss/bc_losses", iteration[1]) + logvector(logger, weighted_pde_losses, "weighted_loss/weighted_pde_losses", + iteration[1]) + logvector(logger, weighted_bc_losses, "weighted_loss/weighted_bc_losses", + iteration[1]) + logscalar(logger, sum_weighted_pde_losses, + "weighted_loss/sum_weighted_pde_losses", iteration[1]) + logscalar(logger, sum_weighted_bc_losses, + "weighted_loss/sum_weighted_bc_losses", iteration[1]) + logscalar(logger, full_weighted_loss, "weighted_loss/full_weighted_loss", + iteration[1]) + logvector(logger, adaloss.pde_loss_weights, "adaptive_loss/pde_loss_weights", + iteration[1]) + logvector(logger, adaloss.bc_loss_weights, "adaptive_loss/bc_loss_weights", + iteration[1]) + end end return full_weighted_loss end - (bc_loss_functions=bc_loss_functions, pde_loss_functions=pde_loss_functions, full_loss_function=loss_function_, - additional_loss_function=additional_loss, flat_initθ=flat_initθ, - inner_pde_loss_functions=_pde_loss_functions, inner_bc_loss_functions=_bc_loss_functions) + (bc_loss_functions = bc_loss_functions, pde_loss_functions = pde_loss_functions, + full_loss_function = loss_function_, + additional_loss_function = additional_loss, flat_initθ = flat_initθ, + inner_pde_loss_functions = _pde_loss_functions, + inner_bc_loss_functions = _bc_loss_functions) end # Convert a PDE problem into an OptimizationProblem function SciMLBase.discretize(pde_system::PDESystem, discretization::PhysicsInformedNN) discretized_functions = discretize_inner_functions(pde_system, discretization) - f = OptimizationFunction(discretized_functions.full_loss_function, Optimization.AutoZygote()) + f = OptimizationFunction(discretized_functions.full_loss_function, + Optimization.AutoZygote()) Optimization.OptimizationProblem(f, discretized_functions.flat_initθ) end diff --git a/src/rode_solve.jl b/src/rode_solve.jl index 82510da994..8cbee80c35 100644 --- a/src/rode_solve.jl +++ b/src/rode_solve.jl @@ -1,4 +1,4 @@ -struct NNRODE{C,W,O,P,K} <: NeuralPDEAlgorithm +struct NNRODE{C, W, O, P, K} <: NeuralPDEAlgorithm chain::C W::W opt::O @@ -6,31 +6,30 @@ struct NNRODE{C,W,O,P,K} <: NeuralPDEAlgorithm autodiff::Bool kwargs::K end -function NNRODE(chain,W,opt=Optim.BFGS(),init_params = nothing;autodiff=false,kwargs...) +function NNRODE(chain, W, opt = Optim.BFGS(), init_params = nothing; autodiff = false, + kwargs...) if init_params === nothing if chain isa FastChain initθ = DiffEqFlux.initial_params(chain) else - initθ,re = Flux.destructure(chain) + initθ, re = Flux.destructure(chain) end else initθ = init_params end - NNRODE(chain,W,opt,initθ,autodiff,kwargs) + NNRODE(chain, W, opt, initθ, autodiff, kwargs) end -function DiffEqBase.solve( - prob::DiffEqBase.AbstractRODEProblem, - alg::NeuralPDEAlgorithm, - args...; - dt, - timeseries_errors = true, - save_everystep=true, - adaptive=false, - abstol = 1f-6, - verbose = false, - maxiters = 100) - +function DiffEqBase.solve(prob::DiffEqBase.AbstractRODEProblem, + alg::NeuralPDEAlgorithm, + args...; + dt, + timeseries_errors = true, + save_everystep = true, + adaptive = false, + abstol = 1.0f-6, + verbose = false, + maxiters = 100) DiffEqBase.isinplace(prob) && error("Only out-of-place methods are allowed!") u0 = prob.u0 @@ -40,8 +39,8 @@ function DiffEqBase.solve( t0 = tspan[1] #hidden layer - chain = alg.chain - opt = alg.opt + chain = alg.chain + opt = alg.opt autodiff = alg.autodiff Wg = alg.W #train points generation @@ -51,55 +50,67 @@ function DiffEqBase.solve( if chain isa FastChain #The phi trial solution if u0 isa Number - phi = (t,W,θ) -> u0 + (t-tspan[1])*first(chain(adapt(DiffEqBase.parameterless_type(θ),[t,W]),θ)) + phi = (t, W, θ) -> u0 + + (t - tspan[1]) * + first(chain(adapt(DiffEqBase.parameterless_type(θ), [t, W]), + θ)) else - phi = (t,W,θ) -> u0 + (t-tspan[1])*chain(adapt(DiffEqBase.parameterless_type(θ),[t,W]),θ) + phi = (t, W, θ) -> u0 + + (t - tspan[1]) * + chain(adapt(DiffEqBase.parameterless_type(θ), [t, W]), θ) end else - _,re = Flux.destructure(chain) + _, re = Flux.destructure(chain) #The phi trial solution if u0 isa Number - phi = (t,W,θ) -> u0 + (t-t0)*first(re(θ)(adapt(DiffEqBase.parameterless_type(θ),[t,W]))) + phi = (t, W, θ) -> u0 + + (t - t0) * + first(re(θ)(adapt(DiffEqBase.parameterless_type(θ), [t, W]))) else - phi = (t,W,θ) -> u0 + (t-t0)*re(θ)(adapt(DiffEqBase.parameterless_type(θ),[t,W])) + phi = (t, W, θ) -> u0 + + (t - t0) * + re(θ)(adapt(DiffEqBase.parameterless_type(θ), [t, W])) end end if autodiff # dfdx = (t,W,θ) -> ForwardDiff.derivative(t->phi(t,θ),t) else - dfdx = (t,W,θ) -> (phi(t+sqrt(eps(t)),W,θ) - phi(t,W,θ))/sqrt(eps(t)) + dfdx = (t, W, θ) -> (phi(t + sqrt(eps(t)), W, θ) - phi(t, W, θ)) / sqrt(eps(t)) end - function inner_loss(t,W,θ) - sum(abs,dfdx(t,W,θ) - f(phi(t,W,θ),p,t,W)) + function inner_loss(t, W, θ) + sum(abs, dfdx(t, W, θ) - f(phi(t, W, θ), p, t, W)) end - Wprob = NoiseProblem(Wg,tspan) - Wsol = solve(Wprob;dt=dt) - W = NoiseGrid(ts , Wsol.W) + Wprob = NoiseProblem(Wg, tspan) + Wsol = solve(Wprob; dt = dt) + W = NoiseGrid(ts, Wsol.W) function loss(θ) - sum(abs2,inner_loss(ts[i],W.W[i],θ) for i in 1:length(ts)) # sum(abs2,phi(tspan[1],θ) - u0) + sum(abs2, inner_loss(ts[i], W.W[i], θ) for i in 1:length(ts)) # sum(abs2,phi(tspan[1],θ) - u0) end - callback = function (p,l) - Wprob = NoiseProblem(Wg,tspan) - Wsol = solve(Wprob;dt=dt) - W = NoiseGrid(ts , Wsol.W) + callback = function (p, l) + Wprob = NoiseProblem(Wg, tspan) + Wsol = solve(Wprob; dt = dt) + W = NoiseGrid(ts, Wsol.W) verbose && println("Current loss is: $l") l < abstol end - res = DiffEqFlux.sciml_train(loss, initθ, opt; cb = callback, maxiters=maxiters, alg.kwargs...) + res = DiffEqFlux.sciml_train(loss, initθ, opt; cb = callback, maxiters = maxiters, + alg.kwargs...) #solutions at timepoints - noiseproblem = NoiseProblem(Wg,tspan) - W = solve(noiseproblem;dt=dt) + noiseproblem = NoiseProblem(Wg, tspan) + W = solve(noiseproblem; dt = dt) if u0 isa Number - u = [(phi(ts[i],W.W[i],res.minimizer)) for i in 1:length(ts)] + u = [(phi(ts[i], W.W[i], res.minimizer)) for i in 1:length(ts)] else - u = [(phi(ts[i],W.W[i],res.minimizer)) for i in 1:length(ts)] + u = [(phi(ts[i], W.W[i], res.minimizer)) for i in 1:length(ts)] end - sol = DiffEqBase.build_solution(prob,alg,ts,u,W = W,calculate_error = false) - DiffEqBase.has_analytic(prob.f) && DiffEqBase.calculate_solution_errors!(sol;timeseries_errors=true,dense_errors=false) + sol = DiffEqBase.build_solution(prob, alg, ts, u, W = W, calculate_error = false) + DiffEqBase.has_analytic(prob.f) && + DiffEqBase.calculate_solution_errors!(sol; timeseries_errors = true, + dense_errors = false) sol end #solve diff --git a/src/stopping_solve.jl b/src/stopping_solve.jl index 19adcc54c9..d7617f272a 100644 --- a/src/stopping_solve.jl +++ b/src/stopping_solve.jl @@ -14,26 +14,26 @@ Arguments: [1]Becker, Sebastian, et al. "Solving high-dimensional optimal stopping problems using deep learning." arXiv preprint arXiv:1908.01602 (2019). """ -struct NNStopping{C,O,S,E} <: NeuralPDEAlgorithm +struct NNStopping{C, O, S, E} <: NeuralPDEAlgorithm chain::C opt::O sdealg::S ensemblealg::E end -NNStopping(chain ; opt=Flux.ADAM(0.1) , sdealg = EM() , ensemblealg = EnsembleThreads()) = NNStopping(chain , opt , sdealg , ensemblealg) - -function DiffEqBase.solve( - prob::SDEProblem, - alg::NNStopping; - abstol = 1f-6, - verbose = false, - maxiters = 300, - trajectories = 1000, - save_everystep = false, - dt, - kwargs... - ) +function NNStopping(chain; opt = Flux.ADAM(0.1), sdealg = EM(), + ensemblealg = EnsembleThreads()) + NNStopping(chain, opt, sdealg, ensemblealg) +end +function DiffEqBase.solve(prob::SDEProblem, + alg::NNStopping; + abstol = 1.0f-6, + verbose = false, + maxiters = 300, + trajectories = 1000, + save_everystep = false, + dt, + kwargs...) tspan = prob.tspan sigma = prob.g μ = prob.f @@ -43,45 +43,47 @@ function DiffEqBase.solve( N = size(ts)[1] T = tspan[2] - m = alg.chain - opt = alg.opt + m = alg.chain + opt = alg.opt sdealg = alg.sdealg ensemblealg = alg.ensemblealg - prob = SDEProblem(μ,sigma,u0,tspan) + prob = SDEProblem(μ, sigma, u0, tspan) ensembleprob = EnsembleProblem(prob) - sim = solve(ensembleprob, sdealg, ensemblealg, dt=dt,trajectories=trajectories,adaptive=false) + sim = solve(ensembleprob, sdealg, ensemblealg, dt = dt, trajectories = trajectories, + adaptive = false) payoff = [] times = [] iter = 0 # for u in sim.u un = [] - function Un(n , X ) + function Un(n, X) if size(un)[1] >= n return un[n] else - if(n == 1) - ans = first(m(X[1])[1]) - un = [ans] - return ans - else - ans = max(first(m(X[n])[n]) , n + 1 - size(ts)[1])*(1 - sum(Un(i , X ) for i in 1:n-1)) - un = vcat( un , ans) - return ans - end - end + if (n == 1) + ans = first(m(X[1])[1]) + un = [ans] + return ans + else + ans = max(first(m(X[n])[n]), n + 1 - size(ts)[1]) * + (1 - sum(Un(i, X) for i in 1:(n - 1))) + un = vcat(un, ans) + return ans + end + end end function loss() reward = 0.00 for u in sim.u X = u.u - reward = reward + sum(Un(i , X )*g(ts[i] , X[i]) for i in 1 : size(ts)[1]) + reward = reward + sum(Un(i, X) * g(ts[i], X[i]) for i in 1:size(ts)[1]) un = [] end return 10000 - reward end - dataset = Iterators.repeated(() , maxiters) + dataset = Iterators.repeated((), maxiters) callback = function () l = loss() @@ -94,21 +96,21 @@ function DiffEqBase.solve( ti = 0 Xt = sim.u[1].u for i in 1:N - un = [] - Usum = Usum + Un(i , Xt) - if Usum >= 1 - Un(i , Xt) + un = [] + Usum = Usum + Un(i, Xt) + if Usum >= 1 - Un(i, Xt) ti = i break - end + end end for u in sim.u X = u.u - price = g(ts[ti] , X[ti]) - payoff = vcat(payoff , price) + price = g(ts[ti], X[ti]) + payoff = vcat(payoff, price) times = vcat(times, ti) iter = iter + 1 # println("SUM : $sump") # println("TIME : $ti") end - sum(payoff)/size(payoff)[1] - end #solve + sum(payoff) / size(payoff)[1] +end #solve diff --git a/src/training_strategies.jl b/src/training_strategies.jl index 9a55807eb4..ebbffc9e9f 100644 --- a/src/training_strategies.jl +++ b/src/training_strategies.jl @@ -11,12 +11,12 @@ end * `points`: number of points in random select training set, * `bcs_points`: number of points in random select training set for boundry conditions (by default, it equals `points`). """ -struct StochasticTraining <:AbstractTrainingStrategy - points:: Int64 - bcs_points:: Int64 +struct StochasticTraining <: AbstractTrainingStrategy + points::Int64 + bcs_points::Int64 end -function StochasticTraining(points;bcs_points = points) +function StochasticTraining(points; bcs_points = points) StochasticTraining(points, bcs_points) end @@ -32,16 +32,18 @@ end For more information look: QuasiMonteCarlo.jl https://github.com/SciML/QuasiMonteCarlo.jl """ -struct QuasiRandomTraining <:AbstractTrainingStrategy - points:: Int64 - bcs_points:: Int64 +struct QuasiRandomTraining <: AbstractTrainingStrategy + points::Int64 + bcs_points::Int64 sampling_alg::QuasiMonteCarlo.SamplingAlgorithm - resampling:: Bool - minibatch:: Int64 + resampling::Bool + minibatch::Int64 end -function QuasiRandomTraining(points;bcs_points = points, sampling_alg = LatinHypercubeSample(),resampling =true, minibatch=0) - QuasiRandomTraining(points,bcs_points,sampling_alg,resampling,minibatch) +function QuasiRandomTraining(points; bcs_points = points, + sampling_alg = LatinHypercubeSample(), resampling = true, + minibatch = 0) + QuasiRandomTraining(points, bcs_points, sampling_alg, resampling, minibatch) end """ @@ -53,7 +55,8 @@ end For more information look: Integrals.jl https://github.com/SciML/Integrals.jl """ -struct QuadratureTraining{Q<:SciMLBase.AbstractIntegralAlgorithm,T} <: AbstractTrainingStrategy +struct QuadratureTraining{Q <: SciMLBase.AbstractIntegralAlgorithm, T} <: + AbstractTrainingStrategy quadrature_alg::Q reltol::T abstol::T @@ -64,4 +67,4 @@ end function QuadratureTraining(; quadrature_alg = CubatureJLh(), reltol = 1e-6, abstol = 1e-3, maxiters = 1_000, batch = 100) QuadratureTraining(quadrature_alg, reltol, abstol, maxiters, batch) -end \ No newline at end of file +end diff --git a/src/transform_inf_integral.jl b/src/transform_inf_integral.jl index 1749fe84cf..d0afb3c91a 100644 --- a/src/transform_inf_integral.jl +++ b/src/transform_inf_integral.jl @@ -1,5 +1,6 @@ -function transform_inf_expr(integrating_depvars, dict_depvar_input, dict_depvars, integrating_variables, transform) +function transform_inf_expr(integrating_depvars, dict_depvar_input, dict_depvars, + integrating_variables, transform) τs = Symbolics.variables(:τ, 1:length(integrating_variables)) τs = Symbol.(τs) dict_transformation_vars = Dict() @@ -25,17 +26,18 @@ function transform_inf_expr(integrating_depvars, dict_depvar_input, dict_depvars dict_depvar_input_[depvar] = ans end - this_eq_pair = Dict(map(intvars -> dict_depvars[intvars] => dict_depvar_input_[intvars], integrating_depvars)) + this_eq_pair = Dict(map(intvars -> dict_depvars[intvars] => dict_depvar_input_[intvars], + integrating_depvars)) this_eq_indvars = unique(vcat(values(this_eq_pair)...)) - return dict_transformation_vars, this_eq_indvars,integrating_var_transformation + return dict_transformation_vars, this_eq_indvars, integrating_var_transformation end function v_inf(t) - return :($t ./ (1 .- $t.^2)) + return :($t ./ (1 .- $t .^ 2)) end -function v_semiinf(t , a , upto_inf) +function v_semiinf(t, a, upto_inf) if a isa Num if upto_inf == true return :($t ./ (1 .- $t)) @@ -51,61 +53,71 @@ function v_semiinf(t , a , upto_inf) end end -function get_inf_transformation_jacobian(integrating_variable, _inf, _semiup, _semilw, _num_semiup, _num_semilw) +function get_inf_transformation_jacobian(integrating_variable, _inf, _semiup, _semilw, + _num_semiup, _num_semilw) j = [] - for var in integrating_variable - if _inf[1] - append!(j, [:((1+$var^2)/(1-$var^2)^2)]) - elseif _semiup[1] || _num_semiup[1] - append!(j, [:(1/(1-$var)^2)]) - elseif _semilw[1] || _num_semilw[1] - append!(j, [:(1/(1+$var)^2)]) - end + for var in integrating_variable + if _inf[1] + append!(j, [:((1 + $var^2) / (1 - $var^2)^2)]) + elseif _semiup[1] || _num_semiup[1] + append!(j, [:(1 / (1 - $var)^2)]) + elseif _semilw[1] || _num_semilw[1] + append!(j, [:(1 / (1 + $var)^2)]) end + end return j end -function transform_inf_integral(lb, ub, integrating_ex, integrating_depvars, dict_depvar_input, dict_depvars, integrating_variable, eltypeθ; dict_transformation_vars = nothing, transformation_vars = nothing) +function transform_inf_integral(lb, ub, integrating_ex, integrating_depvars, + dict_depvar_input, dict_depvars, integrating_variable, + eltypeθ; dict_transformation_vars = nothing, + transformation_vars = nothing) lb_ = Symbolics.tosymbol.(lb) ub_ = Symbolics.tosymbol.(ub) if -Inf in lb_ || Inf in ub_ - if !(integrating_variable isa Array) integrating_variable = [integrating_variable] end lbb = lb_ .=== -Inf ubb = ub_ .=== Inf - _num_semiup = isa.(lb_,Symbol) - _num_semilw = isa.(ub_,Symbol) + _num_semiup = isa.(lb_, Symbol) + _num_semilw = isa.(ub_, Symbol) _none = .!lbb .& .!ubb _inf = lbb .& ubb _semiup = .!lbb .& ubb .& .!_num_semiup - _semilw = lbb .& .!ubb .& .!_num_semilw - + _semilw = lbb .& .!ubb .& .!_num_semilw + function transform_indvars(t, i) if _none[1] return t elseif _inf[1] return v_inf(t) elseif _semiup[1] || _num_semiup[1] - return v_semiinf(t , lb[i] , 1) + return v_semiinf(t, lb[i], 1) elseif _semilw[1] || _num_semilw[1] - return v_semiinf(t , ub[i] , 0) + return v_semiinf(t, ub[i], 0) end end - dict_transformation_vars, transformation_vars, integrating_var_transformation = transform_inf_expr(integrating_depvars, dict_depvar_input, dict_depvars, integrating_variable,transform_indvars) + dict_transformation_vars, transformation_vars, integrating_var_transformation = transform_inf_expr(integrating_depvars, + dict_depvar_input, + dict_depvars, + integrating_variable, + transform_indvars) + + ϵ = 1 / 20 #cbrt(eps(eltypeθ)) - ϵ = 1/20 #cbrt(eps(eltypeθ)) + lb = 0.00 .* _semiup + (-1.00 + ϵ) .* _inf + (-1.00 + ϵ) .* _semilw + _none .* lb + + lb ./ (1 .+ lb) .* _num_semiup + (-1.00 + ϵ) .* _num_semilw + ub = (1.00 - ϵ) .* _semiup + (1.00 - ϵ) .* _inf + 0.00 .* _semilw + _none .* ub + + (1.00 - ϵ) .* _num_semiup + ub ./ (1 .+ ub) .* _num_semilw - lb = 0.00.*_semiup + (-1.00+ϵ).*_inf + (-1.00+ϵ).*_semilw + _none.*lb + lb./(1 .+ lb).*_num_semiup + (-1.00+ϵ).*_num_semilw - ub = (1.00-ϵ).*_semiup + (1.00-ϵ).*_inf + 0.00.*_semilw + _none.*ub + (1.00-ϵ).*_num_semiup + ub./(1 .+ ub).*_num_semilw + j = get_inf_transformation_jacobian(integrating_var_transformation, _inf, _semiup, + _semilw, _num_semiup, _num_semilw) - j = get_inf_transformation_jacobian(integrating_var_transformation, _inf, _semiup, _semilw, _num_semiup, _num_semilw) - integrating_ex = Expr(:call, :*, integrating_ex, j...) end diff --git a/test/IDE_tests.jl b/test/IDE_tests.jl index bcf94a23b9..f0923d6b9b 100644 --- a/test/IDE_tests.jl +++ b/test/IDE_tests.jl @@ -8,7 +8,7 @@ using DomainSets using Random Random.seed!(100) -callback = function (p,l) +callback = function (p, l) println("Current loss is: $l") return false end @@ -19,33 +19,31 @@ println("Integral Tests") @variables i(..) Di = Differential(t) Ii = Integral(t in DomainSets.ClosedInterval(0, t)) -eq = Di(i(t)) + 2*i(t) + 5*Ii(i(t)) ~ 1 -bcs = [i(0.) ~ 0.0] -domains = [t ∈ Interval(0.0,2.0)] -chain = Chain(Dense(1,15,Flux.σ),Dense(15,1)) +eq = Di(i(t)) + 2 * i(t) + 5 * Ii(i(t)) ~ 1 +bcs = [i(0.0) ~ 0.0] +domains = [t ∈ Interval(0.0, 2.0)] +chain = Chain(Dense(1, 15, Flux.σ), Dense(15, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy_ = NeuralPDE.GridTraining(0.1) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = nothing, phi = nothing, - derivative = nothing, - ) -@named pde_system = PDESystem(eq,bcs,domains,[t],[i(t)]) + derivative = nothing) +@named pde_system = PDESystem(eq, bcs, domains, [t], [i(t)]) sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) -prob = NeuralPDE.discretize(pde_system,discretization) -res = Optimization.solve(prob, BFGS(); callback = callback, maxiters=100) +prob = NeuralPDE.discretize(pde_system, discretization) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 100) ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] phi = discretization.phi -analytic_sol_func(t) = 1/2*(exp(-t))*(sin(2*t)) -u_real = [analytic_sol_func(t) for t in ts] -u_predict = [first(phi([t],res.minimizer)) for t in ts] +analytic_sol_func(t) = 1 / 2 * (exp(-t)) * (sin(2 * t)) +u_real = [analytic_sol_func(t) for t in ts] +u_predict = [first(phi([t], res.minimizer)) for t in ts] @test Flux.mse(u_real, u_predict) < 0.001 # plot(ts,u_real) # plot!(ts,u_predict) - ## Simple Integral Test println("Simple Integral Test") @@ -53,27 +51,26 @@ println("Simple Integral Test") @variables u(..) Ix = Integral(x in DomainSets.ClosedInterval(0, x)) # eq = Ix(u(x)) ~ (x^3)/3 -eq = Ix(u(x)*cos(x))~ (x^3)/3 +eq = Ix(u(x) * cos(x)) ~ (x^3) / 3 -bcs = [u(0.) ~ 0.0] -domains = [x ∈ Interval(0.0,1.00)] +bcs = [u(0.0) ~ 0.0] +domains = [x ∈ Interval(0.0, 1.00)] # chain = Chain(Dense(1,15,Flux.σ),Dense(15,1)) -chain = FastChain(FastDense(1,15,Flux.σ),FastDense(15,1)) +chain = FastChain(FastDense(1, 15, Flux.σ), FastDense(15, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy_ = NeuralPDE.GridTraining(0.1) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = initθ, phi = nothing, - derivative = nothing, - ) -@named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)]) -prob = NeuralPDE.discretize(pde_system,discretization) -res = Optimization.solve(prob, BFGS(); callback = callback, maxiters=200) + derivative = nothing) +@named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)]) +prob = NeuralPDE.discretize(pde_system, discretization) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 200) xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] phi = discretization.phi -u_predict = [first(phi([x],res.minimizer)) for x in xs] -u_real = [x^2/cos(x) for x in xs] +u_predict = [first(phi([x], res.minimizer)) for x in xs] +u_real = [x^2 / cos(x) for x in xs] @test Flux.mse(u_real, u_predict) < 0.001 # plot(xs,u_real) @@ -82,32 +79,31 @@ u_real = [x^2/cos(x) for x in xs] #simple multidimensitonal integral test println("simple multidimensitonal integral test") -@parameters x,y +@parameters x, y @variables u(..) Dx = Differential(x) Dy = Differential(y) -Ix = Integral((x,y) in DomainSets.UnitSquare()) -eq = Ix(u(x,y)) ~ 1/3 -bcs = [u(0., 0.) ~ 1, Dx(u(x,y)) ~ -2*x , Dy(u(x ,y)) ~ -2*y ] -domains = [x ∈ Interval(0.0,1.00), y ∈ Interval(0.0,1.00)] -chain = Chain(Dense(2,15,Flux.σ),Dense(15,1)) +Ix = Integral((x, y) in DomainSets.UnitSquare()) +eq = Ix(u(x, y)) ~ 1 / 3 +bcs = [u(0.0, 0.0) ~ 1, Dx(u(x, y)) ~ -2 * x, Dy(u(x, y)) ~ -2 * y] +domains = [x ∈ Interval(0.0, 1.00), y ∈ Interval(0.0, 1.00)] +chain = Chain(Dense(2, 15, Flux.σ), Dense(15, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy_ = NeuralPDE.GridTraining(0.1) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = nothing, phi = nothing, - derivative = nothing, - ) -@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)]) -prob = NeuralPDE.discretize(pde_system,discretization) -res = Optimization.solve(prob, BFGS(); callback = callback, maxiters=100) + derivative = nothing) +@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) +prob = NeuralPDE.discretize(pde_system, discretization) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 100) xs = 0.00:0.01:1.00 ys = 0.00:0.01:1.00 phi = discretization.phi u_real = collect(1 - x^2 - y^2 for y in ys, x in xs); -u_predict = collect(Array(phi([x,y], res.minimizer))[1] for y in ys, x in xs); +u_predict = collect(Array(phi([x, y], res.minimizer))[1] for y in ys, x in xs); @test Flux.mse(u_real, u_predict) < 0.001 # error_ = u_predict .- u_real @@ -116,32 +112,31 @@ u_predict = collect(Array(phi([x,y], res.minimizer))[1] for y in ys, x in xs); # p3 = plot(xs,ys,error_,linetype=:contourf,label = "error") # plot(p1,p2,p3) -@parameters x,y +@parameters x, y @variables u(..) Dx = Differential(x) Dy = Differential(y) -Ix = Integral((x,y) in DomainSets.ProductDomain(UnitInterval(),ClosedInterval(0 ,x))) -eq = Ix(u(x,y)) ~ 5/12 -bcs = [u(0., 0.) ~ 0, Dy(u(x,y)) ~ 2*y , u(x, 0) ~ x ] -domains = [x ∈ Interval(0.0,1.00), y ∈ Interval(0.0,1.00)] -chain = Chain(Dense(2,15,Flux.σ),Dense(15,1)) +Ix = Integral((x, y) in DomainSets.ProductDomain(UnitInterval(), ClosedInterval(0, x))) +eq = Ix(u(x, y)) ~ 5 / 12 +bcs = [u(0.0, 0.0) ~ 0, Dy(u(x, y)) ~ 2 * y, u(x, 0) ~ x] +domains = [x ∈ Interval(0.0, 1.00), y ∈ Interval(0.0, 1.00)] +chain = Chain(Dense(2, 15, Flux.σ), Dense(15, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy_ = NeuralPDE.GridTraining(0.1) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = nothing, phi = nothing, - derivative = nothing, - ) -@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)]) -prob = NeuralPDE.discretize(pde_system,discretization) -res = Optimization.solve(prob, BFGS(); callback = callback, maxiters=100) + derivative = nothing) +@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) +prob = NeuralPDE.discretize(pde_system, discretization) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 100) xs = 0.00:0.01:1.00 ys = 0.00:0.01:1.00 phi = discretization.phi -u_real = collect( x + y^2 for y in ys, x in xs); -u_predict = collect(Array(phi([x,y], res.minimizer))[1] for y in ys, x in xs); +u_real = collect(x + y^2 for y in ys, x in xs); +u_predict = collect(Array(phi([x, y], res.minimizer))[1] for y in ys, x in xs); @test Flux.mse(u_real, u_predict) < 0.01 # error_ = u_predict .- u_real @@ -150,7 +145,6 @@ u_predict = collect(Array(phi([x,y], res.minimizer))[1] for y in ys, x in xs); # p3 = plot(xs,ys,error_,linetype=:contourf,label = "error") # plot(p1,p2,p3) - ## Two variables Integral Test println("Two variables Integral Test") @@ -159,34 +153,33 @@ println("Two variables Integral Test") Dx = Differential(x) Ix = Integral(x in DomainSets.ClosedInterval(1, x)) -eqs = [Ix(u(x)*w(x)) ~ log(abs(x)), - Dx(w(x)) ~ -2/(x^3), - u(x) ~ x ] +eqs = [Ix(u(x) * w(x)) ~ log(abs(x)), + Dx(w(x)) ~ -2 / (x^3), + u(x) ~ x] -bcs = [u(1.) ~ 1.0, w(1.) ~ 1.0] -domains = [x ∈ Interval(1.0,2.0)] +bcs = [u(1.0) ~ 1.0, w(1.0) ~ 1.0] +domains = [x ∈ Interval(1.0, 2.0)] -chains = [FastChain(FastDense(1,15,Flux.σ),FastDense(15,1)) for _ in 1:2] -initθ = map(chain -> Float64.(DiffEqFlux.initial_params(chain)),chains) +chains = [FastChain(FastDense(1, 15, Flux.σ), FastDense(15, 1)) for _ in 1:2] +initθ = map(chain -> Float64.(DiffEqFlux.initial_params(chain)), chains) strategy_ = NeuralPDE.GridTraining(0.1) discretization = NeuralPDE.PhysicsInformedNN(chains, strategy_; - init_params = initθ - ) -@named pde_system = PDESystem(eqs,bcs,domains,[x],[u(x), w(x)]) -prob = NeuralPDE.discretize(pde_system,discretization) -res = Optimization.solve(prob, BFGS(); callback = callback, maxiters=200) + init_params = initθ) +@named pde_system = PDESystem(eqs, bcs, domains, [x], [u(x), w(x)]) +prob = NeuralPDE.discretize(pde_system, discretization) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 200) xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] phi = discretization.phi initθ = discretization.init_params -acum = [0;accumulate(+, length.(initθ))] -sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1] +acum = [0; accumulate(+, length.(initθ))] +sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)] minimizers = [res.minimizer[s] for s in sep] -u_predict = [(phi[1]([x],minimizers[1]))[1] for x in xs] -w_predict = [(phi[2]([x],minimizers[2]))[1] for x in xs] -u_real = [x for x in xs] -w_real = [1/x^2 for x in xs] +u_predict = [(phi[1]([x], minimizers[1]))[1] for x in xs] +w_predict = [(phi[2]([x], minimizers[2]))[1] for x in xs] +u_real = [x for x in xs] +w_real = [1 / x^2 for x in xs] @test Flux.mse(u_real, u_predict) < 0.001 @test Flux.mse(w_real, w_predict) < 0.001 @@ -201,21 +194,22 @@ println("Infinity Integral Test") @variables u(..) I = Integral(x in ClosedInterval(1, x)) Iinf = Integral(x in ClosedInterval(1, Inf)) -eqs = [I(u(x)) ~ Iinf(u(x)) - 1/x] +eqs = [I(u(x)) ~ Iinf(u(x)) - 1 / x] bcs = [u(1) ~ 1] domains = [x ∈ Interval(1.0, 2.0)] chain = FastChain(FastDense(1, 10, Flux.σ), FastDense(10, 1)) initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain)) -discretization = NeuralPDE.PhysicsInformedNN(chain, NeuralPDE.GridTraining(0.1), init_params= initθ) +discretization = NeuralPDE.PhysicsInformedNN(chain, NeuralPDE.GridTraining(0.1), + init_params = initθ) @named pde_system = PDESystem(eqs, bcs, domains, [x], [u(x)]) sym_prob = SciMLBase.symbolic_discretize(pde_system, discretization) prob = SciMLBase.discretize(pde_system, discretization) -res = Optimization.solve(prob, BFGS(); callback =callback, maxiters=200) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 200) xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] phi = discretization.phi -u_predict = [first(phi([x],res.minimizer)) for x in xs] -u_real = [1/x^2 for x in xs] -@test u_real ≈ u_predict rtol = 10^-2 +u_predict = [first(phi([x], res.minimizer)) for x in xs] +u_real = [1 / x^2 for x in xs] +@test u_real≈u_predict rtol=10^-2 # plot(xs,u_real) # plot!(xs,u_predict) @@ -224,21 +218,22 @@ println("Infinity Integral equation Test") @parameters x @variables u(..) I = Integral(x in ClosedInterval(x, Inf)) -eq = I(u(x)) ~ 1/x +eq = I(u(x)) ~ 1 / x domains = [x ∈ Interval(1.0, 2.0)] bcs = [u(1) ~ 1] -chain = FastChain(FastDense(1, 12, Flux.tanh),FastDense(12, 1)) +chain = FastChain(FastDense(1, 12, Flux.tanh), FastDense(12, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) -discretization = NeuralPDE.PhysicsInformedNN(chain, NeuralPDE.GridTraining(0.1), init_params= initθ) +discretization = NeuralPDE.PhysicsInformedNN(chain, NeuralPDE.GridTraining(0.1), + init_params = initθ) @named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)]) sym_prob = SciMLBase.symbolic_discretize(pde_system, discretization) prob = SciMLBase.discretize(pde_system, discretization) prob.f(initθ, nothing) -res = Optimization.solve(prob, BFGS(); callback =callback, maxiters=300) +res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 300) xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] phi = discretization.phi -u_predict = [first(phi([x],res.minimizer)) for x in xs] -u_real = [1/x^2 for x in xs] -@test u_real ≈ u_predict rtol = 10^-2 +u_predict = [first(phi([x], res.minimizer)) for x in xs] +u_real = [1 / x^2 for x in xs] +@test u_real≈u_predict rtol=10^-2 # plot(xs,u_real) # plot!(xs,u_predict) diff --git a/test/NNKolmogorov_tests.jl b/test/NNKolmogorov_tests.jl index d5c2a076d5..12a41a2946 100644 --- a/test/NNKolmogorov_tests.jl +++ b/test/NNKolmogorov_tests.jl @@ -8,19 +8,19 @@ Random.seed!(100) #Using SDEProblem for the Algorithm. # For a diract delta take u0 = Normal(0 , sigma) where sigma --> 0 -u0 = Normal(1.00 , 1.00) -xspan = (-2.0 , 6.0) -tspan = (0.0 , 1.0) -g(u , p , t) = 2.00 -f(u , p , t) = -2.00 +u0 = Normal(1.00, 1.00) +xspan = (-2.0, 6.0) +tspan = (0.0, 1.0) +g(u, p, t) = 2.00 +f(u, p, t) = -2.00 d = 1 sdealg = EM() -prob = SDEProblem(f , g , u0 , (0.0 , 1.0) ; xspan = xspan , d = d) +prob = SDEProblem(f, g, u0, (0.0, 1.0); xspan = xspan, d = d) opt = Flux.ADAM(0.01) -m = Chain(Dense(1, 5, elu),Dense(5, 5, elu) , Dense(5 , 5 , elu) , Dense(5 , 1)) +m = Chain(Dense(1, 5, elu), Dense(5, 5, elu), Dense(5, 5, elu), Dense(5, 1)) ensemblealg = EnsembleThreads() -sol = solve(prob, NNKolmogorov(m,opt , sdealg,ensemblealg) , verbose = true, dt = 0.01, - abstol=1e-10, dx = 0.0001 , trajectories = 100000 , maxiters = 500) +sol = solve(prob, NNKolmogorov(m, opt, sdealg, ensemblealg), verbose = true, dt = 0.01, + abstol = 1e-10, dx = 0.0001, trajectories = 100000, maxiters = 500) # using Plots # @@ -37,11 +37,11 @@ sol = solve(prob, NNKolmogorov(m,opt , sdealg,ensemblealg) , verbose = true, dt # ## The solution is obtained taking the Fourier Transform. -analytical(xi) = pdf.(Normal(3 , sqrt(1.0 + 5.00)) , xi) +analytical(xi) = pdf.(Normal(3, sqrt(1.0 + 5.00)), xi) ##Validation xs = -5:0.00001:5 -x_1 = rand(xs , 1 , 1000) -err_l2 = Flux.mse(analytical(x_1) , m(x_1)) +x_1 = rand(xs, 1, 1000) +err_l2 = Flux.mse(analytical(x_1), m(x_1)) @test err_l2 < 0.01 ## @@ -49,64 +49,65 @@ err_l2 = Flux.mse(analytical(x_1) , m(x_1)) function phi(xi) y = Float64[] for x in xi - y = push!(y , 1.77*x -0.015*x^3) + y = push!(y, 1.77 * x - 0.015 * x^3) end - y = reshape(y , size(xi)[1] , size(xi)[2] ) + y = reshape(y, size(xi)[1], size(xi)[2]) return y end -xspan2 = (-6.0 , 6.0) -tspan2 = (0.0 , 1.0) +xspan2 = (-6.0, 6.0) +tspan2 = (0.0, 1.0) #f = mu and g = sigma -g2(u , p , t) = 0.5*u -f2(u , p , t) = 0.5*0.25*u +g2(u, p, t) = 0.5 * u +f2(u, p, t) = 0.5 * 0.25 * u d2 = 1 sdealg2 = EM() -prob2 = KolmogorovPDEProblem(f2 , g2, phi , xspan2 , tspan2, d2) +prob2 = KolmogorovPDEProblem(f2, g2, phi, xspan2, tspan2, d2) opt2 = Flux.ADAM(0.01) -m2 = Chain(Dense(1, 16, elu) , Dense(16 , 32 , elu),Dense(32 , 16 , elu), Dense(16 , 1)) -sol = solve(prob2, NeuralPDE.NNKolmogorov(m2,opt2 , sdealg2, ensemblealg), verbose = true, dt = 0.01, - dx = 0.0001 , trajectories = 1000 , abstol=1e-6, maxiters = 300) - +m2 = Chain(Dense(1, 16, elu), Dense(16, 32, elu), Dense(32, 16, elu), Dense(16, 1)) +sol = solve(prob2, NeuralPDE.NNKolmogorov(m2, opt2, sdealg2, ensemblealg), verbose = true, + dt = 0.01, + dx = 0.0001, trajectories = 1000, abstol = 1e-6, maxiters = 300) function analytical2(xi) y = Float64[] - a = 1.77*exp(0.5*(0.5)^2*1.0) - b = -0.015*exp(0.5*(0.5*3)^2*1.0) + a = 1.77 * exp(0.5 * (0.5)^2 * 1.0) + b = -0.015 * exp(0.5 * (0.5 * 3)^2 * 1.0) for x in xi - y = push!(y , a*x + b*x^3) + y = push!(y, a * x + b * x^3) end - y = reshape(y , size(xi)[1] , size(xi)[2] ) + y = reshape(y, size(xi)[1], size(xi)[2]) return y end xs2 = -5.00:0.01:5.00 -x_val2 = rand(xs2 , d2 , 50) -errorl2 = Flux.mse(analytical2(x_val2) , m2(x_val2)) +x_val2 = rand(xs2, d2, 50) +errorl2 = Flux.mse(analytical2(x_val2), m2(x_val2)) println("error_l2 = ", errorl2, "\n") @test errorl2 < 0.4 - ##Non-Diagonal Test -f_noise = (du,u,p,t) -> du.=1.01u -g_noise = function (du,u,p,t) - du[1,1] = 0.3u[1] - du[1,2] = 0.6u[1] - du[1,3] = 0.9u[1] - du[1,4] = 0.12u[2] - du[2,1] = 1.2u[1] - du[2,2] = 0.2u[2] - du[2,3] = 0.3u[2] - du[2,4] = 1.8u[2] +f_noise = (du, u, p, t) -> du .= 1.01u +g_noise = function (du, u, p, t) + du[1, 1] = 0.3u[1] + du[1, 2] = 0.6u[1] + du[1, 3] = 0.9u[1] + du[1, 4] = 0.12u[2] + du[2, 1] = 1.2u[1] + du[2, 2] = 0.2u[2] + du[2, 3] = 0.3u[2] + du[2, 4] = 1.8u[2] end -Σ = [1.0 0.3 ; 0.3 1.0] -uo3 = MvNormal([0.0 ; 0.0], Σ) -sdealg3= EM() -xspan3 = (-10.0 , 10.0) -tspan3 = (0.0 , 1.0) +Σ = [1.0 0.3; 0.3 1.0] +uo3 = MvNormal([0.0; 0.0], Σ) +sdealg3 = EM() +xspan3 = (-10.0, 10.0) +tspan3 = (0.0, 1.0) d3 = 2 -prob = SDEProblem(f_noise , g_noise , uo3 , (0.0 , 1.0) ; xspan = xspan3 , d = d3 , noise_rate_prototype=zeros(2,4)) +prob = SDEProblem(f_noise, g_noise, uo3, (0.0, 1.0); xspan = xspan3, d = d3, + noise_rate_prototype = zeros(2, 4)) opt = Flux.ADAM(0.01) -m3 = Chain(Dense(d3, 32, elu) ,Dense(32 , 64 , elu), Dense(64 , 1)) -sol3 = solve(prob, NeuralPDE.NNKolmogorov(m3,opt , sdealg3 , EnsembleThreads()), verbose = true, dt = 0.001, - abstol=1e-6, dx = 0.001, trajectories = 1000,maxiters = 200) +m3 = Chain(Dense(d3, 32, elu), Dense(32, 64, elu), Dense(64, 1)) +sol3 = solve(prob, NeuralPDE.NNKolmogorov(m3, opt, sdealg3, EnsembleThreads()), + verbose = true, dt = 0.001, + abstol = 1e-6, dx = 0.001, trajectories = 1000, maxiters = 200) println("Non-Diagonal test working.") diff --git a/test/NNODE_tests.jl b/test/NNODE_tests.jl index b9958084ff..7754cb216a 100644 --- a/test/NNODE_tests.jl +++ b/test/NNODE_tests.jl @@ -1,70 +1,76 @@ -using Test, Flux, Optim -using Random, NeuralPDE -Random.seed!(100) - -# Run a solve on scalars -linear = (u,p,t) -> cos(2pi*t) -tspan = (0.0f0, 1.0f0) -u0 = 0.0f0 -prob = ODEProblem(linear, u0 ,tspan) -chain = Flux.Chain(Dense(1,5,σ),Dense(5,1)) -opt = Flux.ADAM(0.1, (0.9, 0.95)) -sol = solve(prob, NeuralPDE.NNODE(chain,opt), dt=1/20f0, verbose = true, - abstol=1f-10, maxiters = 200) - -sol = solve(prob, NeuralPDE.NNODE(chain,opt), verbose = true, - abstol=1f-6, maxiters = 200) - -opt = BFGS() -sol = solve(prob, NeuralPDE.NNODE(chain,opt), dt=1/20f0, verbose = true, - abstol=1f-10, maxiters = 200) - -sol = solve(prob, NeuralPDE.NNODE(chain,opt), verbose = true, - abstol=1f-6, maxiters = 200) - -# Run a solve on vectors -linear = (u,p,t) -> [cos(2pi*t)] -tspan = (0.0f0, 1.0f0) -u0 = [0.0f0] -prob = ODEProblem(linear, u0 ,tspan) -chain = Flux.Chain(Dense(1,5,σ),Dense(5,1)) -opt = BFGS() -sol = solve(prob, NeuralPDE.NNODE(chain,opt), dt=1/20f0, abstol=1e-10, - verbose = true, maxiters=200) - -sol = solve(prob, NeuralPDE.NNODE(chain,opt), abstol=1f-6, - verbose = true, maxiters=200) - -@test sol(0.5) isa Vector -@test sol(0.5; idxs = 1) isa Number -@test sol.k isa SciMLBase.OptimizationSolution - -#Example 1 -linear = (u,p,t) -> @. t^3 + 2*t + (t^2)*((1+3*(t^2))/(1+t+(t^3))) - u*(t + ((1+3*(t^2))/(1+t+t^3))) -linear_analytic = (u0,p,t) -> [exp(-(t^2)/2)/(1+t+t^3) + t^2] -prob = ODEProblem(ODEFunction(linear,analytic=linear_analytic),[1f0],(0.0f0,1.0f0)) -chain = Flux.Chain(Dense(1,128,σ),Dense(128,1)) -opt = ADAM(0.01) -sol = solve(prob,NeuralPDE.NNODE(chain,opt),verbose = true, maxiters=400) -@test sol.errors[:l2] < 0.5 - -sol = solve(prob,NeuralPDE.NNODE(chain,opt),verbose = true, maxiters=400, dt = 1/5f0) -@test sol.errors[:l2] < 0.5 - -sol = solve(prob,NeuralPDE.NNODE(chain,opt;strategy = StochasticTraining(100)),verbose = true, maxiters=400, dt = 1/5f0) -@test sol.errors[:l2] < 0.5 - -#Example 2 -linear = (u,p,t) -> -u/5 + exp(-t/5).*cos(t) -linear_analytic = (u0,p,t) -> exp(-t/5)*(u0 + sin(t)) -prob = ODEProblem(ODEFunction(linear,analytic=linear_analytic),0.0f0,(0.0f0,1.0f0)) -chain = Flux.Chain(Dense(1,5,σ),Dense(5,1)) -opt = ADAM(0.1) -sol = solve(prob,NeuralPDE.NNODE(chain,opt),verbose = true, maxiters=400, abstol=1f-8) -@test sol.errors[:l2] < 0.5 - -sol = solve(prob,NeuralPDE.NNODE(chain,opt),verbose = true, maxiters=400, abstol=1f-8, dt = 1/5f0) -@test sol.errors[:l2] < 0.5 - -sol = solve(prob,NeuralPDE.NNODE(chain,opt;strategy = StochasticTraining(100)),verbose = true, maxiters=400, abstol=1f-8, dt = 1/5f0) -@test sol.errors[:l2] < 0.5 \ No newline at end of file +using Test, Flux, Optim +using Random, NeuralPDE +Random.seed!(100) + +# Run a solve on scalars +linear = (u, p, t) -> cos(2pi * t) +tspan = (0.0f0, 1.0f0) +u0 = 0.0f0 +prob = ODEProblem(linear, u0, tspan) +chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) +opt = Flux.ADAM(0.1, (0.9, 0.95)) +sol = solve(prob, NeuralPDE.NNODE(chain, opt), dt = 1 / 20.0f0, verbose = true, + abstol = 1.0f-10, maxiters = 200) + +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, + abstol = 1.0f-6, maxiters = 200) + +opt = BFGS() +sol = solve(prob, NeuralPDE.NNODE(chain, opt), dt = 1 / 20.0f0, verbose = true, + abstol = 1.0f-10, maxiters = 200) + +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, + abstol = 1.0f-6, maxiters = 200) + +# Run a solve on vectors +linear = (u, p, t) -> [cos(2pi * t)] +tspan = (0.0f0, 1.0f0) +u0 = [0.0f0] +prob = ODEProblem(linear, u0, tspan) +chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) +opt = BFGS() +sol = solve(prob, NeuralPDE.NNODE(chain, opt), dt = 1 / 20.0f0, abstol = 1e-10, + verbose = true, maxiters = 200) + +sol = solve(prob, NeuralPDE.NNODE(chain, opt), abstol = 1.0f-6, + verbose = true, maxiters = 200) + +@test sol(0.5) isa Vector +@test sol(0.5; idxs = 1) isa Number +@test sol.k isa SciMLBase.OptimizationSolution + +#Example 1 +linear = (u, p, t) -> @. t^3 + 2 * t + (t^2) * ((1 + 3 * (t^2)) / (1 + t + (t^3))) - + u * (t + ((1 + 3 * (t^2)) / (1 + t + t^3))) +linear_analytic = (u0, p, t) -> [exp(-(t^2) / 2) / (1 + t + t^3) + t^2] +prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), [1.0f0], (0.0f0, 1.0f0)) +chain = Flux.Chain(Dense(1, 128, σ), Dense(128, 1)) +opt = ADAM(0.01) +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, maxiters = 400) +@test sol.errors[:l2] < 0.5 + +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, maxiters = 400, + dt = 1 / 5.0f0) +@test sol.errors[:l2] < 0.5 + +sol = solve(prob, NeuralPDE.NNODE(chain, opt; strategy = StochasticTraining(100)), + verbose = true, maxiters = 400, dt = 1 / 5.0f0) +@test sol.errors[:l2] < 0.5 + +#Example 2 +linear = (u, p, t) -> -u / 5 + exp(-t / 5) .* cos(t) +linear_analytic = (u0, p, t) -> exp(-t / 5) * (u0 + sin(t)) +prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), 0.0f0, (0.0f0, 1.0f0)) +chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) +opt = ADAM(0.1) +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, maxiters = 400, + abstol = 1.0f-8) +@test sol.errors[:l2] < 0.5 + +sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, maxiters = 400, + abstol = 1.0f-8, dt = 1 / 5.0f0) +@test sol.errors[:l2] < 0.5 + +sol = solve(prob, NeuralPDE.NNODE(chain, opt; strategy = StochasticTraining(100)), + verbose = true, maxiters = 400, abstol = 1.0f-8, dt = 1 / 5.0f0) +@test sol.errors[:l2] < 0.5 diff --git a/test/NNPDE_tests.jl b/test/NNPDE_tests.jl index 67b54fc40f..1c9d0d313a 100644 --- a/test/NNPDE_tests.jl +++ b/test/NNPDE_tests.jl @@ -8,7 +8,7 @@ using DomainSets using Random Random.seed!(100) -callback = function (p,l) +callback = function (p, l) println("Current loss is: $l") return false end @@ -21,42 +21,42 @@ function test_ode(strategy_) Dθ = Differential(θ) # 1D ODE - eq = Dθ(u(θ)) ~ θ^3 + 2*θ + (θ^2)*((1+3*(θ^2))/(1+θ+(θ^3))) - u(θ)*(θ + ((1+3*(θ^2))/(1+θ+θ^3))) + eq = Dθ(u(θ)) ~ θ^3 + 2 * θ + (θ^2) * ((1 + 3 * (θ^2)) / (1 + θ + (θ^3))) - + u(θ) * (θ + ((1 + 3 * (θ^2)) / (1 + θ + θ^3))) # Initial and boundary conditions - bcs = [u(0.) ~ 1.0] + bcs = [u(0.0) ~ 1.0] # Space and time domains - domains = [θ ∈ Interval(0.0,1.0)] + domains = [θ ∈ Interval(0.0, 1.0)] # Neural network - chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,1)) + chain = FastChain(FastDense(1, 12, Flux.σ), FastDense(12, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = initθ, phi = nothing, - derivative = nothing, - ) - - @named pde_system = PDESystem(eq,bcs,domains,[θ],[u]) - prob = NeuralPDE.discretize(pde_system,discretization) - sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) - - res = Optimization.solve(prob, ADAM(0.1); maxiters=1000) - prob = remake(prob,u0=res.minimizer) - res = Optimization.solve(prob, ADAM(0.01); maxiters=500) - prob = remake(prob,u0=res.minimizer) - res = Optimization.solve(prob, ADAM(0.001); maxiters=500) + derivative = nothing) + + @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) + prob = NeuralPDE.discretize(pde_system, discretization) + sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) + + res = Optimization.solve(prob, ADAM(0.1); maxiters = 1000) + prob = remake(prob, u0 = res.minimizer) + res = Optimization.solve(prob, ADAM(0.01); maxiters = 500) + prob = remake(prob, u0 = res.minimizer) + res = Optimization.solve(prob, ADAM(0.001); maxiters = 500) phi = discretization.phi - analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2 + analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] - u_real = [analytic_sol_func(t) for t in ts] - u_predict = [first(phi(t,res.minimizer)) for t in ts] + u_real = [analytic_sol_func(t) for t in ts] + u_predict = [first(phi(t, res.minimizer)) for t in ts] - @test u_predict ≈ u_real atol = 0.1 + @test u_predict≈u_real atol=0.1 # using Plots # t_plot = collect(ts) # plot(t_plot ,u_real) @@ -66,86 +66,99 @@ end #TODO There is little meaning in these tests without checking the correctness of the prediction. #TODO I propose to simply remove them. function test_heterogeneous_equation(strategy_) - println("Simple Heterogeneous input PDE, strategy: $(nameof(typeof(strategy_)))") - @parameters x y - @variables p(..) q(..) r(..) s(..) - Dx = Differential(x) - Dy = Differential(y) - - # 2D PDE - eq = p(x) + q(y) + Dx(r(x, y)) + Dy(s(y, x)) ~ 0 - # eq = Dx(p(x)) + Dy(q(y)) + Dx(r(x, y)) + Dy(s(y, x)) + p(x) + q(y) + r(x, y) + s(y, x) ~ 0 - - # Initial and boundary conditions - bcs = [p(1) ~ 0.f0, q(-1) ~ 0.0f0, - r(x, -1) ~ 0.f0, r(1, y) ~ 0.0f0, - s(y, 1) ~ 0.0f0, s(-1, x) ~ 0.0f0] - # bcs = [s(y, 1) ~ 0.0f0] - # Space and time domains - domains = [x ∈ Interval(0.0, 1.0), - y ∈ Interval(0.0, 1.0)] - - # chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) - numhid = 3 - fastchains = [[FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]; - [FastChain(FastDense(2, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]] - discretization = NeuralPDE.PhysicsInformedNN(fastchains, - strategy_) - - @named pde_system = PDESystem(eq, bcs, domains, [x,y], [p(x), q(y), r(x, y), s(y, x)]) - prob = SciMLBase.discretize(pde_system, discretization) - res = Optimization.solve(prob, BFGS(); maxiters=100) + println("Simple Heterogeneous input PDE, strategy: $(nameof(typeof(strategy_)))") + @parameters x y + @variables p(..) q(..) r(..) s(..) + Dx = Differential(x) + Dy = Differential(y) + + # 2D PDE + eq = p(x) + q(y) + Dx(r(x, y)) + Dy(s(y, x)) ~ 0 + # eq = Dx(p(x)) + Dy(q(y)) + Dx(r(x, y)) + Dy(s(y, x)) + p(x) + q(y) + r(x, y) + s(y, x) ~ 0 + + # Initial and boundary conditions + bcs = [p(1) ~ 0.0f0, q(-1) ~ 0.0f0, + r(x, -1) ~ 0.0f0, r(1, y) ~ 0.0f0, + s(y, 1) ~ 0.0f0, s(-1, x) ~ 0.0f0] + # bcs = [s(y, 1) ~ 0.0f0] + # Space and time domains + domains = [x ∈ Interval(0.0, 1.0), + y ∈ Interval(0.0, 1.0)] + + # chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) + numhid = 3 + fastchains = [[FastChain(FastDense(1, numhid, Flux.σ), + FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) + for i in 1:2] + [FastChain(FastDense(2, numhid, Flux.σ), + FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) + for i in 1:2]] + discretization = NeuralPDE.PhysicsInformedNN(fastchains, + strategy_) + + @named pde_system = PDESystem(eq, bcs, domains, [x, y], [p(x), q(y), r(x, y), s(y, x)]) + prob = SciMLBase.discretize(pde_system, discretization) + res = Optimization.solve(prob, BFGS(); maxiters = 100) end ## Heterogeneous system function test_heterogeneous_system(strategy_) - println("Heterogeneous input PDE with derivatives, strategy: $(nameof(typeof(strategy_)))") - @parameters x y - @variables p(..) q(..) - Dx = Differential(x) - Dy = Differential(y) - - # 2D PDE - #TODO Dx(q(y)) = 0 - #TODO so p(x) = 0, q = const is has only trivial solution - eq = p(x) + Dx(q(y)) ~ 0 - - # Initial and boundary conditions - bcs = [p(1) ~ 0.f0, q(-1) ~ 0.0f0] - - # Space and time domains - domains = [x ∈ Interval(0.0, 1.0), - y ∈ Interval(-1.0, 0.0)] - - # chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) - numhid = 3 - fastchains = [[FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]; - [FastChain(FastDense(2, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]] - discretization = NeuralPDE.PhysicsInformedNN(fastchains, - strategy_) - - @named pde_system = PDESystem(eq, bcs, domains, [x,y], [p(x), q(y)]) - prob = SciMLBase.discretize(pde_system, discretization) - res = Optimization.solve(prob, BFGS(); maxiters=100) + println("Heterogeneous input PDE with derivatives, strategy: $(nameof(typeof(strategy_)))") + @parameters x y + @variables p(..) q(..) + Dx = Differential(x) + Dy = Differential(y) + + # 2D PDE + #TODO Dx(q(y)) = 0 + #TODO so p(x) = 0, q = const is has only trivial solution + eq = p(x) + Dx(q(y)) ~ 0 + + # Initial and boundary conditions + bcs = [p(1) ~ 0.0f0, q(-1) ~ 0.0f0] + + # Space and time domains + domains = [x ∈ Interval(0.0, 1.0), + y ∈ Interval(-1.0, 0.0)] + + # chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) + numhid = 3 + fastchains = [[FastChain(FastDense(1, numhid, Flux.σ), + FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) + for i in 1:2] + [FastChain(FastDense(2, numhid, Flux.σ), + FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) + for i in 1:2]] + discretization = NeuralPDE.PhysicsInformedNN(fastchains, + strategy_) + + @named pde_system = PDESystem(eq, bcs, domains, [x, y], [p(x), q(y)]) + prob = SciMLBase.discretize(pde_system, discretization) + res = Optimization.solve(prob, BFGS(); maxiters = 100) end grid_strategy = NeuralPDE.GridTraining(0.1) -quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(), - reltol=1e-3,abstol=1e-3, - maxiters =50, batch=100) -stochastic_strategy = NeuralPDE.StochasticTraining(100; bcs_points= 50) +quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg = CubatureJLh(), + reltol = 1e-3, abstol = 1e-3, + maxiters = 50, batch = 100) +stochastic_strategy = NeuralPDE.StochasticTraining(100; bcs_points = 50) quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100; sampling_alg = LatinHypercubeSample(), - resampling =false, - minibatch = 100 - ) + resampling = false, + minibatch = 100) quasirandom_strategy_resampling = NeuralPDE.QuasiRandomTraining(100; - bcs_points= 50, - sampling_alg = LatticeRuleSample(), - resampling = true, - minibatch = 0) - -strategies = [grid_strategy,stochastic_strategy, quadrature_strategy,quasirandom_strategy,quasirandom_strategy_resampling] + bcs_points = 50, + sampling_alg = LatticeRuleSample(), + resampling = true, + minibatch = 0) + +strategies = [ + grid_strategy, + stochastic_strategy, + quadrature_strategy, + quasirandom_strategy, + quasirandom_strategy_resampling, +] map(strategies) do strategy_ test_ode(strategy_) @@ -160,82 +173,82 @@ end ## Heterogeneous system println("Heterogeneous system") -@parameters x,y,z +@parameters x, y, z @variables u(..), v(..), h(..), p(..) Dz = Differential(z) eqs = [ - u(x,y,z) ~ x+y+z, - v(y,x) ~ x^2 + y^2, + u(x, y, z) ~ x + y + z, + v(y, x) ~ x^2 + y^2, h(z) ~ cos(z), - p(x,z) ~ exp(x)*exp(z), - u(x,y,z) + v(y,x)*Dz(h(z)) - p(x,z) ~ x+y+z - (x^2+y^2)*sin(z) - exp(x)*exp(z) + p(x, z) ~ exp(x) * exp(z), + u(x, y, z) + v(y, x) * Dz(h(z)) - p(x, z) ~ x + y + z - (x^2 + y^2) * sin(z) - + exp(x) * exp(z), ] -bcs = [u(0,0,0) ~ 0.0] +bcs = [u(0, 0, 0) ~ 0.0] domains = [x ∈ Interval(0.0, 1.0), - y ∈ Interval(0.0, 1.0), - z ∈ Interval(0.0, 1.0)] + y ∈ Interval(0.0, 1.0), + z ∈ Interval(0.0, 1.0)] -chain = [FastChain(FastDense(3,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)), - FastChain(FastDense(2,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)), - FastChain(FastDense(1,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)), - FastChain(FastDense(2,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1))] +chain = [ + FastChain(FastDense(3, 12, Flux.tanh), FastDense(12, 12, Flux.tanh), FastDense(12, 1)), + FastChain(FastDense(2, 12, Flux.tanh), FastDense(12, 12, Flux.tanh), FastDense(12, 1)), + FastChain(FastDense(1, 12, Flux.tanh), FastDense(12, 12, Flux.tanh), FastDense(12, 1)), + FastChain(FastDense(2, 12, Flux.tanh), FastDense(12, 12, Flux.tanh), FastDense(12, 1))] initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain)) - grid_strategy = NeuralPDE.GridTraining(0.1) -quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(), - reltol=1e-3,abstol=1e-3, - maxiters =50, batch=100) +quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg = CubatureJLh(), + reltol = 1e-3, abstol = 1e-3, + maxiters = 50, batch = 100) -discretization = NeuralPDE.PhysicsInformedNN(chain,grid_strategy;init_params = initθ) +discretization = NeuralPDE.PhysicsInformedNN(chain, grid_strategy; init_params = initθ) -@named pde_system = PDESystem(eqs,bcs,domains,[x,y,z],[u(x,y,z),v(y,x),h(z),p(x,z)]) +@named pde_system = PDESystem(eqs, bcs, domains, [x, y, z], + [u(x, y, z), v(y, x), h(z), p(x, z)]) -prob = NeuralPDE.discretize(pde_system,discretization) -sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) +prob = NeuralPDE.discretize(pde_system, discretization) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) -callback = function (p,l) +callback = function (p, l) println("Current loss is: $l") return false end -res = Optimization.solve(prob, BFGS();maxiters=1000) +res = Optimization.solve(prob, BFGS(); maxiters = 1000) phi = discretization.phi -analytic_sol_func_ = -[ -(x,y,z) -> x+y+z , -(x,y) -> x^2 + y^2, -(z) -> cos(z), -(x,z) -> exp(x)*exp(z) +analytic_sol_func_ = [ + (x, y, z) -> x + y + z, + (x, y) -> x^2 + y^2, + (z) -> cos(z), + (x, z) -> exp(x) * exp(z), ] -xs,ys,zs = [infimum(d.domain):0.1:supremum(d.domain) for d in domains] +xs, ys, zs = [infimum(d.domain):0.1:supremum(d.domain) for d in domains] -u_real = [analytic_sol_func_[1](x,y,z) for x in xs for y in ys for z in zs] -v_real = [analytic_sol_func_[2](y,x) for y in ys for x in xs ] +u_real = [analytic_sol_func_[1](x, y, z) for x in xs for y in ys for z in zs] +v_real = [analytic_sol_func_[2](y, x) for y in ys for x in xs] h_real = [analytic_sol_func_[3](z) for z in zs] -p_real = [analytic_sol_func_[4](x,z) for x in xs for z in zs] +p_real = [analytic_sol_func_[4](x, z) for x in xs for z in zs] -real_ = [u_real,v_real,h_real,p_real] +real_ = [u_real, v_real, h_real, p_real] initθ = discretization.init_params -acum = [0;accumulate(+, length.(initθ))] -sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1] +acum = [0; accumulate(+, length.(initθ))] +sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)] minimizers = [res.minimizer[s] for s in sep] - -u_predict = [phi[1]([x,y,z],minimizers[1])[1] for x in xs for y in ys for z in zs] -v_predict = [phi[2]([y,x],minimizers[2])[1] for y in ys for x in xs ] -h_predict = [phi[3]([z],minimizers[3])[1] for z in zs] -p_predict = [phi[4]([x,z],minimizers[4])[1] for x in xs for z in zs] -predict = [u_predict,v_predict,h_predict,p_predict] +u_predict = [phi[1]([x, y, z], minimizers[1])[1] for x in xs for y in ys for z in zs] +v_predict = [phi[2]([y, x], minimizers[2])[1] for y in ys for x in xs] +h_predict = [phi[3]([z], minimizers[3])[1] for z in zs] +p_predict = [phi[4]([x, z], minimizers[4])[1] for x in xs for z in zs] +predict = [u_predict, v_predict, h_predict, p_predict] for i in 1:4 - @test predict[i] ≈ real_[i] rtol = 10^-3 + @test predict[i]≈real_[i] rtol=10^-3 end # x_plot = collect(xs) @@ -247,7 +260,6 @@ end # plot(x_plot,y_plot,u_real) # plot!(x_plot,y_plot,u_predict) - ## Example 2, 2D Poisson equation function test_2d_poisson_equation(chain_, strategy_) println("Example 2, 2D Poisson equation, chain: $(nameof(typeof(chain_))), strategy: $(nameof(typeof(strategy_)))") @@ -257,34 +269,36 @@ function test_2d_poisson_equation(chain_, strategy_) 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)) ~ -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)] + 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)] + domains = [x ∈ Interval(0.0, 1.0), + y ∈ Interval(0.0, 1.0)] initθ = Float64.(DiffEqFlux.initial_params(chain_)) discretization = NeuralPDE.PhysicsInformedNN(chain_, strategy_; init_params = initθ) - @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, ADAM(0.1); maxiters=500) + @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, ADAM(0.1); maxiters = 500) phi = discretization.phi - 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) + 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.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))) + 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 = abs.(u_predict .- u_real) - @test u_predict ≈ u_real atol = 2.0 + @test u_predict≈u_real atol=2.0 # p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic"); # p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict"); @@ -292,84 +306,88 @@ function test_2d_poisson_equation(chain_, strategy_) # plot(p1,p2,p3) end -fastchain = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) -fluxchain = Chain(Dense(2,12,Flux.σ),Dense(12,12,Flux.σ),Dense(12,1)) |> f64 +fastchain = FastChain(FastDense(2, 12, Flux.σ), FastDense(12, 12, Flux.σ), FastDense(12, 1)) +fluxchain = Chain(Dense(2, 12, Flux.σ), Dense(12, 12, Flux.σ), Dense(12, 1)) |> f64 chains = [fluxchain, fastchain] for chain in chains test_2d_poisson_equation(chain, grid_strategy) end for strategy_ in strategies - chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) + chain_ = FastChain(FastDense(2, 12, Flux.σ), FastDense(12, 12, Flux.σ), + FastDense(12, 1)) test_2d_poisson_equation(chain_, strategy_) end algs = [CubatureJLp()] #CubatureJLh(), for alg in algs - chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) - strategy_ = NeuralPDE.QuadratureTraining(quadrature_alg = alg,reltol=1e-4,abstol=1e-3,maxiters=30, batch=10) + chain_ = FastChain(FastDense(2, 12, Flux.σ), FastDense(12, 12, Flux.σ), + FastDense(12, 1)) + strategy_ = NeuralPDE.QuadratureTraining(quadrature_alg = alg, reltol = 1e-4, + abstol = 1e-3, maxiters = 30, batch = 10) test_2d_poisson_equation(chain_, strategy_) end - ## Example 3, 3rd-order println("Example 3, 3rd-order ode") @parameters x -@variables u(..) ,Dxu(..) ,Dxxu(..) ,O1(..), O2(..) +@variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..) Dxxx = Differential(x)^3 Dx = Differential(x) # ODE -eq = Dx(Dxxu(x)) ~ cos(pi*x) +eq = Dx(Dxxu(x)) ~ cos(pi * x) # Initial and boundary conditions -bcs_ = [u(0.) ~ 0.0, - u(1.) ~ cos(pi), - Dxu(1.) ~ 1.0] -ep = (cbrt(eps(eltype(Float64))))^2/6 +bcs_ = [u(0.0) ~ 0.0, + u(1.0) ~ cos(pi), + Dxu(1.0) ~ 1.0] +ep = (cbrt(eps(eltype(Float64))))^2 / 6 -der = [Dxu(x) ~ Dx(u(x)) + ep*O1(x) , - Dxxu(x) ~ Dx(Dxu(x))+ ep*O2(x)] +der = [Dxu(x) ~ Dx(u(x)) + ep * O1(x), + Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x)] -bcs = [bcs_;der] +bcs = [bcs_; der] # Space and time domains -domains = [x ∈ Interval(0.0,1.0)] +domains = [x ∈ Interval(0.0, 1.0)] # Neural network -chain = [[FastChain(FastDense(1,12,Flux.tanh),FastDense(12,12,Flux.tanh),FastDense(12,1)) for _ in 1:3]; - [FastChain(FastDense(1,4,Flux.tanh),FastDense(4,1)) for _ in 1:2];] +chain = [[FastChain(FastDense(1, 12, Flux.tanh), FastDense(12, 12, Flux.tanh), + FastDense(12, 1)) for _ in 1:3] + [FastChain(FastDense(1, 4, Flux.tanh), FastDense(4, 1)) for _ in 1:2]] quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100; #points sampling_alg = LatinHypercubeSample()) initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain)) -discretization = NeuralPDE.PhysicsInformedNN(chain,quasirandom_strategy;init_params = initθ) +discretization = NeuralPDE.PhysicsInformedNN(chain, quasirandom_strategy; + init_params = initθ) -@named pde_system = PDESystem(eq,bcs,domains,[x],[u(x),Dxu(x),Dxxu(x),O1(x),O2(x)]) +@named pde_system = PDESystem(eq, bcs, domains, [x], [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)]) -prob = NeuralPDE.discretize(pde_system,discretization) -sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) -discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) +prob = NeuralPDE.discretize(pde_system, discretization) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) +discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) pde_inner_loss_functions = discretized_functions.pde_loss_functions bcs_inner_loss_functions = discretized_functions.bc_loss_functions -cb_ = function (p,l) - println("loss: ", l ) +cb_ = function (p, l) + println("loss: ", l) println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions)) println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions)) return false end -res = Optimization.solve(prob, BFGS(); maxiters=1000) +res = Optimization.solve(prob, BFGS(); maxiters = 1000) phi = discretization.phi[1] -analytic_sol_func(x) = (π*x*(-x+(π^2)*(2*x-3)+1)-sin(π*x))/(π^3) +analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sin(π * x)) / (π^3) xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1] -u_real = [analytic_sol_func(x) for x in xs] -u_predict = [first(phi(x,res.minimizer)) for x in xs] +u_real = [analytic_sol_func(x) for x in xs] +u_predict = [first(phi(x, res.minimizer)) for x in xs] -@test u_predict ≈ u_real atol = 10^-4 +@test u_predict≈u_real atol=10^-4 # x_plot = collect(xs) # plot(x_plot ,u_real) @@ -383,49 +401,49 @@ Dx = Differential(x) Dy = Differential(y) # System of pde -eqs = [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0, - Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0] - # 3*u1(x,0) ~ 2*u2(x,0)] +eqs = [Dx(u1(x, y)) + 4 * Dy(u2(x, y)) ~ 0, + Dx(u2(x, y)) + 9 * Dy(u1(x, y)) ~ 0] +# 3*u1(x,0) ~ 2*u2(x,0)] # Initial and boundary conditions -bcs = [u1(x,0) ~ 2*x, u2(x,0) ~ 3*x] +bcs = [u1(x, 0) ~ 2 * x, u2(x, 0) ~ 3 * x] # Space and time domains -domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)] - +domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] # Neural network -chain1 = FastChain(FastDense(2,15,Flux.tanh),FastDense(15,1)) -chain2 = FastChain(FastDense(2,15,Flux.tanh),FastDense(15,1)) +chain1 = FastChain(FastDense(2, 15, Flux.tanh), FastDense(15, 1)) +chain2 = FastChain(FastDense(2, 15, Flux.tanh), FastDense(15, 1)) -quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(), - reltol=1e-3,abstol=1e-3, - maxiters =50, batch=100) -chain = [chain1,chain2] +quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg = CubatureJLh(), + reltol = 1e-3, abstol = 1e-3, + maxiters = 50, batch = 100) +chain = [chain1, chain2] initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain)) -discretization = NeuralPDE.PhysicsInformedNN(chain,quadrature_strategy; init_params = initθ) +discretization = NeuralPDE.PhysicsInformedNN(chain, quadrature_strategy; + init_params = initθ) -@named pde_system = PDESystem(eqs,bcs,domains,[x,y],[u1(x, y),u2(x, y)]) +@named pde_system = PDESystem(eqs, bcs, domains, [x, y], [u1(x, y), u2(x, y)]) -prob = NeuralPDE.discretize(pde_system,discretization) -sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) +prob = NeuralPDE.discretize(pde_system, discretization) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) -res = Optimization.solve(prob,BFGS(); maxiters=1000) +res = Optimization.solve(prob, BFGS(); maxiters = 1000) phi = discretization.phi -analytic_sol_func(x,y) =[1/3*(6x - y), 1/2*(6x - y)] -xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] -u_real = [[analytic_sol_func(x,y)[i] for x in xs for y in ys] for i in 1:2] +analytic_sol_func(x, y) = [1 / 3 * (6x - y), 1 / 2 * (6x - y)] +xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] +u_real = [[analytic_sol_func(x, y)[i] for x in xs for y in ys] for i in 1:2] initθ = discretization.init_params -acum = [0;accumulate(+, length.(initθ))] -sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1] +acum = [0; accumulate(+, length.(initθ))] +sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)] minimizers = [res.minimizer[s] for s in sep] -u_predict = [[phi[i]([x,y],minimizers[i])[1] for x in xs for y in ys] for i in 1:2] +u_predict = [[phi[i]([x, y], minimizers[i])[1] for x in xs for y in ys] for i in 1:2] -@test u_predict[1] ≈ u_real[1] atol = 0.1 -@test u_predict[2] ≈ u_real[2] atol = 0.1 +@test u_predict[1]≈u_real[1] atol=0.1 +@test u_predict[2]≈u_real[2] atol=0.1 # p1 =plot(xs, ys, u_predict, st=:surface); # p2 = plot(xs, ys, u_real, st=:surface); @@ -441,89 +459,103 @@ Dtt = Differential(t)^2 Dt = Differential(t) #2D PDE -C=1 -eq = Dtt(u(x,t)) ~ C^2*Dxx(u(x,t)) +C = 1 +eq = Dtt(u(x, t)) ~ C^2 * Dxx(u(x, t)) # Initial and boundary conditions -bcs = [u(0,t) ~ 0.,# for all t > 0 - u(1,t) ~ 0.,# for all t > 0 - u(x,0) ~ x*(1. - x), #for all 0 < x < 1 - Dt(u(x,0)) ~ 0. ] #for all 0 < x < 1] +bcs = [u(0, t) ~ 0.0,# for all t > 0 + u(1, t) ~ 0.0,# for all t > 0 + u(x, 0) ~ x * (1.0 - x), #for all 0 < x < 1 + Dt(u(x, 0)) ~ 0.0] #for all 0 < x < 1] # Space and time domains -domains = [x ∈ Interval(0.0,1.0), - t ∈ Interval(0.0,1.0)] +domains = [x ∈ Interval(0.0, 1.0), + t ∈ Interval(0.0, 1.0)] # Neural network -chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1)) +chain = FastChain(FastDense(2, 16, Flux.σ), FastDense(16, 16, Flux.σ), FastDense(16, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) eltypeθ = eltype(initθ) phi = NeuralPDE.Phi(chain) derivative = NeuralPDE.numeric_derivative -indvars = [x,t] +indvars = [x, t] depvars = [u(x, t)] dim = length(domains) -quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(), - reltol=1e-3,abstol=1e-3, - maxiters =50, batch=100) - -integral = NeuralPDE.get_numeric_integral(quadrature_strategy, indvars, depvars, chain isa AbstractArray, derivative) -_pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,phi,derivative,integral, - chain isa AbstractArray,initθ,quadrature_strategy) -_pde_loss_function(rand(2,10), initθ) - -bc_indvars = NeuralPDE.get_argument(bcs,indvars,depvars) -_bc_loss_functions = [NeuralPDE.build_loss_function(bc, indvars, depvars, phi, derivative, integral, - chain isa AbstractArray, initθ, quadrature_strategy, - bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)] -map(loss_f -> loss_f(rand(1,10), initθ),_bc_loss_functions) +quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg = CubatureJLh(), + reltol = 1e-3, abstol = 1e-3, + maxiters = 50, batch = 100) + +integral = NeuralPDE.get_numeric_integral(quadrature_strategy, indvars, depvars, + chain isa AbstractArray, derivative) +_pde_loss_function = NeuralPDE.build_loss_function(eq, indvars, depvars, phi, derivative, + integral, + chain isa AbstractArray, initθ, + quadrature_strategy) +_pde_loss_function(rand(2, 10), initθ) + +bc_indvars = NeuralPDE.get_argument(bcs, indvars, depvars) +_bc_loss_functions = [NeuralPDE.build_loss_function(bc, indvars, depvars, phi, derivative, + integral, + chain isa AbstractArray, initθ, + quadrature_strategy, + bc_indvars = bc_indvar) + for (bc, bc_indvar) in zip(bcs, bc_indvars)] +map(loss_f -> loss_f(rand(1, 10), initθ), _bc_loss_functions) dx = 0.1 -train_sets = NeuralPDE.generate_training_sets(domains,dx,[eq],bcs,eltypeθ,indvars,depvars) -pde_train_set,bcs_train_set = train_sets -pde_bounds, bcs_bounds = NeuralPDE.get_bounds(domains,[eq],bcs,eltypeθ,indvars,depvars,quadrature_strategy) - -lbs,ubs = pde_bounds -pde_loss_functions = [NeuralPDE.get_loss_function(_pde_loss_function, - lbs[1],ubs[1], - eltypeθ, - quadrature_strategy)] +train_sets = NeuralPDE.generate_training_sets(domains, dx, [eq], bcs, eltypeθ, indvars, + depvars) +pde_train_set, bcs_train_set = train_sets +pde_bounds, bcs_bounds = NeuralPDE.get_bounds(domains, [eq], bcs, eltypeθ, indvars, depvars, + quadrature_strategy) + +lbs, ubs = pde_bounds +pde_loss_functions = [ + NeuralPDE.get_loss_function(_pde_loss_function, + lbs[1], ubs[1], + eltypeθ, + quadrature_strategy), +] pde_loss_functions[1](initθ) -lbs,ubs = bcs_bounds -bc_loss_functions = [NeuralPDE.get_loss_function(_loss,lb,ub, - eltypeθ,quadrature_strategy) - for (_loss,lb,ub) in zip(_bc_loss_functions, lbs,ubs)] +lbs, ubs = bcs_bounds +bc_loss_functions = [NeuralPDE.get_loss_function(_loss, lb, ub, + eltypeθ, quadrature_strategy) + for (_loss, lb, ub) in zip(_bc_loss_functions, lbs, ubs)] -map(l->l(initθ) ,bc_loss_functions) +map(l -> l(initθ), bc_loss_functions) -loss_functions = [pde_loss_functions;bc_loss_functions] +loss_functions = [pde_loss_functions; bc_loss_functions] -function loss_function(θ,p) - sum(map(l->l(θ) ,loss_functions)) +function loss_function(θ, p) + sum(map(l -> l(θ), loss_functions)) end f_ = OptimizationFunction(loss_function, Optimization.AutoZygote()) prob = Optimization.OptimizationProblem(f_, initθ) -cb_ = function (p,l) - println("loss: ", l ) +cb_ = function (p, l) + println("loss: ", l) println("losses: ", map(l -> l(p), loss_functions)) return false end -res = Optimization.solve(prob,Optim.BFGS(); maxiters=500,f_abstol=10^-6) +res = Optimization.solve(prob, Optim.BFGS(); maxiters = 500, f_abstol = 10^-6) -xs,ts = [infimum(d.domain):dx:supremum(d.domain) for d in domains] -analytic_sol_func(x,t) = sum([(8/(k^3*pi^3)) * sin(k*pi*x)*cos(C*k*pi*t) for k in 1:2:50000]) +xs, ts = [infimum(d.domain):dx:supremum(d.domain) for d in domains] +function analytic_sol_func(x, t) + sum([(8 / (k^3 * pi^3)) * sin(k * pi * x) * cos(C * k * pi * t) for k in 1:2:50000]) +end -u_predict = reshape([first(phi([x,t],res.minimizer)) for x in xs for t in ts],(length(xs),length(ts))) -u_real = reshape([analytic_sol_func(x,t) for x in xs for t in ts], (length(xs),length(ts))) +u_predict = reshape([first(phi([x, t], res.minimizer)) for x in xs for t in ts], + (length(xs), length(ts))) +u_real = reshape([analytic_sol_func(x, t) for x in xs for t in ts], + (length(xs), length(ts))) -@test u_predict ≈ u_real atol = 0.1 +@test u_predict≈u_real atol=0.1 # diff_u = abs.(u_predict .- u_real) # p1 = plot(xs, ts, u_real, linetype=:contourf,title = "analytic"); @@ -531,7 +563,6 @@ u_real = reshape([analytic_sol_func(x,t) for x in xs for t in ts], (length(xs),l # p3 = plot(xs, ts, diff_u,linetype=:contourf,title = "error"); # plot(p1,p2,p3) - ## Example 6, pde with mixed derivative println("Example 6, pde with mixed derivative") @parameters x y @@ -541,43 +572,47 @@ Dyy = Differential(y)^2 Dx = Differential(x) Dy = Differential(y) -eq = Dxx(u(x,y)) + Dx(Dy(u(x,y))) - 2*Dyy(u(x,y)) ~ -1. +eq = Dxx(u(x, y)) + Dx(Dy(u(x, y))) - 2 * Dyy(u(x, y)) ~ -1.0 # Initial and boundary conditions -bcs = [u(x,0) ~ x, - Dy(u(x,0)) ~ x, - u(x,0) ~ Dy(u(x,0))] +bcs = [u(x, 0) ~ x, + Dy(u(x, 0)) ~ x, + u(x, 0) ~ Dy(u(x, 0))] # Space and time domains -domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)] +domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] quadrature_strategy = NeuralPDE.QuadratureTraining() # Neural network -inner =20 -chain = FastChain(FastDense(2,inner,Flux.tanh),FastDense(inner,inner,Flux.tanh),FastDense(inner,1)) +inner = 20 +chain = FastChain(FastDense(2, inner, Flux.tanh), FastDense(inner, inner, Flux.tanh), + FastDense(inner, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) -discretization = NeuralPDE.PhysicsInformedNN(chain, quadrature_strategy; init_params = initθ) -@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)]) +discretization = NeuralPDE.PhysicsInformedNN(chain, quadrature_strategy; + init_params = initθ) +@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) +prob = NeuralPDE.discretize(pde_system, discretization) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) -res = Optimization.solve(prob,BFGS(); maxiters=1500) +res = Optimization.solve(prob, BFGS(); maxiters = 1500) @show res.original phi = discretization.phi -analytic_sol_func(x,y) = x + x*y +y^2/2 -xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] +analytic_sol_func(x, y) = x + x * y + y^2 / 2 +xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] -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))) +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 = abs.(u_predict .- u_real) -@test u_predict ≈ u_real rtol = 0.1 +@test u_predict≈u_real rtol=0.1 # p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic"); # p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict"); # p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error"); -# plot(p1,p2,p3) \ No newline at end of file +# plot(p1,p2,p3) diff --git a/test/NNPDE_tests_gpu.jl b/test/NNPDE_tests_gpu.jl index 4238c5dbc6..16ae4c79f2 100644 --- a/test/NNPDE_tests_gpu.jl +++ b/test/NNPDE_tests_gpu.jl @@ -8,7 +8,7 @@ import ModelingToolkit: Interval, infimum, supremum using Random Random.seed!(100) -callback = function (p,l) +callback = function (p, l) println("Current loss is: $l") return false end @@ -22,44 +22,44 @@ println("ode") Dθ = Differential(θ) # 1D ODE -eq = Dθ(u(θ)) ~ θ^3 + 2.f0*θ + (θ^2)*((1.f0+3*(θ^2))/(1.f0+θ+(θ^3))) - u(θ)*(θ + ((1.f0+3.f0*(θ^2))/(1.f0+θ+θ^3))) +eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - + u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) # Initial and boundary conditions -bcs = [u(0.) ~ 1.0f0] +bcs = [u(0.0) ~ 1.0f0] # Space and time domains -domains = [θ ∈ Interval(0f0,1f0)] +domains = [θ ∈ Interval(0.0f0, 1.0f0)] # Discretization dt = 0.1f0 # Neural network inner = 20 -chain = Chain(Dense(1,inner,Flux.σ), - Dense(inner,inner,Flux.σ), - Dense(inner,inner,Flux.σ), - Dense(inner,inner,Flux.σ), - Dense(inner,inner,Flux.σ), - Dense(inner,1)) |> gpu +chain = Chain(Dense(1, inner, Flux.σ), + Dense(inner, inner, Flux.σ), + Dense(inner, inner, Flux.σ), + Dense(inner, inner, Flux.σ), + Dense(inner, inner, Flux.σ), + Dense(inner, 1)) |> gpu initθ = DiffEqFlux.initial_params(chain) |> gpu strategy = NeuralPDE.GridTraining(dt) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; - init_params = initθ - ) + init_params = initθ) -@named pde_system = PDESystem(eq,bcs,domains,[θ],[u(θ)]) -prob = NeuralPDE.discretize(pde_system,discretization) -symprob = NeuralPDE.symbolic_discretize(pde_system,discretization) -res = Optimization.solve(prob, ADAM(1e-2); maxiters=2000) +@named pde_system = PDESystem(eq, bcs, domains, [θ], [u(θ)]) +prob = NeuralPDE.discretize(pde_system, discretization) +symprob = NeuralPDE.symbolic_discretize(pde_system, discretization) +res = Optimization.solve(prob, ADAM(1e-2); maxiters = 2000) phi = discretization.phi -analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2 -ts = [infimum(d.domain):dt/10:supremum(d.domain) for d in domains][1] -u_real = [analytic_sol_func(t) for t in ts] -u_predict = [first(Array(phi([t],res.minimizer))) for t in ts] +analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 +ts = [infimum(d.domain):(dt / 10):supremum(d.domain) for d in domains][1] +u_real = [analytic_sol_func(t) for t in ts] +u_predict = [first(Array(phi([t], res.minimizer))) for t in ts] -@test u_predict ≈ u_real atol = 0.2 +@test u_predict≈u_real atol=0.2 # t_plot = collect(ts) # plot(t_plot ,u_real) @@ -72,45 +72,46 @@ println("1D PDE Dirichlet boundary conditions") Dt = Differential(t) Dxx = Differential(x)^2 -eq = Dt(u(t,x)) ~ Dxx(u(t,x)) -bcs = [u(0,x) ~ cos(x), - u(t,0) ~ exp(-t), - u(t,1) ~ exp(-t) * cos(1)] +eq = Dt(u(t, x)) ~ Dxx(u(t, x)) +bcs = [u(0, x) ~ cos(x), + u(t, 0) ~ exp(-t), + u(t, 1) ~ exp(-t) * cos(1)] -domains = [t ∈ Interval(0.0,1.0), - x ∈ Interval(0.0,1.0)] +domains = [t ∈ Interval(0.0, 1.0), + x ∈ Interval(0.0, 1.0)] -@named pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t, x)]) +@named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)]) inner = 30 -chain = FastChain(FastDense(2,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,1))#,(u,p)->gpuones .* u) +chain = FastChain(FastDense(2, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, 1))#,(u,p)->gpuones .* u) strategy = NeuralPDE.StochasticTraining(500) initθ = CuArray(Float64.(DiffEqFlux.initial_params(chain))) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; init_params = initθ) -prob = NeuralPDE.discretize(pdesys,discretization) -symprob = NeuralPDE.symbolic_discretize(pdesys,discretization) +prob = NeuralPDE.discretize(pdesys, discretization) +symprob = NeuralPDE.symbolic_discretize(pdesys, discretization) -res = Optimization.solve(prob, ADAM(0.01);maxiters=1000) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,ADAM(0.001);maxiters=1000) +res = Optimization.solve(prob, ADAM(0.01); maxiters = 1000) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, ADAM(0.001); maxiters = 1000) phi = discretization.phi -u_exact = (t,x) -> exp.(-t) * cos.(x) -ts,xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] -u_predict = reshape([first(Array(phi([t,x],res.minimizer))) for t in ts for x in xs],(length(ts),length(xs))) -u_real = reshape([u_exact(t,x) for t in ts for x in xs ], (length(ts),length(xs))) +u_exact = (t, x) -> exp.(-t) * cos.(x) +ts, xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] +u_predict = reshape([first(Array(phi([t, x], res.minimizer))) for t in ts for x in xs], + (length(ts), length(xs))) +u_real = reshape([u_exact(t, x) for t in ts for x in xs], (length(ts), length(xs))) diff_u = abs.(u_predict .- u_real) -@test u_predict ≈ u_real atol = 1.0 +@test u_predict≈u_real atol=1.0 # p1 = plot(ts, xs, u_real, linetype=:contourf,title = "analytic"); # p2 = plot(ts, xs, u_predict, linetype=:contourf,title = "predict"); @@ -126,24 +127,24 @@ Dx = Differential(x) Dxx = Differential(x)^2 # 1D PDE and boundary conditions -eq = Dt(u(t,x)) ~ Dxx(u(t,x)) -bcs = [u(0,x) ~ cos(x), - Dx(u(t,0)) ~ 0.0, - Dx(u(t,1)) ~ -exp(-t) * sin(1.0)] +eq = Dt(u(t, x)) ~ Dxx(u(t, x)) +bcs = [u(0, x) ~ cos(x), + Dx(u(t, 0)) ~ 0.0, + Dx(u(t, 1)) ~ -exp(-t) * sin(1.0)] # Space and time domains -domains = [t ∈ Interval(0.0,1.0), - x ∈ Interval(0.0,1.0)] +domains = [t ∈ Interval(0.0, 1.0), + x ∈ Interval(0.0, 1.0)] # PDE system -@named pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t, x)]) +@named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)]) inner = 20 -chain = FastChain(FastDense(2,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,1)) +chain = FastChain(FastDense(2, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, 1)) strategy = NeuralPDE.QuasiRandomTraining(500; #points sampling_alg = SobolSample(), @@ -153,21 +154,22 @@ initθ = CuArray(Float64.(DiffEqFlux.initial_params(chain))) discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; init_params = initθ) -prob = NeuralPDE.discretize(pdesys,discretization) -symprob = NeuralPDE.symbolic_discretize(pdesys,discretization) +prob = NeuralPDE.discretize(pdesys, discretization) +symprob = NeuralPDE.symbolic_discretize(pdesys, discretization) -res = Optimization.solve(prob, ADAM(0.1); maxiters=2000) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,ADAM(0.01); maxiters=2000) +res = Optimization.solve(prob, ADAM(0.1); maxiters = 2000) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, ADAM(0.01); maxiters = 2000) phi = discretization.phi -u_exact = (t,x) -> exp(-t) * cos(x) -ts,xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] -u_predict = reshape([first(Array(phi([t,x],res.minimizer))) for t in ts for x in xs],(length(ts),length(xs))) -u_real = reshape([u_exact(t,x) for t in ts for x in xs ], (length(ts),length(xs))) +u_exact = (t, x) -> exp(-t) * cos(x) +ts, xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains] +u_predict = reshape([first(Array(phi([t, x], res.minimizer))) for t in ts for x in xs], + (length(ts), length(xs))) +u_real = reshape([u_exact(t, x) for t in ts for x in xs], (length(ts), length(xs))) diff_u = abs.(u_predict .- u_real) -@test u_predict ≈ u_real atol = 1.0 +@test u_predict≈u_real atol=1.0 # p1 = plot(ts, xs, u_real, linetype=:contourf,title = "analytic"); # p2 = plot(ts, xs, u_predict, linetype=:contourf,title = "predict"); @@ -181,7 +183,7 @@ println("2D PDE") Dxx = Differential(x)^2 Dyy = Differential(y)^2 Dt = Differential(t) -t_min= 0.0 +t_min = 0.0 t_max = 2.0 x_min = 0.0 x_max = 2.0 @@ -189,28 +191,28 @@ y_min = 0.0 y_max = 2.0 # 3D PDE -eq = Dt(u(t,x,y)) ~ Dxx(u(t,x,y)) + Dyy(u(t,x,y)) +eq = Dt(u(t, x, y)) ~ Dxx(u(t, x, y)) + Dyy(u(t, x, y)) -analytic_sol_func(t,x,y) = exp(x+y)*cos(x+y+4t) +analytic_sol_func(t, x, y) = exp(x + y) * cos(x + y + 4t) # Initial and boundary conditions -bcs = [u(t_min,x,y) ~ analytic_sol_func(t_min,x,y), - u(t,x_min,y) ~ analytic_sol_func(t,x_min,y), - u(t,x_max,y) ~ analytic_sol_func(t,x_max,y), - u(t,x,y_min) ~ analytic_sol_func(t,x,y_min), - u(t,x,y_max) ~ analytic_sol_func(t,x,y_max)] +bcs = [u(t_min, x, y) ~ analytic_sol_func(t_min, x, y), + u(t, x_min, y) ~ analytic_sol_func(t, x_min, y), + u(t, x_max, y) ~ analytic_sol_func(t, x_max, y), + u(t, x, y_min) ~ analytic_sol_func(t, x, y_min), + u(t, x, y_max) ~ analytic_sol_func(t, x, y_max)] # Space and time domains -domains = [t ∈ Interval(t_min,t_max), - x ∈ Interval(x_min,x_max), - y ∈ Interval(y_min,y_max)] +domains = [t ∈ Interval(t_min, t_max), + x ∈ Interval(x_min, x_max), + y ∈ Interval(y_min, y_max)] # Neural network inner = 25 -chain = FastChain(FastDense(3,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,inner,Flux.σ), - FastDense(inner,1))#,(u,p)->gpuones .* u) +chain = FastChain(FastDense(3, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, inner, Flux.σ), + FastDense(inner, 1))#,(u,p)->gpuones .* u) initθ = CuArray(Float64.(DiffEqFlux.initial_params(chain))) @@ -219,21 +221,22 @@ discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; init_params = initθ) -@named pde_system = PDESystem(eq,bcs,domains,[t,x,y],[u(t,x,y)]) -prob = NeuralPDE.discretize(pde_system,discretization) -symprob = NeuralPDE.symbolic_discretize(pde_system,discretization) +@named pde_system = PDESystem(eq, bcs, domains, [t, x, y], [u(t, x, y)]) +prob = NeuralPDE.discretize(pde_system, discretization) +symprob = NeuralPDE.symbolic_discretize(pde_system, discretization) -res = Optimization.solve(prob,ADAM(0.01);maxiters=2500) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,ADAM(0.001);maxiters=2500) +res = Optimization.solve(prob, ADAM(0.01); maxiters = 2500) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, ADAM(0.001); maxiters = 2500) @show res.original phi = discretization.phi -ts,xs,ys = [infimum(d.domain):0.1:supremum(d.domain) for d in domains] -u_real = [analytic_sol_func(t,x,y) for t in ts for x in xs for y in ys] -u_predict = [first(Array(phi([t, x, y], res.minimizer))) for t in ts for x in xs for y in ys] +ts, xs, ys = [infimum(d.domain):0.1:supremum(d.domain) for d in domains] +u_real = [analytic_sol_func(t, x, y) for t in ts for x in xs for y in ys] +u_predict = [first(Array(phi([t, x, y], res.minimizer))) for t in ts for x in xs + for y in ys] -@test u_predict ≈ u_real rtol = 0.2 +@test u_predict≈u_real rtol=0.2 # using Plots # using Printf diff --git a/test/NNParamKolmogorov_tests.jl b/test/NNParamKolmogorov_tests.jl index d2c17bd392..1381010e3d 100644 --- a/test/NNParamKolmogorov_tests.jl +++ b/test/NNParamKolmogorov_tests.jl @@ -4,56 +4,59 @@ using StochasticDiffEq using LinearAlgebra d = 1 -m = Chain(Dense(3 , 16, elu),Dense(16, 16, elu) , Dense(16 , 5 , elu) , Dense(5 , 1)) +m = Chain(Dense(3, 16, elu), Dense(16, 16, elu), Dense(16, 5, elu), Dense(5, 1)) ensemblealg = EnsembleThreads() γ_mu_prototype = nothing -γ_sigma_prototype = zeros(d , d , 1) +γ_sigma_prototype = zeros(d, d, 1) γ_phi_prototype = nothing sdealg = EM() -tspan = (0.00 , 1.00) +tspan = (0.00, 1.00) trajectories = 10000 -function phi(x , y_phi) - x.^2 +function phi(x, y_phi) + x .^ 2 end -g(x , γ_sigma) = γ_sigma[: , : , 1] -f(x , γ_mu_1 ,γ_mu_2 ) = [0.00] -xspan = (0.00 , 3.00) -y_domain = KolmogorovParamDomain((0.00 , 2.00) , (0.00 , 0.00) , (0.00 , 0.00) ) +g(x, γ_sigma) = γ_sigma[:, :, 1] +f(x, γ_mu_1, γ_mu_2) = [0.00] +xspan = (0.00, 3.00) +y_domain = KolmogorovParamDomain((0.00, 2.00), (0.00, 0.00), (0.00, 0.00)) dt = 0.01 dx = 0.01 dy = 0.01 opt = Flux.ADAM(1e-2) -prob = ParamKolmogorovPDEProblem(f , g , phi , xspan , tspan , d , y_domain ; Y_sigma_prototype = γ_sigma_prototype) -sol = solve(prob, NNParamKolmogorov(m,opt , sdealg,ensemblealg) , verbose = true, dt = 0.01, - abstol=1e-10, dx = 0.01 , trajectories = trajectories , maxiters = 400 , use_gpu = false) - -x_test = rand(xspan[1]:dx:xspan[2] , d , 1 , 1000) -t_test = rand(tspan[1]:dt:tspan[2] , 1 , 1000) -γ_sigma_test = rand(0.3:dy:0.3 , d , d, 1, 1000) +prob = ParamKolmogorovPDEProblem(f, g, phi, xspan, tspan, d, y_domain; + Y_sigma_prototype = γ_sigma_prototype) +sol = solve(prob, NNParamKolmogorov(m, opt, sdealg, ensemblealg), verbose = true, dt = 0.01, + abstol = 1e-10, dx = 0.01, trajectories = trajectories, maxiters = 400, + use_gpu = false) +x_test = rand(xspan[1]:dx:xspan[2], d, 1, 1000) +t_test = rand(tspan[1]:dt:tspan[2], 1, 1000) +γ_sigma_test = rand(0.3:dy:0.3, d, d, 1, 1000) function getTestData(x_test, t_test, γ_sigma_test) - X_test = reshape([], 3 ,0) - for i in 1:length(t_test) - K = vcat(t_test[i] , x_test[: , : , i] , reshape(γ_sigma_test[: , : , : , i] , d^2*(size(γ_sigma_prototype))[3] , 1)) - X_test = hcat(X_test , K) - end - return X_test + X_test = reshape([], 3, 0) + for i in 1:length(t_test) + K = vcat(t_test[i], x_test[:, :, i], + reshape(γ_sigma_test[:, :, :, i], d^2 * (size(γ_sigma_prototype))[3], 1)) + X_test = hcat(X_test, K) + end + return X_test end X_test = getTestData(x_test, t_test, γ_sigma_test) -function analytical(x , t , y) - #return sum(x.^2) + t*tr(y*y) for multidimensional - return x^2 + t*(y*y) +function analytical(x, t, y) + #return sum(x.^2) + t*tr(y*y) for multidimensional + return x^2 + t * (y * y) end function y_analytical(X) - y_test = reshape([], 1 ,0) - for i in 1:length(t_test) - y_test = hcat(y_test , analytical(X[: , i][2] , X[: , i][1] , γ_sigma_test[: , : , 1, i][1] ) ) - end - return y_test + y_test = reshape([], 1, 0) + for i in 1:length(t_test) + y_test = hcat(y_test, + analytical(X[:, i][2], X[:, i][1], γ_sigma_test[:, :, 1, i][1])) + end + return y_test end y_test = y_analytical(X_test) diff --git a/test/NNRODE_tests.jl b/test/NNRODE_tests.jl index 412ee20a2b..c834cb49ab 100644 --- a/test/NNRODE_tests.jl +++ b/test/NNRODE_tests.jl @@ -1,39 +1,40 @@ using Flux, OptimizationFlux, StochasticDiffEq, DiffEqNoiseProcess, Optim, Test -using NeuralPDE +using NeuralPDE using Random Random.seed!(100) println("Test Case 1") -linear = (u,p,t,W) -> 2u*sin(W) +linear = (u, p, t, W) -> 2u * sin(W) tspan = (0.00f0, 1.00f0) u0 = 1.0f0 -dt = 1/50f0 -W = WienerProcess(0.0,0.0,nothing) -prob = RODEProblem(linear, u0 ,tspan, noise = W) -chain = Flux.Chain(Dense(2,8,relu),Dense(8,16 , relu),Dense(16,1)) +dt = 1 / 50.0f0 +W = WienerProcess(0.0, 0.0, nothing) +prob = RODEProblem(linear, u0, tspan, noise = W) +chain = Flux.Chain(Dense(2, 8, relu), Dense(8, 16, relu), Dense(16, 1)) opt = ADAM(1e-4) -sol = solve(prob, NeuralPDE.NNRODE(chain,W,opt), dt=dt, verbose = true, - abstol=1e-10, maxiters = 3000) +sol = solve(prob, NeuralPDE.NNRODE(chain, W, opt), dt = dt, verbose = true, + abstol = 1e-10, maxiters = 3000) W2 = NoiseWrapper(sol.W) -prob1 = RODEProblem(linear , u0 , tspan , noise = W2) -sol2 = solve(prob1,RandomEM(),dt=dt) -err = Flux.mse(sol.u , sol2.u) +prob1 = RODEProblem(linear, u0, tspan, noise = W2) +sol2 = solve(prob1, RandomEM(), dt = dt) +err = Flux.mse(sol.u, sol2.u) @test err < 0.3 println("Test Case 2") -linear = (u,p,t,W) -> t^3 + 2*t + (t^2)*((1+3*(t^2))/(1+t+(t^3))) - u*(t + ((1+3*(t^2))/(1+t+t^3))) + 5*W +linear = (u, p, t, W) -> t^3 + 2 * t + (t^2) * ((1 + 3 * (t^2)) / (1 + t + (t^3))) - + u * (t + ((1 + 3 * (t^2)) / (1 + t + t^3))) + 5 * W tspan = (0.00f0, 1.00f0) u0 = 1.0f0 -dt = 1/100f0 -W = WienerProcess(0.0,0.0,nothing) -prob = RODEProblem(linear, u0 ,tspan, noise = W) -chain = Flux.Chain(Dense(2,32,sigmoid), Dense(32,32,sigmoid) , Dense(32,1)) +dt = 1 / 100.0f0 +W = WienerProcess(0.0, 0.0, nothing) +prob = RODEProblem(linear, u0, tspan, noise = W) +chain = Flux.Chain(Dense(2, 32, sigmoid), Dense(32, 32, sigmoid), Dense(32, 1)) opt = ADAM(1e-3) -sol = solve(prob, NeuralPDE.NNRODE(chain,W,opt), dt=dt, verbose = true, - abstol=1e-10, maxiters = 2000) +sol = solve(prob, NeuralPDE.NNRODE(chain, W, opt), dt = dt, verbose = true, + abstol = 1e-10, maxiters = 2000) W2 = NoiseWrapper(sol.W) -prob1 = RODEProblem(linear , u0 , tspan , noise = W2) -sol2 = solve(prob1,RandomEM(),dt=dt) -err = Flux.mse(sol.u , sol2.u) +prob1 = RODEProblem(linear, u0, tspan, noise = W2) +sol2 = solve(prob1, RandomEM(), dt = dt) +err = Flux.mse(sol.u, sol2.u) @test err < 0.4 diff --git a/test/Stopping_tests.jl b/test/Stopping_tests.jl index b033d83480..0e875ad99e 100644 --- a/test/Stopping_tests.jl +++ b/test/Stopping_tests.jl @@ -1,4 +1,4 @@ -using Test, Flux , StochasticDiffEq , LinearAlgebra +using Test, Flux, StochasticDiffEq, LinearAlgebra using NeuralPDE using Random @@ -8,43 +8,44 @@ d = 1 r = 0.04f0 beta = 0.2f0 T = 1 -u0 = fill(80.00 , d , 1) +u0 = fill(80.00, d, 1) sdealg = EM() ensemblealg = EnsembleThreads() -f(du,u,p,t) = (du .= r*u) -sigma(du,u,p,t) = (du .= Diagonal(beta*u)) -tspan = (0.0 , 1.0) +f(du, u, p, t) = (du .= r * u) +sigma(du, u, p, t) = (du .= Diagonal(beta * u)) +tspan = (0.0, 1.0) N = 50 -dt = tspan[2]/49 +dt = tspan[2] / 49 K = 100.00 -function g(t , x) - return exp(-r*t)*(max(K - maximum(x) , 0)) +function g(t, x) + return exp(-r * t) * (max(K - maximum(x), 0)) end -prob = SDEProblem(f , sigma , u0 , tspan ; g = g) +prob = SDEProblem(f, sigma, u0, tspan; g = g) opt = Flux.ADAM(0.1) -m = Chain(Dense(d , 5, tanh), Dense(5, 16 , tanh) , Dense(16 , N ), softmax) -sol = solve(prob, NeuralPDE.NNStopping( m, opt , sdealg , ensemblealg), verbose = true, dt = dt, - abstol=1e-6, maxiters = 20 , trajectories = 200) +m = Chain(Dense(d, 5, tanh), Dense(5, 16, tanh), Dense(16, N), softmax) +sol = solve(prob, NeuralPDE.NNStopping(m, opt, sdealg, ensemblealg), verbose = true, + dt = dt, + abstol = 1e-6, maxiters = 20, trajectories = 200) ##Analytical Binomial Tree approach for American Options -function BinomialTreeAM1D(S0 , N , r , beta) - V = zeros(N+1) - dT = T/N - u = exp(beta*sqrt(dT)) - d = 1/u - S_T = [S0*(u^j)* (d^(N-j)) for j in 0:N] - a = exp(r*dT) - p = (a - d)/(u - d) +function BinomialTreeAM1D(S0, N, r, beta) + V = zeros(N + 1) + dT = T / N + u = exp(beta * sqrt(dT)) + d = 1 / u + S_T = [S0 * (u^j) * (d^(N - j)) for j in 0:N] + a = exp(r * dT) + p = (a - d) / (u - d) q = 1.0 - p - V = [max(K - x , 0) for x in S_T] - for i in N-1:-1:0 - V[1:end-1] = exp(-r*dT).*(p*V[2:end] + q*V[1:end-1]) - S_T = S_T*u - V = [max(K - S_T[i] , V[i]) for i in 1:size(S_T)[1]] + V = [max(K - x, 0) for x in S_T] + for i in (N - 1):-1:0 + V[1:(end - 1)] = exp(-r * dT) .* (p * V[2:end] + q * V[1:(end - 1)]) + S_T = S_T * u + V = [max(K - S_T[i], V[i]) for i in 1:size(S_T)[1]] end return V[1] end -real_sol = BinomialTreeAM1D(u0[1] , N , r , beta) +real_sol = BinomialTreeAM1D(u0[1], N, r, beta) error = abs(sol - real_sol) @test error < 0.5 diff --git a/test/adaptive_loss_tests.jl b/test/adaptive_loss_tests.jl index cbf22bae22..1f7713ba21 100644 --- a/test/adaptive_loss_tests.jl +++ b/test/adaptive_loss_tests.jl @@ -5,19 +5,22 @@ import ModelingToolkit: Interval, infimum, supremum using DomainSets using Random -nonadaptive_loss = NeuralPDE.NonAdaptiveLoss(pde_loss_weights=1, bc_loss_weights=1) -gradnormadaptive_loss = NeuralPDE.GradientScaleAdaptiveLoss(100, pde_loss_weights=1e3, bc_loss_weights=1) -adaptive_loss = NeuralPDE.MiniMaxAdaptiveLoss(100; pde_loss_weights=1, bc_loss_weights=1) -adaptive_losses = [nonadaptive_loss, gradnormadaptive_loss,adaptive_loss] -maxiters=4000 -seed=60 +nonadaptive_loss = NeuralPDE.NonAdaptiveLoss(pde_loss_weights = 1, bc_loss_weights = 1) +gradnormadaptive_loss = NeuralPDE.GradientScaleAdaptiveLoss(100, pde_loss_weights = 1e3, + bc_loss_weights = 1) +adaptive_loss = NeuralPDE.MiniMaxAdaptiveLoss(100; pde_loss_weights = 1, + bc_loss_weights = 1) +adaptive_losses = [nonadaptive_loss, gradnormadaptive_loss, adaptive_loss] +maxiters = 4000 +seed = 60 ## 2D Poisson equation -function test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed=60, maxiters=4000) +function test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed = 60, maxiters = 4000) Random.seed!(seed) hid = 40 - chain_ = FastChain(FastDense(2,hid,Flux.σ),FastDense(hid,hid,Flux.σ),FastDense(hid,1)) - strategy_ = NeuralPDE.StochasticTraining(256) + chain_ = FastChain(FastDense(2, hid, Flux.σ), FastDense(hid, hid, Flux.σ), + FastDense(hid, 1)) + strategy_ = NeuralPDE.StochasticTraining(256) @info "adaptive reweighting test outdir:, maxiters: $(maxiters), 2D Poisson equation, adaptive_loss: $(nameof(typeof(adaptive_loss))) " @parameters x y @variables u(..) @@ -25,14 +28,14 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed=60, maxiters 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)) ~ -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)] + 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)] + domains = [x ∈ Interval(0.0, 1.0), + y ∈ Interval(0.0, 1.0)] initθ = Float64.(DiffEqFlux.initial_params(chain_)) iteration = [0] @@ -41,29 +44,29 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed=60, maxiters init_params = initθ, adaptive_loss = adaptive_loss, logger = nothing, - iteration=iteration) + iteration = iteration) - - @named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)]) - prob = NeuralPDE.discretize(pde_system,discretization) + @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) + prob = NeuralPDE.discretize(pde_system, discretization) phi = discretization.phi - sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) - + sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) - 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_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys))) + 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_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys], + (length(xs), length(ys))) - callback = function (p,l) + callback = function (p, l) iteration[1] += 1 if iteration[1] % 100 == 0 @info "Current loss is: $l, iteration is $(iteration[1])" end return false end - res = Optimization.solve(prob, ADAM(0.03); maxiters=maxiters, callback =callback) + res = Optimization.solve(prob, ADAM(0.03); maxiters = maxiters, callback = callback) - u_predict = reshape([first(phi([x,y],res.minimizer)) 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))) diff_u = abs.(u_predict .- u_real) total_diff = sum(diff_u) total_u = sum(abs.(u_real)) @@ -73,14 +76,15 @@ function test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed=60, maxiters #p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict"); #p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error"); #(plot=plot(p1,p2,p3), error=total_diff, total_diff_rel=total_diff_rel) - (error=total_diff, total_diff_rel=total_diff_rel) + (error = total_diff, total_diff_rel = total_diff_rel) end - - @info "testing that the adaptive loss methods roughly succeed" -test_2d_poisson_equation_adaptive_loss_no_logs_run_seediters(adaptive_loss) = test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed=seed, maxiters=maxiters) -error_results_no_logs = map(test_2d_poisson_equation_adaptive_loss_no_logs_run_seediters, adaptive_losses) +function test_2d_poisson_equation_adaptive_loss_no_logs_run_seediters(adaptive_loss) + test_2d_poisson_equation_adaptive_loss(adaptive_loss; seed = seed, maxiters = maxiters) +end +error_results_no_logs = map(test_2d_poisson_equation_adaptive_loss_no_logs_run_seediters, + adaptive_losses) # accuracy tests @show error_results_no_logs[1][:total_diff_rel] diff --git a/test/additional_loss_tests.jl b/test/additional_loss_tests.jl index af79b097aa..591d67a359 100644 --- a/test/additional_loss_tests.jl +++ b/test/additional_loss_tests.jl @@ -22,31 +22,31 @@ _σ = 0.5 dx = 0.01 # here we use normalization condition: dx*p(x) ~ 1, in order to get non-zero solution. #(α - 3*β*x^2)*p(x) + (α*x - β*x^3)*Dx(p(x)) ~ (_σ^2/2)*Dxx(p(x)) -eq = [Dx((α*x - β*x^3)*p(x)) ~ (_σ^2/2)*Dxx(p(x))] +eq = [Dx((α * x - β * x^3) * p(x)) ~ (_σ^2 / 2) * Dxx(p(x))] x_0 = -2.2 x_end = 2.2 # Initial and boundary conditions -bcs = [p(x_0) ~ 0. ,p(x_end) ~ 0.] +bcs = [p(x_0) ~ 0.0, p(x_end) ~ 0.0] # Space and time domains -domains = [x ∈ Interval(-2.2,2.2)] +domains = [x ∈ Interval(-2.2, 2.2)] # Neural network inn = 18 -chain = FastChain(FastDense(1,inn,Flux.σ), - FastDense(inn,inn,Flux.σ), - FastDense(inn,inn,Flux.σ), - FastDense(inn,1)) +chain = FastChain(FastDense(1, inn, Flux.σ), + FastDense(inn, inn, Flux.σ), + FastDense(inn, inn, Flux.σ), + FastDense(inn, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) lb = [x_0] ub = [x_end] -function norm_loss_function(phi,θ,p) - function inner_f(x,θ) - dx*phi(x, θ) .- 1 +function norm_loss_function(phi, θ, p) + function inner_f(x, θ) + dx * phi(x, θ) .- 1 end prob = IntegralProblem(inner_f, lb, ub, θ) - norm2 = solve(prob, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters =10); + norm2 = solve(prob, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters = 10) abs(norm2[1]) end # norm_loss_function(phi,initθ,nothing) @@ -54,36 +54,36 @@ end discretization = NeuralPDE.PhysicsInformedNN(chain, NeuralPDE.GridTraining(dx); init_params = initθ, - additional_loss=norm_loss_function) + additional_loss = norm_loss_function) -@named pde_system = PDESystem(eq,bcs,domains,[x],[p(x)]) -prob = NeuralPDE.discretize(pde_system,discretization) -discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) +@named pde_system = PDESystem(eq, bcs, domains, [x], [p(x)]) +prob = NeuralPDE.discretize(pde_system, discretization) +discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) pde_inner_loss_functions = discretized_functions.pde_loss_functions bcs_inner_loss_functions = discretized_functions.bc_loss_functions phi = discretization.phi -cb_ = function (p,l) - println("loss: ", l ) +cb_ = function (p, l) + println("loss: ", l) println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions)) println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions)) - println("additional_loss: ", norm_loss_function(phi,p,nothing)) + println("additional_loss: ", norm_loss_function(phi, p, nothing)) return false end -res = Optimization.solve(prob,LBFGS(),maxiters=400) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,BFGS(),maxiters=2000) +res = Optimization.solve(prob, LBFGS(), maxiters = 400) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, BFGS(), maxiters = 2000) C = 142.88418699042 -analytic_sol_func(x) = C*exp((1/(2*_σ^2))*(2*α*x^2 - β*x^4)) +analytic_sol_func(x) = C * exp((1 / (2 * _σ^2)) * (2 * α * x^2 - β * x^4)) xs = [infimum(d.domain):dx:supremum(d.domain) for d in domains][1] -u_real = [analytic_sol_func(x) for x in xs] -u_predict = [first(phi(x,res.u)) for x in xs] +u_real = [analytic_sol_func(x) for x in xs] +u_predict = [first(phi(x, res.u)) for x in xs] -@test u_predict ≈ u_real rtol = 1e-3 +@test u_predict≈u_real rtol=1e-3 # plot(xs ,u_real, label = "analytic") # plot!(xs ,u_predict, label = "predict") @@ -92,75 +92,75 @@ u_predict = [first(phi(x,res.u)) for x in xs] println("Example 8, Lorenz System") Random.seed!(1234) -@parameters t ,σ_ ,β, ρ +@parameters t, σ_, β, ρ @variables x(..), y(..), z(..) Dt = Differential(t) -eqs = [Dt(x(t)) ~ σ_*(y(t) - x(t)), - Dt(y(t)) ~ x(t)*(ρ - z(t)) - y(t), - Dt(z(t)) ~ x(t)*y(t) - β*z(t)] - +eqs = [Dt(x(t)) ~ σ_ * (y(t) - x(t)), + Dt(y(t)) ~ x(t) * (ρ - z(t)) - y(t), + Dt(z(t)) ~ x(t) * y(t) - β * z(t)] bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0] -domains = [t ∈ Interval(0.0,1.0)] +domains = [t ∈ Interval(0.0, 1.0)] dt = 0.05 input_ = length(domains) n = 12 -chain = [FastChain(FastDense(input_,n,Flux.tanh),FastDense(n,n,Flux.σ),FastDense(n,1)) for _ in 1:3] +chain = [FastChain(FastDense(input_, n, Flux.tanh), FastDense(n, n, Flux.σ), + FastDense(n, 1)) for _ in 1:3] #Generate Data -function lorenz!(du,u,p,t) - du[1] = 10.0*(u[2]-u[1]) - du[2] = u[1]*(28.0-u[3]) - u[2] - du[3] = u[1]*u[2] - (8/3)*u[3] +function lorenz!(du, u, p, t) + du[1] = 10.0 * (u[2] - u[1]) + du[2] = u[1] * (28.0 - u[3]) - u[2] + du[3] = u[1] * u[2] - (8 / 3) * u[3] end -u0 = [1.0;0.0;0.0] -tspan = (0.0,1.0) -prob = ODEProblem(lorenz!,u0,tspan) -sol = solve(prob, Tsit5(), dt=0.1) +u0 = [1.0; 0.0; 0.0] +tspan = (0.0, 1.0) +prob = ODEProblem(lorenz!, u0, tspan) +sol = solve(prob, Tsit5(), dt = 0.1) ts = [infimum(d.domain):dt:supremum(d.domain) for d in domains][1] function getData(sol) data = [] us = hcat(sol(ts).u...) ts_ = hcat(sol(ts).t...) - return [us,ts_] + return [us, ts_] end data = getData(sol) #Additional Loss Function initθs = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain)) -acum = [0;accumulate(+, length.(initθs))] -sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1] -(u_ , t_) = data +acum = [0; accumulate(+, length.(initθs))] +sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)] +(u_, t_) = data len = length(data[2]) -function additional_loss(phi, θ , p) - return sum(sum(abs2, phi[i](t_ , θ[sep[i]]) .- u_[[i], :])/len for i in 1:1:3) +function additional_loss(phi, θ, p) + return sum(sum(abs2, phi[i](t_, θ[sep[i]]) .- u_[[i], :]) / len for i in 1:1:3) end discretization = NeuralPDE.PhysicsInformedNN(chain, NeuralPDE.GridTraining(dt); - init_params =initθs, - param_estim=true, - additional_loss=additional_loss) -testθ =reduce(vcat,initθs) + init_params = initθs, + param_estim = true, + additional_loss = additional_loss) +testθ = reduce(vcat, initθs) additional_loss(discretization.phi, testθ, nothing) -@named pde_system = PDESystem(eqs,bcs,domains, - [t],[x(t), y(t), z(t)],[σ_, ρ, β], - defaults=Dict([p => 1.0 for p in [σ_, ρ, β]])) -prob = NeuralPDE.discretize(pde_system,discretization) -sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) -discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) -discretized_functions.full_loss_function([testθ;ones(3)], Float64[]) +@named pde_system = PDESystem(eqs, bcs, domains, + [t], [x(t), y(t), z(t)], [σ_, ρ, β], + defaults = Dict([p => 1.0 for p in [σ_, ρ, β]])) +prob = NeuralPDE.discretize(pde_system, discretization) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) +discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) +discretized_functions.full_loss_function([testθ; ones(3)], Float64[]) -res = Optimization.solve(prob, Optim.BFGS(); maxiters=6000) -p_ = res.minimizer[end-2:end] +res = Optimization.solve(prob, Optim.BFGS(); maxiters = 6000) +p_ = res.minimizer[(end - 2):end] @test sum(abs2, p_[1] - 10.00) < 0.1 @test sum(abs2, p_[2] - 28.00) < 0.1 -@test sum(abs2, p_[3] - (8/3)) < 0.1 +@test sum(abs2, p_[3] - (8 / 3)) < 0.1 #Plotting the system # initθ = discretization.init_params # acum = [0;accumulate(+, length.(initθ))] @@ -180,11 +180,11 @@ eq = [u(0) ~ u(0)] bc = [u(0) ~ u(0)] x0 = 0 x_end = pi -dx =pi/10 -domain = [x ∈ Interval(x0,x_end)] +dx = pi / 10 +domain = [x ∈ Interval(x0, x_end)] -hidden =10 -chain = FastChain(FastDense(1,hidden, Flux.tanh), +hidden = 10 +chain = FastChain(FastDense(1, hidden, Flux.tanh), FastDense(hidden, hidden, Flux.sin), FastDense(hidden, hidden, Flux.tanh), FastDense(hidden, 1)) @@ -193,33 +193,33 @@ initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy = NeuralPDE.GridTraining(dx) xs = collect(x0:dx:x_end)' -aproxf_(x) = @. cos(pi*x) -data =aproxf_(xs) +aproxf_(x) = @. cos(pi * x) +data = aproxf_(xs) -function additional_loss_(phi, θ , p) - sum(abs2,phi(xs,θ) .- data) +function additional_loss_(phi, θ, p) + sum(abs2, phi(xs, θ) .- data) end -discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; - initial_params=initθ, - additional_loss=additional_loss_) +discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; + initial_params = initθ, + additional_loss = additional_loss_) phi = discretization.phi phi(xs, initθ) -additional_loss_(phi, initθ , nothing) +additional_loss_(phi, initθ, nothing) -@named pde_system = PDESystem(eq,bc,domain,[x],[u(x)]) -prob = NeuralPDE.discretize(pde_system,discretization) +@named pde_system = PDESystem(eq, bc, domain, [x], [u(x)]) +prob = NeuralPDE.discretize(pde_system, discretization) -res = Optimization.solve(prob,ADAM(0.01),maxiters=500) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,BFGS(),maxiters=500) +res = Optimization.solve(prob, ADAM(0.01), maxiters = 500) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, BFGS(), maxiters = 500) -@test phi(xs,res.u) ≈ aproxf_(xs) rtol = 0.01 +@test phi(xs, res.u)≈aproxf_(xs) rtol=0.01 # xs_ = xs' # plot(xs_,data') # plot!(xs_, phi(xs,res.u)') # func(x,y) = -20.0 * exp(-0.2 * sqrt(0.5 * (x^2 + y^2))) - exp(0.5 * (cos(2 * pi * x) + cos(2 * pi * y))) + e + 20 -# func(x,y) = -abs(sin(x) * cos(y) * exp(abs(1 - (sqrt(x^2 + y^2)/pi)))) \ No newline at end of file +# func(x,y) = -abs(sin(x) * cos(y) * exp(abs(1 - (sqrt(x^2 + y^2)/pi)))) diff --git a/test/direct_function_tests.jl b/test/direct_function_tests.jl index 789452380f..4d06824350 100644 --- a/test/direct_function_tests.jl +++ b/test/direct_function_tests.jl @@ -16,34 +16,34 @@ println("Approximation of function 1D") func(x) = @. 2 + abs(x - 0.5) eq = [u(x) ~ func(x)] -bc = [u(0)~u(0)] +bc = [u(0) ~ u(0)] x0 = 0 x_end = 2 -dx= 0.001 -domain = [x ∈ Interval(x0,x_end)] +dx = 0.001 +domain = [x ∈ Interval(x0, x_end)] xs = collect(x0:dx:x_end) func_s = func(xs) -hidden =10 -chain = FastChain(FastDense(1,hidden, Flux.tanh), - FastDense(hidden, hidden, Flux.tanh), - FastDense(hidden, 1)) +hidden = 10 +chain = FastChain(FastDense(1, hidden, Flux.tanh), + FastDense(hidden, hidden, Flux.tanh), + FastDense(hidden, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy = NeuralPDE.GridTraining(0.01) -discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; initial_params=initθ) -@named pde_system = PDESystem(eq,bc,domain,[x],[u(x)]) -prob = NeuralPDE.discretize(pde_system,discretization) -sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) +discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; initial_params = initθ) +@named pde_system = PDESystem(eq, bc, domain, [x], [u(x)]) +prob = NeuralPDE.discretize(pde_system, discretization) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) -res = Optimization.solve(prob,ADAM(0.05),maxiters=1000) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,BFGS(initial_stepnorm = 0.01),maxiters=500) +res = Optimization.solve(prob, ADAM(0.05), maxiters = 1000) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, BFGS(initial_stepnorm = 0.01), maxiters = 500) -@test discretization.phi(xs',res.u) ≈ func(xs') rtol = 0.01 +@test discretization.phi(xs', res.u)≈func(xs') rtol=0.01 # plot(xs,func(xs)) # plot!(xs, discretization.phi(xs',res.u)') @@ -53,16 +53,16 @@ println("Approximation of function 1D 2") @parameters x @variables u(..) -func(x) = @. cos(5pi*x)*x +func(x) = @. cos(5pi * x) * x eq = [u(x) ~ func(x)] -bc = [u(0)~u(0)] +bc = [u(0) ~ u(0)] x0 = 0 x_end = 4 -domain = [x ∈ Interval(x0,x_end)] +domain = [x ∈ Interval(x0, x_end)] -hidden =20 -chain = FastChain(FastDense(1,hidden, Flux.sin), +hidden = 20 +chain = FastChain(FastDense(1, hidden, Flux.sin), FastDense(hidden, hidden, Flux.sin), FastDense(hidden, hidden, Flux.sin), FastDense(hidden, 1)) @@ -70,20 +70,20 @@ initθ = DiffEqFlux.initial_params(chain) strategy = NeuralPDE.GridTraining(0.01) -discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; initial_params=initθ) -@named pde_system = PDESystem(eq,bc,domain,[x],[u(x)]) -sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization) -prob = NeuralPDE.discretize(pde_system,discretization) +discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; initial_params = initθ) +@named pde_system = PDESystem(eq, bc, domain, [x], [u(x)]) +sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization) +prob = NeuralPDE.discretize(pde_system, discretization) -res = Optimization.solve(prob,ADAM(0.01),maxiters=500) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,BFGS(),maxiters=1000) +res = Optimization.solve(prob, ADAM(0.01), maxiters = 500) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, BFGS(), maxiters = 1000) -dx= 0.01 +dx = 0.01 xs = collect(x0:dx:x_end) func_s = func(xs) -@test discretization.phi(xs',res.u) ≈ func(xs') rtol = 0.01 +@test discretization.phi(xs', res.u)≈func(xs') rtol=0.01 # plot(xs,func(xs)) # plot!(xs, discretization.phi(xs',res.u)') @@ -91,11 +91,11 @@ func_s = func(xs) ## Approximation of function 2D println("Approximation of function 2D") -@parameters x,y +@parameters x, y @variables u(..) -func(x,y) = -cos(x) * cos(y) * exp(-((x - pi)^2 + (y - pi)^2)) -eq = [u(x,y) ~ func(x,y)] -bc = [u(0,0) ~ u(0,0)] +func(x, y) = -cos(x) * cos(y) * exp(-((x - pi)^2 + (y - pi)^2)) +eq = [u(x, y) ~ func(x, y)] +bc = [u(0, 0) ~ u(0, 0)] x0 = -10 x_end = 10 @@ -105,35 +105,36 @@ d = 0.4 domain = [x ∈ Interval(x0, x_end), y ∈ Interval(y0, y_end)] -hidden =25 -chain = FastChain(FastDense(2,hidden, Flux.tanh), +hidden = 25 +chain = FastChain(FastDense(2, hidden, Flux.tanh), FastDense(hidden, hidden, Flux.tanh), FastDense(hidden, hidden, Flux.tanh), FastDense(hidden, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) strategy = NeuralPDE.GridTraining(d) -discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; initial_params=initθ) -@named pde_system = PDESystem(eq,bc,domain,[x,y],[u(x, y)]) -prob = NeuralPDE.discretize(pde_system,discretization) -symprob = NeuralPDE.symbolic_discretize(pde_system,discretization) -discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) +discretization = NeuralPDE.PhysicsInformedNN(chain, strategy; initial_params = initθ) +@named pde_system = PDESystem(eq, bc, domain, [x, y], [u(x, y)]) +prob = NeuralPDE.discretize(pde_system, discretization) +symprob = NeuralPDE.symbolic_discretize(pde_system, discretization) +discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) discretized_functions.full_loss_function(initθ, Float64[]) -res = Optimization.solve(prob,ADAM(0.01),maxiters=500) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,BFGS(),maxiters=1000) -prob = remake(prob,u0=res.minimizer) -res = Optimization.solve(prob,BFGS(),maxiters=500) +res = Optimization.solve(prob, ADAM(0.01), maxiters = 500) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, BFGS(), maxiters = 1000) +prob = remake(prob, u0 = res.minimizer) +res = Optimization.solve(prob, BFGS(), maxiters = 500) phi = discretization.phi xs = collect(x0:0.1:x_end) ys = collect(y0:0.1:y_end) -u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys))) -u_real = reshape([func(x,y) 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([func(x, y) for x in xs for y in ys], (length(xs), length(ys))) diff_u = abs.(u_predict .- u_real) -@test u_predict ≈ u_real rtol = 0.05 +@test u_predict≈u_real rtol=0.05 # p1 = plot(xs, ys, u_real, st=:surface,title = "analytic"); # p2 = plot(xs, ys, u_predict, st=:surface,title = "predict"); diff --git a/test/forward_tests.jl b/test/forward_tests.jl index fa383f7746..02b01d5748 100644 --- a/test/forward_tests.jl +++ b/test/forward_tests.jl @@ -9,107 +9,105 @@ import ModelingToolkit: Interval @variables u(..) Dx = Differential(x) - eq = Dx(u(x)) ~ 0. - bcs = [u(0.) ~ u(0.)] - domains = [x ∈ Interval(0.0,1.0)] - chain = FastChain((x,p) -> x.^2) + eq = Dx(u(x)) ~ 0.0 + bcs = [u(0.0) ~ u(0.0)] + domains = [x ∈ Interval(0.0, 1.0)] + chain = FastChain((x, p) -> x .^ 2) - chain([1],Float64[]) + chain([1], Float64[]) strategy_ = NeuralPDE.GridTraining(0.1) - discretization = NeuralPDE.PhysicsInformedNN(chain,strategy_;init_params = Float64[]) - @named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)]) - prob = NeuralPDE.discretize(pde_system,discretization) - discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) + discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = Float64[]) + @named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)]) + prob = NeuralPDE.discretize(pde_system, discretization) + discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) eqs = pde_system.eqs bcs = pde_system.bcs domains = pde_system.domain dx = strategy_.dx eltypeθ = eltype(discretized_functions.flat_initθ) - depvars,indvars,dict_indvars,dict_depvars, dict_depvar_input = NeuralPDE.get_vars(pde_system.ivs, - pde_system.dvs) + depvars, indvars, dict_indvars, dict_depvars, dict_depvar_input = NeuralPDE.get_vars(pde_system.ivs, + pde_system.dvs) - train_sets = generate_training_sets(domains,dx,eqs,bcs,eltypeθ, - dict_indvars,dict_depvars) + train_sets = generate_training_sets(domains, dx, eqs, bcs, eltypeθ, + dict_indvars, dict_depvars) pde_train_sets, bcs_train_sets = train_sets - pde_train_sets = NeuralPDE.adapt(eltypeθ,pde_train_sets)[1] + pde_train_sets = NeuralPDE.adapt(eltypeθ, pde_train_sets)[1] - train_data =pde_train_sets + train_data = pde_train_sets pde_loss_function = discretized_functions.inner_pde_loss_functions[1] - dudx(x) = @. 2*x - @test pde_loss_function(train_data, Float64[]) ≈ dudx(train_data) rtol = 1e-8 + dudx(x) = @. 2 * x + @test pde_loss_function(train_data, Float64[])≈dudx(train_data) rtol=1e-8 end @testset "derivatives" begin - chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1)) + chain = FastChain(FastDense(2, 16, Flux.σ), FastDense(16, 16, Flux.σ), FastDense(16, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain)) eltypeθ = eltype(initθ) phi = NeuralPDE.Phi(chain) derivative = NeuralPDE.numeric_derivative - u_ = (cord, θ, phi)->sum(phi(cord, θ)) + u_ = (cord, θ, phi) -> sum(phi(cord, θ)) - phi([1,2], initθ) + phi([1, 2], initθ) phi_ = (p) -> phi(p, initθ)[1] - dphi = Zygote.gradient(phi_,[1.,2.]) + dphi = Zygote.gradient(phi_, [1.0, 2.0]) - function get_ε(dim, der_num,eltypeθ) + function get_ε(dim, der_num, eltypeθ) epsilon = cbrt(eps(eltypeθ)) ε = zeros(eltypeθ, dim) ε[der_num] = epsilon ε end - eps_x = get_ε(2, 1,Float64) - eps_y = get_ε(2, 2,Float64) + eps_x = get_ε(2, 1, Float64) + eps_y = get_ε(2, 2, Float64) - dphi_x = derivative(phi,u_,[1.,2.],[eps_x],1,initθ) - dphi_y = derivative(phi,u_,[1.,2.],[eps_y],1,initθ) + dphi_x = derivative(phi, u_, [1.0, 2.0], [eps_x], 1, initθ) + dphi_y = derivative(phi, u_, [1.0, 2.0], [eps_y], 1, initθ) #first order derivatives - @test isapprox(dphi[1][1], dphi_x, atol=1e-8) - @test isapprox(dphi[1][2], dphi_y, atol=1e-8) + @test isapprox(dphi[1][1], dphi_x, atol = 1e-8) + @test isapprox(dphi[1][2], dphi_y, atol = 1e-8) - dphi_x = derivative(phi,u_,[1.,2.],[[ 0.0049215667, 0.0]],1,initθ) - dphi_y = derivative(phi,u_,[1.,2.],[[0.0, 0.0049215667]],1,initθ) + dphi_x = derivative(phi, u_, [1.0, 2.0], [[0.0049215667, 0.0]], 1, initθ) + dphi_y = derivative(phi, u_, [1.0, 2.0], [[0.0, 0.0049215667]], 1, initθ) - hess_phi = Zygote.hessian(phi_,[1,2]) + hess_phi = Zygote.hessian(phi_, [1, 2]) - dphi_xx = derivative(phi,u_,[1.,2.],[eps_x,eps_x],2,initθ) - dphi_xy = derivative(phi,u_,[1.,2.],[eps_x,eps_y],2,initθ) - dphi_yy = derivative(phi,u_,[1.,2.],[eps_y,eps_y],2,initθ) + dphi_xx = derivative(phi, u_, [1.0, 2.0], [eps_x, eps_x], 2, initθ) + dphi_xy = derivative(phi, u_, [1.0, 2.0], [eps_x, eps_y], 2, initθ) + dphi_yy = derivative(phi, u_, [1.0, 2.0], [eps_y, eps_y], 2, initθ) #second order derivatives - @test isapprox(hess_phi[1], dphi_xx, atol=1e-5) - @test isapprox(hess_phi[2], dphi_xy, atol=1e-5) - @test isapprox(hess_phi[4], dphi_yy, atol=1e-5) + @test isapprox(hess_phi[1], dphi_xx, atol = 1e-5) + @test isapprox(hess_phi[2], dphi_xy, atol = 1e-5) + @test isapprox(hess_phi[4], dphi_yy, atol = 1e-5) end - @testset "Integral" begin #semi-infinite intervals @parameters x @variables u(..) I = Integral(x in ClosedInterval(0, Inf)) eq = I(u(x)) ~ 0 - bcs = [u(1.) ~ exp(1)/(exp(2) + 3)] + bcs = [u(1.0) ~ exp(1) / (exp(2) + 3)] domains = [x ∈ Interval(1.0, 2.0)] - chain = FastChain((x,p) -> exp.(x) ./ (exp.(2 .*x) .+ 3)) - chain([1],Float64[]) + chain = FastChain((x, p) -> exp.(x) ./ (exp.(2 .* x) .+ 3)) + chain([1], Float64[]) strategy_ = NeuralPDE.GridTraining(0.1) - discretization = NeuralPDE.PhysicsInformedNN(chain,strategy_;init_params = Float64[]) - @named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)]) + discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = Float64[]) + @named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)]) sym_prob = SciMLBase.symbolic_discretize(pde_system, discretization) - prob = NeuralPDE.discretize(pde_system,discretization) - discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) - inner_loss =discretized_functions.inner_pde_loss_functions[1] - exact_u = π/(3*sqrt(3)) - @test inner_loss(ones(1,1), Float64[])[1] ≈ exact_u rtol = 1e-5 - + prob = NeuralPDE.discretize(pde_system, discretization) + discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) + inner_loss = discretized_functions.inner_pde_loss_functions[1] + exact_u = π / (3 * sqrt(3)) + @test inner_loss(ones(1, 1), Float64[])[1]≈exact_u rtol=1e-5 #infinite intervals @parameters x @@ -118,15 +116,15 @@ end eqs = I(u(x)) ~ 0 domains = [x ∈ Interval(1.0, 2.0)] bcs = [u(1) ~ u(1)] - chain = FastChain((x,p) -> x .* exp.(-x .^2)) - chain([1],Float64[]) + chain = FastChain((x, p) -> x .* exp.(-x .^ 2)) + chain([1], Float64[]) - discretization = NeuralPDE.PhysicsInformedNN(chain,strategy_;init_params = Float64[]) + discretization = NeuralPDE.PhysicsInformedNN(chain, strategy_; init_params = Float64[]) @named pde_system = PDESystem(eqs, bcs, domains, [x], [u(x)]) sym_prob = SciMLBase.symbolic_discretize(pde_system, discretization) prob = SciMLBase.discretize(pde_system, discretization) - discretized_functions = NeuralPDE.discretize_inner_functions(pde_system,discretization) - inner_loss =discretized_functions.inner_pde_loss_functions[1] + discretized_functions = NeuralPDE.discretize_inner_functions(pde_system, discretization) + inner_loss = discretized_functions.inner_pde_loss_functions[1] exact_u = 0 - @test inner_loss(ones(1,1), Float64[])[1] ≈ exact_u rtol = 1e-9 + @test inner_loss(ones(1, 1), Float64[])[1]≈exact_u rtol=1e-9 end diff --git a/test/neural_adapter_tests.jl b/test/neural_adapter_tests.jl index a68d011b1b..1d671cfdcc 100644 --- a/test/neural_adapter_tests.jl +++ b/test/neural_adapter_tests.jl @@ -8,7 +8,7 @@ using Statistics using Random Random.seed!(100) -callback = function (p,l) +callback = function (p, l) println("Current loss is: $l") return false end @@ -21,94 +21,98 @@ 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) +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)] +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) +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 = Flux.tanh -chain1 = Chain(Dense(2,inner,af), - Dense(inner,inner,af), - Dense(inner,1)) +chain1 = Chain(Dense(2, inner, af), + Dense(inner, inner, af), + Dense(inner, 1)) initθ = Float64.(DiffEqFlux.initial_params(chain1)) discretization = NeuralPDE.PhysicsInformedNN(chain1, quadrature_strategy; init_params = initθ) -@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=2000) +@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 = 2000) phi = discretization.phi inner_ = 10 af = Flux.tanh -chain2 = FastChain(FastDense(2,inner_,af), - FastDense(inner_,inner_,af), - FastDense(inner_,inner_,af), - FastDense(inner_,1)) +chain2 = FastChain(FastDense(2, inner_, af), + FastDense(inner_, inner_, af), + FastDense(inner_, inner_, af), + FastDense(inner_, 1)) -initθ2 =Float64.(DiffEqFlux.initial_params(chain2)) +initθ2 = Float64.(DiffEqFlux.initial_params(chain2)) -function loss(cord,θ) - chain2(cord,θ) .- phi(cord,res.minimizer) +function loss(cord, θ) + chain2(cord, θ) .- phi(cord, res.minimizer) end grid_strategy = NeuralPDE.GridTraining(0.05) -quadrature_strategy = NeuralPDE.QuadratureTraining(reltol=1e-5,abstol=1e-5, - maxiters =50, batch=100) +quadrature_strategy = NeuralPDE.QuadratureTraining(reltol = 1e-5, abstol = 1e-5, + maxiters = 50, batch = 100) stochastic_strategy = NeuralPDE.StochasticTraining(400) -quasirandom_strategy = NeuralPDE.QuasiRandomTraining(400,resampling =false,minibatch = 200) +quasirandom_strategy = NeuralPDE.QuasiRandomTraining(400, resampling = false, + minibatch = 200) quasirandom_strategy_resampling = NeuralPDE.QuasiRandomTraining(250) -strategies1 = [grid_strategy,quadrature_strategy] +strategies1 = [grid_strategy, quadrature_strategy] reses_1 = map(strategies1) do strategy_ println("Neural adapter Poisson equation, strategy: $(nameof(typeof(strategy_)))") - prob_ = NeuralPDE.neural_adapter(loss,initθ2,pde_system, strategy_) - res_ = Optimization.solve(prob_, ADAM(0.01); maxiters=8000) - prob_ = remake(prob_,u0=res_.minimizer) - res_ = Optimization.solve(prob_, BFGS(); maxiters=200) + prob_ = NeuralPDE.neural_adapter(loss, initθ2, pde_system, strategy_) + res_ = Optimization.solve(prob_, ADAM(0.01); maxiters = 8000) + prob_ = remake(prob_, u0 = res_.minimizer) + res_ = Optimization.solve(prob_, BFGS(); maxiters = 200) end -strategies2 = [stochastic_strategy,quasirandom_strategy]# quasirandom_strategy_resampling] +strategies2 = [stochastic_strategy, quasirandom_strategy]# quasirandom_strategy_resampling] reses_2 = map(strategies2) do strategy_ println("Neural adapter Poisson equation, strategy: $(nameof(typeof(strategy_)))") - prob_ = NeuralPDE.neural_adapter(loss,initθ2,pde_system, strategy_) - res_ = Optimization.solve(prob_, ADAM(0.01); maxiters=8000) - prob_ = remake(prob_,u0=res_.minimizer) - res_ = Optimization.solve(prob_, BFGS(); maxiters=200) + prob_ = NeuralPDE.neural_adapter(loss, initθ2, pde_system, strategy_) + res_ = Optimization.solve(prob_, ADAM(0.01); maxiters = 8000) + prob_ = remake(prob_, u0 = res_.minimizer) + res_ = Optimization.solve(prob_, BFGS(); maxiters = 200) end -reses_ = [reses_1;reses_2;] +reses_ = [reses_1; reses_2] discretizations = map(res_ -> NeuralPDE.PhysicsInformedNN(chain2, grid_strategy; - init_params = res_.minimizer),reses_) + init_params = res_.minimizer), + reses_) -probs = map(discret -> NeuralPDE.discretize(pde_system,discret),discretizations) -phis = map(discret -> discret.phi ,discretizations) +probs = map(discret -> NeuralPDE.discretize(pde_system, discret), discretizations) +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) -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.minimizer)) 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_predicts = map(zip(phis,reses_)) do (phi_,res_) - reshape([first(phi_([x,y],res_.minimizer)) for x in xs for y in ys],(length(xs),length(ys))) +u_predicts = map(zip(phis, reses_)) do (phi_, res_) + reshape([first(phi_([x, y], res_.minimizer)) for x in xs for y in ys], + (length(xs), length(ys))) end -u_real = reshape([analytic_sol_func(x,y) 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))) -@test_broken u_predict ≈ u_real atol = 1e-3 -@test_broken u_predicts[1] ≈ u_real atol = 1e-2 +@test_broken u_predict≈u_real atol=1e-3 +@test_broken u_predicts[1]≈u_real atol=1e-2 map(u_predicts[2:end]) do upred - @test_broken upred ≈ u_real atol = 1e-2 + @test_broken upred≈u_real atol=1e-2 end #using Plots @@ -122,7 +126,6 @@ end # p6 = plot(xs, ys, diff_u_,linetype=:contourf,title = "error_"); # plot(p2,p1,p5,p3,p6) - ## Example, 2D Poisson equation, domain decomposition println("Example, 2D Poisson equation, domain decomposition") @parameters x y @@ -130,10 +133,10 @@ println("Example, 2D Poisson equation, domain decomposition") 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)) ~ -sin(pi * x) * sin(pi * 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) ~ -sin(pi * 1) * sin(pi * y), + u(x, 0) ~ 0.0, u(x, 1) ~ -sin(pi * x) * sin(pi * 1)] # Space x_0 = 0.0 @@ -141,38 +144,39 @@ x_end = 1.0 x_domain = Interval(x_0, x_end) y_domain = Interval(0.0, 1.0) domains = [x ∈ x_domain, - y ∈ y_domain] + y ∈ y_domain] count_decomp = 10 # Neural network af = Flux.tanh inner = 12 -chains = [FastChain(FastDense(2, inner, af), FastDense(inner, inner, af), FastDense(inner, 1)) for _ in 1:count_decomp] +chains = [FastChain(FastDense(2, inner, af), FastDense(inner, inner, af), + FastDense(inner, 1)) for _ in 1:count_decomp] initθs = map(c -> Float64.(c), DiffEqFlux.initial_params.(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] +xs_ = infimum(x_domain):(1 / count_decomp):supremum(x_domain) +xs_domain = [(xs_[i], xs_[i + 1]) for i in 1:(length(xs_) - 1)] domains_map = map(xs_domain) do (xs_dom) x_domain_ = Interval(xs_dom...) domains_ = [x ∈ x_domain_, - y ∈ y_domain] + y ∈ y_domain] end -analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2) -function create_bcs(x_domain_,phi_bound) - x_0, x_e = x_domain_.left, x_domain_.right +analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * 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)] + 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)] 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)] + 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)] bcs end @@ -183,22 +187,23 @@ pde_system_map = [] for i in 1:count_decomp println("decomposition $i") domains_ = domains_map[i] - phi_in(cord) = phis[i-1](cord,reses[i-1].minimizer) + phi_in(cord) = phis[i - 1](cord, reses[i - 1].minimizer) # phi_bound(x,y) = if (x isa Matrix) phi_in(vcat(x, fill(y,size(x)))) else phi_in(vcat(fill(x,size(y)),y)) end - phi_bound(x,y) = phi_in(vcat(x,y)) - @register_symbolic phi_bound(x,y) + phi_bound(x, y) = phi_in(vcat(x, y)) + @register_symbolic phi_bound(x, y) global phi_bound - Base.Broadcast.broadcasted(::typeof(phi_bound), x,y) = phi_bound(x,y) + Base.Broadcast.broadcasted(::typeof(phi_bound), x, y) = phi_bound(x, y) bcs_ = create_bcs(domains_[1].domain, phi_bound) - @named pde_system_ = PDESystem(eq, bcs_, domains_, [x, y], [u(x,y)]) - push!(pde_system_map,pde_system_) - strategy = NeuralPDE.GridTraining([0.1/count_decomp, 0.1]) + @named pde_system_ = PDESystem(eq, bcs_, domains_, [x, y], [u(x, y)]) + push!(pde_system_map, pde_system_) + strategy = NeuralPDE.GridTraining([0.1 / count_decomp, 0.1]) - discretization = NeuralPDE.PhysicsInformedNN(chains[i], strategy; init_params=initθs[i]) + discretization = NeuralPDE.PhysicsInformedNN(chains[i], strategy; + init_params = initθs[i]) - prob = NeuralPDE.discretize(pde_system_,discretization) - symprob = NeuralPDE.symbolic_discretize(pde_system_,discretization) - res_ = Optimization.solve(prob, BFGS(), maxiters=1500) + prob = NeuralPDE.discretize(pde_system_, discretization) + symprob = NeuralPDE.symbolic_discretize(pde_system_, discretization) + res_ = Optimization.solve(prob, BFGS(), maxiters = 1500) @show res_.minimum phi = discretization.phi push!(reses, res_) @@ -224,61 +229,64 @@ function compose_result(dx) xs_ = infimum(x_domain):dx:supremum(x_domain) xs = collect(xs_) function index_of_interval(x_) - for (i,x_domain) in enumerate(xs_domain) - if x_<= x_domain[2] && x_>= x_domain[1] + for (i, x_domain) in enumerate(xs_domain) + if x_ <= x_domain[2] && x_ >= x_domain[1] return i end end end for x_ in xs i = index_of_interval(x_) - u_predict_sub = [first(phis[i]([x_,y],reses[i].minimizer)) for y in ys] - u_real_sub = [analytic_sol_func(x_,y) for y in ys] + u_predict_sub = [first(phis[i]([x_, y], reses[i].minimizer)) for y in ys] + u_real_sub = [analytic_sol_func(x_, y) for y in ys] diff_u_sub = u_predict_sub .- u_real_sub - append!(u_predict_array,u_predict_sub) - append!(diff_u_array,diff_u_sub) + append!(u_predict_array, u_predict_sub) + append!(diff_u_array, diff_u_sub) end - xs,ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains] - u_predict = reshape(u_predict_array,(length(xs),length(ys))) - diff_u = reshape(diff_u_array, (length(xs),length(ys))) + xs, ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains] + u_predict = reshape(u_predict_array, (length(xs), length(ys))) + diff_u = reshape(diff_u_array, (length(xs), length(ys))) u_predict, diff_u end -dx= 0.01 +dx = 0.01 u_predict, diff_u = compose_result(dx) - inner_ = 18 af = Flux.tanh -chain2 = FastChain(FastDense(2,inner_,af), - FastDense(inner_,inner_,af), - FastDense(inner_,inner_,af), - FastDense(inner_,inner_,af), - FastDense(inner_,1)) +chain2 = FastChain(FastDense(2, inner_, af), + FastDense(inner_, inner_, af), + FastDense(inner_, inner_, af), + FastDense(inner_, inner_, af), + FastDense(inner_, 1)) -initθ2 =Float64.(DiffEqFlux.initial_params(chain2)) +initθ2 = Float64.(DiffEqFlux.initial_params(chain2)) -@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x,y)]) +@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) losses = map(1:count_decomp) do i - loss(cord,θ) = chain2(cord,θ) .- phis[i](cord,reses[i].minimizer) + loss(cord, θ) = chain2(cord, θ) .- phis[i](cord, reses[i].minimizer) end -prob_ = NeuralPDE.neural_adapter(losses,initθ2, pde_system_map,NeuralPDE.GridTraining([0.1/count_decomp,0.1])) -res_ = Optimization.solve(prob_, BFGS(); maxiters=2000) +prob_ = NeuralPDE.neural_adapter(losses, initθ2, pde_system_map, + NeuralPDE.GridTraining([0.1 / count_decomp, 0.1])) +res_ = Optimization.solve(prob_, BFGS(); maxiters = 2000) @show res_.minimum -prob_ = NeuralPDE.neural_adapter(losses,res_.minimizer, pde_system_map, NeuralPDE.GridTraining(0.01)) -res_ = Optimization.solve(prob_, BFGS(); maxiters=1000) +prob_ = NeuralPDE.neural_adapter(losses, res_.minimizer, pde_system_map, + NeuralPDE.GridTraining(0.01)) +res_ = Optimization.solve(prob_, BFGS(); maxiters = 1000) @show res_.minimum phi_ = NeuralPDE.Phi(chain2) -xs,ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains] -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))) +xs, ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains] +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 -@test u_predict ≈ u_real rtol = 0.1 -@test u_predict_ ≈ u_real rtol = 0.1 +@test u_predict≈u_real rtol=0.1 +@test u_predict_≈u_real rtol=0.1 # p1 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict 1"); # p2 = plot(xs, ys, u_predict_,linetype=:contourf,title = "predict 2"); diff --git a/test/runtests.jl b/test/runtests.jl index 2235e54c03..2103936c89 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,56 +3,56 @@ using SafeTestsets const GROUP = get(ENV, "GROUP", "All") -const is_APPVEYOR = Sys.iswindows() && haskey(ENV,"APPVEYOR") +const is_APPVEYOR = Sys.iswindows() && haskey(ENV, "APPVEYOR") -const is_TRAVIS = haskey(ENV,"TRAVIS") +const is_TRAVIS = haskey(ENV, "TRAVIS") -const is_CI = haskey(ENV,"CI") +const is_CI = haskey(ENV, "CI") function dev_subpkg(subpkg) subpkg_path = joinpath(dirname(@__DIR__), "lib", subpkg) - Pkg.develop(PackageSpec(path=subpkg_path)) + Pkg.develop(PackageSpec(path = subpkg_path)) end @time begin - if GROUP == "All" || GROUP == "NNODE" - @time @safetestset "NNODE" begin include("NNODE_tests.jl") end - end - if GROUP == "All" || GROUP == "NNPDE1" - @time @safetestset "NNPDE" begin include("NNPDE_tests.jl") end - end - if GROUP == "All" || GROUP == "NNPDE2" - @time @safetestset "Additional Loss" begin include("additional_loss_tests.jl") end - @time @safetestset "Direction Function Approximation" begin include("direct_function_tests.jl") end - end - if GROUP == "All" || GROUP == "NeuralAdapter" - @time @safetestset "NeuralAdapter" begin include("neural_adapter_tests.jl") end - end - if GROUP == "All" || GROUP == "IntegroDiff" - @time @safetestset "IntegroDiff" begin include("IDE_tests.jl") end - end - if GROUP == "All" || GROUP == "AdaptiveLoss" - @time @safetestset "AdaptiveLoss" begin include("adaptive_loss_tests.jl") end - end - if GROUP == "All" || GROUP == "NNKOLMOGOROV" - @time @safetestset "NNKolmogorov" begin include("NNKolmogorov_tests.jl") end - end - if GROUP == "All" || GROUP == "NNSTOPPINGTIME" - @time @safetestset "NNStopping" begin include("Stopping_tests.jl") end - end - if GROUP == "All" || GROUP == "NNRODE" + if GROUP == "All" || GROUP == "NNODE" + @time @safetestset "NNODE" begin include("NNODE_tests.jl") end + end + if GROUP == "All" || GROUP == "NNPDE1" + @time @safetestset "NNPDE" begin include("NNPDE_tests.jl") end + end + if GROUP == "All" || GROUP == "NNPDE2" + @time @safetestset "Additional Loss" begin include("additional_loss_tests.jl") end + @time @safetestset "Direction Function Approximation" begin include("direct_function_tests.jl") end + end + if GROUP == "All" || GROUP == "NeuralAdapter" + @time @safetestset "NeuralAdapter" begin include("neural_adapter_tests.jl") end + end + if GROUP == "All" || GROUP == "IntegroDiff" + @time @safetestset "IntegroDiff" begin include("IDE_tests.jl") end + end + if GROUP == "All" || GROUP == "AdaptiveLoss" + @time @safetestset "AdaptiveLoss" begin include("adaptive_loss_tests.jl") end + end + if GROUP == "All" || GROUP == "NNKOLMOGOROV" + @time @safetestset "NNKolmogorov" begin include("NNKolmogorov_tests.jl") end + end + if GROUP == "All" || GROUP == "NNSTOPPINGTIME" + @time @safetestset "NNStopping" begin include("Stopping_tests.jl") end + end + if GROUP == "All" || GROUP == "NNRODE" @time @safetestset "NNRODE" begin include("NNRODE_tests.jl") end @time @safetestset "NNParamKolmogorov" begin include("NNParamKolmogorov_tests.jl") end - end - if GROUP == "All" || GROUP == "Forward" + end + if GROUP == "All" || GROUP == "Forward" @time @safetestset "Forward" begin include("forward_tests.jl") end - end - if GROUP == "All" || GROUP == "Logging" + end + if GROUP == "All" || GROUP == "Logging" dev_subpkg("NeuralPDELogging") subpkg_path = joinpath(dirname(@__DIR__), "lib", "NeuralPDELogging") - Pkg.test(PackageSpec(name="NeuralPDELogging", path=subpkg_path)) - end - if !is_APPVEYOR && GROUP == "GPU" - @safetestset "NNPDE_gpu" begin include("NNPDE_tests_gpu.jl") end - end + Pkg.test(PackageSpec(name = "NeuralPDELogging", path = subpkg_path)) + end + if !is_APPVEYOR && GROUP == "GPU" + @safetestset "NNPDE_gpu" begin include("NNPDE_tests_gpu.jl") end + end end