Merge pull request #2768 from AayushSabharwal/as/mtkize-names

ChrisRackauckas · web-flow · commit a285fc29ddb3 · 2024-06-07T08:16:11.000-04:00
feat: allow specifying variable names in `modelingtoolkitize`
diff --git a/src/systems/diffeqs/modelingtoolkitize.jl b/src/systems/diffeqs/modelingtoolkitize.jl
@@ -3,29 +3,44 @@ $(TYPEDSIGNATURES)
 
 Generate `ODESystem`, dependent variables, and parameters from an `ODEProblem`.
 """
-function modelingtoolkitize(prob::DiffEqBase.ODEProblem; kwargs...)
+function modelingtoolkitize(
+        prob::DiffEqBase.ODEProblem; u_names = nothing, p_names = nothing, kwargs...)
     prob.f isa DiffEqBase.AbstractParameterizedFunction &&
         return prob.f.sys
-    @parameters t
-
+    t = t_nounits
     p = prob.p
     has_p = !(p isa Union{DiffEqBase.NullParameters, Nothing})
 
-    _vars = define_vars(prob.u0, t)
+    if u_names !== nothing
+        varnames_length_check(prob.u0, u_names; is_unknowns = true)
+        _vars = [_defvar(name)(t) for name in u_names]
+    elseif SciMLBase.has_sys(prob.f)
+        varnames = getname.(variable_symbols(prob.f.sys))
+        varidxs = variable_index.((prob.f.sys,), varnames)
+        invpermute!(varnames, varidxs)
+        _vars = [_defvar(name)(t) for name in varnames]
+    else
+        _vars = define_vars(prob.u0, t)
+    end
 
     vars = prob.u0 isa Number ? _vars : ArrayInterface.restructure(prob.u0, _vars)
     params = if has_p
-        _params = define_params(p)
+        if p_names === nothing && SciMLBase.has_sys(prob.f)
+            p_names = Dict(parameter_index(prob.f.sys, sym) => sym
+            for sym in parameter_symbols(prob.f.sys))
+        end
+        _params = define_params(p, p_names)
         p isa Number ? _params[1] :
-        (p isa Tuple || p isa NamedTuple || p isa AbstractDict ? _params :
+        (p isa Tuple || p isa NamedTuple || p isa AbstractDict || p isa MTKParameters ?
+         _params :
          ArrayInterface.restructure(p, _params))
     else
         []
     end
 
     var_set = Set(vars)
 
-    D = Differential(t)
+    D = D_nounits
     mm = prob.f.mass_matrix
 
     if mm === I
@@ -70,6 +85,8 @@ function modelingtoolkitize(prob::DiffEqBase.ODEProblem; kwargs...)
     default_p = if has_p
         if prob.p isa AbstractDict
             Dict(v => prob.p[k] for (k, v) in pairs(_params))
+        elseif prob.p isa MTKParameters
+            Dict(params .=> reduce(vcat, prob.p))
         else
             Dict(params .=> vec(collect(prob.p)))
         end
@@ -125,44 +142,96 @@ function Base.showerror(io::IO, e::ModelingtoolkitizeParametersNotSupportedError
     println(io, e.type)
 end
 
-function define_params(p)
+function varnames_length_check(vars, names; is_unknowns = false)
+    if length(names) != length(vars)
+        throw(ArgumentError("""
+            Number of $(is_unknowns ? "unknowns" : "parameters") ($(length(vars))) \
+            does not match number of names ($(length(names))).
+        """))
+    end
+end
+
+function define_params(p, _ = nothing)
     throw(ModelingtoolkitizeParametersNotSupportedError(typeof(p)))
 end
 
-function define_params(p::AbstractArray)
-    [toparam(variable(:α, i)) for i in eachindex(p)]
+function define_params(p::AbstractArray, names = nothing)
+    if names === nothing
+        [toparam(variable(:α, i)) for i in eachindex(p)]
+    else
+        varnames_length_check(p, names)
+        [toparam(variable(names[i])) for i in eachindex(p)]
+    end
 end
 
-function define_params(p::Number)
-    [toparam(variable(:α))]
+function define_params(p::Number, names = nothing)
+    if names === nothing
+        [toparam(variable(:α))]
+    elseif names isa Union{AbstractArray, AbstractDict}
+        varnames_length_check(p, names)
+        [toparam(variable(names[i])) for i in eachindex(p)]
+    else
+        [toparam(variable(names))]
+    end
 end
 
-function define_params(p::AbstractDict)
-    OrderedDict(k => toparam(variable(:α, i)) for (i, k) in zip(1:length(p), keys(p)))
+function define_params(p::AbstractDict, names = nothing)
+    if names === nothing
+        OrderedDict(k => toparam(variable(:α, i)) for (i, k) in zip(1:length(p), keys(p)))
+    else
+        varnames_length_check(p, names)
+        OrderedDict(k => toparam(variable(names[k])) for k in keys(p))
+    end
 end
 
-function define_params(p::Union{SLArray, LArray})
-    [toparam(variable(x)) for x in LabelledArrays.symnames(typeof(p))]
+function define_params(p::Union{SLArray, LArray}, names = nothing)
+    if names === nothing
+        [toparam(variable(x)) for x in LabelledArrays.symnames(typeof(p))]
+    else
+        varnames_length_check(p, names)
+        [toparam(variable(names[i])) for i in eachindex(p)]
+    end
 end
 
-function define_params(p::Tuple)
-    tuple((toparam(variable(:α, i)) for i in eachindex(p))...)
+function define_params(p::Tuple, names = nothing)
+    if names === nothing
+        tuple((toparam(variable(:α, i)) for i in eachindex(p))...)
+    else
+        varnames_length_check(p, names)
+        tuple((toparam(variable(names[i])) for i in eachindex(p))...)
+    end
 end
 
-function define_params(p::NamedTuple)
-    NamedTuple(x => toparam(variable(x)) for x in keys(p))
+function define_params(p::NamedTuple, names = nothing)
+    if names === nothing
+        NamedTuple(x => toparam(variable(x)) for x in keys(p))
+    else
+        varnames_length_check(p, names)
+        NamedTuple(x => toparam(variable(names[x])) for x in keys(p))
+    end
 end
 
-function define_params(p::MTKParameters)
-    bufs = (p...,)
-    i = 1
-    ps = []
-    for buf in bufs
-        for _ in buf
-            push!(ps, toparam(variable(:α, i)))
+function define_params(p::MTKParameters, names = nothing)
+    if names === nothing
+        bufs = (p...,)
+        i = 1
+        ps = []
+        for buf in bufs
+            for _ in buf
+                push!(
+                    ps,
+                    if names === nothing
+                        toparam(variable(:α, i))
+                    else
+                        toparam(variable(names[i]))
+                    end
+                )
+            end
         end
+        return identity.(ps)
+    else
+        return collect(values(names))
     end
-    return identity.(ps)
 end
 
 """
