better BPINN formulation, improvements to include dataset domain points

.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

@@ -10,6 +10,11 @@ on:
       - master
     paths-ignore:
       - 'docs/**'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref_name != github.event.repository.default_branch || github.ref != 'refs/tags/v*' }}
+
 jobs:
   test:
     runs-on: ubuntu-latest
@@ -23,14 +28,13 @@ jobs:
           - NNPDE1
           - NNPDE2
           - AdaptiveLoss
-          - Logging
           - Forward
           - DGM
           - NNODE
           - NeuralAdapter
           - IntegroDiff
         version:
-          - "1"
+          - "1.10"
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -55,7 +59,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/SpellCheck.yml

@@ -10,4 +10,4 @@ jobs:
       - name: Checkout Actions Repository
         uses: actions/checkout@v4
       - name: Check spelling
-        uses: crate-ci/typos@v1.25.0 
+        uses: crate-ci/typos@v1.26.0 

.github/workflows/Tests.yml

@@ -23,14 +23,17 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        version:
+          - "1"
+          - "lts"
+          - "pre"
         group:
           - "QA"
           - "ODEBPINN"
           - "PDEBPINN"
           - "NNPDE1"
           - "NNPDE2"
           - "AdaptiveLoss"
-          - "Logging"
           - "Forward"
           - "DGM"
           - "NNODE"
@@ -39,5 +42,7 @@ jobs:
     uses: "SciML/.github/.github/workflows/tests.yml@v1"
     with:
       group: "${{ matrix.group }}"
-      coverage-directories: "src,lib/NeuralPDELogging/src"
+      julia-version: "${{ matrix.version }}"
+      coverage-directories: "src,ext"
+      julia-runtest-depwarn: "yes"  # TensorBoardLogger has a global depwarn
     secrets: "inherit"

Project.toml

@@ -1,100 +1,133 @@
 name = "NeuralPDE"
 uuid = "315f7962-48a3-4962-8226-d0f33b1235f0"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "5.16.0"
+version = "5.17.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"
+Aqua = "0.8.9"
+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"
+Hwloc = "3.3.0"
+Integrals = "4.5"
+InteractiveUtils = "<0.0.1, 1"
+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"
+Printf = "1.10"
 QuasiMonteCarlo = "0.3.2"
 Random = "1"
+ReTestItems = "1.29.0"
+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"
+Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 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"
+ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+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", "Hwloc", "InteractiveUtils", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "ReTestItems", "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/make.jl

@@ -18,5 +18,4 @@ makedocs(sitename = "NeuralPDE.jl",
         canonical = "https://docs.sciml.ai/NeuralPDE/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)

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)