[xtalax]

Build PDE solution wrappers array-wise

Summary

On a system compiled with mtkcompile, building the PDE solution wrapper grew much faster than the grid: 1.5 s at 1,521 states, 5.6 s at 3,481 and 22.0 s at 6,241 for a 2D heat equation. This PR reads each grid point that is an unknown by its state index. The remaining observed points, such as eliminated boundary points, are evaluated together with one observed function. With that change the same three cases wrap in 0.12 s, 0.26 s and 0.20 s.

The DAEProblem path was already fast on master and is unchanged here. The 70 s that TinyFDTD measured at 19k states was on code from before #701. On master, the staggered 2D TM DAEProblem wraps in 0.07 s at 19,683 states, and TinyFDTD measured 0.043 s at 19,220.

Cause

PDETimeSeriesSolution first tried array_observed_solution for each field. On a compiled system the field's observed equation is u(t) ~ array_literal(u[1, 1], u[2, 1], ...), with one statement per grid point. Evaluating it calls the observed function generated for it, and the first call compiles that function. At 41 × 41 points the compile took 1.43 s (the second call takes 0.7 ms), and its cost grows faster than the number of grid points. A profile of the first wrap at 1,521 states puts 684 of 748 samples inside that generated function.

Change

• Unknowns by index. grid_indices looks up every grid point with SymbolicIndexingInterface.variable_index, which is O(1) through the index cache. Grid points that are unknowns are read with sol[i, :].
• Observed points together. grid_observed_values evaluates the rest with a single getsym call, instead of building and compiling one observed function per element. At 1,521 states (160 observed boundary points), single-element observed calls cost about 8 ms each. One batched call covering all 160 took 0.13 s.
• Fully observed fields. A field with no unknowns at all, such as a reduced-order reconstruction, still goes through array_observed_solution. That is the case Reconstruct array observables once per PDE field #676 added it for, and its test still passes.
• PDENoTimeSolution uses the same two helpers.

fd1f976c added the batched observed evaluation, and its message credits it with the 22 s cost. That was wrong: the cost came from the array observed function, which ae3cb73d routes around. The batching still matters, because the boundary points now go through it.

Timings

SciMLBase.wrap_sol on a raw solution with 11 saves, after a warm-up wrap of a small system (Julia 1.12.6, ModelingToolkit 11.45.3).

System	States	master (53855c9)	this PR
2D heat, ODEProblem(mtkcompile(sys))	1,521	1.51 s	0.12 s
2D heat, ODEProblem(mtkcompile(sys))	3,481	5.59 s	0.26 s
2D heat, ODEProblem(mtkcompile(sys))	6,241	21.99 s	0.20 s
staggered 2D TM, DAEProblem	5,043	0.01 s	0.01 s

I did not go past 6,241 states on the compiled path. There, mtkcompile itself takes 468 s to build the problem, so the wrapper is no longer the bottleneck.

Tests

"Test 10" and "Test 11" in test/Sol_Interface/solution_interface.jl compare every stored grid value of the wrapper with the raw solution's value for that discrete variable. They cover the DAEProblem path, where every grid point is an unknown; the compiled ODE path with observed boundary points and a time-dependent boundary value; a staggered Yee layout; and a steady-state NonlinearProblem. They also check the interpolant at the grid points. They check against the raw solution, not against master's output, so they guard the routing against regressions; they do not reproduce a slowdown.

Local runs on ae3cb73d (GROUP=<g> julia --project -e 'using Pkg; Pkg.test()', Julia 1.12.6). Every group exits 0, and the broken counts are the existing @test_broken markers:

• Sol_Interface: 431/431, including Reconstruct array observables once per PDE field #676's array-observables testset.
• Components: passed.
• DAE: 22/22.
• Stationary: 9/9.
• Wave_Eq_Staggered: 305 pass, 5 broken.
• Pseudospectral: 70/70 and 174/174.
• Complex: 2/2.
• Discretization: 3618/3618 and 56/56.

Runic is clean.

Overlap

This overlaps #716 (gpu-solution-wrap), which also edits PDETimeSeriesSolution and PDENoTimeSolution, to read from a host copy of the solution. They should compose, since the element reads here would simply read from #716's host copy, but I have not tested the combination. Whichever merges second needs a small rebase in timedep.jl and timeindep.jl.

🤖 Posted by an AI agent · harness: Claude Code 2.1.296 · model: claude-opus-5-5
Review: unreviewed
Session: 644ed0ff-c8c9-4cc3-b4c6-f326231a4030