Merge pull request #955 from nicholaskl97/stateful-lux-fix

ChrisRackauckas · web-flow · commit 0d6524547e0d · 2025-06-27T07:42:18.000Z
Adds initial states option to `Phi` and `PhysicsInformedNN`
diff --git a/Project.toml b/Project.toml
@@ -57,6 +57,7 @@ Adapt = "4"
 AdvancedHMC = "0.6.1, 0.7"
 Aqua = "0.8.9"
 ArrayInterface = "7.11"
+Boltz = "1"
 CUDA = "5.5.2"
 ChainRulesCore = "1.24"
 ComponentArrays = "0.15.16"
@@ -114,6 +115,7 @@ julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Boltz = "4544d5e4-abc5-4dea-817f-29e4c205d9c8"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
@@ -134,4 +136,4 @@ TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "FastGaussQuadrature", "Flux", "Hwloc", "InteractiveUtils", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "ReTestItems", "StochasticDiffEq", "TensorBoardLogger", "Test"]
+test = ["Aqua", "Boltz", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "FastGaussQuadrature", "Flux", "Hwloc", "InteractiveUtils", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "ReTestItems", "StochasticDiffEq", "TensorBoardLogger", "Test"]
diff --git a/README.md b/README.md
@@ -47,7 +47,7 @@ the documentation, which contains the unreleased features.
 
 ```julia
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers
-import ModelingToolkit: Interval, infimum, supremum
+import DomainSets: Interval, infimum, supremum
 
 @parameters x y
 @variables u(..)
diff --git a/docs/src/developer/debugging.md b/docs/src/developer/debugging.md
@@ -7,7 +7,7 @@ PDE solvers.
 
 ```julia
 using NeuralPDE, ModelingToolkit, Flux, Zygote
-import ModelingToolkit: Interval, infimum, supremum
+import DomainSets: Interval, infimum, supremum
 # 2d wave equation, neumann boundary condition
 @parameters x, t
 @variables u(..)
diff --git a/docs/src/examples/3rd.md b/docs/src/examples/3rd.md
@@ -17,7 +17,7 @@ We will use physics-informed neural networks.
 ```@example 3rdDerivative
 using NeuralPDE, Lux, ModelingToolkit
 using Optimization, OptimizationOptimJL, OptimizationOptimisers
-import ModelingToolkit: Interval, infimum, supremum
+import DomainSets: Interval, infimum, supremum
 
 @parameters x
 @variables u(..)
diff --git a/docs/src/examples/heterogeneous.md b/docs/src/examples/heterogeneous.md
@@ -11,7 +11,7 @@ Here, we write an arbitrary heterogeneous system:
 
 ```@example heterogeneous
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
-import ModelingToolkit: Interval
+import DomainSets: Interval
 
 @parameters x y
 @variables p(..) q(..) r(..) s(..)
diff --git a/docs/src/examples/linear_parabolic.md b/docs/src/examples/linear_parabolic.md
@@ -28,7 +28,7 @@ with a physics-informed neural network.
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers,
       OptimizationOptimJL, LineSearches
 using Plots
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..), w(..)
diff --git a/docs/src/examples/nonlinear_elliptic.md b/docs/src/examples/nonlinear_elliptic.md
@@ -29,7 +29,7 @@ This is done using a derivative neural network approximation.
 ```@example nonlinear_elliptic
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, Roots
 using Plots
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters x, y
 Dx = Differential(x)
diff --git a/docs/src/examples/nonlinear_hyperbolic.md b/docs/src/examples/nonlinear_hyperbolic.md
@@ -37,7 +37,7 @@ using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, Roots,
       LineSearches
 using SpecialFunctions
 using Plots
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..), w(..)
diff --git a/docs/src/examples/wave.md b/docs/src/examples/wave.md
@@ -17,7 +17,7 @@ Further, the solution of this equation with the given boundary conditions is pre
 
 ```@example wave
 using NeuralPDE, Lux, Optimization, OptimizationOptimJL
-using ModelingToolkit: Interval
+using DomainSets: Interval
 
 @parameters t, x
 @variables u(..)
@@ -99,7 +99,7 @@ with grid discretization `dx = 0.05` and physics-informed neural networks. Here,
 ```@example wave2
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
 using Plots, Printf
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..) Dxu(..) Dtu(..) O1(..) O2(..)
diff --git a/docs/src/tutorials/constraints.md b/docs/src/tutorials/constraints.md
@@ -23,7 +23,7 @@ with Physics-Informed Neural Networks.
 ```@example fokkerplank
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches
 using Integrals, Cubature
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 # the example is taken from this article https://arxiv.org/abs/1910.10503
 @parameters x
 @variables p(..)
