Skip to content

Lower computed initialization maps via symbolic tracing - #516

Closed
ChrisRackauckas wants to merge 3 commits into
masterfrom
lower-computed-init-maps
Closed

Lower computed initialization maps via symbolic tracing#516
ChrisRackauckas wants to merge 3 commits into
masterfrom
lower-computed-init-maps

Conversation

@ChrisRackauckas

@ChrisRackauckas ChrisRackauckas commented Aug 29, 2026

Copy link
Copy Markdown
Member

Summary

  • Trace ModelingToolkit initialization state/parameter maps with Symbolics variables instead of integer index tokens. Map outputs classify per entry: a bare traced variable stays a static gather recipe (InitializationSourceIndex), a literal becomes an InitializationConstant, and a computed symbolic expression compiles through Symbolics.build_function into a generated function (InitializationScalarExpression / InitializationArrayExpression) evaluated on-device against the same flat source vector. Arrays with any computed entry compile as one static-output function.
  • RuntimeGeneratedFunctions.drop_expr strips the stored body::Expr so compiled recipes are isbits and fit inside GPU kernel problems (the same mechanism MTKBase codegen uses).
  • Symbolics and RuntimeGeneratedFunctions join the ModelingToolkitBaseExt trigger list — both always load with ModelingToolkitBase, so extension activation is unchanged.
  • Raw HomotopyProblem initialization now fails with a clear ArgumentError instead of a MethodError, and the tutorial documents the homotopy(actual, simplified) → actual rewrite (exact by the operator's own numeric semantics: homotopy(a, s) = a) that yields an equivalent supported initialization, with a runnable recipe.
  • Tutorial restriction list updated: the "fallback getters are not yet lowered" limitation is lifted; value-dependent branching in maps remains unsupported (untraceable).

The bug this also fixes

The integer-token trace was not merely limited — it was unsound. A computed map output like 2·token + 1 is integer-valued; when it happens to land in 1:source_count it is indistinguishable from a legitimate index, and initialization silently gathers wrong values while reporting success. Reproducer (plain ODE, initialization_eqs = [mx^3 + mx ~ 2, my ~ 2mx + 1], correct init mx = 1, my = 3):

master (silent wrong answer):

COMPAT OK: true
(sol.u[1]).u[1] = [0.0, 1.0]      # wrong, retcode Success

this PR:

COMPAT OK: true
(sol.u[1]).u[1] = [3.0, 1.0]      # correct (state order [my, mx])
compat umap: InitializationStateMap
compat pmap: InitializationParameterMap

Whether master errors or silently mis-initializes depends only on where the aliased value lands: the same equations under different variable naming produce the out-of-range case, which errors. The new Computed initialization maps testset hits that shape on master:

Computed initialization maps: Error During Test ...
  ModelingToolkit initialization maps must copy numeric values from the ODE or initialization problem.
   [2] source_recipe(index::Float64, source_count::Int64)

With this PR both new testsets pass:

Computed initialization maps     |    4      4  35.9s
Homotopy-stripped initialization |    3      3   4.0s

Testing

Review notes: the map-tracing rewrite touches every recipe path, so the pendulum + Test-3 suites doubling as regression coverage for pure-gather maps is the main safety net. Judgment call worth scrutiny: arrays with any computed entry compile the whole array into one generated function (bare entries become direct indexing) rather than mixing per-element recipes.


This PR should be ignored until reviewed by @ChrisRackauckas.

🤖 Generated with Claude Code

Note on Documentation CI: .github/workflows/GPU.yml gates the Documentation job with if: github.event_name == 'push' || !github.event.pull_request.draft, so it shows as skipping while this PR is a draft and only runs once the PR is marked ready for review (or after merge — which is how #511's missing_docs failure escaped to master). Local validation of the edited page is in the Verification section; the full local build cannot go green on this machine (Pascal code 801 on CUDA example pages).

Trace MTK initialization state and parameter maps with Symbolics variables
instead of integer index tokens. Entries copied from a source slot stay static
gather recipes; literal outputs become constants; computed entries (for example
ODE states that are observed variables of the torn initialization system)
compile through `build_function` into isbits generated functions evaluated
against the same source vector on the device.

The integer-token scheme could not represent computed entries, and worse, an
integer-valued computed output could alias a valid source index and initialize
with wrong values while reporting success. The symbolic trace removes that
aliasing by construction.

Raw homotopy initialization problems now fail with a clear error, and the
tutorial documents rewriting `homotopy(actual, simplified)` nodes to `actual`
to obtain an equivalent supported initialization.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Older ModelingToolkit versions in the downgrade CI lower homotopy nodes
without producing a HomotopyProblem, so there is no rejection to assert.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

# Conflicts:
#	docs/src/tutorials/modelingtoolkit.md
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant