Merge pull request #783 from sathvikbhagavan/sb/compats

build: compat bumps to latest versions

.buildkite/pipeline.yml

@@ -2,7 +2,7 @@ steps:
   - label: "GPU"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.9"
+          version: "1"
       - JuliaCI/julia-test#v1:
            coverage: false # 1000x slowdown
     agents:
@@ -20,7 +20,7 @@ steps:
     matrix:
       setup:
         version:
-          - "1.9"
+          - "1"
         group:
           - "NNODE"
           - "NeuralAdapter"
@@ -44,7 +44,7 @@ steps:
   - label: "Documentation"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.9"
+          version: "1"
     command: |
       julia --project=docs -e '
         println("--- :julia: Instantiating project")

.github/workflows/CI.yml

@@ -25,7 +25,7 @@ jobs:
           - Logging
           - Forward
         version:
-          - "1.9"
+          - "1"
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

.github/workflows/Downstream.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        julia-version: [1,1.6]
+        julia-version: [1]
         os: [ubuntu-latest]
         package:
           - {user: SciML, repo: PDESystemLibrary.jl, group: NeuralPDE}

.github/workflows/Invalidations.yml

@@ -18,7 +18,7 @@ jobs:
     steps:
       - uses: julia-actions/setup-julia@v1
         with:
-          version: "1.9"
+          version: "1"
       - uses: actions/checkout@v4
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-invalidations@v1

Project.toml

@@ -9,6 +9,7 @@ AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -18,7 +19,6 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
-IntegralsCubature = "c31f79ba-6e32-46d4-a52f-182a8ac42a54"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
@@ -30,7 +30,6 @@ Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -48,16 +47,17 @@ ArrayInterface = "6, 7"
 CUDA = "4"
 ChainRulesCore = "1"
 ComponentArrays = "0.13.2, 0.14, 0.15"
+Cubature = "1.5"
 DiffEqBase = "6"
 DiffEqNoiseProcess = "5.1"
 Distributions = "0.23, 0.24, 0.25"
 DocStringExtensions = "0.8, 0.9"
-DomainSets = "0.6"
+DomainSets = "0.6, 0.7"
 Flux = "0.13, 0.14"
 ForwardDiff = "0.10"
 Functors = "0.4"
-Integrals = "3.1"
-IntegralsCubature = "=0.2.2"
+Integrals = "4"
+LineSearches = "7.2"
 LogDensityProblems = "2"
 Lux = "0.4, 0.5"
 MCMCChains = "6"
@@ -67,10 +67,9 @@ Optim = ">= 1.7.8"
 Optimisers = "0.2, 0.3"
 Optimization = "3"
 QuasiMonteCarlo = "0.3.2"
-RecursiveArrayTools = "2.31"
 Reexport = "1.0"
 RuntimeGeneratedFunctions = "0.5"
-SciMLBase = "1.91, 2"
+SciMLBase = "2"
 Statistics = "1"
 StochasticDiffEq = "6.13"
 SymbolicUtils = "1"
@@ -81,12 +80,15 @@ julia = "1.6"
 
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [targets]
-test = ["Test", "CUDA", "SafeTestsets", "OptimizationOptimisers", "OptimizationOptimJL", "Pkg", "OrdinaryDiffEq"]
+test = ["Test", "CUDA", "SafeTestsets", "OptimizationOptimisers", "OptimizationOptimJL", "Pkg", "OrdinaryDiffEq", "LineSearches", "cuDNN", "LuxCUDA"]

docs/Project.toml

@@ -1,12 +1,12 @@
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
+Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
-IntegralsCubature = "c31f79ba-6e32-46d4-a52f-182a8ac42a54"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
@@ -24,12 +24,13 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
 AdvancedHMC = "0.5"
+Cubature = "1.5"
 DiffEqBase = "6.106"
+Distributions = "0.23, 0.24, 0.25"
 Documenter = "1"
-DomainSets = "0.6"
+DomainSets = "0.6, 0.7"
 Flux = "0.13, 0.14"
-Integrals = "3.3"
-IntegralsCubature = "=0.2.2"
+Integrals = "4"
 Lux = "0.4, 0.5"
 ModelingToolkit = "8.33"
 MonteCarloMeasurements = "1"
@@ -40,6 +41,6 @@ OptimizationOptimisers = "0.1"
 OptimizationPolyalgorithms = "0.1"
 OrdinaryDiffEq = "6.31"
 Plots = "1.36"
-QuasiMonteCarlo = "0.2"
+QuasiMonteCarlo = "0.3"
 Roots = "2.0"
 SpecialFunctions = "2.1"

docs/make.jl

@@ -3,6 +3,7 @@ using Documenter, NeuralPDE
 cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml", force = true)
 cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true)
 
+ENV["JULIA_DEBUG"] = "Documenter"
 ENV["GKSwstype"] = "100"
 using Plots
 

docs/src/manual/bpinns.md