diff --git a/docs/src/tutorials/derivative_neural_network.md b/docs/src/tutorials/derivative_neural_network.md
@@ -54,7 +54,7 @@ using the second numeric derivative `Dt(Dtu1(t,x))`.
 ```@example derivativenn
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers,
       OptimizationOptimJL, LineSearches, Plots
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters t, x
 Dt = Differential(t)
diff --git a/docs/src/tutorials/dgm.md b/docs/src/tutorials/dgm.md
@@ -54,7 +54,7 @@ u(t, 1) & = 0
 using NeuralPDE
 using ModelingToolkit, Optimization, OptimizationOptimisers
 using Distributions
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 using MethodOfLines, OrdinaryDiffEq
 using Plots
 
diff --git a/docs/src/tutorials/gpu.md b/docs/src/tutorials/gpu.md
@@ -45,7 +45,7 @@ In total, this looks like:
 using NeuralPDE, Lux, LuxCUDA, Random, ComponentArrays
 using Optimization
 using OptimizationOptimisers
-import ModelingToolkit: Interval
+import DomainSets: Interval
 using Plots
 using Printf
 
diff --git a/docs/src/tutorials/integro_diff.md b/docs/src/tutorials/integro_diff.md
@@ -44,7 +44,7 @@ u(0) = 0
 
 ```@example integro
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, DomainSets
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 using Plots
 
 @parameters t
diff --git a/docs/src/tutorials/low_level.md b/docs/src/tutorials/low_level.md
@@ -14,7 +14,7 @@ with Physics-Informed Neural Networks. Here is an example of using the low-level
 
 ```@example low_level
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..)
diff --git a/docs/src/tutorials/low_level_2.md b/docs/src/tutorials/low_level_2.md
@@ -28,7 +28,8 @@ With Bayesian Physics-Informed Neural Networks, here is an example of using `Bay
 
 ```@example low_level_2
 using NeuralPDE, Lux, ModelingToolkit, LinearAlgebra, AdvancedHMC
-import ModelingToolkit: Interval, infimum, supremum, Distributions
+import ModelingToolkit: Distributions
+import DomainSets: Interval, infimum, supremum
 using Plots, MonteCarloMeasurements
 
 @parameters x, t, α
diff --git a/docs/src/tutorials/neural_adapter.md b/docs/src/tutorials/neural_adapter.md
@@ -15,7 +15,7 @@ Using the example of 2D Poisson equation, it is shown how, using the method neur
 
 ```@example neural_adapter
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 using Random, ComponentArrays
 
 @parameters x y
diff --git a/docs/src/tutorials/param_estim.md b/docs/src/tutorials/param_estim.md
@@ -17,7 +17,7 @@ We start by defining the problem,
 ```@example param_estim
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, OrdinaryDiffEq,
       Plots, LineSearches
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 @parameters t, σ_, β, ρ
 @variables x(..), y(..), z(..)
 Dt = Differential(t)
diff --git a/docs/src/tutorials/pdesystem.md b/docs/src/tutorials/pdesystem.md
@@ -29,7 +29,7 @@ Using physics-informed neural networks.
 
 ```@example poisson
 using NeuralPDE, Lux, Optimization, OptimizationOptimJL, LineSearches, Plots
-using ModelingToolkit: Interval
+using DomainSets: Interval
 
 @parameters x y
 @variables u(..)
@@ -91,7 +91,7 @@ The ModelingToolkit PDE interface for this example looks like this:
 
 ```@example poisson
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
-using ModelingToolkit: Interval
+using DomainSets: Interval
 using Plots
 
 @parameters x y
diff --git a/docs/src/tutorials/systems.md b/docs/src/tutorials/systems.md
@@ -37,7 +37,7 @@ with physics-informed neural networks.
 ```@example system
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches,
       OptimizationOptimisers
-using ModelingToolkit: Interval, infimum, supremum
+using DomainSets: Interval, infimum, supremum
 
 @parameters t, x
 @variables u1(..), u2(..), u3(..)
