Merge pull request #921 from SciML/sb/fix_CI

Fix CI

Author: ChrisRackauckas

src/NeuralPDE.jl

@@ -29,7 +29,7 @@ using Reexport: @reexport
 using RuntimeGeneratedFunctions: RuntimeGeneratedFunctions, @RuntimeGeneratedFunction
 using SciMLBase: SciMLBase, BatchIntegralFunction, IntegralProblem, NoiseProblem,
                  OptimizationFunction, OptimizationProblem, ReturnCode, discretize,
-                 isinplace, solve, symbolic_discretize
+                 isinplace, solve, symbolic_discretize, ODEProblem, ODESolution
 using Statistics: Statistics, mean
 using QuasiMonteCarlo: QuasiMonteCarlo, LatinHypercubeSample
 using WeightInitializers: glorot_uniform, zeros32

src/PDE_BPINN.jl

@@ -20,7 +20,8 @@ function LogDensityProblems.logdensity(ltd::PDELogTargetDensity, θ)
                priorlogpdf(ltd, θ) + L2LossData(ltd, θ)
     else
         return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) +
-               priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd)
+               priorlogpdf(ltd, θ) + L2LossData(ltd, θ) +
+               ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd)
     end
 end
 
@@ -57,11 +58,11 @@ function get_lossy(pinnrep, dataset, Dict_differentials)
     # each sub vector has dataset's indvar coord's datafree_colloc_loss_function, n_subvectors = n_rows_dataset(or n_indvar_coords_dataset)
     # zip each colloc equation with args for each build_loss call per equation vector
     data_colloc_loss_functions = [[build_loss_function(pinnrep, eq, pde_indvar)
-                                       for (eq, pde_indvar, integration_indvar) in zip(
-                                          colloc_equation,
-                                          pinnrep.pde_indvars,
-                                          pinnrep.pde_integration_vars)]
-                                      for colloc_equation in colloc_equations]
+                                   for (eq, pde_indvar, integration_indvar) in zip(
+                                      colloc_equation,
+                                      pinnrep.pde_indvars,
+                                      pinnrep.pde_integration_vars)]
+                                  for colloc_equation in colloc_equations]
 
     return data_colloc_loss_functions
 end

src/ode_solve.jl

@@ -281,6 +281,12 @@ function (f::NNODEInterpolation)(t::Vector, idxs, ::Type{Val{0}}, p, continuity)
     return DiffEqArray([out[idxs, i] for i in axes(out, 2)], t)
 end
 
+function (sol::ODESolution{T, N, U, U2, D, T2, R, D2, P, A})(
+        t::AbstractVector{<:Number}, ::Type{deriv}, idxs::Nothing,
+        continuity) where {T, N, U, U2, D, T2, R, D2, P, A <: NNODE, deriv}
+    sol.interp(t, idxs, deriv, sol.prob.p, continuity)
+end
+
 SciMLBase.interp_summary(::NNODEInterpolation) = "Trained neural network interpolation"
 SciMLBase.allowscomplex(::NNODE) = true
 

src/pino_ode_solve.jl

