refactor: trigger build with latest Lux (#882)

* chore: update to latest Lux

* test: run tests enabling type-instability checks

* chore: bump minimum versions

* adapt fix

* undo NNPDE_tests.jl

* Revert "adapt fix"

This reverts commit 71c7071.

* Array{eltypeθ}

* Revert "Array{eltypeθ}"

This reverts commit 6d6d670.

* convert

* add convert topost  adapt

* fix ε_

* fix

* fix convert

* update convert, adapt

* update project.toml

* update Project

* update Project

* update Project

* update

* fix

* typos

* fix: dependencies

* test: merge qa testing

* chore: run formatter

* fix: BPINN ODE testing

* fix: update minimum versions

* refactor: update DGM implementation

* test: mark weighted training test as broken

* refactor: remove junk boilerplate

* fix: element type handling

* fix: incorrect DGM architecture

* refactor: rearrange exports

* test: run logging with non-error depwarn

* fix: forward tests

* fix: downgrade testing

* refactor: cleanup NNODE

* refactor: use explicit imports

* refactor: cleanup NNDAE

* feat: bring back NNRODE

* refactor: cleanup PINN code

* fix: eltype conversions in IntegroDiff

* refactor: cleanup neural adapter code

* refactor: bayesian PINN ODEs

* fix: missing NNRODE tests

* fix: try fixing more tests

* fix: different device handling

* fix: Bayesian PINN

* test: try reducing maxiters

* refactor: more cleanup of neural adapter

* docs: update compat

* docs: cleanup

* refactor: cleanup of deps a bit

* fix: allow scalar for number types

* fix: neural adapter tests

* fix: final round of cleanup

* fix: remove incorrect NNRODE implementation

* refactor: remove NeuralPDELogging in-favor of extension (#901)

---------

Co-authored-by: KirillZubov <[email protected]>

.JuliaFormatter.toml

@@ -1,2 +1,3 @@
 style = "sciml"
-format_markdown = true
+format_markdown = true
+annotate_untyped_fields_with_any = false

.buildkite/pipeline.yml

@@ -1,15 +1,15 @@
 steps:
-  - label: "GPU"
+  - label: "CUDA"
     plugins:
       - JuliaCI/julia#v1:
           version: "1"
       - JuliaCI/julia-test#v1:
-           coverage: false # 1000x slowdown
+           coverage: true
     agents:
       queue: "juliagpu"
       cuda: "*"
     env:
-      GROUP: 'GPU'
+      GROUP: 'CUDA'
       JULIA_PKG_SERVER: "" # it often struggles with our large artifacts
       # SECRET_CODECOV_TOKEN: "..."
     timeout_in_minutes: 240

.github/workflows/CompatHelper.yml

@@ -23,4 +23,4 @@ jobs:
       - name: CompatHelper.main()
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: julia -e 'using CompatHelper; CompatHelper.main(;subdirs=["", "docs", "lib/NeuralPDELogging"])'
+        run: julia -e 'using CompatHelper; CompatHelper.main(;subdirs=["", "docs"])'

.github/workflows/Downgrade.yml

@@ -30,7 +30,7 @@ jobs:
           - NeuralAdapter
           - IntegroDiff
         version:
-          - "1"
+          - "1.10"
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -55,7 +55,7 @@ jobs:
           GROUP: ${{ matrix.group }}
       - uses: julia-actions/julia-processcoverage@v1
         with:
-          directories: src,lib/NeuralPDELogging/src
+          directories: src,ext
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info

.github/workflows/Tests.yml

@@ -23,6 +23,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        version:
+          - "1.10"
         group:
           - "QA"
           - "ODEBPINN"
@@ -39,5 +41,6 @@ jobs:
     uses: "SciML/.github/.github/workflows/tests.yml@v1"
     with:
       group: "${{ matrix.group }}"
-      coverage-directories: "src,lib/NeuralPDELogging/src"
+      coverage-directories: "src,ext"
+      julia-version: "${{ matrix.version }}"
     secrets: "inherit"

Project.toml

@@ -4,97 +4,128 @@ authors = ["Chris Rackauckas <[email protected]>"]
 version = "5.16.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
-DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
-Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 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"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
-UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
+
+[extensions]
+NeuralPDETensorBoardLoggerExt = "TensorBoardLogger"
+
 [compat]
+ADTypes = "1.9.0"
 Adapt = "4"
 AdvancedHMC = "0.6.1"
 Aqua = "0.8"
-ArrayInterface = "7.9"
-CUDA = "5.3"
+ArrayInterface = "7.11"
+CUDA = "5.5.2"
 ChainRulesCore = "1.24"
-ComponentArrays = "0.15.14"
+ComponentArrays = "0.15.16"
+ConcreteStructs = "0.2.3"
 Cubature = "1.5"
 DiffEqNoiseProcess = "5.20"
 Distributions = "0.25.107"
 DocStringExtensions = "0.9.3"
-DomainSets = "0.6, 0.7"
-Flux = "0.14.11"
+DomainSets = "0.7"
+ExplicitImports = "1.10.1"
+Flux = "0.14.22"
 ForwardDiff = "0.10.36"
-Functors = "0.4.10"
-Integrals = "4.4"
-LineSearches = "7.2"
-LinearAlgebra = "1"
+Functors = "0.4.12"
+Integrals = "4.5"
+IntervalSets = "0.7.10"
+LineSearches = "7.3"
+LinearAlgebra = "1.10"
 LogDensityProblems = "2"
-Lux = "0.5.58"
-LuxCUDA = "0.3.2"
+Lux = "1.1.0"
+LuxCUDA = "0.3.3"
+LuxCore = "1.0.1"
+LuxLib = "1.3.2"
 MCMCChains = "6"
-MethodOfLines = "0.11"
-ModelingToolkit = "9.9"
+MLDataDevices = "1.2.0"
+MethodOfLines = "0.11.6"
+ModelingToolkit = "9.46"
 MonteCarloMeasurements = "1.1"
-Optim = "1.7.8"
-Optimization = "3.24, 4"
-OptimizationOptimJL = "0.2.1"
-OptimizationOptimisers = "0.2.1, 0.3"
-OrdinaryDiffEq = "6.74"
-Pkg = "1"
+Optimisers = "0.3.3"
+Optimization = "4"
+OptimizationOptimJL = "0.4"
+OptimizationOptimisers = "0.3"
+OrdinaryDiffEq = "6.87"
+Pkg = "1.10"
+Printf = "1.10"
 QuasiMonteCarlo = "0.3.2"
 Random = "1"
+RecursiveArrayTools = "3.27.0"
 Reexport = "1.2"
 RuntimeGeneratedFunctions = "0.5.12"
 SafeTestsets = "0.1"
-SciMLBase = "2.28"
+SciMLBase = "2.56"
 Statistics = "1.10"
-SymbolicUtils = "1.5, 2, 3"
-Symbolics = "5.27.1, 6"
-Test = "1"
-UnPack = "1"
-Zygote = "0.6.69"
+StochasticDiffEq = "6.69.1"
+SymbolicIndexingInterface = "0.3.31"
+SymbolicUtils = "3.7.2"
+Symbolics = "6.14"
+TensorBoardLogger = "0.1.24"
+Test = "1.10"
+WeightInitializers = "1.0.3"
+Zygote = "0.6.71"
 julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
+ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
+TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "CUDA", "SafeTestsets", "OptimizationOptimJL", "Pkg", "OrdinaryDiffEq", "LineSearches", "LuxCUDA", "Flux", "MethodOfLines"]
+test = ["Aqua", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "SafeTestsets", "StochasticDiffEq", "TensorBoardLogger", "Test"]

docs/Project.toml

@@ -35,20 +35,20 @@ DiffEqBase = "6.148"
 Distributions = "0.25.107"
 Documenter = "1"
 DomainSets = "0.6, 0.7"
-Flux = "0.14.11"
+Flux = "0.14.17"
 Integrals = "4"
 LineSearches = "7.2"
-Lux = "0.5.22"
+Lux = "1"
 LuxCUDA = "0.3.2"
 MethodOfLines = "0.11"
 ModelingToolkit = "9.7"
 MonteCarloMeasurements = "1"
-NeuralPDE = "5.14"
-Optimization = "3.24, 4"
-OptimizationOptimJL = "0.2.1, 0.3, 0.4"
-OptimizationOptimisers = "0.2.1, 0.3"
-OptimizationPolyalgorithms = "0.2"
-OrdinaryDiffEq = "6.74"
+NeuralPDE = "5"
+Optimization = "4"
+OptimizationOptimJL = "0.4"
+OptimizationOptimisers = "0.3"
+OptimizationPolyalgorithms = "0.3"
+OrdinaryDiffEq = "6.87"
 Plots = "1.36"
 QuasiMonteCarlo = "0.3.2"
 Random = "1"

docs/src/examples/3rd.md

@@ -36,18 +36,18 @@ bcs = [u(0.0) ~ 0.0,
 domains = [x ∈ Interval(0.0, 1.0)]
 
 # Neural network
-chain = Lux.Chain(Dense(1, 8, Lux.σ), Dense(8, 1))
+chain = Chain(Dense(1, 8, σ), Dense(8, 1))
 
 discretization = PhysicsInformedNN(chain, QuasiRandomTraining(20))
 @named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)])
 prob = discretize(pde_system, discretization)
 
 callback = function (p, l)
