Merge pull request #153 from saschatimme/st/variable-arrays

Add easy possibility to create array of variables

Author: ChrisRackauckas

README.md

@@ -109,6 +109,36 @@ f = @eval eval(nlsys_func)
 f2 = (du,u) -> f(du,u,(10.0,26.0,2.33))
 ```
 
+### Example: Arrays of variables
+
+Sometimes it is convenient to define arrays of variables to model things like `x₁,…,x₃`.
+The `@variables` and `@parameters` macros support this with the following syntax:
+
+```julia
+julia> @variables x[1:3];
+julia> x
+3-element Array{Operation,1}:
+ x[1]()
+ x[2]()
+ x[3]()
+
+# support for arbitrary ranges and tensors
+julia> @variables y[2:3,1:5:6];
+julia> y
+2×2 Array{Operation,2}:
+ y[2,1]()  y[2,6]()
+ y[3,1]()  y[3,6]()
+
+
+# also works for dependent variables
+julia> @parameters t; @variables z[1:3](t);
+julia> z
+3-element Array{Operation,1}:
+ z[1](t())
+ z[2](t())
+ z[3](t())
+ ```
+
 ## Core Principles
 
 The core idea behind ModelingToolkit.jl is that mathematical equations require
@@ -145,7 +175,9 @@ context-aware single variable of the IR. Its fields are described as follows:
 
 For example, the following code defines an independent variable `t`, a parameter
 `α`, a function parameter `σ`, a variable `x` which depends on `t`, a variable
-`y` with no dependents, and a variable `z` which depends on `t`, `α`, and `x(t)`.
+`y` with no dependents, a variable `z` which depends on `t`, `α`, and `x(t)`
+and a parameters `β₁` and `β₂`.
+
 
 ```julia
 t = Variable(:t; known = true)()  # independent variables are treated as known
@@ -155,19 +187,21 @@ w = Variable(:w; known = false)   # unknown, left uncalled
 x = Variable(:x; known = false)(t)  # unknown, depends on `t`
 y = Variable(:y; known = false)()   # unknown, no dependents
 z = Variable(:z; known = false)(t, α, x)  # unknown, multiple arguments
+β₁ = Variable(:β, 1; known = true)() # with index 1
+β₂ = Variable(:β, 2; known = true)() # with index 2
 
-expr = x + y^α + σ(3) * (z - t) - w(t - 1)
+expr = β₁ * x + y^α + σ(3) * (z - t) - β₂ * w(t - 1)
 ```
 
 We can rewrite this more concisely using macros. Note the difference between
 including and excluding empty parentheses. When in call format, variables are
 aliased to the given call, allowing implicit use of dependents for convenience.
 
 ```julia
-@parameters t α σ(..)
+@parameters t α σ(..) β[1:2]
 @variables w(..) x(t) y() z(t, α, x)
 
-expr = x + y^α + σ(3) * (z - t) - w(t - 1)
+expr = β[1] * x + y^α + σ(3) * (z - t) - β[2] * w(t - 1)
 ```
 
 Note that `@parameters` and `@variables` implicitly add `()` to values that
@@ -276,7 +310,7 @@ is accessible via a function-based interface. This means that all macros are
 syntactic sugar in some form. For example, the variable construction:
 
 ```julia
