diff --git a/docs/pages.jl b/docs/pages.jl index 054b01f9e..31b18ec31 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -22,7 +22,8 @@ pages = ["index.md", "examples/heterogeneous.md", "examples/linear_parabolic.md", "examples/nonlinear_elliptic.md", - "examples/nonlinear_hyperbolic.md"], + "examples/nonlinear_hyperbolic.md", + "examples/complex.md"], "Manual" => Any["manual/ode.md", "manual/dae.md", "manual/pinns.md", diff --git a/docs/src/examples/complex.md b/docs/src/examples/complex.md new file mode 100644 index 000000000..0e1456e41 --- /dev/null +++ b/docs/src/examples/complex.md @@ -0,0 +1,95 @@ +# Complex Equations with PINNs + +NeuralPDE supports training PINNs with complex differential equations. This example will demonstrate how to use it for [`NNODE`](@ref). Let us consider a system of [bloch equations](https://en.wikipedia.org/wiki/Bloch_equations). Note [`QuadratureTraining`](@ref) cannot be used with complex equations due to current limitations of computing quadratures. + +As the input to this neural network is time which is real, we need to initialize the parameters of the neural network with complex values for it to output and train with complex values. + +```@example complex +using Random, NeuralPDE +using OrdinaryDiffEq +using Lux, OptimizationOptimisers +using Plots +rng = Random.default_rng() +Random.seed!(100) + +function bloch_equations(u, p, t) + Ω, Δ, Γ = p + γ = Γ / 2 + ρ₁₁, ρ₂₂, ρ₁₂, ρ₂₁ = u + d̢ρ = [im * Ω * (ρ₁₂ - ρ₂₁) + Γ * ρ₂₂; + -im * Ω * (ρ₁₂ - ρ₂₁) - Γ * ρ₂₂; + -(γ + im * Δ) * ρ₁₂ - im * Ω * (ρ₂₂ - ρ₁₁); + conj(-(γ + im * Δ) * ρ₁₂ - im * Ω * (ρ₂₂ - ρ₁₁))] + return d̢ρ +end + +u0 = zeros(ComplexF64, 4) +u0[1] = 1.0 +time_span = (0.0, 2.0) +parameters = [2.0, 0.0, 1.0] + +problem = ODEProblem(bloch_equations, u0, time_span, parameters) + +chain = Lux.Chain( + Lux.Dense(1, 16, tanh; init_weight = (rng, a...) -> Lux.kaiming_normal(rng, ComplexF64, a...)) , + Lux.Dense(16, 4; init_weight = (rng, a...) -> Lux.kaiming_normal(rng, ComplexF64, a...)) + ) +ps, st = Lux.setup(rng, chain) + +opt = OptimizationOptimisers.Adam(0.01) +ground_truth = solve(problem, Tsit5(), saveat = 0.01) +alg = NNODE(chain, opt, ps; strategy = StochasticTraining(500)) +sol = solve(problem, alg, verbose = false, maxiters = 5000, saveat = 0.01) +``` + +Now, lets plot the predictions. + +`u1`: + +```@example complex +plot(sol.t, real.(reduce(hcat, sol.u)[1, :])); +plot!(ground_truth.t, real.(reduce(hcat, ground_truth.u)[1, :])) +``` + +```@example complex +plot(sol.t, imag.(reduce(hcat, sol.u)[1, :])); +plot!(ground_truth.t, imag.(reduce(hcat, ground_truth.u)[1, :])) +``` + +`u2`: + +```@example complex +plot(sol.t, real.(reduce(hcat, sol.u)[2, :])); +plot!(ground_truth.t, real.(reduce(hcat, ground_truth.u)[2, :])) +``` + +```@example complex +plot(sol.t, imag.(reduce(hcat, sol.u)[2, :])); +plot!(ground_truth.t, imag.(reduce(hcat, ground_truth.u)[2, :])) +``` + +`u3`: + +```@example complex +plot(sol.t, real.(reduce(hcat, sol.u)[3, :])); +plot!(ground_truth.t, real.(reduce(hcat, ground_truth.u)[3, :])) +``` + +```@example complex +plot(sol.t, imag.(reduce(hcat, sol.u)[3, :])); +plot!(ground_truth.t, imag.(reduce(hcat, ground_truth.u)[3, :])) +``` + +`u4`: + +```@example complex +plot(sol.t, real.(reduce(hcat, sol.u)[4, :])); +plot!(ground_truth.t, real.(reduce(hcat, ground_truth.u)[4, :])) +``` + +```@example complex +plot(sol.t, imag.(reduce(hcat, sol.u)[4, :])); +plot!(ground_truth.t, imag.(reduce(hcat, ground_truth.u)[4, :])) +``` + +We can see it is able to learn the real parts of `u1`, `u2` and imaginary parts of `u3`, `u4`. diff --git a/docs/src/tutorials/neural_adapter.md b/docs/src/tutorials/neural_adapter.md index 93f0dd036..56d3a0679 100644 --- a/docs/src/tutorials/neural_adapter.md +++ b/docs/src/tutorials/neural_adapter.md @@ -69,7 +69,7 @@ function loss(cord, θ) ch2 .- phi(cord, res.u) end -strategy = NeuralPDE.QuadratureTraining(; reltol = 1e-6) +strategy = NeuralPDE.QuadratureTraining(; reltol = 1e-6, abstol = 1e-3) prob_ = NeuralPDE.neural_adapter(loss, init_params2, pde_system, strategy) res_ = Optimization.solve(prob_, OptimizationOptimisers.Adam(5e-3); maxiters = 10000) @@ -173,7 +173,7 @@ for i in 1:count_decomp 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.QuadratureTraining(; reltol = 1e-6) + strategy = NeuralPDE.QuadratureTraining(; reltol = 1e-6, abstol = 1e-3) discretization = NeuralPDE.PhysicsInformedNN(chains[i], strategy; init_params = init_params[i]) @@ -243,10 +243,10 @@ callback = function (p, l) end prob_ = NeuralPDE.neural_adapter(losses, init_params2, pde_system_map, - NeuralPDE.QuadratureTraining(; reltol = 1e-6)) + NeuralPDE.QuadratureTraining(; reltol = 1e-6, abstol = 1e-3)) res_ = Optimization.solve(prob_, OptimizationOptimisers.Adam(5e-3); maxiters = 5000) prob_ = NeuralPDE.neural_adapter(losses, res_.u, pde_system_map, - NeuralPDE.QuadratureTraining(; reltol = 1e-6)) + NeuralPDE.QuadratureTraining(; reltol = 1e-6, abstol = 1e-3)) res_ = Optimization.solve(prob_, OptimizationOptimisers.Adam(5e-3); maxiters = 5000) phi_ = PhysicsInformedNN(chain2, strategy; init_params = res_.u).phi diff --git a/src/ode_solve.jl b/src/ode_solve.jl index 8af6a708d..eb5ae942e 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -326,6 +326,7 @@ function (f::NNODEInterpolation)(t::Vector, idxs, ::Type{Val{0}}, p, continuity) end SciMLBase.interp_summary(::NNODEInterpolation) = "Trained neural network interpolation" +SciMLBase.allowscomplex(::NNODE) = true function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, alg::NNODE, @@ -357,6 +358,8 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, !(chain isa Lux.AbstractExplicitLayer) && error("Only Lux.AbstractExplicitLayer neural networks are supported") phi, init_params = generate_phi_θ(chain, t0, u0, init_params) + ((eltype(eltype(init_params).types[1]) <: Complex || eltype(eltype(init_params).types[2]) <: Complex) && alg.strategy isa QuadratureTraining) && + error("QuadratureTraining cannot be used with complex parameters. Use other strategies.") init_params = if alg.param_estim ComponentArrays.ComponentArray(; depvar = ComponentArrays.ComponentArray(init_params), p = prob.p) diff --git a/test/NNODE_tests.jl b/test/NNODE_tests.jl index 1e2ba3c05..72bb47f81 100644 --- a/test/NNODE_tests.jl +++ b/test/NNODE_tests.jl @@ -5,6 +5,7 @@ import Lux, OptimizationOptimisers, OptimizationOptimJL using Flux using LineSearches +rng = Random.default_rng() Random.seed!(100) @testset "Scalar" begin @@ -250,6 +251,45 @@ end @test reduce(hcat, sol.u)≈u_ atol=1e-2 end +@testset "Complex Numbers" begin + function bloch_equations(u, p, t) + Ω, Δ, Γ = p + γ = Γ / 2 + ρ₁₁, ρ₂₂, ρ₁₂, ρ₂₁ = u + d̢ρ = [im * Ω * (ρ₁₂ - ρ₂₁) + Γ * ρ₂₂; + -im * Ω * (ρ₁₂ - ρ₂₁) - Γ * ρ₂₂; + -(γ + im * Δ) * ρ₁₂ - im * Ω * (ρ₂₂ - ρ₁₁); + conj(-(γ + im * Δ) * ρ₁₂ - im * Ω * (ρ₂₂ - ρ₁₁))] + return d̢ρ + end + + u0 = zeros(ComplexF64, 4) + u0[1] = 1 + time_span = (0.0, 2.0) + parameters = [2.0, 0.0, 1.0] + + problem = ODEProblem(bloch_equations, u0, time_span, parameters) + + chain = Lux.Chain( + Lux.Dense(1, 16, tanh; init_weight = (rng, a...) -> Lux.kaiming_normal(rng, ComplexF64, a...)) , + Lux.Dense(16, 4; init_weight = (rng, a...) -> Lux.kaiming_normal(rng, ComplexF64, a...)) + ) + ps, st = Lux.setup(rng, chain) + + opt = OptimizationOptimisers.Adam(0.01) + ground_truth = solve(problem, Tsit5(), saveat = 0.01) + strategies = [StochasticTraining(500), GridTraining(0.01), WeightedIntervalTraining([0.1, 0.4, 0.4, 0.1], 500)] + + @testset "$(nameof(typeof(strategy)))" for strategy in strategies + alg = NNODE(chain, opt, ps; strategy) + sol = solve(problem, alg, verbose = false, maxiters = 5000, saveat = 0.01) + @test sol.u ≈ ground_truth.u rtol=1e-1 + end + + alg = NNODE(chain, opt, ps; strategy = QuadratureTraining()) + @test_throws ErrorException solve(problem, alg, verbose = false, maxiters = 5000, saveat = 0.01) +end + @testset "Translating from Flux" begin println("Translating from Flux") linear = (u, p, t) -> cos(2pi * t)