Initialization in ModelingToolkit

Specific Domains Modelling & Simulations
1

Questions, etc.

  1. If any of the functionality discussed below already exists, please let me know!
  2. Does anyone else find the ideas useful? Or am I the only one that struggles with initialization?
  3. If anyone with better programming skills manage to develop better versions of the functionality below, with better testing of input arguments, etc., and with correction of the mistakes I certainly have done, please do so!

Motivation

I am testing a “helper” function for initialization of ODEProblems which at the time is named get_initialization_maps (should probably be re-named to get_mtkmaps or something). The two positional arguments are:

  1. csys - the mtkcompile’d symbolic model, and
  2. sys - the uncompiled, complete symbolic model

Here is the use:

u0_map, guess_map = get_initialization_maps(csys, sys)

prob = ODEProblem(csys, u0_map, (0,10.0); guesses = guess_map)

The motivation for this helper function is that I find it somewhat tricky to ensure complete u0_map and guess_map when doing this manually, in particular when the sys is set up from some library.

The function is listed towards the end of this topic.

Working principle

I assume that mtkcompile generates an index 1 DAE M\frac{\mathrm{d}u}{\mathrm{d}t} = f(u;t,p) with u being the unknowns of csys, which can be simplified to:

\begin{align} M_\mathrm{d}\frac{\mathrm{d}u_\mathrm{d}}{\mathrm{d}t} = g(u;t,p) \\ 0 = h(u; t,p) \end{align}

Here, matrix M_\mathrm{d} is non-singular, so I could have assumed it was I.

For proper initialization, I need to specify initial maps u0_map for u_\mathrm{d} — and no other variables — and guess maps for u_\mathrm{a} = u\setminus u_\mathrm{d}. However, there is no harm in providing guess maps for differential unknowns, so I can just as well provide guess_map for all unknowns u.

Function get_initialization_maps (developed with the help of Google AI Studio…) does the following:

  1. Parses the equations of csys and finds the differential variables u_\mathrm{d} – to be used in u0_map.
  2. Reads the unknowns of csys for use in guess_map.
  3. Parses csys and finds all user-provided initial values and guess values in the underlying code, both for unknowns and observed variables.
  4. Compares unknowns + observed variables of csys and unknowns of sys to spot any variables introduced in the index reduction process – these will be given initial/guess values 0.0.
  5. For u0_map, pairs all observed variables + algebraic variables that may happen to have been given initial values or guesses in underlying code to nothing, thus effectively removing them.
  6. For u0_map, sets differential variables to provided initial values (first priority), or guess values (if no initial values are provided), or missing if neither initial values nor guess values have been provided.
  7. For guess_map, sets the unknowns equal to provided initial values (first priority, if any have been given) or guess variables, or missing.

If any variables in u0_map or guess_map have gotten value missing, they need to be manually updated with a value that removes the missing value. I have a function update_map for that (should probably be named set_mtkmap or set_mtkmaps??).

Good and Bad

OK, the above strategy and my code (given below) works for a relatively simple case of my library-in-the-works. To be more precise, it works if all variables are scalar variables.

If some variables are symbolic vectors, I struggle. Presumably, the problem lies in that in sys, the symbolic vectors are still vectors, while in csys, the vectors have been “flattened” into elements of the vector??

I show a modified get_initialization_maps at the bottom of this topic, with a screen shot of how it fails.

Working functions

Here is my function get_initialization_maps:

"""
    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

And here is the update_map function that I use to (i) check whether there are missing values (only one argument), and (ii) replace a missing mapping.

"""
    update_map(base_map, patch_map=nothing)

Audit or repair a ModelingToolkit initialization map. 
Supports both Vector{Pair} and Dict as input, and returns Vector{Pair}.

- If one argument: Returns a vector of variables currently mapped to `missing`.
- If two arguments: Replaces `missing` or existing values with those in `patch_map`, 
  respecting 'nothing' (follower) constraints and existing system schema.
"""
function update_map(base_map, patch_map=nothing)
    # Mode 1: Audit
    if patch_map === nothing
        pairs = (base_map isa AbstractDict) ? collect(base_map) : base_map
        return [p for p in pairs if p.second === missing]
    end

    # Mode 2: Repair
    current_dict = Dict(base_map)
    patch_dict = Dict(patch_map)
    
    # 1. Validation
    for (var, val) in patch_dict
        if !haskey(current_dict, var)
            error("Update Failed: Variable '\$var' does not exist in the map.")
        end
        if current_dict[var] === nothing
            error("Update Failed: Variable '\$var' is a 'nothing' (follower) and cannot be fixed manually.")
        end
    end

    # 2. Construction (Preserving order if base_map is a Vector)
    if base_map isa AbstractVector
        # 'map' here correctly refers to the Julia Base.map function
        return map(base_map) do (var, val)
            haskey(patch_dict, var) ? (var => patch_dict[var]) : (var => val)
        end
    else
        # If input was a Dict, we return a Vector{Pair} to be compatible with ODESystem/ODEProblem
        return [v => (haskey(patch_dict, v) ? patch_dict[v] : current_dict[v]) 
                for v in keys(current_dict)]
    end
end

I should mention that I have a “twin” function to get_initialization_maps where the second argument is not sys, but instead the solution of a successful run. That function picks out the end values of the solution and insert these in u0_map and guess_map. Useful for, e.g., initializing the system at steady state, and for linearization.

Non-working function

The following modification get_initialization_maps is meant to generalize the previous function and allow for variables that are vectors. But it does not work properly.

"""
    get_initialization_maps(csys, sys=nothing)

