MTK-based Parser for Physics-Informed Neural Networks (PINNs MVP)

[ajatshatru01]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

This PR introduces the MVP implementation of an experimental, ModelingToolkit-based symbolic parser for PINNs in NeuralPDE.jl.

The primary goal of this project is to bypass the old, manual AST-transformation-based residual parsing system and replace it with a modern, compilation-centric approach using Symbolics.jl and ModelingToolkit.jl features.

Key Features Implemented

• Accessor Compliance: Accesses the MTK PDESystem cleanly using standard accessors (ModelingToolkit.get_eqs, get_bcs, get_domain, etc.).

• Symbolic NN Substitutions: Leverages ModelingToolkitNeuralNets.SymbolicNeuralNetwork and @parameters to define callable NN/DNN wrappers, allowing neural networks to be substituted directly into symbolic equations.

• Compiled Residuals: Uses Symbolics.build_function to compile equations into high-performance, point-wise numerical residual functions.

• Training Strategy Compatibility: Implements _wrap_as_datafree to wrap compiled residuals into the (cord::Matrix, θ) -> Matrix{1, N} signature, bridging the parser output directly to NeuralPDE’s existing training strategies (e.g., GridTraining, StochasticTraining, QuasiRandomTraining).

• Differentiable Loss: Integrates with Lux/Zygote to ensure the compiled loss function is fully differentiable with respect to model parameters.

Codebase Changes

1. [NEW] symbolic_pinn_parser.jl
   Contains the core parser implementation:

• Structural representations: SymbolicPINNSystem, SymbolicPINNDerivativeWrapper, and SymbolicPINNNeuralSpec.

• Substitution logic: symbolic_pinn_residual searches for dependent variables and derivative terms and replaces them with symbolic NN calls.

• Compiled functions: build_symbolic_pinn_loss acts as the main entry point to parse a system and return the loss functions, sampled collocation points, and wrapped datafree functions.

2. [MODIFY] NeuralPDE.jl

• Included the new parser source file symbolic_pinn_parser.jl.

• Imported SymbolicNeuralNetwork from ModelingToolkitNeuralNets.

3. [MODIFY] Project.toml

• Added ModelingToolkitNeuralNets to dependencies and ReTestItems to test/extras.

4. [NEW] symbolic_pinn_parser_heat_eq_tests.jl
   Added a comprehensive test suite covering:

• Basic MVP Functionality: Verifies correct parsing, substitution, residual construction, and boundary condition evaluation for the 1D Heat Equation.

• Datafree Loss Format Compatibility: Confirms that datafree loss wrappers return the standard (1, N) matrix size expected by training strategies.

• AD Compatibility: Verifies that Zygote successfully differentiates the constructed loss with respect to parameters.

To verify:- run

julia> using ReTestItems

julia> runtests("test/symbolic_pinn_parser_heat_eq_tests.jl")

I had run the test on my device and the result is as follows:-

Add any other context about the problem here.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Introduces an experimental "symbolic PINN parser" MVP that lowers a ModelingToolkit PDESystem (heat equation) into symbolic residuals backed by SymbolicNeuralNetwork, compiles them into callable PDE/BC residual functions, and exposes a simple mean-square loss plus (cord, θ) "datafree" wrappers compatible with NeuralPDE training strategies.

Changes:

• Adds src/symbolic_pinn_parser.jl implementing parse_pde_system, symbolic residual construction with same-direction derivative handling, compilation, collocation sampling, and a build_symbolic_pinn_loss entry point.
• Wires the new file and ModelingToolkitNeuralNets.SymbolicNeuralNetwork dependency into src/NeuralPDE.jl and Project.toml.
• Adds test/symbolic_pinn_parser_heat_eq_tests.jl covering parsing, datafree loss matrix shape, and Zygote gradient compatibility.

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 4 comments.

File	Description
src/symbolic_pinn_parser.jl	New parser/compiler producing symbolic PINN residuals and loss functions.
src/NeuralPDE.jl	Imports SymbolicNeuralNetwork and includes the new parser file.
Project.toml	Adds ModelingToolkitNeuralNets and ReTestItems dependencies.
test/symbolic_pinn_parser_heat_eq_tests.jl	Tests for parsing, datafree format, and Zygote gradients.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sathvikbhagavan]

Can you post an example of the generated loss function for the heat equation that it produces? It might be good to add tests to check its correctness later.

[ChrisRackauckas]

Using the pattern matching and rules system is probably more robust here. See @rule

[AayushSabharwal]

@rule is bad and should not be used. It is slow and has a bunch of sharp edges. I have ideas for a much more robust alternative but haven't found the time to flesh it out yet.

[ajatshatru01]

Using the pattern matching and rules system is probably more robust here. See @rule

as highlighted by @AayushSabharwal @rule has significant overhead per node, nevertheless, my recursive tree-walking is also fragile. I explored a bit and found about prewalk, which is provided by SymbolicUtils.

@AayushSabharwal single pass prewalk substitution works for this ??

[AayushSabharwal]

Prewalk is fine

[ajatshatru01]

ok will implement it by this week

[ajatshatru01]

@AayushSabharwal so i checked the SymbolicUtlis provided Rewriters.Prewalk, i dont think it explicitly stops recursing once a match is found and thus cant be used in my use-case, i am thinking of implementing my own version of prewalk which stops recursion once matched
what's your thoughts on this?

[AayushSabharwal]

What exactly do you need to do - what is the input and what is the desired output?

[ajatshatru01]

ah i got it, the inbuilt prewalk is sufficient, i just needed to wrap my matcher logic

[ajatshatru01]

Can you post an example of the generated loss function for the heat equation that it produces? It might be good to add tests to check its correctness later.

I tested the symbolic PINN loss generation for the 1D heat equation:

The generated symbolic residuals are:

I also verified the generated expressions by substituting the analytical solution exp(-pi^2 * input[2]) * sin(pi * input[1]) and all PDE and BC residuals evaluate to zero

These residuals are then automatically lowered into executable loss functions via _compiled_residual, _wrap_as_datafree, and _mean_square, yielding pde_loss(theta), bc_loss(theta) and loss(theta)

@sathvikbhagavan do you specifically require the loss expression cause then i will just add a function to my parser mvp such that every PDE automatically gets a symbolic loss expression generated from the parser output

[sathvikbhagavan]

do you specifically require the loss expression cause then i will just add a function to my parser mvp such that every PDE automatically gets a symbolic loss expression generated from the parser output

Yeah, that would be good

[ajatshatru01]

do you specifically require the loss expression cause then i will just add a function to my parser mvp such that every PDE automatically gets a symbolic loss expression generated from the parser output

Yeah, that would be good

hey, really sorry for the late reply, i am a bit busy this week due to a college related program, however i am currently implementing prewalk substitution, once that gets done and i integrate the mvp with discretize api, i will rerun the loss generation script and provide you with the generated loss expression, its just that this week is a lot hectic, really sorry about this

[ajatshatru01]

Yeah, that would be good

julia> include("train_pinn_mvp.jl")
(pde_coords, bc1_coords, bc2_coords, bc3_coords, theta)->begin
N_pde = size(pde_coords, 2)
pde_loss = sum(1:N_pde) do j
pt = #= C:\Users\ajats\Sciml\NeuralPDE.jl\train_pinn_mvp.jl:47 =# @view(pde_coords[:, j])
abs2((function (x, t, NN, DNN, p)
#= C:\Users\ajats.julia\packages\Symbolics\AmMIj\src\build_function.jl:143 =# @inbounds begin
begin
(+)(()(-1, (getindex)(DNN((SymbolicUtils.array_literal)((2,), x, t), p, [1, 1]), 1)), (getindex)(DNN((SymbolicUtils.array_literal)((2,), x, t), p, [2]), 1))
end
end
end)(pt..., theta))
end / N_pde
N_bc1 = size(bc1_coords, 2)
bc1_loss = sum(1:N_bc1) do j
pt = #= C:\Users\ajats\Sciml\NeuralPDE.jl\train_pinn_mvp.jl:54 =# @view(bc1_coords[:, j])
abs2((function (x, t, NN, DNN, p)
#= C:\Users\ajats.julia\packages\Symbolics\AmMIj\src\build_function.jl:143 =# @inbounds begin
begin
(getindex)(NN((SymbolicUtils.array_literal)((2,), 0.0, t), p), 1)
end
end
end)(pt..., theta))
end / N_bc1
N_bc2 = size(bc2_coords, 2)
bc2_loss = sum(1:N_bc2) do j
pt = #= C:\Users\ajats\Sciml\NeuralPDE.jl\train_pinn_mvp.jl:61 =# @view(bc2_coords[:, j])
abs2((function (x, t, NN, DNN, p)
#= C:\Users\ajats.julia\packages\Symbolics\AmMIj\src\build_function.jl:143 =# @inbounds begin
begin
(getindex)(NN((SymbolicUtils.array_literal)((2,), 1.0, t), p), 1)
end
end
end)(pt..., theta))
end / N_bc2
N_bc3 = size(bc3_coords, 2)
bc3_loss = sum(1:N_bc3) do j
pt = #= C:\Users\ajats\Sciml\NeuralPDE.jl\train_pinn_mvp.jl:68 =# @view(bc3_coords[:, j])
abs2((function (x, t, NN, DNN, p)
#= C:\Users\ajats.julia\packages\Symbolics\AmMIj\src\build_function.jl:143 =# @inbounds begin
begin
(+)((getindex)(NN((SymbolicUtils.array_literal)((2,), x, 0.0), p), 1), ()(-1, NaNMath.sin((*)(x, (identity)(π)))))
end
end
end)(pt..., theta))
end / N_bc3
return pde_loss + bc1_loss + bc2_loss + bc3_loss
end

