diff --git a/demo_symbolic_expression.txt b/demo_symbolic_expression.txt new file mode 100644 index 000000000..f70c616de --- /dev/null +++ b/demo_symbolic_expression.txt @@ -0,0 +1,11 @@ +# Compact symbolic PINN loss template +# NN_j(...) are registered neural outputs (not expanded layer-by-layer). + +PDE residual terms: +(Differential(t, 1)(NN_1(t, x)) + Differential(x, 1)(NN_1(t, x)))^2 + +Boundary residual terms: +(NN_1(t, x) - sin(6.283185307179586x))^2 + (NN_1(t, x) - sin(-6.283185307179586t))^2 + (-sin(6.283185307179586(1.0 - t)) + NN_1(t, x))^2 + +Loss template: +sum_over_10_point_grid(PDE residual terms) + 15.0 * (Boundary residual terms) diff --git a/demo_symbolic_parser.jl b/demo_symbolic_parser.jl new file mode 100644 index 000000000..4f437f509 --- /dev/null +++ b/demo_symbolic_parser.jl @@ -0,0 +1,75 @@ +import Pkg +Pkg.activate(temp=true) +Pkg.develop(path=".") +Pkg.add(["ModelingToolkit", "Lux", "DomainSets", "Optimization", "OptimizationOptimisers", "ComponentArrays", "Plots", "ModelingToolkitNeuralNets"]) + +using NeuralPDE +using ModelingToolkit +using DomainSets +using ComponentArrays +using Optimization +using OptimizationOptimisers +using Lux +using Plots +using Random + +Random.seed!(42) + +println("==== Symbolic PINN Parser Demo (MVP) ====") + +# Advection Equation MVP problem: +# ∂u/∂t + 1.0 * ∂u/∂x = 0 +# u(0, x) = sin(2pi * x) +# True solution: u(t, x) = sin(2pi * (x - t)) + +@parameters t x +@variables u(..) +Dt = Differential(t) +Dx = Differential(x) + +eq = Dt(u(t, x)) + 1.0 * Dx(u(t, x)) ~ 0 + +bcs = [ + u(0.0, x) ~ sin(2pi * x), + u(t, 0.0) ~ sin(-2pi * t), + u(t, 1.0) ~ sin(2pi * (1.0 - t)) +] + +domains = [ + t ∈ Interval(0.0, 1.0), + x ∈ Interval(0.0, 1.0) +] + +@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t, x)]) + +# Use the newly ported MVP parser +loss_func, p_init, chain, st = build_pinn_loss( + pde_system; + width=6, + depth=1, + activation=tanh, + n_points=10, + bc_weight=15.0, + symbolic_expression_path="demo_symbolic_expression.txt" +) + +# Convert to Optimization problem +obj = OptimizationFunction((theta, _) -> loss_func(theta), Optimization.AutoForwardDiff()) +prob = OptimizationProblem(obj, collect(p_init), nothing) + +println("Starting training for Advection Equation via Symbolic Loss Function...") +res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.01), maxiters=1500) +println("Training complete. Final objective = ", res.objective) + +# Evaluate and print comparison +theta_final = ComponentArray(res.u, getaxes(p_init)) + +xs = collect(range(0.0, 1.0, length=10)) +u_pred = [first(Lux.apply(chain, [0.5, xi], theta_final, st)[1]) for xi in xs] +u_true = [sin(2pi * (xi - 0.5)) for xi in xs] + +println("\n=== Results ===") +println("Predictions: ", round.(u_pred, digits=4)) +println("Exact: ", round.(u_true, digits=4)) +println("Max error: ", round(maximum(abs.(u_pred .- u_true)), digits=4)) +println("\nSymbolic PINN parser executed successfully!\n") diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 3fc56bfdc..13df5d026 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -92,6 +92,7 @@ include("PDE_BPINN.jl") include("dgm.jl") include("NN_SDE_solve.jl") include("NN_SDE_weaksolve.jl") +include("symbolic_pinn_parser.jl") export PINOODE export NNODE, NNDAE @@ -99,6 +100,7 @@ export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde export NNSDE export SDEPINN export PhysicsInformedNN, discretize +export build_pinn_loss export BPINNsolution, BayesianPINN export DeepGalerkin diff --git a/src/adaptive_losses.jl b/src/adaptive_losses.jl index dbf86c158..80bce38dc 100644 --- a/src/adaptive_losses.jl +++ b/src/adaptive_losses.jl @@ -1,4 +1,4 @@ -abstract type AbstractAdaptiveLoss end +# AbstractAdaptiveLoss is defined in pinn_types.jl # Utils vectorify(x::Vector, ::Type{T}) where {T <: Real} = T.(x) diff --git a/src/pinn_types.jl b/src/pinn_types.jl index fd910a56b..abf3a6d4c 100644 --- a/src/pinn_types.jl +++ b/src/pinn_types.jl @@ -83,9 +83,13 @@ methodology. the vector. * `param_estim`: whether the parameters of the differential equation should be included in the values sent to the `additional_loss` function. Defaults to `false`. -* `logger`: ?? needs docs -* `log_options`: ?? why is this separate from the logger? -* `iteration`: used to control the iteration counter??? +* `logger`: a logger object (e.g. from TensorBoardLogger.jl) for recording training + metrics. Defaults to `nothing` (no logging). +* `log_options`: a `LogOptions` struct controlling logging behavior such as + `log_frequency`. Separate from `logger` to allow configuring logging behavior + independently of the logging backend. +* `iteration`: a `Ref{Int}` or `Vector{Int}` for controlling the iteration counter. + If `nothing` (default), an internal counter is used with self-incrementing enabled. * `kwargs`: Extra keyword arguments which are splatted to the `OptimizationProblem` on `solve`. """ @@ -190,182 +194,126 @@ function BayesianPINN(args...; dataset = nothing, kwargs...) end """ -`PINNRepresentation`` + AbstractAdaptiveLoss + +An abstract type for adaptive loss weighting strategies used in PINN training. +Concrete subtypes dynamically adjust the relative importance of different loss +components (PDE, boundary conditions, additional) during optimization. + +See also: [`NonAdaptiveLoss`](@ref), [`GradientScaleAdaptiveLoss`](@ref), +[`MiniMaxAdaptiveLoss`](@ref). +""" +abstract type AbstractAdaptiveLoss end + +""" + PINNRepresentation An internal representation of a physics-informed neural network (PINN). This is the struct used internally and returned for introspection by `symbolic_discretize`. ## Fields -$(FIELDS) +- `eqs`: The equations of the PDE. +- `bcs`: The boundary condition equations. +- `domains`: The domains for each of the independent variables. +- `eq_params`: The symbolic parameters of the PDE system, or `SciMLBase.NullParameters()` if none. +- `defaults`: The default values (initial conditions) dictionary from the PDESystem. +- `default_p`: The numeric values of the PDE parameters extracted from `defaults`, or `nothing` + if there are no parameters. +- `param_estim`: Whether parameters of the differential equation are estimated alongside the + neural network weights. +- `additional_loss`: The `additional_loss` function as provided by the user. +- `adaloss`: The adaptive loss function, an `AbstractAdaptiveLoss`. +- `depvars`: The dependent variables of the system as a `Vector{Symbol}`. +- `indvars`: The independent variables of the system as a `Vector{Symbol}`. +- `dict_indvars`: A `Dict{Symbol, Int}` mapping independent variable symbols to their integer indices. +- `dict_depvars`: A `Dict{Symbol, Int}` mapping dependent variable symbols to their integer indices. +- `dict_depvar_input`: A `Dict{Symbol, Vector{Symbol}}` mapping each dependent variable symbol to + the vector of independent variable symbols that it depends on. +- `logger`: The logger as provided by the user (e.g. from TensorBoardLogger.jl), or `nothing`. +- `multioutput`: Whether there are multiple outputs, i.e. a system of PDEs. +- `iteration`: The iteration counter used inside the cost function. +- `init_params`: The initial parameters as provided by the user. If the PDE is a system of PDEs, this + will be an array of arrays. If Lux.jl is used, then this is an array of ComponentArrays. +- `flat_init_params`: The initial parameters as a flattened array. This is the array that is used in the + construction of the OptimizationProblem. If a Lux.jl neural network is used, then this + flattened form is a `ComponentArray`. If the equation is a system of equations, then + `flat_init_params.depvar.x` are the parameters for the neural network corresponding + to the dependent variable `x`, and i.e. if `depvar[i] == :x` then for `phi[i]`. + If `param_estim = true`, then `flat_init_params.p` are the parameters and + `flat_init_params.depvar.x` are the neural network parameters, so + `flat_init_params.depvar.x` would be the parameters of the neural network for the + dependent variable `x` if it's a system. +- `phi`: The representation of the test function of the PDE solution. +- `derivative`: The function used for computing the derivative. +- `strategy`: The training strategy as provided by the user, an `AbstractTrainingStrategy`. +- `pde_indvars`: The independent variables used in each PDE equation. For `QuadratureTraining`, + these are the arguments; for other strategies, these are the variables. +- `bc_indvars`: The independent variables used in each boundary condition. For `QuadratureTraining`, + these are the arguments; for other strategies, these are the variables. +- `pde_integration_vars`: The integration variables found in each PDE equation (for integral terms). +- `bc_integration_vars`: The integration variables found in each boundary condition (for integral terms). +- `integral`: The numeric integration function used for integral terms in the PDE, or `nothing` + before initialization. +- `symbolic_pde_loss_functions`: The PDE loss functions as represented in Julia AST. +- `symbolic_bc_loss_functions`: The boundary condition loss functions as represented in Julia AST. +- `loss_functions`: The `PINNLossFunctions`, i.e. the generated loss functions. """ mutable struct PINNRepresentation - """ - The equations of the PDE - """ eqs::Any - """ - The boundary condition equations - """ bcs::Any - """ - The domains for each of the independent variables - """ domains::Any - """ - ??? - """ eq_params::Any - """ - ??? - """ defaults::Any - """ - ??? - """ default_p::Any - """ - Whether parameters are to be appended to the `additional_loss` - """ - param_estim::Any - """ - The `additional_loss` function as provided by the user - """ + param_estim::Bool additional_loss::Any - """ - The adaptive loss function - """ - adaloss::Any - """ - The dependent variables of the system - """ - depvars::Any - """ - The independent variables of the system - """ - indvars::Any - """ - A dictionary form of the independent variables. Define the structure ??? - """ - dict_indvars::Any - """ - A dictionary form of the dependent variables. Define the structure ??? - """ - dict_depvars::Any - """ - ??? - """ - dict_depvar_input::Any - """ - The logger as provided by the user - """ + adaloss::AbstractAdaptiveLoss + depvars::Vector{Symbol} + indvars::Vector{Symbol} + dict_indvars::Dict{Symbol, Int} + dict_depvars::Dict{Symbol, Int} + dict_depvar_input::Dict{Symbol, Vector{Symbol}} logger::Any - """ - Whether there are multiple outputs, i.e. a system of PDEs - """ multioutput::Bool - """ - The iteration counter used inside the cost function - """ iteration::Any - """ - The initial parameters as provided by the user. If the PDE is a system of PDEs, this - will be an array of arrays. If Lux.jl is used, then this is an array of ComponentArrays. - """ init_params::Any - """ - The initial parameters as a flattened array. This is the array that is used in the - construction of the OptimizationProblem. If a Lux.jl neural network is used, then this - flattened form is a `ComponentArray`. If the equation is a system of equations, then - `flat_init_params.depvar.x` are the parameters for the neural network corresponding - to the dependent variable `x`, and i.e. if `depvar[i] == :x` then for `phi[i]`. - If `param_estim = true`, then `flat_init_params.p` are the parameters and - `flat_init_params.depvar.x` are the neural network parameters, so - `flat_init_params.depvar.x` would be the parameters of the neural network for the - dependent variable `x` if it's a system. - """ flat_init_params::Any - """ - The representation of the test function of the PDE solution - """ phi::Any - """ - The function used for computing the derivative - """ derivative::Any - """ - The training strategy as provided by the user - """ strategy::AbstractTrainingStrategy - """ - ??? - """ pde_indvars::Any - """ - ??? - """ bc_indvars::Any - """ - ??? - """ pde_integration_vars::Any - """ - ??? - """ bc_integration_vars::Any - """ - ??? - """ integral::Any - """ - The PDE loss functions as represented in Julia AST - """ symbolic_pde_loss_functions::Any - """ - The boundary condition loss functions as represented in Julia AST - """ symbolic_bc_loss_functions::Any - """ - The PINNLossFunctions, i.e. the generated loss functions - """ loss_functions::Any end """ -`PINNLossFunctions`` + PINNLossFunctions -The generated functions from the PINNRepresentation +The generated functions from the `PINNRepresentation`. ## Fields -$(FIELDS) +- `bc_loss_functions`: The boundary condition loss functions. +- `pde_loss_functions`: The PDE loss functions. +- `full_loss_function`: The full loss function, combining the PDE and boundary condition loss + functions. This is the loss function that is used by the optimizer. +- `additional_loss_function`: The wrapped `additional_loss`, as pieced together for the optimizer. +- `datafree_pde_loss_functions`: The pre-data version of the PDE loss function. +- `datafree_bc_loss_functions`: The pre-data version of the BC loss function. """ -struct PINNLossFunctions - """ - The boundary condition loss functions - """ - bc_loss_functions::Any - """ - The PDE loss functions - """ - pde_loss_functions::Any - """ - The full loss function, combining the PDE and boundary condition loss functions. - This is the loss function that is used by the optimizer. - """ - full_loss_function::Any - """ - The wrapped `additional_loss`, as pieced together for the optimizer. - """ - additional_loss_function::Any - """ - The pre-data version of the PDE loss function - """ - datafree_pde_loss_functions::Any - """ - The pre-data version of the BC loss function - """ - datafree_bc_loss_functions::Any +@concrete struct PINNLossFunctions + bc_loss_functions + pde_loss_functions + full_loss_function + additional_loss_function + datafree_pde_loss_functions + datafree_bc_loss_functions end get_u() = (cord, θ, phi) -> phi(cord, θ) diff --git a/src/symbolic_pinn_parser.jl b/src/symbolic_pinn_parser.jl new file mode 100644 index 000000000..a23a95425 --- /dev/null +++ b/src/symbolic_pinn_parser.jl @@ -0,0 +1,197 @@ +# Symbolic PINN Parser MVP + +function _build_domain_grids(domains; n_points=20) + grids = Dict{Any, Vector{Float64}}() + for dom in domains + iv = dom.variables + interval = dom.domain + lo = DomainSets.leftendpoint(interval) + hi = DomainSets.rightendpoint(interval) + grids[iv] = collect(range(Float64(lo), Float64(hi), length=n_points)) + end + return grids +end + +function _replace_dv_calls(expr, dvs, ivs, nn_out_sym) + dv_op_names = [string(Symbolics.operation(Symbolics.unwrap(dv))) for dv in dvs] + expr_unwrapped = Symbolics.unwrap(expr) + + pw = SymbolicUtils.Postwalk(x -> begin + if SymbolicUtils.istree(x) + op = Symbolics.operation(x) + op_name = string(op) + idx = findfirst(==(op_name), dv_op_names) + if idx !== nothing + args = collect(Symbolics.arguments(x)) + coord_map = Dict{Any, Any}() + for i in 1:min(length(ivs), length(args)) + coord_map[ivs[i]] = args[i] + end + nn_at_args = substitute(nn_out_sym[idx], coord_map) + return Symbolics.unwrap(nn_at_args) + end + end + return x + end) + + return Num(pw(expr_unwrapped)) +end + +function _build_compact_symbolic_loss(eqs, bcs, dvs, ivs; n_points=20, bc_weight=10.0) + iv_list = join(string.(ivs), ", ") + + pde_terms = String[] + for eq in eqs + res_str = string(eq.lhs - eq.rhs) + for (j, dv) in enumerate(dvs) + dv_op = split(string(dv), "(")[1] + call_pat = Regex("\\b" * dv_op * "\\([^\\)]*\\)") + res_str = replace(res_str, call_pat => "NN_$(j)($(iv_list))") + end + push!(pde_terms, "(" * res_str * ")^2") + end + + bc_terms = String[] + for bc in bcs + res_bc_str = string(bc.lhs - bc.rhs) + for (j, dv) in enumerate(dvs) + dv_op = split(string(dv), "(")[1] + call_pat = Regex("\\b" * dv_op * "\\([^\\)]*\\)") + res_bc_str = replace(res_bc_str, call_pat => "NN_$(j)($(iv_list))") + end + push!(bc_terms, "(" * res_bc_str * ")^2") + end + + io = IOBuffer() + println(io, "# Compact symbolic PINN loss template") + println(io, "# NN_j(...) are registered neural outputs (not expanded layer-by-layer).") + println(io) + + println(io, "PDE residual terms:") + if isempty(pde_terms) + println(io, "0") + else + println(io, join(pde_terms, " + ")) + end + println(io) + + println(io, "Boundary residual terms:") + if isempty(bc_terms) + println(io, "0") + else + println(io, join(bc_terms, " + ")) + end + println(io) + + println(io, "Loss template:") + println( + io, + "sum_over_" * string(n_points) * "_point_grid(PDE residual terms) + " * + string(bc_weight) * " * (Boundary residual terms)" + ) + + return String(take!(io)) +end + +function build_pinn_loss( + pde_system, + chain = nothing; + width=16, + depth=2, + activation=tanh, + n_points=20, + bc_weight=10.0, + show_symbolic_expression=true, + symbolic_expression_path="pinn_symbolic_expression.txt", + symbolic_expression_style=:compact, + rng=Random.default_rng() +) + eqs = collect(pde_system.eqs) + bcs = collect(pde_system.bcs) + dvs = collect(pde_system.dvs) + ivs = collect(pde_system.ivs) + doms = collect(pde_system.domain) + + if chain === nothing + # We manually construct a fully connected network as fallback + # If ModelingToolkitNeuralNets is not available we could build manually via Lux + chain = Lux.Chain(Lux.Dense(length(ivs), width, activation), Lux.Dense(width, length(dvs))) + end + + p_init, st = Lux.setup(rng, chain) + p_ca = ComponentArray(p_init) + @variables p_sym[1:length(p_ca)] + p_sym_ca = ComponentArray(p_sym, getaxes(p_ca)) + + nn_out_sym, _ = Lux.apply(chain, ivs, p_sym_ca, st) + + grids = _build_domain_grids(doms; n_points=n_points) + + # 1. Compile PDE residuals into functions + compiled_res_funcs = [] + for eq in eqs + res = eq.lhs - eq.rhs + res = _replace_dv_calls(res, dvs, ivs, nn_out_sym) + res = expand_derivatives(res) + + # Compile to a function taking (p, ivs...) + bf = Symbolics.build_function(res, p_sym, ivs..., expression=Val(false)) + push!(compiled_res_funcs, bf isa Tuple ? bf[1] : bf) + end + + # 2. Compile Boundary Condition residuals + compiled_bc_funcs = [] + for bc in bcs + res_bc = bc.lhs - bc.rhs + res_bc = _replace_dv_calls(res_bc, dvs, ivs, nn_out_sym) + res_bc = expand_derivatives(res_bc) + + # Evaluate boundaries across the grid + bf_bc = Symbolics.build_function(res_bc, p_sym, ivs..., expression=Val(false)) + push!(compiled_bc_funcs, bf_bc isa Tuple ? bf_bc[1] : bf_bc) + end + + # Grid arrays + iv_vecs = [grids[iv] for iv in ivs] + + if show_symbolic_expression + output_path = isabspath(symbolic_expression_path) ? symbolic_expression_path : joinpath(pwd(), symbolic_expression_path) + + expr_text = if symbolic_expression_style == :expanded + "Expanded symbolic expression tracking full grid is disabled in Lazy Grid Sum mode." + elseif symbolic_expression_style == :compact + _build_compact_symbolic_loss(eqs, bcs, dvs, ivs; n_points=n_points, bc_weight=bc_weight) + else + error("symbolic_expression_style must be :compact or :expanded") + end + + open(output_path, "w") do io + print(io, expr_text) + end + println("Symbolic PINN loss expression (" * String(symbolic_expression_style) * ") saved to: " * output_path) + end + + loss_func = (p) -> begin + loss = zero(eltype(p)) + + # PDE Residuals over grid + for cpt in Iterators.product(iv_vecs...) + for rf in compiled_res_funcs + r = rf(p, cpt...) + loss += r^2 + end + end + + # Boundary Conditions over the grid + for cpt in Iterators.product(iv_vecs...) + for bcf in compiled_bc_funcs + r_bc = bcf(p, cpt...) + loss += bc_weight * r_bc^2 + end + end + + return loss + end + + return loss_func, p_ca, chain, st +end diff --git a/test/pinn_types_stability_tests.jl b/test/pinn_types_stability_tests.jl new file mode 100644 index 000000000..5274ad8d3 --- /dev/null +++ b/test/pinn_types_stability_tests.jl @@ -0,0 +1,36 @@ +@testitem "PINNRepresentation Type Stability Smoke Test" tags = [:nnpde1] begin + using NeuralPDE, ModelingToolkit, Lux, DomainSets, Optimization, OptimizationOptimisers + + @parameters x + @variables u(..) + Dx = Differential(x) + + # 1D ODE: u'(x) = cos(x) + eq = Dx(u(x)) ~ cos(x) + bcs = [u(0.0) ~ 0.0] + domains = [x ∈ Interval(0.0, Float64(π))] + + # Neural Network Architecture + chain = Chain(Dense(1, 8, tanh), Dense(8, 1)) + + # Discretization using Grid Training + discretization = PhysicsInformedNN(chain, GridTraining(0.1)) + @named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)]) + + # Convert PDESystem to OptimizationProblem + prob = discretize(pde_system, discretization) + + # Run optimization + res = solve(prob, Adam(0.01); maxiters = 50) + + # Verify predictions + phi = discretization.phi + u_pred = [first(phi([xi], res.u)) for xi in 0.0:0.5:Float64(π)] + u_real = [sin(xi) for xi in 0.0:0.5:Float64(π)] + + max_error = maximum(abs.(u_pred .- u_real)) + + # We do a loose bound check just to verify the pipeline doesn't error + # and the gradients are propagating reasonably + @test max_error < 1.0 +end