diff --git a/src/systems/nonlinear/modelingtoolkitize.jl b/src/systems/nonlinear/modelingtoolkitize.jl
@@ -3,17 +3,34 @@ $(TYPEDSIGNATURES)
 
 Generate `NonlinearSystem`, dependent variables, and parameters from an `NonlinearProblem`.
 """
-function modelingtoolkitize(prob::NonlinearProblem; kwargs...)
+function modelingtoolkitize(
+        prob::NonlinearProblem; u_names = nothing, p_names = nothing, kwargs...)
     p = prob.p
     has_p = !(p isa Union{DiffEqBase.NullParameters, Nothing})
 
-    _vars = reshape([variable(:x, i) for i in eachindex(prob.u0)], size(prob.u0))
+    if u_names !== nothing
+        varnames_length_check(prob.u0, u_names; is_unknowns = true)
+        _vars = [variable(name) for name in u_names]
+    elseif SciMLBase.has_sys(prob.f)
+        varnames = getname.(variable_symbols(prob.f.sys))
+        varidxs = variable_index.((prob.f.sys,), varnames)
+        invpermute!(varnames, varidxs)
+        _vars = [variable(name) for name in varnames]
+    else
+        _vars = [variable(:x, i) for i in eachindex(prob.u0)]
+    end
+    _vars = reshape(_vars, size(prob.u0))
 
     vars = prob.u0 isa Number ? _vars : ArrayInterface.restructure(prob.u0, _vars)
     params = if has_p
-        _params = define_params(p)
+        if p_names === nothing && SciMLBase.has_sys(prob.f)
+            p_names = Dict(parameter_index(prob.f.sys, sym) => sym
+            for sym in parameter_symbols(prob.f.sys))
+        end
+        _params = define_params(p, p_names)
         p isa Number ? _params[1] :
-        (p isa Tuple || p isa NamedTuple ? _params :
+        (p isa Tuple || p isa NamedTuple || p isa AbstractDict || p isa MTKParameters ?
+         _params :
          ArrayInterface.restructure(p, _params))
     else
         []
@@ -29,14 +46,25 @@ function modelingtoolkitize(prob::NonlinearProblem; kwargs...)
     eqs = vcat([0.0 ~ rhs[i] for i in 1:length(out_def)]...)
 
     sts = vec(collect(vars))
-
+    _params = params
+    params = values(params)
     params = if params isa Number || (params isa Array && ndims(params) == 0)
         [params[1]]
     else
         vec(collect(params))
     end
     default_u0 = Dict(sts .=> vec(collect(prob.u0)))
-    default_p = has_p ? Dict(params .=> vec(collect(prob.p))) : Dict()
+    default_p = if has_p
+        if prob.p isa AbstractDict
+            Dict(v => prob.p[k] for (k, v) in pairs(_params))
+        elseif prob.p isa MTKParameters
+            Dict(params .=> reduce(vcat, prob.p))
+        else
+            Dict(params .=> vec(collect(prob.p)))
+        end
+    else
+        Dict()
+    end
 
     de = NonlinearSystem(eqs, sts, params,
         defaults = merge(default_u0, default_p);
diff --git a/src/systems/optimization/modelingtoolkitize.jl b/src/systems/optimization/modelingtoolkitize.jl
@@ -3,25 +3,70 @@ $(TYPEDSIGNATURES)
 
 Generate `OptimizationSystem`, dependent variables, and parameters from an `OptimizationProblem`.
 """