ran a script to generate the entire loss expression for a 1D heat equation

[ajatshatru01]

i tried to get it into an understandable form and it represents this:-

I believe everything is correct with respect to the loss function being generated

[ajatshatru01]

@ChrisRackauckas i have added the loss function above
Also this is the AST for the generated loss function
Expr
head: Symbol ->
args: Array{Any}((2,))
1: Expr
head: Symbol tuple
args: Array{Any}((5,))
1: Symbol pde_coords
2: Symbol bc1_coords
3: Symbol bc2_coords
4: Symbol bc3_coords
5: Symbol theta
2: Expr
head: Symbol block
args: Array{Any}((9,))
1: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol N_pde
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol size
2: Symbol pde_coords
3: Int64 2
2: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol pde_loss
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol /
2: Expr
3: Symbol N_pde
3: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol N_bc1
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol size
2: Symbol bc1_coords
3: Int64 2
4: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol bc1_loss
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol /
2: Expr
3: Symbol N_bc1
5: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol N_bc2
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol size
2: Symbol bc2_coords
3: Int64 2
6: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol bc2_loss
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol /
2: Expr
3: Symbol N_bc2
7: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol N_bc3
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol size
2: Symbol bc3_coords
3: Int64 2
8: Expr
head: Symbol =
args: Array{Any}((2,))
1: Symbol bc3_loss
2: Expr
head: Symbol call
args: Array{Any}((3,))
1: Symbol /
2: Expr
3: Symbol N_bc3
9: Expr
head: Symbol return
args: Array{Any}((1,))
1: Expr
head: Symbol call
args: Array{Any}((5,))
1: Symbol +
2: Symbol pde_loss
3: Symbol bc1_loss
4: Symbol bc2_loss
5: Symbol bc3_loss

[sathvikbhagavan]

@ChrisRackauckas i have added the loss function above Also this is the AST for the generated loss function Expr head: Symbol -> args: Array{Any}((2,)) 1: Expr head: Symbol tuple args: Array{Any}((5,)) 1: Symbol pde_coords 2: Symbol bc1_coords 3: Symbol bc2_coords 4: Symbol bc3_coords 5: Symbol theta 2: Expr head: Symbol block args: Array{Any}((9,)) 1: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_pde 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol pde_coords 3: Int64 2 2: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol pde_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_pde 3: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc1 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc1_coords 3: Int64 2 4: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc1_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc1 5: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc2 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc2_coords 3: Int64 2 6: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc2_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc2 7: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc3 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc3_coords 3: Int64 2 8: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc3_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc3 9: Expr head: Symbol return args: Array{Any}((1,)) 1: Expr head: Symbol call args: Array{Any}((5,)) 1: Symbol + 2: Symbol pde_loss 3: Symbol bc1_loss 4: Symbol bc2_loss 5: Symbol bc3_loss

Can't this be printed as a symbolic expression?

[ajatshatru01]

@ChrisRackauckas i have added the loss function above Also this is the AST for the generated loss function Expr head: Symbol -> args: Array{Any}((2,)) 1: Expr head: Symbol tuple args: Array{Any}((5,)) 1: Symbol pde_coords 2: Symbol bc1_coords 3: Symbol bc2_coords 4: Symbol bc3_coords 5: Symbol theta 2: Expr head: Symbol block args: Array{Any}((9,)) 1: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_pde 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol pde_coords 3: Int64 2 2: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol pde_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_pde 3: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc1 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc1_coords 3: Int64 2 4: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc1_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc1 5: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc2 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc2_coords 3: Int64 2 6: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc2_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc2 7: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc3 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc3_coords 3: Int64 2 8: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc3_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc3 9: Expr head: Symbol return args: Array{Any}((1,)) 1: Expr head: Symbol call args: Array{Any}((5,)) 1: Symbol + 2: Symbol pde_loss 3: Symbol bc1_loss 4: Symbol bc2_loss 5: Symbol bc3_loss

Can't this be printed as a symbolic expression?

i have only printed symbolic residuals
will try this and send u the results

[ajatshatru01]

@sathvikbhagavan can you also pls check the integration with discretize.jl , for now i have made it such that we can pass a boolean when writing a script and set it to true to invoke my parser and when set to false it invokes the default parser

[sathvikbhagavan]

@sathvikbhagavan can you also pls check the integration with discretize.jl , for now i have made it such that we can pass a boolean when writing a script and set it to true to invoke my parser and when set to false it invokes the default parser

I think we should just remove the old one and have a breaking release once the parser is finished and tested completely. @ChrisRackauckas, your thoughts?

[ajatshatru01]

@sathvikbhagavan can you also pls check the integration with discretize.jl , for now i have made it such that we can pass a boolean when writing a script and set it to true to invoke my parser and when set to false it invokes the default parser

I think we should just remove the old one and have a breaking release once the parser is finished and tested completely. @ChrisRackauckas, your thoughts?

Yea i get it, i have just opted for this till i am 100% sure that my parser is working and good
i even ran my parser on a 1D and 4D equations to test convergence wrt default method and i was giving comparable results, somewhat better too

[ajatshatru01]

@ChrisRackauckas i have added the loss function above Also this is the AST for the generated loss function Expr head: Symbol -> args: Array{Any}((2,)) 1: Expr head: Symbol tuple args: Array{Any}((5,)) 1: Symbol pde_coords 2: Symbol bc1_coords 3: Symbol bc2_coords 4: Symbol bc3_coords 5: Symbol theta 2: Expr head: Symbol block args: Array{Any}((9,)) 1: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_pde 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol pde_coords 3: Int64 2 2: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol pde_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_pde 3: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc1 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc1_coords 3: Int64 2 4: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc1_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc1 5: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc2 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc2_coords 3: Int64 2 6: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc2_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc2 7: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol N_bc3 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol size 2: Symbol bc3_coords 3: Int64 2 8: Expr head: Symbol = args: Array{Any}((2,)) 1: Symbol bc3_loss 2: Expr head: Symbol call args: Array{Any}((3,)) 1: Symbol / 2: Expr 3: Symbol N_bc3 9: Expr head: Symbol return args: Array{Any}((1,)) 1: Expr head: Symbol call args: Array{Any}((5,)) 1: Symbol + 2: Symbol pde_loss 3: Symbol bc1_loss 4: Symbol bc2_loss 5: Symbol bc3_loss

Can't this be printed as a symbolic expression?

i have only printed symbolic residuals will try this and send u the results