Generate initialization maps (`u0_map` and `guess_map`) for a compiled 
ModelingToolkit system (v11.x). 

This version explicitly handles the "scalar-equals-vector" error by using 
regex-based index extraction. It identifies scalarized unknowns (e.g., `mp[1]`), 
retrieves the parent vector from metadata, and extracts the specific element.

### Arguments
- `csys`: The compiled or simplified `ODESystem`.
- `sys`: (Optional) The original uncompiled `ODESystem` used to identify 
  variables introduced during index reduction.

### Returns
- `u0_map::Vector{Pair}`: Targets for the `u0` argument in `ODEProblem`.
- `guess_map::Vector{Pair}`: Seeds for the `guesses` argument in `ODEProblem`.
"""
function get_initialization_maps(csys, sys=nothing)
    # 1. Aggregate Metadata
    # We merge defaults and guesses from the compiled system.
    combined_meta = merge(
        ModelingToolkit.default_values(csys),
        ModelingToolkit.get_guesses(csys)
    )

    # 2. Helper: Canonical Normalization
    # Strips time, namespaces, port aliases (ax), and parentheses.
    # Does NOT strip indices yet, as we need them for extraction.
    function normalize_full(v)
        s = string(v)
        s = replace(s, r"\(t\)" => "")
        s = replace(s, "₊" => ".")
        s = replace(s, ".ax." => ".")
        s = replace(s, r"^ax\." => "")
        s = replace(s, r"[\(\)]" => "")
        return s
    end

    # Pre-calculate normalized string map for metadata
    # Metadata keys are usually array-level (no indices)
    meta_strings = Dict()
    for (k, v) in combined_meta
        # Normalize the key and ensure no index is present in the metadata key
        nk = replace(normalize_full(k), r"\[\d+\]" => "")
        meta_strings[nk] = v
    end
    
    # Identify original variables for the 'added' check
    sys_norms = sys === nothing ? Set{String}() : Set(replace(normalize_full(u), r"\[\d+\]" => "") for u in unknowns(sys))

    # 3. Helper: Value Resolution Logic
    function resolve_value(v)
        v_str = normalize_full(v)
        
        # Extract index: "name[idx]"
        idx_match = match(r"\[(\d+)\]", v_str)
        
        # Get the parent name by stripping the index
        v_parent = replace(v_str, r"\[\d+\]" => "")
        
        if haskey(meta_strings, v_parent)
            val = meta_strings[v_parent]
            
            if idx_match !== nothing
                # If the unknown is an indexed element, extract the specific scalar
                idx = parse(Int, idx_match.captures[1])
                if val isa AbstractArray && idx <= length(val)
                    return val[idx]
                else
                    # If metadata is already scalar or index is OOB, return as is
                    return val
                end
            end
            # No index in unknown, return metadata value (scalar or array)
            return val
        end
        return missing
    end

    # 4. Structural Analysis
    eqs = equations(csys)
    diff_vars_list = [arguments(eq.lhs)[1] for eq in eqs if ModelingToolkit.isdifferential(eq.lhs)]
    diff_vars_set = Set(Symbolics.unwrap.(diff_vars_list))
    csys_unks = unknowns(csys)
    
    # 5. Map Construction
    u0_map = Pair{Any, Any}[]
    guess_map = Pair{Any, Any}[]

    for v in csys_unks
        val = resolve_value(v)
        
        # Default to 0.0 for variables added by the compiler (e.g. dummy derivatives)
        if val === missing
            v_parent = replace(normalize_full(v), r"\[\d+\]" => "")
            val = (sys !== nothing && v_parent ∉ sys_norms) ? 0.0 : missing
        end

        push!(guess_map, v => val)

        if Symbolics.unwrap(v) ∈ diff_vars_set
            push!(u0_map, v => val)
        else
            # Algebraic variables are 'nothing' in u0_map
            push!(u0_map, v => nothing)
        end
    end

    # Observed variables are always 'nothing' in u0_map
    for v in observed(csys)
        push!(u0_map, v.lhs => nothing)
    end

    return u0_map, guess_map
end

This function sets “flattened” vector elements equal to the entire vector, instead of the corresponding element of the vector…

image
image895×175 35.3 KB