diff --git a/src/discretize.jl b/src/discretize.jl
@@ -88,7 +88,7 @@ function build_symbolic_loss_function(pinnrep::PINNRepresentation, eqs;
         params_symbols = Symbol[]
         expr_params = Expr[]
         for (i, eq_param) in enumerate(eq_params)
-            push!(expr_params, :(ArrayInterface.allowed_getindex(p, ($i):($i))))
+            push!(expr_params, :(ArrayInterface.allowed_getindex(p, ($i))))
             push!(params_symbols, Symbol(:($eq_param)))
         end
         params_eq = Expr(:(=), build_expr(:tuple, params_symbols),
diff --git a/src/pinn_types.jl b/src/pinn_types.jl
@@ -32,18 +32,18 @@ Fields:
     smodel <: StatefulLuxLayer
 end
 
-function Phi(layer::AbstractLuxLayer)
-    return Phi(StatefulLuxLayer{true}(
-        layer, nothing, initialstates(Random.default_rng(), layer)))
+function Phi(layer::AbstractLuxLayer; init_states = nothing)
+    init_states === nothing && (init_states = initialstates(Random.default_rng(), layer))
+    return Phi(StatefulLuxLayer{true}(layer, nothing, init_states))
 end
 
 (f::Phi)(x::Number, θ) = only(cdev(f([x], θ)))
 
 (f::Phi)(x::AbstractArray, θ) = f.smodel(safe_get_device(θ)(x), θ)
 
 """
-    PhysicsInformedNN(chain, strategy; init_params = nothing, phi = nothing,
-                      param_estim = false, additional_loss = nothing,
+    PhysicsInformedNN(chain, strategy; init_params = nothing, init_states = nothing,
+                      phi = nothing, param_estim = false, additional_loss = nothing,
                       adaptive_loss = nothing, logger = nothing, log_options = LogOptions(),
                       iteration = nothing, kwargs...)
 
@@ -66,6 +66,9 @@ methodology.
 * `init_params`: the initial parameters of the neural networks. If `init_params` is not
   given, then the neural network default parameters are used. Note that for Lux, the default
   will convert to Float64.
+* `init_states`: the initial states of the neural networks. If `init_states` is not
+  given, then the neural network default states are used. Note that for Lux, the default
+  will convert to Float64.
 * `phi`: a trial solution, specified as `phi(x,p)` where `x` is the coordinates vector for
   the dependent variable and `p` are the weights of the phi function (generally the weights
   of the neural network defining `phi`). By default, this is generated from the `chain`.
@@ -90,6 +93,7 @@ methodology.
     chain <: Union{AbstractLuxLayer, AbstractArray{<:AbstractLuxLayer}}
     strategy <: Union{Nothing, AbstractTrainingStrategy}
     init_params
+    init_states
     phi <: Union{Phi, AbstractArray{<:Phi}}
     derivative
     param_estim
@@ -104,10 +108,10 @@ methodology.
 end
 
 function PhysicsInformedNN(
-        chain, strategy; init_params = nothing, derivative = nothing, param_estim = false,
-        phi::Union{Nothing, Phi, AbstractArray{<:Phi}} = nothing, additional_loss = nothing,
-        adaptive_loss = nothing, logger = nothing, log_options = LogOptions(),
-        iteration = nothing, kwargs...)
+        chain, strategy; init_params = nothing, init_states = nothing, derivative = nothing,
+        param_estim = false, phi::Union{Nothing, Phi, AbstractArray{<:Phi}} = nothing,
+        additional_loss = nothing, adaptive_loss = nothing, logger = nothing,
+        log_options = LogOptions(), iteration = nothing, kwargs...)
     multioutput = chain isa AbstractArray
     if multioutput
         chain = map(chain) do cᵢ
@@ -118,7 +122,11 @@ function PhysicsInformedNN(
         chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain))
     end
 
-    phi = phi === nothing ? (multioutput ? map(Phi, chain) : Phi(chain)) : phi
+    phi = phi === nothing ?
+        (multioutput ?
+            map(x -> Phi(x[1], x[2]), zip(chain, init_states)) :
+            Phi(chain; init_states)) :
+        phi
 
     derivative = ifelse(derivative === nothing, numeric_derivative, derivative)
 
@@ -133,9 +141,9 @@ function PhysicsInformedNN(
         self_increment = true
     end
 
-    return PhysicsInformedNN(chain, strategy, init_params, phi, derivative, param_estim,
-        additional_loss, adaptive_loss, logger, log_options, iteration, self_increment,
-        multioutput, kwargs)
+    return PhysicsInformedNN(chain, strategy, init_params, init_states, phi, derivative,
+        param_estim, additional_loss, adaptive_loss, logger, log_options, iteration,
+        self_increment, multioutput, kwargs)
 end
 
 """
diff --git a/test/BPINN_PDE_tests.jl b/test/BPINN_PDE_tests.jl
@@ -2,7 +2,7 @@
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -37,7 +37,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -80,7 +80,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -136,7 +136,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -186,7 +186,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
     import Flux
 
     Random.seed!(100)
@@ -228,7 +228,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -289,7 +289,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -356,7 +356,7 @@ end
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
-    import ModelingToolkit: Interval, infimum, supremum
+    import DomainSets: Interval, infimum, supremum
 
     Random.seed!(100)
 
@@ -459,7 +459,7 @@ end
         [α],
         defaults = Dict([α => 2.0]))
 
-    # neccesarry for loss function construction (involves Operator masking) 
+    # neccesarry for loss function construction (involves Operator masking)
     eqs = pde_system.eqs
     Dict_differentials = Dict()
     exps = toexpr.(eqs)
diff --git a/test/IDE_tests.jl b/test/IDE_tests.jl
diff --git a/test/NNPDE_cuda_tests.jl b/test/NNPDE_cuda_tests.jl
diff --git a/test/NNPDE_tests.jl b/test/NNPDE_tests.jl
diff --git a/test/adaptive_loss_tests.jl b/test/adaptive_loss_tests.jl
diff --git a/test/additional_loss_tests.jl b/test/additional_loss_tests.jl
diff --git a/test/dgm_tests.jl b/test/dgm_tests.jl
diff --git a/test/direct_function_tests.jl b/test/direct_function_tests.jl
diff --git a/test/forward_tests.jl b/test/forward_tests.jl
diff --git a/test/neural_adapter_tests.jl b/test/neural_adapter_tests.jl