-    println("Current loss is: $l")
+    (p.iter % 500 == 0 || p.iter == 2000) && println("Current loss is: $l")
     return false
 end
 
-res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.01); maxiters = 2000)
+res = solve(prob, OptimizationOptimisers.Adam(0.01); maxiters = 2000, callback)
 phi = discretization.phi
 ```
 

docs/src/examples/complex.md

@@ -5,10 +5,7 @@ NeuralPDE supports training PINNs with complex differential equations. This exam
 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
+using Random, NeuralPDE, OrdinaryDiffEq, Lux, OptimizationOptimisers, Plots
 rng = Random.default_rng()
 Random.seed!(100)
 
@@ -30,11 +27,9 @@ 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...))
+chain = Chain(
+    Dense(1, 16, tanh; init_weight = kaiming_normal(ComplexF64)),
+    Dense(16, 4; init_weight = kaiming_normal(ComplexF64))
 )
 ps, st = Lux.setup(rng, chain)
 

docs/src/examples/heterogeneous.md

@@ -31,11 +31,11 @@ domains = [x ∈ Interval(0.0, 1.0),
     y ∈ Interval(0.0, 1.0)]
 
 numhid = 3
-chains = [[Lux.Chain(Dense(1, numhid, Lux.σ), Dense(numhid, numhid, Lux.σ),
-               Dense(numhid, 1)) for i in 1:2]
-          [Lux.Chain(Dense(2, numhid, Lux.σ), Dense(numhid, numhid, Lux.σ),
-               Dense(numhid, 1)) for i in 1:2]]
-discretization = NeuralPDE.PhysicsInformedNN(chains, QuadratureTraining())
+chains = [[Chain(Dense(1, numhid, σ), Dense(numhid, numhid, σ), Dense(numhid, 1))
+           for i in 1:2]
+          [Chain(Dense(2, numhid, σ), Dense(numhid, numhid, σ), Dense(numhid, 1))
+           for i in 1:2]]
+discretization = PhysicsInformedNN(chains, QuadratureTraining())
 
 @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)

docs/src/examples/ks.md

@@ -53,14 +53,13 @@ bcs = [u(x, 0) ~ u_analytic(x, 0),
     Dx(u(10, t)) ~ du(10, t)]
 
 # Space and time domains
-domains = [x ∈ Interval(-10.0, 10.0),
-    t ∈ Interval(0.0, 1.0)]
+domains = [x ∈ Interval(-10.0, 10.0), t ∈ Interval(0.0, 1.0)]
 # Discretization
 dx = 0.4;
 dt = 0.2;
 
 # Neural network
-chain = Lux.Chain(Dense(2, 12, Lux.σ), Dense(12, 12, Lux.σ), Dense(12, 1))
+chain = Chain(Dense(2, 12, σ), Dense(12, 12, σ), Dense(12, 1))
 
 discretization = PhysicsInformedNN(chain, GridTraining([dx, dt]))
 @named pde_system = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])
@@ -72,7 +71,7 @@ callback = function (p, l)
 end
 
 opt = OptimizationOptimJL.BFGS()
-res = Optimization.solve(prob, opt; maxiters = 2000)
+res = Optimization.solve(prob, opt; maxiters = 2000, callback)
 phi = discretization.phi
 ```
 

docs/src/examples/linear_parabolic.md

@@ -70,7 +70,7 @@ domains = [x ∈ Interval(0.0, 1.0),
 # Neural network
 input_ = length(domains)
 n = 15
-chain = [Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, 1)) for _ in 1:2]
+chain = [Chain(Dense(input_, n, σ), Dense(n, n, σ), Dense(n, 1)) for _ in 1:2]
 
 strategy = StochasticTraining(500)
 discretization = PhysicsInformedNN(chain, strategy)
@@ -82,18 +82,17 @@ sym_prob = symbolic_discretize(pdesystem, discretization)
 pde_inner_loss_functions = sym_prob.loss_functions.pde_loss_functions
 bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
-global iteration = 0
 callback = function (p, l)
-    if iteration % 10 == 0
+    if p.iter % 500 == 0
+        println("iter: ", p.iter)
         println("loss: ", l)
         println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
         println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
     end
-    global iteration += 1
     return false
 end
 
-res = Optimization.solve(prob, OptimizationOptimisers.Adam(1e-2); maxiters = 10000)
+res = solve(prob, OptimizationOptimisers.Adam(1e-2); maxiters = 5000, callback)
 
 phi = discretization.phi