@@ -14,8 +14,11 @@ NeuralPDE.ahmc_bayesian_pinn_pde
 
 ## `symbolic_discretize` for `BayesianPINN` and lower level interface.
 
-```@docs
+```@docs; canonical=false
 SciMLBase.symbolic_discretize(::PDESystem, ::NeuralPDE.AbstractPINN)
+```
+
+```@docs
 NeuralPDE.BPINNstats
 NeuralPDE.BPINNsolution
 ```

docs/src/tutorials/constraints.md

@@ -22,7 +22,7 @@ with Physics-Informed Neural Networks.
 
 ```@example fokkerplank
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
-using Integrals, IntegralsCubature
+using Integrals, Cubature
 import ModelingToolkit: Interval, infimum, supremum
 # the example is taken from this article https://arxiv.org/abs/1910.10503
 @parameters x

src/BPINN_ode.jl

@@ -150,7 +150,6 @@ BPINN Solution contains the original solution from AdvancedHMC.jl sampling(BPINN
 > estimated_nn_params - Probabilistic Estimate of NN params from sampled weights,biases
 > estimated_de_params - Probabilistic Estimate of DE params from sampled unknown DE paramters
 """
-
 struct BPINNsolution{O <: BPINNstats, E, NP, OP, P}
     original::O
     ensemblesol::E

src/NeuralPDE.jl

@@ -11,7 +11,7 @@ using Reexport, Statistics
 using Zygote, ForwardDiff, Random, Distributions
 using Adapt, DiffEqNoiseProcess, StochasticDiffEq
 using Optimization
-using Integrals, IntegralsCubature
+using Integrals, Cubature
 using QuasiMonteCarlo
 using RuntimeGeneratedFunctions
 using SciMLBase
@@ -31,7 +31,6 @@ import ModelingToolkit: Interval, infimum, supremum #,Ball
 import SciMLBase: @add_kwonly, parameterless_type
 import Optimisers
 import UnPack: @unpack
-import RecursiveArrayTools
 import ChainRulesCore, Flux, Lux, ComponentArrays
 import ChainRulesCore: @non_differentiable
 

src/ode_solve.jl

@@ -272,8 +272,8 @@ end
 
 function generate_loss(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, batch)
     ts = tspan[1]:(strategy.dx):tspan[2]
-
     # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken
+    autodiff && throw(ArgumentError("autodiff not supported for GridTraining."))
     function loss(θ, _)
         if batch
             sum(abs2, inner_loss(phi, f, autodiff, ts, θ, p))
@@ -287,6 +287,7 @@ end
 function generate_loss(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p,
                        batch)
     # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken
+    autodiff && throw(ArgumentError("autodiff not supported for StochasticTraining."))
     function loss(θ, _)
         ts = adapt(parameterless_type(θ),
                    [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)])
@@ -302,6 +303,7 @@ end
 
 function generate_loss(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p,
                        batch)
+    autodiff && throw(ArgumentError("autodiff not supported for WeightedIntervalTraining."))
     minT = tspan[1]
     maxT = tspan[2]
 

src/training_strategies.jl

@@ -320,9 +320,10 @@ function get_loss_function(loss_function, lb, ub, eltypeθ, strategy::Quadrature
             # mean(abs2,loss_(x,θ), dims=2)
             # size_x = fill(size(x)[2],(1,1))
             x = adapt(parameterless_type(ComponentArrays.getdata(θ)), x)
-            sum(abs2, loss_(x, θ), dims = 2) #./ size_x
+            sum(abs2, view(loss_(x, θ), 1, :), dims = 2) #./ size_x
         end
-        prob = IntegralProblem(integrand, lb, ub, θ, batch = strategy.batch, nout = 1)
+        integral_function = BatchIntegralFunction(integrand, max_batch = strategy.batch)
+        prob = IntegralProblem(integral_function, lb, ub, θ)
         solve(prob,
               strategy.quadrature_alg,
               reltol = strategy.reltol,

test/NNODE_tests.jl

@@ -10,14 +10,14 @@ 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))
+chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) |> Flux.f64
 luxchain = Lux.Chain(Lux.Dense(1, 5, Lux.σ), Lux.Dense(5, 1))
 opt = OptimizationOptimisers.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)
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(chain, opt; autodiff = true), dt = 1 / 20.0f0,
+@test_throws ArgumentError solve(prob, NeuralPDE.NNODE(chain, opt; autodiff = true), dt = 1 / 20.0f0,
                        verbose = true, abstol = 1.0f-10, maxiters = 200)
 
 sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true,
@@ -26,7 +26,7 @@ sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true,
 sol = solve(prob, NeuralPDE.NNODE(luxchain, opt), dt = 1 / 20.0f0, verbose = true,
             abstol = 1.0f-10, maxiters = 200)
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(luxchain, opt; autodiff = true),
+@test_throws ArgumentError solve(prob, NeuralPDE.NNODE(luxchain, opt; autodiff = true),
                        dt = 1 / 20.0f0,
                        verbose = true, abstol = 1.0f-10, maxiters = 200)
 
@@ -51,14 +51,14 @@ 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))
+chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) |> Flux.f64
 luxchain = Lux.Chain(Lux.Dense(1, 5, σ), Lux.Dense(5, 1))
 
 opt = OptimizationOptimJL.BFGS()
 sol = solve(prob, NeuralPDE.NNODE(chain, opt), dt = 1 / 20.0f0, abstol = 1e-10,
             verbose = true, maxiters = 200)
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(chain, opt; autodiff = true), dt = 1 / 20.0f0,
+@test_throws ArgumentError solve(prob, NeuralPDE.NNODE(chain, opt; autodiff = true), dt = 1 / 20.0f0,
                        abstol = 1e-10, verbose = true, maxiters = 200)
 
 sol = solve(prob, NeuralPDE.NNODE(chain, opt), abstol = 1.0f-6,
@@ -67,7 +67,7 @@ sol = solve(prob, NeuralPDE.NNODE(chain, opt), abstol = 1.0f-6,
 sol = solve(prob, NeuralPDE.NNODE(luxchain, opt), dt = 1 / 20.0f0, abstol = 1e-10,
             verbose = true, maxiters = 200)
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(luxchain, opt; autodiff = true),
+@test_throws ArgumentError solve(prob, NeuralPDE.NNODE(luxchain, opt; autodiff = true),
                        dt = 1 / 20.0f0,
                        abstol = 1e-10, verbose = true, maxiters = 200)
 
@@ -83,20 +83,20 @@ 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))
+chain = Flux.Chain(Dense(1, 128, σ), Dense(128, 1)) |> Flux.f64
 luxchain = Lux.Chain(Lux.Dense(1, 128, σ), Lux.Dense(128, 1))
 opt = OptimizationOptimisers.Adam(0.01)
 
 sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, maxiters = 400)
 @test sol.errors[:l2] < 0.5
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(chain, opt; batch = true), verbose = true,
+@test_throws AssertionError solve(prob, NeuralPDE.NNODE(chain, opt; batch = true), verbose = true,
                        maxiters = 400)
 
 sol = solve(prob, NeuralPDE.NNODE(luxchain, opt), verbose = true, maxiters = 400)
 @test sol.errors[:l2] < 0.5
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(luxchain, opt; batch = true), verbose = true,
+@test_throws AssertionError solve(prob, NeuralPDE.NNODE(luxchain, opt; batch = true), verbose = true,
                        maxiters = 400)
 
 sol = solve(prob,
@@ -142,23 +142,23 @@ sol = solve(prob, NeuralPDE.NNODE(luxchain, opt; batch = true), verbose = true,
 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))
+chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) |> Flux.f64
 luxchain = Lux.Chain(Lux.Dense(1, 5, σ), Lux.Dense(5, 1))
 
 opt = OptimizationOptimisers.Adam(0.1)
 sol = solve(prob, NeuralPDE.NNODE(chain, opt), verbose = true, maxiters = 400,
             abstol = 1.0f-8)
 @test sol.errors[:l2] < 0.5
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(chain, opt; batch = true), verbose = true,
+@test_throws AssertionError solve(prob, NeuralPDE.NNODE(chain, opt; batch = true), verbose = true,
                        maxiters = 400,
                        abstol = 1.0f-8)
 
 sol = solve(prob, NeuralPDE.NNODE(luxchain, opt), verbose = true, maxiters = 400,
             abstol = 1.0f-8)
 @test sol.errors[:l2] < 0.5
 
-@test_throws Any solve(prob, NeuralPDE.NNODE(luxchain, opt; batch = true), verbose = true,
+@test_throws AssertionError solve(prob, NeuralPDE.NNODE(luxchain, opt; batch = true), verbose = true,
                        maxiters = 400,
                        abstol = 1.0f-8)
 
@@ -214,7 +214,7 @@ u0 = [0.0f0, -1.0f0 / 2pi]
 linear_analytic = (u0, p, t) -> [sin(2pi * t) / 2pi, -cos(2pi * t) / 2pi]
 odefunction = ODEFunction(linear, analytic = linear_analytic)
 prob = ODEProblem(odefunction, u0, tspan)
-chain = Flux.Chain(Dense(1, 10, σ), Dense(10, 2))
+chain = Flux.Chain(Dense(1, 10, σ), Dense(10, 2)) |> Flux.f64
 opt = OptimizationOptimisers.Adam(0.1)
 alg = NeuralPDE.NNODE(chain, opt; autodiff = false)
 
@@ -313,7 +313,7 @@ sol1 = solve(prob, alg1, verbose = true, abstol = 1.0f-10, maxiters = 200)
 @test sol1.errors[:l2] < 0.5
 
 # StochasticTraining(Flux Chain)
-chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1))
+chain = Flux.Chain(Dense(1, 5, σ), Dense(5, 1)) |> Flux.f64
 
 (u_, t_) = (u_analytical(ts), ts)
 function additional_loss(phi, θ)