feat: allow specifying variable names in modelingtoolkitize

Author: AayushSabharwal

src/systems/diffeqs/modelingtoolkitize.jl

@@ -3,19 +3,22 @@ $(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
+        _vars = [_defvar(name)(t) for name in u_names]
+    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)
+        _params = define_params(p, p_names)
         p isa Number ? _params[1] :
         (p isa Tuple || p isa NamedTuple || p isa AbstractDict ? _params :
          ArrayInterface.restructure(p, _params))
@@ -25,7 +28,7 @@ function modelingtoolkitize(prob::DiffEqBase.ODEProblem; kwargs...)
 
     var_set = Set(vars)
 
-    D = Differential(t)
+    D = D_nounits
     mm = prob.f.mass_matrix
 
     if mm === I
@@ -125,41 +128,89 @@ function Base.showerror(io::IO, e::ModelingtoolkitizeParametersNotSupportedError
     println(io, e.type)
 end
 
-function define_params(p)
+function varnames_length_check(vars, names)
+        if length(names) == length(p)
+            throw(ArgumentError("""
+                Number of parameters ($(length(p))) 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(name)) for name in names]
+    end
 end
 
-function define_params(p::Number)
-    [toparam(variable(:α))]
+function define_params(p::Number, names = nothing)
+    if names === nothing
+        [toparam(variable(:α))]
+    elseif names isa AbstractArray
+        varnames_length_check(p, names)
+        [toparam(variable(name)) for name in names]
+    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(name)) for name in names]
+    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(name)) for name in names))
+    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)
+function define_params(p::MTKParameters, names = nothing)
     bufs = (p...,)
     i = 1
     ps = []
     for buf in bufs
         for _ in buf
-            push!(ps, toparam(variable(:α, i)))
+            push!(
+                ps,
+                if names === nothing
+                    toparam(variable(:α, i))
+                else
+                    toparam(variable(names[i]))
+                end
+            )
         end
     end
     return identity.(ps)