@sathvikbhagavan so i generated the symbolic loss expression for the same pde over some concrete collocation points pde_pts = [[0.25, 0.5], [0.75, 0.5]]
bc1_pts = [[0.0, 0.5]]
bc2_pts = [[1.0, 0.5]]
bc3_pts = [[0.5, 0.0]]
this is the loss expression:-
abs2((NN(SymbolicUtils.array_literal((2,), 1.0, 0.5), p))[1]) + (abs2(-(DNN(SymbolicUtils.array_literal((2,), 0.75, 0.5), p, [1, 1]))[1] + (DNN(SymbolicUtils.array_literal((2,), 0.75, 0.5), p, [2]))[1]) + abs2(-(DNN(SymbolicUtils.array_literal((2,), 0.25, 0.5), p, [1, 1]))[1] + (DNN(SymbolicUtils.array_literal((2,), 0.25, 0.5), p, [2]))[1])) / 2 + abs2((NN(SymbolicUtils.array_literal((2,), 0.0, 0.5), p))[1]) + abs2(-sin(0.5π) + (NN(SymbolicUtils.array_literal((2,), 0.5, 0.0), p))[1])

it is the same loss as before but with symbolic abstraction

[ajatshatru01]

@sathvikbhagavan so i had implemented exact AD in my parser, however i tested the training time with my parser and the legacy path that is currently being implemented and mine was about 3 times slower, the reason being that legacy path uses Finite-Difference (FD) approximations for spatial derivatives instead of Exact AD, uses stencils and thus there is a single pass zygote AD, whereas in my parser point-by-point Zygote.gradient calls are executed which takes a lot of time.

I am thinking of keep both, the finite difference approximation for faster training and exact ad for better precision and let the users decide which to use, like we can use a bool to toggle between fd approximation and exact ad

[ajatshatru01]

@ChrisRackauckas @sathvikbhagavan @AayushSabharwal i am sticking to finite difference approximations after feedbacks from Chris, since applying exact AD support led to extremely slow training time
now since this parser is supposed to replace the current legacy parser, i have decided to add all the existing supports to my parser
I have identified 4 potential improvements-
A. Support for Integro-Differential Equations (IDEs)
B. Parameter Estimation
C. Bayesian PINNs
D. verify and test compatibility with all training strategies

Any improvement feedbacks are welcome

[sathvikbhagavan]

Cool. For now, lets have finite difference based. But I think eventually we should also have the option of doing AD for the derivatives. @ChrisRackauckas, is it possible to make it faster than three times slow time @ajatshatru01 mentions or is it a limitation of using Zygote + Forward AD?

[ajatshatru01]

Cool. For now, lets have finite difference based. But I think eventually we should also have the option of doing AD for the derivatives. @ChrisRackauckas, is it possible to make it faster than three times slow time @ajatshatru01 mentions or is it a limitation of using Zygote + Forward AD?

yea i think being able to add AD support while maintaining the fast training time can be a separate project goal, i will look into that once this project is done

[ajatshatru01]

@sathvikbhagavan @ChrisRackauckas @AayushSabharwal eally sorry for not attending the meet, i am having a strong fever
as for my weekly progress

1. implemented equation-parameter support in symbolic PINN parser path
2. added integral and integro-differential support to symbolic parser

do review them, the changes are in the last two commits

[ajatshatru01]

Just a clarification, what is the legacy path? Is that the old parser? I think we should just replace the old parser completely with this new one.

@sathvikbhagavan yes it the old parser which to be replaced by this new one, hence why i am migrating all functionality from the old to new

[AayushSabharwal]

While the symbolic manipulation generally looks like it'll do the right thing, it's very type-unstable and thus could be much faster. I'm not sure how much of that is in scope for this project. I've highlighted the major pain points here, but generally ensuring that all containers (arrays, dictionaries, structs, etc.) have concrete types will go a long way.

Perhaps more important than the symbolic manipulation code is SymbolicPINNResidual. To me, it looks like the argument construction will not infer the lengths/types of all the tuples it builds, which can be catastrophic for inference. I imagine this residual function is called in a hot loop, in which case it likely contributes more to overall runtime than the symbolic manipulation.

[AayushSabharwal]

Is it possible to give all fields a concrete type?

[ajatshatru01]

I can rewrite the structs using type parameters so that every instance constructed has concrete, fully inferred types
That works??

[AayushSabharwal]

That works but it's still suboptimal. There's no contract on what the fields are, so

• Many edge cases have to be handled downstream. E.g. a list of variables might be Vector{Symbolics.SymbolicT} or Vector{Num} or Vector{Any} and all three need handling. If the type is restricted, downstream code is simplified.
• Type parameters force a lot of recompilation. All functions which take this struct are recompiled for every combination of type parameters.

[ajatshatru01]

okk

[AayushSabharwal]

Just a gut feeling, but it might be possible to narrow down the number of edge cases here if we can have concrete types for the structs above.

[ajatshatru01]

Even if we give SymbolicPINNSystem concrete types, op_domain_variables is read dynamically from the user's symbolic equation AST. Because MTK is extremely loose with how it parses coordinate variables:
We cannot predict which of the 4 forms (scalar, tuple, vector, or symbolic call) MTK will return.
Thus, the if-else block on lines 82–90 acts as a normalization layer to convert whatever MTK gives us into a uniform Vector{Symbolics.Num} before we do index lookups. I believe this is unavoidable because the types are determined by ModelingToolkit's internals rather than our own structs.

[AayushSabharwal]

ModelingToolkit stores all expressions as Symbolics.SymbolicT. As far as I can see, op_domain_variables comes from op.domain.variables where op isa Integral. This field is concretely typed as Symbolics.SymbolicT so you know exactly the type that is input to this function. Can you give examples of how this path gets different inputs?

[ajatshatru01]

yes they do store all expressions under the Symbolic type hierarchy, but the shape and structure of the symbolic AST layout returned by op.domain.variables changes completely depending on whether the integration is 1D or multi-dimensional.

Here are the two cases:

Case 1: 1D Integration (e.g., Integral(t in ClosedInterval(0, 1)))
Value of variables: t (a single symbol).
AST Form: A basic leaf node. SymbolicUtils.iscall(variables) returns false.

Case 2: Multi-D Integration (e.g., Integral((x, y) in ProductDomain(...)))
Value of variables: tuple(x, y) (a symbolic call representation of a tuple).
AST Form: A call node. SymbolicUtils.iscall(variables) returns true (with operation tuple and arguments [x, y]).

Even though the input container is always a Symbolic subtype, the downstream code needs a flat Julia list of the integrating variables (Vector{Symbolic}) to identify their indices in the system.Because tuple(x, y) is a symbolic call, we cannot simply iterate over it like a vector. We have to:

• Detect if it's a symbolic tuple call (operation(unwrapped_vars) === tuple).

• Retrieve its arguments ([x, y]) using SymbolicUtils.arguments(unwrapped_vars).

• Fall back to wrapping a single symbol [t] in a vector if it is not a call.

[AayushSabharwal]

Isn't this quadratic in time complexity?

[ajatshatru01]

we can make it linear complexity by constructing dict or set, but wont the allocation overhead be way more than the time complexity then? dimensions are small so wont that makes the linear search is the fastest, allocation-free way to execute this check?

[AayushSabharwal]

Yep, quadratic only matters if this will have large inputs. If that is not the case, then this is fine.

[AayushSabharwal]

Suggested change

	derivative_vars = Any[]
	derivative_vars = Symbolics.SymbolicT[]

All symbolic variables after unwrap have this type, so any buffers storing them can have a concrete type. This can improve the speed of symbolic processing code significantly.

[ajatshatru01]

got it

[AayushSabharwal]

Suggested change

	args = collect(SymbolicUtils.arguments(deriv_info.call))
	args = SymbolicUtils.arguments(deriv_info.call)

This collect seems unnecessary. arguments returns a read-only buffer, so accidental mutation will error and it will still behave like a normal vector.

[ajatshatru01]

ok, i have replaced it everywhere

[AayushSabharwal]

Note that ntuple (and tuples in general) is really only worth it if the length and type of each element is inferrable.

[AayushSabharwal]

Continuing the above comment, this function will likely infer as ::Any since different branches return tuples of different lengths and types. If eq_param_count in the above struct is made a type parameter and this function then dispatches on ::Val{eq_param_count}, it will infer better.

[AayushSabharwal]

If these argument tuples are not inferred, f.compiled is a dynamic dispatch every single call and will likely infer as returning ::Any. This will cause the dispatch itself to be slow, and will poison inference for all code that calls this residual function.

[AayushSabharwal]

Might be worth checking if this infers due to Iterators.product(axes...) not being inferrable, though it likely should.

[AayushSabharwal]

Similar to a previous comment, SymbolicUtils.query accomplishes precisely this.

[ajatshatru01]

