Merge pull request #3224 from AayushSabharwal/as/parse-string

feat: add function to parse variable from string

docs/src/basics/FAQ.md

@@ -82,6 +82,16 @@ parameter_index(sys, sym)
 Note that while the variable index will be an integer, the parameter index is a struct of
 type `ParameterIndex` whose internals should not be relied upon.
 
+## Can I index with strings?
+
+Strings are not considered symbolic variables, and thus cannot directly be used for symbolic
+indexing. However, ModelingToolkit does provide a method to parse the string representation of
+a variable, given the system in which that variable exists.
+
+```@docs
+ModelingToolkit.parse_variable
+```
+
 ## Transforming value maps to arrays
 
 ModelingToolkit.jl allows (and recommends) input maps like `[x => 2.0, y => 3.0]`

src/systems/abstractsystem.jl

@@ -3302,6 +3302,96 @@ function dump_unknowns(sys::AbstractSystem)
     end
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Return the variable in `sys` referred to by its string representation `str`.
+Roughly supports the following CFG:
+
+```
+varname                  = "D(" varname ")" | "Differential(" iv ")(" varname ")" | arrvar | maybe_dummy_var
+arrvar                   = maybe_dummy_var "[idxs...]"
+idxs                     = int | int "," idxs
+maybe_dummy_var          = namespacedvar | namespacedvar "(" iv ")" |
+                           namespacedvar "(" iv ")" "ˍ" ts | namespacedvar "ˍ" ts |
+                           namespacedvar "ˍ" ts "(" iv ")"
+ts                       = iv | iv ts
+namespacedvar            = ident "₊" namespacedvar | ident "." namespacedvar | ident
+```
+
+Where `iv` is the independent variable, `int` is an integer and `ident` is an identifier.
+"""
+function parse_variable(sys::AbstractSystem, str::AbstractString)
+    iv = has_iv(sys) ? string(getname(get_iv(sys))) : nothing
+
+    # I'd write a regex to validate `str`, but https://xkcd.com/1171/
+    str = strip(str)
+    derivative_level = 0
+    while ((cond1 = startswith(str, "D(")) || startswith(str, "Differential(")) && endswith(str, ")")
+        if cond1
+            derivative_level += 1
+            str = _string_view_inner(str, 2, 1)
+            continue
+        end
+        _tmpstr = _string_view_inner(str, 13, 1)
+        if !startswith(_tmpstr, "$iv)(")
+            throw(ArgumentError("Expected differential with respect to independent variable $iv in $str"))
+        end
+        derivative_level += 1
+        str = _string_view_inner(_tmpstr, length(iv) + 2, 0)
+    end
+
+    arr_idxs = nothing
+    if endswith(str, ']')
+        open_idx = only(findfirst('[', str))
+        idxs_range = nextind(str, open_idx):prevind(str, lastindex(str))
+        idxs_str = view(str, idxs_range)
+        str = view(str, firstindex(str):prevind(str, open_idx))
+        arr_idxs = map(Base.Fix1(parse, Int), eachsplit(idxs_str, ","))
+    end
+
+    if iv !== nothing && endswith(str, "($iv)")
+        str = _string_view_inner(str, 0, 2 + length(iv))
+    end
+
+    dummyderivative_level = 0
+    if iv !== nothing && (dd_idx = findfirst('ˍ', str)) !== nothing
+        t_idx = findnext(iv, str, dd_idx)
+        while t_idx !== nothing
+            dummyderivative_level += 1
+            t_idx = findnext(iv, str, nextind(str, last(t_idx)))
+        end
+        str = view(str, firstindex(str):prevind(str, dd_idx))
+    end
+
+    if iv !== nothing && endswith(str, "($iv)")
+        str = _string_view_inner(str, 0, 2 + length(iv))
+    end
+
+    cur = sys
+    for ident in eachsplit(str, ('.', NAMESPACE_SEPARATOR))
+        ident = Symbol(ident)
+        hasproperty(cur, ident) ||
+            throw(ArgumentError("System $(nameof(cur)) does not have a subsystem/variable named $(ident)"))
+        cur = getproperty(cur, ident)
+    end
+
+    if arr_idxs !== nothing
+        cur = cur[arr_idxs...]
+    end
+
+    for i in 1:(derivative_level + dummyderivative_level)
+        cur = Differential(get_iv(sys))(cur)
+    end
+
+    return cur
+end
+
+function _string_view_inner(str, startoffset, endoffset)
+    view(str,
+        nextind(str, firstindex(str), startoffset):prevind(str, lastindex(str), endoffset))
+end
+
 ### Functions for accessing algebraic/differential equations in systems ###
 
 """

test/variable_utils.jl

@@ -1,6 +1,7 @@
 using ModelingToolkit, Test
-using ModelingToolkit: value, vars
+using ModelingToolkit: value, vars, parse_variable
 using SymbolicUtils: <ₑ
+
 @parameters α β δ
 expr = (((1 / β - 1) + δ) / α)^(1 / (α - 1))
 ref = sort([β, δ, α], lt = <ₑ)
@@ -41,3 +42,106 @@ ts = collect_ivs([eq])
     res = vars(fn([x, y], z))
     @test length(res) == 3
 end