-function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem; kwargs...)
+function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem;
+        u_names = nothing, p_names = nothing, kwargs...)
     num_cons = isnothing(prob.lcons) ? 0 : length(prob.lcons)
     if prob.p isa Tuple || prob.p isa NamedTuple
         p = [x for x in prob.p]
     else
         p = prob.p
     end
-
-    vars = ArrayInterface.restructure(prob.u0,
-        [variable(:x, i) for i in eachindex(prob.u0)])
-    params = if p isa DiffEqBase.NullParameters
-        []
-    elseif p isa MTKParameters
-        [variable(:α, i) for i in eachindex(vcat(p...))]
+    has_p = !(p isa Union{DiffEqBase.NullParameters, Nothing})
+    if u_names !== nothing
+        varnames_length_check(prob.u0, u_names; is_unknowns = true)
+        _vars = [variable(name) for name in u_names]
+    elseif SciMLBase.has_sys(prob.f)
+        varnames = getname.(variable_symbols(prob.f.sys))
+        varidxs = variable_index.((prob.f.sys,), varnames)
+        invpermute!(varnames, varidxs)
+        _vars = [variable(name) for name in varnames]
+        if prob.f.sys isa OptimizationSystem
+            for (i, sym) in enumerate(variable_symbols(prob.f.sys))
+                if hasbounds(sym)
+                    _vars[i] = Symbolics.setmetadata(
+                        _vars[i], VariableBounds, getbounds(sym))
+                end
+            end
+        end
+    else
+        _vars = [variable(:x, i) for i in eachindex(prob.u0)]
+    end
+    _vars = reshape(_vars, size(prob.u0))
+    vars = ArrayInterface.restructure(prob.u0, _vars)
+    params = if has_p
+        if p_names === nothing && SciMLBase.has_sys(prob.f)
+            p_names = Dict(parameter_index(prob.f.sys, sym) => sym
+            for sym in parameter_symbols(prob.f.sys))
+        end
+        if p isa MTKParameters
+            old_to_new = Dict()
+            for sym in parameter_symbols(prob)
+                idx = parameter_index(prob, sym)
+                old_to_new[unwrap(sym)] = unwrap(p_names[idx])
+            end
+            order = reorder_parameters(prob.f.sys, full_parameters(prob.f.sys))
+            for arr in order
+                for i in eachindex(arr)
+                    arr[i] = old_to_new[arr[i]]
+                end
+            end
+            _params = order
+        else
+            _params = define_params(p, p_names)
+        end
+        p isa Number ? _params[1] :
+        (p isa Tuple || p isa NamedTuple || p isa AbstractDict || p isa MTKParameters ?
+         _params :
+         ArrayInterface.restructure(p, _params))
     else
-        ArrayInterface.restructure(p, [variable(:α, i) for i in eachindex(p)])
+        []
     end
 
-    eqs = prob.f(vars, params)
+    if p isa MTKParameters
+        eqs = prob.f(vars, params...)
+    else
+        eqs = prob.f(vars, params)
+    end
 
     if DiffEqBase.isinplace(prob) && !isnothing(prob.f.cons)
         lhs = Array{Num}(undef, num_cons)
@@ -58,10 +103,32 @@ function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem; kwargs...)
     else
         cons = []
     end
+    params = values(params)
+    params = if params isa Number || (params isa Array && ndims(params) == 0)
+        [params[1]]
+    elseif p isa MTKParameters
+        reduce(vcat, params)
+    else
+        vec(collect(params))
+    end
 
-    de = OptimizationSystem(eqs, vec(vars), vec(toparam.(params));
+    sts = vec(collect(vars))
+    default_u0 = Dict(sts .=> vec(collect(prob.u0)))
+    default_p = if has_p
+        if prob.p isa AbstractDict
+            Dict(v => prob.p[k] for (k, v) in pairs(_params))
+        elseif prob.p isa MTKParameters
+            Dict(params .=> reduce(vcat, prob.p))
+        else
+            Dict(params .=> vec(collect(prob.p)))
+        end
+    else
+        Dict()
+    end
+    de = OptimizationSystem(eqs, sts, params;
         name = gensym(:MTKizedOpt),
         constraints = cons,
+        defaults = merge(default_u0, default_p),
         kwargs...)
     de
 end
diff --git a/test/modelingtoolkitize.jl b/test/modelingtoolkitize.jl