@AayushSabharwal implemented all the optimisations you suggested, will commit it as soon as i am done with making the data-driven bpinn support possible with the parser. I am more focused on migrating all the supports from old parser to new as of now and will do a through insppection for optimisation later on.

[ChrisRackauckas]

(p) -> abs2((NN(SymbolicUtils.array_literal((2,), 1.0, 0.5), p))[1]) + (abs2(-(DNN(SymbolicUtils.array_literal((2,), 0.75, 0.5), p, [1, 1]))[1] + (DNN(SymbolicUtils.array_literal((2,), 0.75, 0.5), p, [2]))[1]) + abs2(-(DNN(SymbolicUtils.array_literal((2,), 0.25, 0.5), p, [1, 1]))[1] + (DNN(SymbolicUtils.array_literal((2,), 0.25, 0.5), p, [2]))[1])) / 2 + abs2((NN(SymbolicUtils.array_literal((2,), 0.0, 0.5), p))[1]) + abs2(-sin(0.5π) + (NN(SymbolicUtils.array_literal((2,), 0.5, 0.0), p))[1])

pde_pts = [[0.25, 0.5], [0.75, 0.5]]
bc1_pts = [[0.0, 0.5]]
bc2_pts = [[1.0, 0.5]]
bc3_pts = [[0.5, 0.0]]
theta = [pde_pts; bc1_pts; bc2_pts; bc3_pts]