-@parameters t σ ρ β
+@parameters t σ ρ β[1:3]
 @variables x(t) y(t) z(t)
 @derivatives D'~t
 ```
@@ -287,7 +321,7 @@ is syntactic sugar for:
 t = Variable(:t; known = true)()
 σ = Variable(:σ; known = true)
 ρ = Variable(:ρ; known = true)()
-β = Variable(:β; known = true)()
+β = [Variable(:β, i; known = true)() for i in 1:3]
 x = Variable(:x)(t)
 y = Variable(:y)(t)
 z = Variable(:z)(t)

src/variables.jl

@@ -19,6 +19,11 @@ struct Variable <: Function
     known::Bool
     Variable(name; known = false) = new(name, known)
 end
+function Variable(name, indices...; known = false)
+    var_name = Symbol("$name[$(join(indices, ","))]")
+    Variable(var_name; known=known)
+end
+
 (x::Variable)(args...) = Operation(x, collect(Expression, args))
 
 Base.isequal(x::Variable, y::Variable) = (x.name, x.known) == (y.name, y.known)
@@ -74,29 +79,64 @@ function _parse_vars(macroname, known, x)
     x = flatten_expr!(x)
     for _var in x
         iscall = isa(_var, Expr) && _var.head == :call
-        issym    = _var isa Symbol
-        @assert iscall || issym "@$macroname expects a tuple of expressions or an expression of a tuple (`@$macroname x y z(t)` or `@$macroname x, y, z(t)`)"
+        isarray = isa(_var, Expr) && _var.head == :ref
+        issym  = _var isa Symbol
+        @assert iscall || isarray || issym "@$macroname expects a tuple of expressions or an expression of a tuple (`@$macroname x y z(t) v[1:3] w[1:2,1:4]` or `@$macroname x, y, z(t) v[1:3] w[1:2,1:4]`)"
 
         if iscall
-            var_name = _var.args[1]
-            if _var.args[end] == :..
-                expr = :($var_name = $Variable($(Meta.quot(var_name)); known = $known))
-            else
-                expr = :($var_name = $Variable($(Meta.quot(var_name)); known = $known)($(_var.args[2:end]...)))
-            end
+            var_name, expr = _construct_vars(_var.args[1], known, _var.args[2:end])
         else
-            # Implicit 0-args call
-            var_name = _var
-            expr = :($var_name = $Variable($(Meta.quot(var_name)); known = $known)())
+            var_name, expr = _construct_vars(_var, known, nothing)
         end
-
         push!(var_names, var_name)
         push!(ex.args, expr)
     end
     push!(ex.args, build_expr(:tuple, var_names))
     return ex
 end
 
+function _construct_vars(_var, known, call_args)
+    issym  = _var isa Symbol
+    isarray = isa(_var, Expr) && _var.head == :ref
+    if isarray
+        var_name = _var.args[1]
+        indices = _var.args[2:end]
+        expr = _construct_array_vars(var_name, known, call_args, indices...)
+    else
+        # Implicit 0-args call
+        var_name = _var
+        expr = _construct_var(var_name, known, call_args)
+    end
+    var_name, :($var_name = $expr)
+end
+
+function _construct_var(var_name, known, call_args)
+    if call_args === nothing
+        :(Variable($(Meta.quot(var_name)); known = $known)())
+    elseif call_args[end] == :..
+        :(Variable($(Meta.quot(var_name)); known = $known))
+    else
+        :(Variable($(Meta.quot(var_name)); known = $known)($(call_args...)))
+    end
+end
+
+function _construct_var(var_name, known, call_args, ind)
+    if call_args === nothing
+        :(Variable($(Meta.quot(var_name)), $ind...; known = $known)())
+    elseif call_args[end] == :..
+        :(Variable($(Meta.quot(var_name)), $ind...; known = $known))
+    else
+        :(Variable($(Meta.quot(var_name)), $ind...; known = $known)($(call_args...)))
+    end
+end
+
+
+function _construct_array_vars(var_name, known, call_args, indices...)
+    :(map(Iterators.product($(indices...))) do ind
+        $(_construct_var(var_name, known, call_args, :ind))
+    end)
+end
+
 
 """
 $(SIGNATURES)

test/variable_parsing.jl

@@ -34,3 +34,25 @@ D1 = Differential(t)
 @test @macroexpand(@variables x, y, z(t)) == @macroexpand(@variables x y z(t))
 
 convert(Expression, :($x == 0 ? $y : $x))
+
+# Test array expressions
+@parameters begin
+    t[1:2]
+    s[1:2:4,1:2]
+end
+@parameters σ[1:2](..)
+
+t1 = [Variable(:t, 1; known = true)(),
+      Variable(:t, 2; known = true)()]
+s1 = [Variable(:s, 1, 1; known = true)() Variable(:s, 1, 2; known = true)()
+      Variable(:s, 3, 1; known = true)() Variable(:s, 3, 2; known = true)()]
+σ1 = [Variable(:σ, 1; known = true),
+      Variable(:σ, 2; known = true)]
+@test isequal(t1, t)
+@test isequal(s1, s)
+@test isequal(σ1, σ)
+
+@parameters t
+@variables x[1:2](t)
+x1 = [Variable(:x, 1)(t), Variable(:x, 2)(t)]
+@test isequal(x1, x)