Initialization maps for ODEProblem?

[B-LIE]

To me, a major problem with using MTK is to find good initialization maps for u0 and guess values.

The above problem is particularly a problem when I try to compose a system from library components (I am attempting to do a small library...), where one does not know in advance what will end up as differential equation and algebraic equations. As a motivating example, consider a rigid pipe with an inelastic fluid. The fluid mass in the pipe is constant, and I use momentum balance to compute the flow rate. Because of constant mass, if I connect two pipes in series, their mass flow rates must be the same, so I can not choose initial momentum independently.

So I have ended up specifying guess values for all variables in the component constructor. This way, I do not overspecify the initialization. At the same time, I have not specified any initial "states", so the solver will give a warning about an underspecified system and that it will use some least squares routine based on guess values. I don't like this particularly, either.

Desired solution

I assume that the compiled symbolic model (csys) is an index 1 DAE. Then I need an initial value (u0) for all differential unknowns, and I need a guess value for all algebraic unknowns. Since it doesn't harm to give guess values in addition to initial values for the unknowns, it is fine to give guess values for all unknowns.

So I have a working function that parses the equations(csys) and picks out differential variables form the unknowns. Then it checks whether there are initial values or guess values specified for these differential variables, and chooses the specified initial value if both are given, and the guess value is made initial value if only guess value exists. If none exists, the value is set to missing. The missing values must be manually added. Also, all non-differential variables are set to nothing in u0 just to cancel out any variables that have been given a value, but would lead to conflicting initialization.

The function does the same for the guess map, but only contains the unknowns.

Because index reduction may add new variables (e.g., time derivatives of original variables in the original symbolic model sys prior to compilation), I assume I can find these by checking for the set difference between csys: unknowns UNION observed variables, and sys: unknowns. I set these extra variables to zero initial value.

The above strategy has worked well in limited testing of systems with scalar symbolic variables. However, the function fails when I try to use it for a system with vector symbolic variables. As an example of a system with vector symbolic variable: a pipe with elastic fluid.

Describe alternatives you’ve considered

I haven't really tried other solutions. It is possible that MTK already has this functionality, but I have not been able to find it. Also, my code is a result of interaction with Google AI Studio and lots of iterations. But "we" have not been able to make the code work for symbolic vector variables. Parts of the problem is that in sys (the un-simplified symbolic model), the vector variables are still vectors, while in csys, they seem to have been "flattened" into scalars.

Additional context

Here is my working function for scalars.

"""
    get_initialization_maps(csys, sys=nothing)

Generate initialization maps (`u0_map` and `guess_map`) for a compiled 
ModelingToolkit system (v11.x), optionally auditing against the uncompiled 
system to handle index reduction.

### Strategy
1. **Differential Variables (The Anchors):**
   - If added by `mtkcompile` (index reduction), set to `0.0` for a neutral start.
   - If original, assigned the value from `initial_conditions` (via `default_values`). 
     Falls back to `guesses` metadata. If both are missing, marked as `missing`.
2. **Algebraic & Observed Variables (The Followers):**
   - **Target Map (`u0_map`):** Set to `nothing` to allow the equations to dictate 
     values, preventing over-specification.
   - **Seed Map (`guess_map`):** Assigned values from `guesses` (priority) or 
     `initial_conditions`. If neither exist, marked as `missing`.

### Arguments
- `csys`: The compiled or simplified `ODESystem` (e.g., from `mtkcompile`).
- `sys`: (Optional) The original uncompiled `ODESystem` used to detect 
  index-reduced variables.

### Returns
- `u0_map`: Vector of pairs for the `u0` argument in `ODEProblem` (Targets).
- `guess_map`: Vector of pairs for the `guesses` argument (Seeds).
"""
function get_initialization_maps(csys, sys=nothing)
    # 1. Structural Analysis via Equation Inspection
    eqs = equations(csys)
    
    # Identify differential variables (differentiated in the compiled system)
    diff_vars_list = [arguments(eq.lhs)[1] for eq in eqs if ModelingToolkit.isdifferential(eq.lhs)]
    diff_vars_set = Set(diff_vars_list)
    
    # Identify all unknowns in the compiled system
    csys_unks = unknowns(csys)
    
    # Identify variables added by MTK (e.g., during index reduction)
    added_vars_set = if sys === nothing
        Set()
    else
        sys_unks_set = Set(unknowns(sys))
        Set([u for u in csys_unks if u ∉ sys_unks_set])
    end
    
    # Algebraic unknowns: unknowns in csys that are not differential
    alg_vars = [u for u in csys_unks if u ∉ diff_vars_set]
    
    # Observed variables (Simplified out of the solver state)
    obs_vars = [v.lhs for v in observed(csys)]
    
    # 2. Metadata Extraction
    # ic_defaults = 'initial_conditions' keyword values
    # explicit_guesses = '[guess = ...]' or 'guesses' keyword values
    ic_defaults = ModelingToolkit.default_values(csys)
    explicit_guesses = ModelingToolkit.get_guesses(csys)

    # 3. Build u0_map (The Targets)
    u0_map = Pair{Any, Any}[]

    for v in diff_vars_list
        if v ∈ added_vars_set
            push!(u0_map, v => 0.0)
        else
            # Precedence: Initial Condition > Guess
            val = get(ic_defaults, v, get(explicit_guesses, v, missing))
            push!(u0_map, v => val)
        end
    end

    for v in [alg_vars; obs_vars]
        push!(u0_map, v => nothing)
    end

    # 4. Build guess_map (The Seeds)
    # Returned as Vector{Pair} for consistency and order preservation
    guess_map = Pair{Any, Any}[]
    
    for v in csys_unks
        if v ∈ added_vars_set
            push!(guess_map, v => 0.0)
        else
            # Precedence: Explicit Guess > Initial Condition
            val = get(explicit_guesses, v, get(ic_defaults, v, missing))
            push!(guess_map, v => val)
        end
    end

    return u0_map, guess_map
end

This does not work for @variables that are defined as vectors. Also: note that my function is not of professional quality, and can be improved. The name is not super good either (is get_mtkmaps better?).