(p, theta) -> sum(abs2((NN(SymbolicUtils.array_literal((2,), x), p)), theta)

[ChrisRackauckas]

QuadratureStrategy

theta = PINNParameter(points, domain)

Integrals.jl for the quadrature, 1D uses QuadGK 2D - 8D uses Cubature, 9D+ uses Cuba QMC, Integrals.jl default solver should handle that so use solve(prob) form.

(p, theta) -> solve(IntegralProblem(x->(abs2((NN(SymbolicUtils.array_literal((2,), x), p))), theta.domain))

[ChrisRackauckas]

Just aim for strategies that have a fixed number of points (gridstrategy, randomstrategy, quasirandomstrategy) for the first merge. Then after review and merge, follow up with quadraturestrategy. Then follow up with BPINN tests (shouldn't require many code changes) and Integro-Differential terms.

[ajatshatru01]

(p) -> abs2((NN(SymbolicUtils.array_literal((2,), 1.0, 0.5), p))[1]) + (abs2(-(DNN(SymbolicUtils.array_literal((2,), 0.75, 0.5), p, [1, 1]))[1] + (DNN(SymbolicUtils.array_literal((2,), 0.75, 0.5), p, [2]))[1]) + abs2(-(DNN(SymbolicUtils.array_literal((2,), 0.25, 0.5), p, [1, 1]))[1] + (DNN(SymbolicUtils.array_literal((2,), 0.25, 0.5), p, [2]))[1])) / 2 + abs2((NN(SymbolicUtils.array_literal((2,), 0.0, 0.5), p))[1]) + abs2(-sin(0.5π) + (NN(SymbolicUtils.array_literal((2,), 0.5, 0.0), p))[1])

pde_pts = [[0.25, 0.5], [0.75, 0.5]]
bc1_pts = [[0.0, 0.5]]
bc2_pts = [[1.0, 0.5]]
bc3_pts = [[0.5, 0.0]]
theta = [pde_pts; bc1_pts; bc2_pts; bc3_pts]

(p, theta) -> sum(abs2((NN(SymbolicUtils.array_literal((2,), x), p)), theta)

so as of now we are baking specific grid points into the symbolic equation AST. Instead, we should keep the coordinate points as variables in the symbolic representation, compile a function of the form (p, theta) -> ..., and pass the coordinates (theta) dynamically during the training loop.

[sathvikbhagavan]

Yes, all strategies sample points every iteration except GridTraining (which is not a very useful strategy anyway)

[ajatshatru01]

@ChrisRackauckas @sathvikbhagavan compiled residual function is coordinate-free, ran the tests for gridstrategy, randomstrategy, quasirandomstrategy all of which passed, do review and merge
since i have done much of the work on integral and bpinns too, i would ask you to merge this pr as a whole as of now and i will be submitting further prs with tweaks to it

[ajatshatru01]

ran an additional script to verify the correctness:-
"""
Demonstrates that the symbolic PINN parser produces coordinate-free compiled
residual functions. Collocation coordinates are never baked into the symbolic
AST — they are injected at runtime by the training strategy on every iteration.

PDE: ∂ₜu = ∂ₓₓu on [0,1]×[0,1]
BCs: u(0,t)=0, u(1,t)=0, u(x,0)=sin(πx)
"""

using NeuralPDE, ModelingToolkit, DomainSets, Lux, Symbolics, Zygote
using QuasiMonteCarlo
import DomainSets: Interval

sep  = "=" ^ 72
dash = "─" ^ 72

# ── System ────────────────────────────────────────────────────────────────────
@parameters x t
@variables u(..)
Dt = Differential(t);  Dxx = Differential(x)^2

eq  = Dt(u(x, t)) ~ Dxx(u(x, t))
bcs = [u(0.0, t) ~ 0.0, u(1.0, t) ~ 0.0, u(x, 0.0) ~ sin(pi * x)]
domains = [x ∈ Interval(0.0, 1.0), t ∈ Interval(0.0, 1.0)]
@named heat_sys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])

chain = Lux.Chain(Lux.Dense(2, 4, tanh), Lux.Dense(4, 1))

sym = NeuralPDE.build_symbolic_pinn_loss(heat_sys, chain; n_interior = 2, n_bc = 2)
θ₀  = sym.theta0
pde_fn = sym.datafree_pde_loss_functions[1]

# ── 1. Symbolic residual expressions ─────────────────────────────────────────
println(sep)
println("1. SYMBOLIC RESIDUAL EXPRESSIONS")
println("   x, t appear as free symbols — no collocation coordinates baked in.")
println(sep)
println()
println("PDE  ∂ₜu − ∂ₓₓu:")
println("  ", sym.pde_residuals[1])
println()
println("BC1  u(0,t) = 0:      ", sym.bc_residuals[1])
println("BC2  u(1,t) = 0:      ", sym.bc_residuals[2])
println("BC3  u(x,0) = sin(πx):", sym.bc_residuals[3])

# ── 2. Compiled function signature ───────────────────────────────────────────
println()
println(sep)
println("2. COMPILED RESIDUAL SIGNATURE  f(cord::Matrix, θ) → (1, N) Matrix")
println("   cord is the (D×N) coordinate matrix passed by the strategy at runtime.")
println(sep)
println()

cord_a = rand(2, 6)
cord_b = rand(2, 6)
println("  f(cord_a, θ₀) = ", pde_fn(cord_a, θ₀))
println("  f(cord_b, θ₀) = ", pde_fn(cord_b, θ₀))
println()
println("  Same fn, different coord matrices → different outputs: ",
        pde_fn(cord_a, θ₀) != pde_fn(cord_b, θ₀))

# ── 3. Strategy-level proof ───────────────────────────────────────────────────
println()
println(sep)
println("3. STRATEGY-LEVEL PROOF")
println("   Same compiled f(cord, θ). Only the coordinate matrix changes.")
println(sep)

# GridTraining
println()
println(dash)
println("  GridTraining(0.1)")
println(dash)
prob = discretize(heat_sys, PhysicsInformedNN(chain, GridTraining(0.1); symbolic_parser = true))
g1, g2 = prob.f(prob.u0, nothing), prob.f(prob.u0, nothing)
println("  call 1: ", g1)
println("  call 2: ", g2)
println("  deterministic (fixed grid, same cord reused): ", g1 == g2)

# StochasticTraining
println()
println(dash)
println("  StochasticTraining(30)")
println(dash)
prob = discretize(heat_sys, PhysicsInformedNN(chain, StochasticTraining(30); symbolic_parser = true))
s1, s2, s3 = prob.f(prob.u0, nothing), prob.f(prob.u0, nothing), prob.f(prob.u0, nothing)
println("  call 1: ", s1)
println("  call 2: ", s2)
println("  call 3: ", s3)
println("  resamples each call (fresh rand(D,N) cord): ", s1 != s2 && s2 != s3)

# QuasiRandomTraining
println()
println(dash)
println("  QuasiRandomTraining(30, resampling=true)")
println(dash)
prob = discretize(heat_sys, PhysicsInformedNN(
    chain,
    QuasiRandomTraining(30; sampling_alg = LatinHypercubeSample(), resampling = true);
    symbolic_parser = true
))
q1, q2, q3 = prob.f(prob.u0, nothing), prob.f(prob.u0, nothing), prob.f(prob.u0, nothing)
println("  call 1: ", q1)
println("  call 2: ", q2)
println("  call 3: ", q3)
println("  resamples each call (fresh QMC cord):       ", q1 != q2 && q2 != q3)

# ── 4. Gradient ───────────────────────────────────────────────────────────────
println()
println(sep)
println("4. ZYGOTE GRADIENT  ∇θ L(θ)  — differentiable end-to-end")
println(sep)
println()
grad = Zygote.gradient(sym.loss, θ₀)
println("  ∇θ exists     : ", grad !== nothing)
println("  all finite    : ", all(isfinite, grad[1]))
println("  ‖∇θ‖₂         = ", round(sqrt(sum(abs2, grad[1])); digits = 6))
println()


the results:- 
julia> include("print_symbolic_loss.jl")
========================================================================
1. SYMBOLIC RESIDUAL EXPRESSIONS
   x, t appear as free symbols — no collocation coordinates baked in.
========================================================================

PDE  ∂ₜu − ∂ₓₓu:
  (-(NN(SymbolicUtils.array_literal((2,), x, -6.055454452393343e-6 + t), p))[1] + (NN(SymbolicUtils.array_literal((2,), x, 6.055454452393343e-6 + t), p))[1]) / 1.2110908904786686e-5 + (-(NN(SymbolicUtils.array_literal((2,), -0.0001220703125 + x, t), p))[1] - (NN(SymbolicUtils.array_literal((2,), 0.0001220703125 + x, t), p))[1] + 2(NN(SymbolicUtils.array_literal((2,), x, t), p))[1]) / 1.4901161193847656e-8

BC1  u(0,t) = 0:      (NN(SymbolicUtils.array_literal((2,), 0.0, t), p))[1]
BC2  u(1,t) = 0:      (NN(SymbolicUtils.array_literal((2,), 1.0, t), p))[1]
BC3  u(x,0) = sin(πx):(NN(SymbolicUtils.array_literal((2,), x, 0.0), p))[1] - sin(π*x)

========================================================================
2. COMPILED RESIDUAL SIGNATURE  f(cord::Matrix, θ) → (1, N) Matrix
   cord is the (D×N) coordinate matrix passed by the strategy at runtime.
========================================================================

  f(cord_a, θ₀) = [-0.3790013296851279 -1.9565315848112543 -1.3269908182084345 -0.43534367247458833 -1.9982956986226101 -0.8044579281433774]
  f(cord_b, θ₀) = [-0.2655453251403499 -1.5211624165102517 -0.3604337031627938 -2.079655116554596 -0.5906811323718157 -1.0906198714735775]

  Same fn, different coord matrices → different outputs: true

========================================================================
3. STRATEGY-LEVEL PROOF
   Same compiled f(cord, θ). Only the coordinate matrix changes.
========================================================================

────────────────────────────────────────────────────────────────────────
  GridTraining(0.1)
────────────────────────────────────────────────────────────────────────
  call 1: 2.857348457213691
  call 2: 2.857348457213691
  deterministic (fixed grid, same cord reused): true

────────────────────────────────────────────────────────────────────────
  StochasticTraining(30)
────────────────────────────────────────────────────────────────────────
  call 1: 2.6731926252811586
  call 2: 2.5495762665568895
  call 3: 2.5725803327055647
  resamples each call (fresh rand(D,N) cord): true

────────────────────────────────────────────────────────────────────────
  QuasiRandomTraining(30, resampling=true)
────────────────────────────────────────────────────────────────────────
  call 1: 1.5765115477074356
  call 2: 1.418271049600992
  call 3: 1.5799330450684184
  resamples each call (fresh QMC cord):       true

========================================================================
4. ZYGOTE GRADIENT  ∇θ L(θ)  — differentiable end-to-end
========================================================================

  ∇θ exists     : true
  all finite    : true
  ‖∇θ‖₂         = 5.929817

[AayushSabharwal]

It's probably worth replacing this NamedTuple with a struct so that integrand_info can be a concretely typed vector instead of Vector{Any}.

[AayushSabharwal]

Splatting into a vector is very difficult for the Julia compiler. Prefer using vcat or array concatenation syntax [lb_args; ub_args].

[AayushSabharwal]

Suggested change

	return SymbolicUtils.term(_solve_pinn_integral, call_args...; type = Real, shape = SymbolicUtils.ShapeVecT())
	return Symbolics.STerm(_solve_pinn_integral, call_args; type = Real, shape = SymbolicUtils.ShapeVecT())

This avoids the splat.

[AayushSabharwal]

Splatting a vector is slow

[AayushSabharwal]

Suggested change

	if SymbolicUtils.iscall(ex) && any(dv_op -> isequal(SymbolicUtils.operation(ex), dv_op), dv_ops)
	if SymbolicUtils.iscall(ex) && any(dv_op -> isequal(SymbolicUtils.operation(ex), dv_op)::Bool, dv_ops)::Bool

This is better for performance. operation(ex) infers as ::Any, so isequal is a dynamic dispatch and won't infer as ::Bool. We can make this even better if we know what sort of operations dv_ops contains. Another improvement is to make dv_ops a Set if it is large enough.

[sathvikbhagavan]

Can you train for more iterations to check if the solution converges? As its a small system, huge number of iterations won't be needed. Do it for all integration tests. Also, have a couple more systems other than heat equation.

[sathvikbhagavan]

I am of the opinion to completely remove the legacy path, not have a fallback as the new parser is much better than the old one and have a breaking change. @ChrisRackauckas, is that ok or do we need to keep this?

[ajatshatru01]

yea i will deprecate the legacy path once the new parser is fully verified and merged, it will be mostly deleting the code sections of legacy path then

[sathvikbhagavan]

Can you check the correctness of it, not just finiteness? Same applies to other places

[sathvikbhagavan]

Can we have a way to pass weights (optional) to these two losses? As its one of the basic techniques used in PINN training. These can be fixed or potentially trainable. We can assume to be fixed for now.

[ajatshatru01]

@AayushSabharwal did all the improvements you asked for
@sathvikbhagavan added and tested the tests you had asked for, along with adding weights to losses support
additionally, ran another script to verify the loss convergence and solution correctness

using Pkg
Pkg.activate(".")

using NeuralPDE, ModelingToolkit, DomainSets, Lux, Optimization, OptimizationOptimisers, Random, Printf
import DomainSets: Interval, ClosedInterval

function print_report(name, init_loss, final_loss, max_err, mean_err)
    println("="^65)
    println(" SYSTEM: $name")
    println("="^65)
    @printf("  Initial Loss   : %.6f\n", init_loss)
    @printf("  Final Loss     : %.6f  (Loss Reduction: %.2fx)\n", final_loss, init_loss / max(final_loss, 1e-12))
    if max_err !== nothing
        @printf("  Max Abs Error  : %.6f\n", max_err)
        @printf("  Mean Abs Error : %.6f\n", mean_err)
        println("  Status         : SUCCESS (Converged & Solution Verified)")
    else
        println("  Status         : SUCCESS (Converged)")
    end
    println()
end

println("\n" * "#"^65)
println(" SYMBOLIC PINN PARSER - CONVERGENCE & ACCURACY VERIFICATION")
println("#"^65 * "\n")

# -------------------------------------------------------------------
# System 1: 1D Heat Equation
# -------------------------------------------------------------------
Random.seed!(100)
@parameters x t
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2

eq  = Dt(u(x, t)) ~ Dxx(u(x, t))
bcs = [
    u(0.0, t) ~ 0.0,
    u(1.0, t) ~ 0.0,
    u(x, 0.0) ~ sin(pi * x),
]
domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
@named heat_sys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])

chain1 = Lux.Chain(Lux.Dense(2, 12, tanh), Lux.Dense(12, 1))
discretization1 = PhysicsInformedNN(chain1, GridTraining(0.1); symbolic_parser = true)
prob1 = discretize(heat_sys, discretization1)

init_loss1 = prob1.f(prob1.u0, nothing)
sol1 = solve(prob1, OptimizationOptimisers.Adam(0.02); maxiters = 600)
final_loss1 = prob1.f(sol1.u, nothing)

phi1 = discretization1.phi
xs = 0.1:0.1:0.9
ts = 0.1:0.1:0.9
u_pred1 = [first(phi1([x, t], sol1.u)) for x in xs for t in ts]
u_true1 = [exp(-pi^2 * t) * sin(pi * x) for x in xs for t in ts]
errs1 = abs.(u_pred1 .- u_true1)

print_report("1D Heat Equation (PDE)", init_loss1, final_loss1, maximum(errs1), sum(errs1)/length(errs1))

# -------------------------------------------------------------------
# System 2: 2D Poisson Equation
# -------------------------------------------------------------------
Random.seed!(100)
@parameters x y
@variables u2(..)
Dxx2 = Differential(x)^2
Dyy2 = Differential(y)^2

eq2 = Dxx2(u2(x, y)) + Dyy2(u2(x, y)) ~ -sin(pi * x) * sin(pi * y)
bcs2 = [
    u2(0.0, y) ~ 0.0,
    u2(1.0, y) ~ 0.0,
    u2(x, 0.0) ~ 0.0,
    u2(x, 1.0) ~ 0.0,
]
domains2 = [x in Interval(0.0, 1.0), y in Interval(0.0, 1.0)]
@named poisson_sys = PDESystem(eq2, bcs2, domains2, [x, y], [u2(x, y)])

chain2 = Lux.Chain(Lux.Dense(2, 12, tanh), Lux.Dense(12, 12, tanh), Lux.Dense(12, 1))
discretization2 = PhysicsInformedNN(chain2, GridTraining(0.1); symbolic_parser = true)
prob2 = discretize(poisson_sys, discretization2)

init_loss2 = prob2.f(prob2.u0, nothing)
sol2 = solve(prob2, OptimizationOptimisers.Adam(0.02); maxiters = 600)
final_loss2 = prob2.f(sol2.u, nothing)

phi2 = discretization2.phi
xs2 = 0.1:0.1:0.9
ys2 = 0.1:0.1:0.9
analytic_poisson(x, y) = sin(pi * x) * sin(pi * y) / (2 * pi^2)
u_pred2 = [first(phi2([x, y], sol2.u)) for x in xs2 for y in ys2]
u_true2 = [analytic_poisson(x, y) for x in xs2 for y in ys2]
errs2 = abs.(u_pred2 .- u_true2)

print_report("2D Poisson Equation (Elliptic PDE)", init_loss2, final_loss2, maximum(errs2), sum(errs2)/length(errs2))

# -------------------------------------------------------------------
# System 3: 1D Wave Equation
# -------------------------------------------------------------------
Random.seed!(100)
@parameters x t
@variables u3(..)
Dxx3 = Differential(x)^2
Dtt3 = Differential(t)^2
Dt3  = Differential(t)

C = 1.0
eq3 = Dtt3(u3(x, t)) ~ C^2 * Dxx3(u3(x, t))
bcs3 = [
    u3(0.0, t) ~ 0.0,
    u3(1.0, t) ~ 0.0,
    u3(x, 0.0) ~ sin(pi * x),
    Dt3(u3(x, 0.0)) ~ 0.0,
]
domains3 = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
@named wave_sys = PDESystem(eq3, bcs3, domains3, [x, t], [u3(x, t)])

chain3 = Lux.Chain(Lux.Dense(2, 12, tanh), Lux.Dense(12, 12, tanh), Lux.Dense(12, 1))
discretization3 = PhysicsInformedNN(chain3, GridTraining(0.1); symbolic_parser = true)
prob3 = discretize(wave_sys, discretization3)

init_loss3 = prob3.f(prob3.u0, nothing)
sol3 = solve(prob3, OptimizationOptimisers.Adam(0.01); maxiters = 1000)
final_loss3 = prob3.f(sol3.u, nothing)

phi3 = discretization3.phi
xs3 = 0.2:0.1:0.8
ts3 = 0.2:0.1:0.8
analytic_wave(x, t) = sin(pi * x) * cos(pi * t)
u_pred3 = [first(phi3([x, t], sol3.u)) for x in xs3 for t in ts3]
u_true3 = [analytic_wave(x, t) for x in xs3 for t in ts3]
errs3 = abs.(u_pred3 .- u_true3)

print_report("1D Wave Equation (Hyperbolic PDE)", init_loss3, final_loss3, maximum(errs3), sum(errs3)/length(errs3))

# -------------------------------------------------------------------
# System 4: 1D Heterogeneous ODE
# -------------------------------------------------------------------
Random.seed!(100)
@parameters θ
@variables u4(..)
Dθ = Differential(θ)

eq4 = Dθ(u4(θ)) ~ θ^3 + 2.0f0 * θ +
    (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) -
    u4(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3)))

bcs4 = [u4(0.0) ~ 1.0f0]
domains4 = [θ ∈ Interval(0.0f0, 1.0f0)]

chain4 = Lux.Chain(Lux.Dense(1, 12, σ), Lux.Dense(12, 1))
discretization4 = PhysicsInformedNN(chain4, GridTraining(0.05); symbolic_parser = true)
@named ode_sys = PDESystem(eq4, bcs4, domains4, [θ], [u4(θ)])

prob4 = discretize(ode_sys, discretization4)
init_loss4 = prob4.f(prob4.u0, nothing)
sol4 = solve(prob4, OptimizationOptimisers.Adam(0.02); maxiters = 500)
final_loss4 = prob4.f(sol4.u, nothing)

phi4 = discretization4.phi
analytic_ode(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2
ts4 = 0.1:0.1:0.9
u_pred4 = [first(phi4([t], sol4.u)) for t in ts4]
u_true4 = [analytic_ode(t) for t in ts4]
errs4 = abs.(u_pred4 .- u_true4)

print_report("1D Non-linear ODE", init_loss4, final_loss4, maximum(errs4), sum(errs4)/length(errs4))

# -------------------------------------------------------------------
# System 5: Integro-Differential Equation (IDE)
# -------------------------------------------------------------------
Random.seed!(100)
@parameters t
@variables i5(..)
Di5 = Differential(t)
Ii5 = Integral(t in ClosedInterval(0, 2))
eq5 = Di5(i5(t)) + 2 * i5(t) + 5 * Ii5(i5(t)) ~ 1
bcs5 = [i5(0.0) ~ 0.0]
domains5 = [t ∈ Interval(0.0, 2.0)]

chain5 = Lux.Chain(Lux.Dense(1, 8, σ), Lux.Dense(8, 1))
@named ide_sys = PDESystem(eq5, bcs5, domains5, [t], [i5(t)])

discretization5 = PhysicsInformedNN(chain5, GridTraining(0.25); symbolic_parser = true)
prob5 = discretize(ide_sys, discretization5)
init_loss5 = prob5.f(prob5.u0, nothing)
sol5 = solve(prob5, OptimizationOptimisers.Adam(0.01); maxiters = 300)
final_loss5 = prob5.f(sol5.u, nothing)

# -------------------------------------------------------------------
# System 6: 3D Poisson Equation (Higher-Dimensional PDE)
# -------------------------------------------------------------------
Random.seed!(100)
@parameters x y z
@variables u6(..)
Dxx6 = Differential(x)^2
Dyy6 = Differential(y)^2
Dzz6 = Differential(z)^2

eq6 = Dxx6(u6(x, y, z)) + Dyy6(u6(x, y, z)) + Dzz6(u6(x, y, z)) ~ -3.0f0 * (pi^2) * sin(pi * x) * sin(pi * y) * sin(pi * z)
bcs6 = [
    u6(0.0, y, z) ~ 0.0, u6(1.0, y, z) ~ 0.0,
    u6(x, 0.0, z) ~ 0.0, u6(x, 1.0, z) ~ 0.0,
    u6(x, y, 0.0) ~ 0.0, u6(x, y, 1.0) ~ 0.0,
]
domains6 = [x in Interval(0.0, 1.0), y in Interval(0.0, 1.0), z in Interval(0.0, 1.0)]
@named poisson3d_sys = PDESystem(eq6, bcs6, domains6, [x, y, z], [u6(x, y, z)])

chain6 = Lux.Chain(Lux.Dense(3, 16, tanh), Lux.Dense(16, 16, tanh), Lux.Dense(16, 1))
discretization6 = PhysicsInformedNN(chain6, GridTraining(0.2); symbolic_parser = true)
prob6 = discretize(poisson3d_sys, discretization6)

init_loss6 = prob6.f(prob6.u0, nothing)
sol6 = solve(prob6, OptimizationOptimisers.Adam(0.02); maxiters = 600)
final_loss6 = prob6.f(sol6.u, nothing)

phi6 = discretization6.phi
xs6 = 0.2:0.3:0.8
ys6 = 0.2:0.3:0.8
zs6 = 0.2:0.3:0.8
analytic_sol3d(x, y, z) = sin(pi * x) * sin(pi * y) * sin(pi * z)
u_pred6 = [first(phi6([x, y, z], sol6.u)) for x in xs6 for y in ys6 for z in zs6]
u_true6 = [analytic_sol3d(x, y, z) for x in xs6 for y in ys6 for z in zs6]
errs6 = abs.(u_pred6 .- u_true6)

print_report("3D Poisson Equation (3D Spatial PDE)", init_loss6, final_loss6, maximum(errs6), sum(errs6)/length(errs6))

println("#"^65)
println(" ALL SYSTEMS VERIFIED SUCCESSFULLY!")
println("#"^65 * "\n")

results:-

julia> include("verify_symbolic_convergence.jl")
  Activating project at `C:\Users\ajats\Sciml\NeuralPDE.jl`

#################################################################
 SYMBOLIC PINN PARSER - CONVERGENCE & ACCURACY VERIFICATION
#################################################################

=================================================================
 SYSTEM: 1D Heat Equation (PDE)
=================================================================
  Initial Loss   : 1.332178
  Final Loss     : 0.060450  (Loss Reduction: 22.04x)
  Max Abs Error  : 0.103236
  Mean Abs Error : 0.030847
  Status         : SUCCESS (Converged & Solution Verified)

=================================================================
 SYSTEM: 2D Poisson Equation (Elliptic PDE)
=================================================================
  Initial Loss   : 3.407725
  Final Loss     : 0.001870  (Loss Reduction: 1822.19x)
  Max Abs Error  : 0.009839
  Mean Abs Error : 0.004260
  Status         : SUCCESS (Converged & Solution Verified)

=================================================================
 SYSTEM: 1D Wave Equation (Hyperbolic PDE)
=================================================================
  Initial Loss   : 1.923722
  Final Loss     : 0.001912  (Loss Reduction: 1006.20x)
  Max Abs Error  : 0.359575
  Mean Abs Error : 0.226117
  Status         : SUCCESS (Converged & Solution Verified)

=================================================================
 SYSTEM: 1D Non-linear ODE
=================================================================
  Initial Loss   : 16.347634
  Final Loss     : 0.007419  (Loss Reduction: 2203.45x)
  Max Abs Error  : 0.015439
  Mean Abs Error : 0.008036
  Status         : SUCCESS (Converged & Solution Verified)

=================================================================
 SYSTEM: 3D Poisson Equation (3D Spatial PDE)
=================================================================
  Initial Loss   : 76.348611
  Final Loss     : 0.087446  (Loss Reduction: 873.09x)
  Max Abs Error  : 0.122697
  Mean Abs Error : 0.063464
  Status         : SUCCESS (Converged & Solution Verified)

#################################################################
 ALL SYSTEMS VERIFIED SUCCESSFULLY!
#################################################################

[AayushSabharwal]

This looks mostly okay now. There are still some type-instabilities, but it's hard to say whether they're easily fixable, or how much they'll affect performance. Ideally, code that generates functions should only really be type-unstable due to the type of the RuntimeGeneratedFunction. I think this is acceptable for an initial version, and the best way to tackle any further instabilities is with proper profiling.

[ajatshatru01]

@sathvikbhagavan @ChrisRackauckas deleted the code for legacy parser and made the symbolic parser the default path
i also did a functionality testing with the new parser

@testitem "Symbolic PINN parser: Legacy Parity - Parameterized 1D Heat Equation" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux, Optimization, OptimizationOptimisers
    using Test
    import DomainSets: Interval

    @parameters x t a
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    # Heat equation with PDE parameter 'a'
    eq = Dt(u(x, t)) ~ a * Dxx(u(x, t))
    bcs = [
        u(0.0, t) ~ 0.0,
        u(1.0, t) ~ 0.0,
        u(x, 0.0) ~ sin(pi * x),
    ]
    domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
    @named pde_sys = PDESystem(
        eq,
        bcs,
        domains,
        [x, t],
        [u(x, t)],
        [a],
        initial_conditions = Dict([a => 1.0])
    )

    chain = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))
    discretization = PhysicsInformedNN(chain, GridTraining(0.25))

    prob = discretize(pde_sys, discretization)
    @test prob isa OptimizationProblem
    @test prob.u0 !== nothing

    res = prob.f(prob.u0, nothing)
    @test isfinite(res)
    @test res >= 0
end

@testitem "Symbolic PINN parser: Legacy Parity - Training Strategy Parity" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux, QuasiMonteCarlo, Integrals
    using Test
    import DomainSets: Interval

    @parameters x t
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    eq = Dt(u(x, t)) ~ Dxx(u(x, t))
    bcs = [
        u(0.0, t) ~ 0.0,
        u(1.0, t) ~ 0.0,
        u(x, 0.0) ~ sin(pi * x),
    ]
    domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
    @named sys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])

    chain = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))

    # Test 1: GridTraining
    prob_grid = discretize(sys, PhysicsInformedNN(chain, GridTraining(0.25)))
    @test isfinite(prob_grid.f(prob_grid.u0, nothing))

    # Test 2: StochasticTraining
    prob_stoch = discretize(sys, PhysicsInformedNN(chain, StochasticTraining(30)))
    @test isfinite(prob_stoch.f(prob_stoch.u0, nothing))

    # Test 3: QuasiRandomTraining
    qr_strat = QuasiRandomTraining(30; sampling_alg = LatinHypercubeSample())
    prob_qr = discretize(sys, PhysicsInformedNN(chain, qr_strat))
    @test isfinite(prob_qr.f(prob_qr.u0, nothing))

    # Test 4: QuadratureTraining
    quad_strat = QuadratureTraining(quadrature_alg = CubatureJLh(), reltol = 1e-3, abstol = 1e-3, maxiters = 50)
    prob_quad = discretize(sys, PhysicsInformedNN(chain, quad_strat))
    @test isfinite(prob_quad.f(prob_quad.u0, nothing))
end

@testitem "Symbolic PINN parser: Legacy Parity - Adaptive Losses & Additional Loss" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux
    using Test
    import DomainSets: Interval

    @parameters x t
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    eq = Dt(u(x, t)) ~ Dxx(u(x, t))
    bcs = [
        u(0.0, t) ~ 0.0,
        u(1.0, t) ~ 0.0,
        u(x, 0.0) ~ sin(pi * x),
    ]
    domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
    @named sys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])

    chain = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))

    # Additional loss function
    add_loss(phi, θ, p) = sum(abs2, θ) * 1e-4

    # Test with MiniMaxAdaptiveLoss & additional_loss
    disc = PhysicsInformedNN(
        chain, GridTraining(0.25);
        adaptive_loss = MiniMaxAdaptiveLoss(10),
        additional_loss = add_loss
    )

    prob = discretize(sys, disc)
    @test prob isa OptimizationProblem
    @test isfinite(prob.f(prob.u0, nothing))
end

@testitem "Symbolic PINN parser: Legacy Parity - Derivative BCs (Neumann / Robin)" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux
    using Test
    import DomainSets: Interval

    @parameters x t
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    eq = Dt(u(x, t)) ~ Dxx(u(x, t))
    # Neumann derivative boundary conditions
    bcs = [
        Dx(u(0.0, t)) ~ 0.0,
        Dx(u(1.0, t)) ~ 0.0,
        u(x, 0.0) ~ cos(pi * x),
    ]
    domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
    @named neumann_sys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])

    chain = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))
    prob = discretize(neumann_sys, PhysicsInformedNN(chain, GridTraining(0.25)))

    @test prob isa OptimizationProblem
    @test isfinite(prob.f(prob.u0, nothing))
end