@@ -260,7 +260,10 @@ p,t = rand(1, 50, 10), rand(1, 50, 10)
 interp(p, t)
 ```
 """
-function (f::PINOODEInterpolation)(p::AbstractArray, t::AbstractArray)
+(f::PINOODEInterpolation)(p, t) = f(t, nothing, Val{0}, p, nothing)
+
+function (f::PINOODEInterpolation)(
+        t::AbstractArray, ::Nothing, ::Type{Val{0}}, p::AbstractArray, continuity)
     if f.phi.model isa DeepONet
         f.phi((p, t), f.θ)
     elseif f.phi.model isa Chain
@@ -273,7 +276,8 @@ function (f::PINOODEInterpolation)(p::AbstractArray, t::AbstractArray)
     end
 end
 
-function (f::PINOODEInterpolation)(p::AbstractArray, t::Number)
+function (f::PINOODEInterpolation)(
+        t::Number, ::Nothing, ::Type{Val{0}}, p::AbstractArray, continuity)
     if f.phi.model isa DeepONet
         t_ = [t]
         f.phi((p, t_), f.θ)
@@ -288,8 +292,16 @@ end
 SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation"
 SciMLBase.allowscomplex(::PINOODE) = true
 
-function (sol::SciMLBase.AbstractODESolution)(t::Union{Number, AbstractArray})
-    sol.interp(sol.prob.p, t)
+function (sol::ODESolution{T, N, U, U2, D, T2, R, D2, P, A})(
+        t::AbstractArray, ::Type{deriv}, idxs::Nothing,
+        continuity) where {T, N, U, U2, D, T2, R, D2, P, A <: PINOODE, deriv}
+    sol.interp(t, idxs, deriv, sol.prob.p, continuity)
+end
+
+function (sol::ODESolution{T, N, U, U2, D, T2, R, D2, P, A})(
+        t::AbstractVector{<:Number}, ::Type{deriv}, idxs::Nothing,
+        continuity) where {T, N, U, U2, D, T2, R, D2, P, A <: PINOODE, deriv}
+    sol.interp(t, idxs, deriv, sol.prob.p, continuity)
 end
 
 function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,

src/training_strategies.jl

@@ -97,10 +97,11 @@ function merge_strategy_with_loglikelihood_function(pinnrep::PINNRepresentation,
     return pde_loss_functions, bc_loss_functions
 end
 
-function get_points_loss_functions(loss_function, train_set, eltypeθ, strategy::GridTraining;
+function get_points_loss_functions(
+        loss_function, train_set, eltypeθ, strategy::GridTraining;
         τ = nothing)
-        # loss_function length is number of all points loss is being evaluated upon
-        # train sets rows are for each indvar, cols are coordinates (row_1,row_2,..row_n) at which loss evaluated
+    # loss_function length is number of all points loss is being evaluated upon
+    # train sets rows are for each indvar, cols are coordinates (row_1,row_2,..row_n) at which loss evaluated
     function loss(θ, std)
         logpdf(
             MvNormal(loss_function(train_set, θ)[1, :],

test/BPINN_PDE_tests.jl

@@ -526,4 +526,4 @@ end
     α = 1
     @test abs(param_new - α) < 0.2 * α
     @test abs(param_new - α) < abs(param_old - α)
-end
+end

test/BPINN_tests.jl

@@ -426,4 +426,4 @@ end
              mean(abs, u[2, :] .- pmean(sol_pestim2.ensemblesol[2]))
 
     @test Loss_1 > Loss_2
-end
+end

test/NNODE_tests.jl

@@ -142,7 +142,7 @@ end
 end
 
 @testitem "Training Strategy: Others" tags=[:nnode] begin
-    using OrdinaryDiffEq, Random, Lux, Optimisers
+    using OrdinaryDiffEq, Random, Lux, Optimisers, Integrals
 
     Random.seed!(100)
 
@@ -165,13 +165,14 @@ end
     @testset "$(nameof(typeof(strategy)))" for strategy in [
         GridTraining(0.01),
         StochasticTraining(1000),
-        QuadratureTraining(reltol = 1e-3, abstol = 1e-6, maxiters = 50, batch = 100)
+        QuadratureTraining(reltol = 1e-3, abstol = 1e-6, maxiters = 50,
+            batch = 100, quadrature_alg = QuadGKJL())
     ]
         alg = NNODE(luxchain, opt; additional_loss, strategy)
         @test begin
             sol = solve(prob, alg; verbose = false, maxiters = 500, abstol = 1e-6)
             sol.errors[:l2] < 0.5
-        end broken=(strategy isa QuadratureTraining)
+        end
     end
 end
 

test/NNPDE_tests.jl

@@ -187,10 +187,10 @@ end
         phi = discretization.phi
 
         xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
-        analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2)
+        analytic = (x, y) -> (sinpi(x) * sinpi(y)) / (2pi^2)
 
         u_predict = [first(phi([x, y], res.u)) for x in xs for y in ys]
-        u_real = [analytic_sol_func(x, y) for x in xs for y in ys]
+        u_real = [analytic(x, y) for x in xs for y in ys]
 
         @test u_predict≈u_real atol=2.0
     end
@@ -394,6 +394,8 @@ end
 @testitem "PDE VI: PDE with mixed derivative" tags=[:nnpde1] setup=[NNPDE1TestSetup] begin
     using Lux, Random, Optimisers, DomainSets, Cubature, QuasiMonteCarlo, Integrals
     import ModelingToolkit: Interval, infimum, supremum
+    using OptimizationOptimJL: BFGS
+    using LineSearches: BackTracking
 
     @parameters x y
     @variables u(..)
@@ -414,15 +416,15 @@ end
     # Space and time domains
     domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
 
-    strategy = StochasticTraining(1024)
-    inner = 20
-    chain = Chain(Dense(2, inner, tanh), Dense(inner, inner, tanh), Dense(inner, 1))
+    strategy = StochasticTraining(2048)
+    inner = 32
+    chain = Chain(Dense(2, inner, sigmoid), Dense(inner, inner, sigmoid), Dense(inner, 1))
 
     discretization = PhysicsInformedNN(chain, strategy)
     @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
 
     prob = discretize(pde_system, discretization)
-    res = solve(prob, Adam(0.01); maxiters = 5000, callback)
+    res = solve(prob, BFGS(); maxiters = 500, callback)
     phi = discretization.phi
 
     analytic_sol_func(x, y) = x + x * y + y^2 / 2

test/PINO_ode_tests.jl

@@ -231,7 +231,7 @@ end
         Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast),
             Dense(10 => 10, Lux.tanh_fast)))
 
-    u = rand(2, 50)
+    u = rand(3, 50)
     v = rand(1, 40, 1)
     θ, st = Lux.setup(Random.default_rng(), deeponet)
     c = deeponet((u, v), θ, st)[1]