+
+@testset "parse_variable with iv: $iv" for iv in [t, only(@independent_variables tt)]
+    D = Differential(iv)
+    function Lorenz(; name)
+        @variables begin
+            x(iv)
+            y(iv)
+            z(iv)
+        end
+        @parameters begin
+            σ
+            ρ
+            β
+        end
+        sys = ODESystem(
+            [D(D(x)) ~ σ * (y - x)
+             D(y) ~ x * (ρ - z) - y
+             D(z) ~ x * y - β * z], iv; name)
+    end
+    function ArrSys(; name)
+        @variables begin
+            x(iv)[1:2]
+        end
+        @parameters begin
+            p[1:2, 1:2]
+        end
+        sys = ODESystem([D(D(x)) ~ p * x], iv; name)
+    end
+    function Outer(; name)
+        @named 😄 = Lorenz()
+        @named arr = ArrSys()
+        sys = ODESystem(Equation[], iv; name, systems = [😄, arr])
+    end
+
+    @mtkbuild sys = Outer()
+    for (str, var) in [
+        # unicode system, scalar variable
+        ("😄.x", sys.😄.x),
+        ("😄.x($iv)", sys.😄.x),
+        ("😄₊x", sys.😄.x),
+        ("😄₊x($iv)", sys.😄.x),
+        # derivative
+        ("D(😄.x)", D(sys.😄.x)),
+        ("D(😄.x($iv))", D(sys.😄.x)),
+        ("D(😄₊x)", D(sys.😄.x)),
+        ("D(😄₊x($iv))", D(sys.😄.x)),
+        ("Differential($iv)(😄.x)", D(sys.😄.x)),
+        ("Differential($iv)(😄.x($iv))", D(sys.😄.x)),
+        ("Differential($iv)(😄₊x)", D(sys.😄.x)),
+        ("Differential($iv)(😄₊x($iv))", D(sys.😄.x)),
+        # other derivative
+        ("😄.xˍ$iv", D(sys.😄.x)),
+        ("😄.x($iv)ˍ$iv", D(sys.😄.x)),
+        ("😄₊xˍ$iv", D(sys.😄.x)),
+        ("😄₊x($iv)ˍ$iv", D(sys.😄.x)),
+        # scalar parameter
+        ("😄.σ", sys.😄.σ),
+        ("😄₊σ", sys.😄.σ),
+        # array variable
+        ("arr.x", sys.arr.x),
+        ("arr₊x", sys.arr.x),
+        ("arr.x($iv)", sys.arr.x),
+        ("arr₊x($iv)", sys.arr.x),
+        # getindex
+        ("arr.x[1]", sys.arr.x[1]),
+        ("arr₊x[1]", sys.arr.x[1]),
+        ("arr.x($iv)[1]", sys.arr.x[1]),
+        ("arr₊x($iv)[1]", sys.arr.x[1]),
+        # derivative
+        ("D(arr.x($iv))", D(sys.arr.x)),
+        ("D(arr₊x($iv))", D(sys.arr.x)),
+        ("D(arr.x[1])", D(sys.arr.x[1])),
+        ("D(arr₊x[1])", D(sys.arr.x[1])),
+        ("D(arr.x($iv)[1])", D(sys.arr.x[1])),
+        ("D(arr₊x($iv)[1])", D(sys.arr.x[1])),
+        ("Differential($iv)(arr.x($iv))", D(sys.arr.x)),
+        ("Differential($iv)(arr₊x($iv))", D(sys.arr.x)),
+        ("Differential($iv)(arr.x[1])", D(sys.arr.x[1])),
+        ("Differential($iv)(arr₊x[1])", D(sys.arr.x[1])),
+        ("Differential($iv)(arr.x($iv)[1])", D(sys.arr.x[1])),
+        ("Differential($iv)(arr₊x($iv)[1])", D(sys.arr.x[1])),
+        # other derivative
+        ("arr.xˍ$iv", D(sys.arr.x)),
+        ("arr₊xˍ$iv", D(sys.arr.x)),
+        ("arr.xˍ$iv($iv)", D(sys.arr.x)),
+        ("arr₊xˍ$iv($iv)", D(sys.arr.x)),
+        ("arr.xˍ$iv[1]", D(sys.arr.x[1])),
+        ("arr₊xˍ$iv[1]", D(sys.arr.x[1])),
+        ("arr.xˍ$iv($iv)[1]", D(sys.arr.x[1])),
+        ("arr₊xˍ$iv($iv)[1]", D(sys.arr.x[1])),
+        ("arr.x($iv)ˍ$iv", D(sys.arr.x)),
+        ("arr₊x($iv)ˍ$iv", D(sys.arr.x)),
+        ("arr.x($iv)ˍ$iv[1]", D(sys.arr.x[1])),
+        ("arr₊x($iv)ˍ$iv[1]", D(sys.arr.x[1])),
+        # array parameter
+        ("arr.p", sys.arr.p),
+        ("arr₊p", sys.arr.p),
+        ("arr.p[1, 2]", sys.arr.p[1, 2]),
+        ("arr₊p[1, 2]", sys.arr.p[1, 2])
+    ]
+        @test isequal(parse_variable(sys, str), var)
+    end
+end