@testitem "Symbolic PINN parser: Legacy Parity - Initial Value ODE Systems" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux
    using Test
    import DomainSets: Interval

    @parameters t
    @variables u(..)
    Dt = Differential(t)

    # Pure initial value ODE system
    eq = Dt(u(t)) ~ -u(t)
    bcs = [u(0.0) ~ 1.0]
    domains = [t in Interval(0.0, 1.0)]
    @named ode_sys = PDESystem(eq, bcs, domains, [t], [u(t)])

    chain = Lux.Chain(Lux.Dense(1, 8, tanh), Lux.Dense(8, 1))
    prob = discretize(ode_sys, PhysicsInformedNN(chain, GridTraining(0.1)))

    @test prob isa OptimizationProblem
    @test isfinite(prob.f(prob.u0, nothing))
end

@testitem "Symbolic PINN parser: Legacy Parity - Bayesian PINN (BPINN)" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux
    using Test
    import DomainSets: Interval

    @parameters x t
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    eq = Dt(u(x, t)) ~ Dxx(u(x, t))
    bcs = [
        u(0.0, t) ~ 0.0,
        u(1.0, t) ~ 0.0,
        u(x, 0.0) ~ sin(pi * x),
    ]
    domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
    @named sys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])

    chain = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))
    bpinn_disc = BayesianPINN([chain], GridTraining(0.2))

    sym_rep = symbolic_discretize(sys, bpinn_disc)
    @test sym_rep !== nothing
end

@testitem "Symbolic PINN parser: Legacy Parity - Integro-Differential Equations (IDE)" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux, Symbolics
    using Test
    import DomainSets: Interval

    @parameters x
    @variables u(..)
    Ix = Integral(x in Interval(0.0, x))

    # Volterra integro-differential equation
    eq = Ix(u(x)) ~ x^2 / 2
    bcs = [u(0.0) ~ 0.0]
    domains = [x in Interval(0.0, 1.0)]
    @named ide_sys = PDESystem(eq, bcs, domains, [x], [u(x)])

    chain = Lux.Chain(Lux.Dense(1, 8, tanh), Lux.Dense(8, 1))
    prob = discretize(ide_sys, PhysicsInformedNN(chain, GridTraining(0.1)))

    @test prob isa OptimizationProblem
    @test isfinite(prob.f(prob.u0, nothing))
end

@testitem "Symbolic PINN parser: Legacy Parity - Coupled Multi-Variable PDESystems" tags = [:symbolicpinn] begin
    using NeuralPDE, ModelingToolkit, DomainSets, Lux
    using Test
    import DomainSets: Interval

    @parameters x t
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)

    # Coupled system of PDEs: ∂u/∂t = ∂v/∂x, ∂v/∂t = ∂u/∂x
    eqs = [
        Dt(u(x, t)) ~ Dx(v(x, t)),
        Dt(v(x, t)) ~ Dx(u(x, t))
    ]
    bcs = [
        u(0.0, t) ~ 0.0,
        v(0.0, t) ~ 0.0,
        u(x, 0.0) ~ sin(pi * x),
        v(x, 0.0) ~ cos(pi * x)
    ]
    domains = [x in Interval(0.0, 1.0), t in Interval(0.0, 1.0)]
    @named coupled_sys = PDESystem(eqs, bcs, domains, [x, t], [u(x, t), v(x, t)])

    chain1 = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))
    chain2 = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))
    prob = discretize(coupled_sys, PhysicsInformedNN([chain1, chain2], GridTraining(0.2)))

    @test prob isa OptimizationProblem
    @test isfinite(prob.f(prob.u0, nothing))
end

julia> include("test/symbolic_pinn_parser_legacy_parity_tests.jl")
21:42:28 | START test item "Symbolic PINN parser: Legacy Parity - Parameterized 1D Heat Equation" at test\symbolic_pinn_parser_legacy_parity_tests.jl:1
Test Summary:                                                        | Pass  Total  Time
Symbolic PINN parser: Legacy Parity - Parameterized 1D Heat Equation |    4      4  4.3s
21:42:32 | DONE test item "Symbolic PINN parser: Legacy Parity - Parameterized 1D Heat Equation" 4.3 secs (98.3% compile, 2.8% recompile, 4.3% GC), 4.25 M allocs (217.713 MB)
21:42:32 | START test item "Symbolic PINN parser: Legacy Parity - Training Strategy Parity" at test\symbolic_pinn_parser_legacy_parity_tests.jl:41
Test Summary:                                                  | Pass  Total   Time
Symbolic PINN parser: Legacy Parity - Training Strategy Parity |    4      4  31.2s
21:43:03 | DONE test item "Symbolic PINN parser: Legacy Parity - Training Strategy Parity" 31.2 secs (99.5% compile, 2.2% GC), 23.31 M allocs (1.198 GB)
21:43:03 | START test item "Symbolic PINN parser: Legacy Parity - Adaptive Losses & Additional Loss" at test\symbolic_pinn_parser_legacy_parity_tests.jl:81
Test Summary:                                                           | Pass  Total  Time
Symbolic PINN parser: Legacy Parity - Adaptive Losses & Additional Loss |    2      2  2.6s
21:43:06 | DONE test item "Symbolic PINN parser: Legacy Parity - Adaptive Losses & Additional Loss" 2.6 secs (95.8% compile, 11.7% GC), 633.52 K allocs (33.507 MB)
21:43:06 | START test item "Symbolic PINN parser: Legacy Parity - Derivative BCs (Neumann / Robin)" at test\symbolic_pinn_parser_legacy_parity_tests.jl:117
Test Summary:                                                          | Pass  Total  Time
Symbolic PINN parser: Legacy Parity - Derivative BCs (Neumann / Robin) |    2      2  2.3s
21:43:08 | DONE test item "Symbolic PINN parser: Legacy Parity - Derivative BCs (Neumann / Robin)" 2.3 secs (95.1% compile), 1.51 M allocs (77.189 MB)
21:43:08 | START test item "Symbolic PINN parser: Legacy Parity - Initial Value ODE Systems" at test\symbolic_pinn_parser_legacy_parity_tests.jl:145
Test Summary:                                                   | Pass  Total  Time
Symbolic PINN parser: Legacy Parity - Initial Value ODE Systems |    2      2  6.2s
21:43:15 | DONE test item "Symbolic PINN parser: Legacy Parity - Initial Value ODE Systems" 6.2 secs (98.4% compile, 0.5% recompile), 4.19 M allocs (215.073 MB)
21:43:15 | START test item "Symbolic PINN parser: Legacy Parity - Bayesian PINN (BPINN)" at test\symbolic_pinn_parser_legacy_parity_tests.jl:167
Test Summary:                                               | Pass  Total  Time
Symbolic PINN parser: Legacy Parity - Bayesian PINN (BPINN) |    1      1  6.2s
21:43:21 | DONE test item "Symbolic PINN parser: Legacy Parity - Bayesian PINN (BPINN)" 6.2 secs (98.2% compile, 4.1% GC), 3.65 M allocs (191.219 MB)
21:43:21 | START test item "Symbolic PINN parser: Legacy Parity - Integro-Differential Equations (IDE)" at test\symbolic_pinn_parser_legacy_parity_tests.jl:193
Test Summary:                                                              | Pass  Total  Time
Symbolic PINN parser: Legacy Parity - Integro-Differential Equations (IDE) |    2      2  5.5s
21:43:26 | DONE test item "Symbolic PINN parser: Legacy Parity - Integro-Differential Equations (IDE)" 5.5 secs (97.8% compile, 14.8% recompile, 4.4% GC), 3.99 M allocs (207.387 MB)
21:43:26 | START test item "Symbolic PINN parser: Legacy Parity - Coupled Multi-Variable PDESystems" at test\symbolic_pinn_parser_legacy_parity_tests.jl:215
Test Summary:                                                           | Pass  Total   Time
Symbolic PINN parser: Legacy Parity - Coupled Multi-Variable PDESystems |    2      2  10.1s
21:43:37 | DONE test item "Symbolic PINN parser: Legacy Parity - Coupled Multi-Variable PDESystems" 10.1 secs (98.7% compile), 4.55 M allocs (234.392 MB)

[DhairyaLGandhi]

This looks like a bug bear. Do we need separate scalar and array functions here @AayushSabharwal ? Wouldn't the hash cons'd IR have naturally amortized the cost of the NN? Why do we need to manually handle it?

[AayushSabharwal]

Yep, it should be handled by CSE. Though the evaluated nn_vals isn't actually used or passed anywhere, which feels like dead code at best and a bug at worst.

[ajatshatru01]

so should i remove the manual mat_fn closure and let CSE handle it itself? @AayushSabharwal