Additions from algebraic GPU work

Major:
- Added "convex_evaluator" and "all_evaluators" functions
- Added methods for apply_transform with an Equation or Vector{Equation} as the input instead of an ODESystem
- Added pull_vars to extract variable names from Symbolics expressions
- Added shrink_eqs to shrink vectors of equations via substitution

Minor:
- Added a fix for variables like x[1] that are not dynamic functions of other variables such as "t".
- Added support for ^2 operation
- Corrected mid_expr to now work properly with Symbolics

Author: RXGottlieb

Manifest.toml

@@ -0,0 +1,15 @@
+name = "SourceCodeMcCormick"
+uuid = "a7283dc5-4ecf-47fb-a95b-1412723fc960"
+authors = ["Robert Gottlieb <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
+
+[compat]
+IfElse = "0.1.0 - 1.1.1"
+ModelingToolkit = "~8"
+SymbolicUtils = "~0.19"

Project.toml

@@ -3,6 +3,9 @@ module SourceCodeMcCormick
 
 using ModelingToolkit
 using SymbolicUtils.Code
+using IfElse
+using DocStringExtensions
+
 # TODO: Need to import Assignment and other stuff probably
 # Check out the functionality in ModelingToolkit.jl and Symbolics.jl
 
@@ -14,7 +17,9 @@ function transform_rule end
 
 export McCormickIntervalTransform
 
-export apply_transform, extract_terms, genvar, genparam, get_name
+export apply_transform, extract_terms, genvar, genparam, get_name,
+        factor!, binarize!, pull_vars, shrink_eqs, convex_evaluator,
+        all_evaluators
 
 include(joinpath(@__DIR__, "interval", "interval.jl"))
 include(joinpath(@__DIR__, "relaxation", "relaxation.jl"))

src/SourceCodeMcCormick.jl

@@ -22,21 +22,23 @@ end
 function var_names(::IntervalTransform, s::Real)
     return s, s
 end
-function var_names(::IntervalTransform, s::Term{Real, Nothing}) #Any terms like "Differential"
-    if length(s.arguments)>1
-        error("Multiple arguments not supported.")
-    end
-    if typeof(s.arguments[1])<:Term #then it has args
+function var_names(::IntervalTransform, s::Term{Real, Nothing}) #Any terms like "Differential", or "x[1]" (NOT x[1](t))
+    if typeof(s.arguments[1])<:Term #then it has typical args like "x", "y", ...
         args = Symbol[]
         for i in s.arguments[1].arguments
             push!(args, get_name(i))
         end
         var = get_name(s.arguments[1])
         var_lo = genvar(Symbol(string(var)*"_lo"), args)
         var_hi = genvar(Symbol(string(var)*"_hi"), args)
-    elseif typeof(s.arguments[1])<:Sym #Then it has no args
-        var_lo = genparam(Symbol(string(s.arguments[1].name)*"_lo"))
-        var_hi = genparam(Symbol(string(s.arguments[1].name)*"_hi"))
+    elseif typeof(s.arguments[1])<:Sym #Then it has no typical args, i.e., x[1] has args Any[x, 1]
+        if length(s.arguments)==1
+            var_lo = genparam(Symbol(string(s.arguments[1].name)*"_lo"))
+            var_hi = genparam(Symbol(string(s.arguments[1].name)*"_hi"))
+        else
+            var_lo = genparam(Symbol(string(s.arguments[1].name)*"_"*string(s.arguments[2])*"_lo"))
+            var_hi = genparam(Symbol(string(s.arguments[1].name)*"_"*string(s.arguments[2])*"_hi"))
+        end
     else
         error("Type of argument invalid")
     end
@@ -74,18 +76,13 @@ function translate_initial_conditions(::IntervalTransform, prob::ODESystem, new_
 end
 
 
-
-# function var_names(::IntervalTransform, s::Number)
-#     sL = s
-#     sU = s
-#     sL, sU
-# end
-
 # Helper functions for navigating SymbolicUtils structures
 get_name(x::Sym{SymbolicUtils.FnType{Tuple{Any}, Real}, Nothing}) = x.name
 
 """
-Takes x[1,1] returns :x_1_1
+    get_name
+
+Take a Symbolic-type object such as `x[1,1]` and return a symbol like `:x_1_1`.
 """
 function get_name(s::Term{SymbolicUtils.FnType{Tuple, Real}, Nothing})
     d = s.arguments

src/interval/interval.jl

@@ -1,6 +1,4 @@
 
-using IfElse
-
 # Transformation rules should have the general form listed below:
 # Check IntervalArithmetic for valid interval bound rules (we won't be able 
 # to do correctly rounded stuff on GPUs but otherwise the operators should work out)
@@ -13,6 +11,11 @@ function transform_rule(::IntervalTransform, nothing, yL, yU, xL, xU)
     ru = Equation(yU, xU)
     return rl, ru
 end
+function transform_rule(::IntervalTransform, ::typeof(getindex), yL, yU, xL, xU)
+    rl = Equation(yL, xL)
+    ru = Equation(yU, xU)
+    return rl, ru
+end
 function transform_rule(::IntervalTransform, ::typeof(exp), yL, yU, xL, xU)
     rl = Equation(yL, exp(xL))
     ru = Equation(yU, exp(xU))
@@ -63,5 +66,12 @@ function transform_rule(::IntervalTransform, ::typeof(max), zL, zU, xL, xU, yL,
     ru = Equation(zU, max(xU, yU))
     return rl, ru
 end
+function transform_rule(::IntervalTransform, ::typeof(^), zL, zU, xL, xU, yL, yU)
+    ~((typeof(yL) <: Int) || (typeof(yL) <: AbstractFloat)) && error("Symbolic exponents not currently supported.")
+    ~(yL == 2) && error("Exponents besides 2 not currently supported")
+    rl = Equation(zL, max(min(xU, 0.0), xL)^2)
+    ru = Equation(zU, IfElse.ifelse(xU < 0.0, xL, IfElse.ifelse(xL > 0, xU, IfElse.ifelse(abs(xL) >= abs(xU), xL, xU)))^2)
+    return rl, ru
+end
 
 # TODO: /, ^

src/interval/rules.jl

@@ -19,9 +19,6 @@ function var_names(::McCormickTransform, s::Real)
     return s, s
 end
 function var_names(::McCormickTransform, s::Term{Real, Nothing}) #Any terms like "Differential"
-    if length(s.arguments)>1
-        error("Multiple arguments not supported.")
-    end
     if typeof(s.arguments[1])<:Term #then it has args
         args = Symbol[]
         for i in s.arguments[1].arguments
@@ -31,8 +28,13 @@ function var_names(::McCormickTransform, s::Term{Real, Nothing}) #Any terms like
         var_cv = genvar(Symbol(string(var)*"_cv"), args)
         var_cc = genvar(Symbol(string(var)*"_cc"), args)
     elseif typeof(s.arguments[1])<:Sym #Then it has no args
-        var_cv = genparam(Symbol(string(s.arguments[1].name)*"_cv"))
-        var_cc = genparam(Symbol(string(s.arguments[1].name)*"_cc"))
+        if length(s.arguments)==1
+            var_cv = genparam(Symbol(string(s.arguments[1].name)*"_cv"))
+            var_cc = genparam(Symbol(string(s.arguments[1].name)*"_cc"))
+        else
+            var_cv = genparam(Symbol(string(s.arguments[1].name)*"_"*string(s.arguments[2])*"_cv"))
+            var_cc = genparam(Symbol(string(s.arguments[1].name)*"_"*string(s.arguments[2])*"_cc"))
+        end
     else
         error("Type of argument invalid")
     end
@@ -105,7 +107,7 @@ end
 line_expr(x, xL, xU, zL, zU) = IfElse.ifelse(zU > zL, (zL*(xU - x) + zU*(x - xL))/(xU - xL), zU)
 
 # A symbolic way of computing the mid of three numbers (returns IfElse block)
-mid_expr(a, b, c) = IfElse.ifelse((a < b) && (b < c), y, IfElse.ifelse((c < b) && (b < a), b,
-                    IfElse.ifelse((b < a) && (a < c), x, IfElse.ifelse((c < a) && (a < b), a, c))))
+mid_expr(a, b, c) = IfElse.ifelse(a < b, IfElse.ifelse(b < c, b, IfElse.ifelse(c < a, a, c)),
+                        IfElse.ifelse(c < b, b, IfElse.ifelse(a < c, a, c)))
 
 include(joinpath(@__DIR__, "rules.jl"))

src/relaxation/relaxation.jl

@@ -1,5 +1,3 @@
-    
-using IfElse
 
 # Note: McCormick transform is only the convex and concave portions of the transformation,
 # so that the interval transforms (which are the same as for a regular Interval Transform)
@@ -24,6 +22,11 @@ end
 #=
 Unitary Rules
 =#
+function transform_rule(::McCormickTransform, ::typeof(getindex), yL, yU, ycv, ycc, xL, xU, xcv, xcc)
+    rcv = Equation(ycv, xcv)
+    rcc = Equation(ycc, xcc)
+    return rcv, rcc
+end
 function transform_rule(::McCormickTransform, ::typeof(exp), yL, yU, ycv, ycc, xL, xU, xcv, xcc)
     mcv = mid_expr(xcv, xcc, xL)
     mcc = mid_expr(xcv, xcc, xU)
@@ -154,6 +157,17 @@ function transform_rule(::McCormickTransform, ::typeof(max), zL, zU, zcv, zcc, x
     return rcv, rcc
 end
 
+function transform_rule(::McCormickTransform, ::typeof(^), zL, zU, zcv, zcc, xL, xU, xcv, xcc, yL, yU, ycv, ycc)
+    ~((typeof(yL) <: Int) || (typeof(yL) <: AbstractFloat)) && error("Symbolic exponents not currently supported.")
+    ~(yL == 2) && error("Exponents besides 2 not currently supported")
+    mcv = mid_expr(xcv, xcc, max(min(xU, 0.0), xL))
+    mcc = mid_expr(xcv, xcc, IfElse.ifelse(xU < 0.0, xL, IfElse.ifelse(xL > 0, xU, IfElse.ifelse(abs(xL) >= abs(xU), xL, xU))))
+    rcv = Equation(zcv, mcv^2)
+    rcc = Equation(zcc, (xL+xU)*mcc - xU*xL)
+    return rcv, rcc
+end
+
+
 #=
 TODO: Add other operators. It's probably helpful to break the McCormick overload and McCormick + Interval Outputs
 into separate transform_rules since the coupling for the ODEs are one directional and potentially useful.

src/relaxation/rules.jl

@@ -1,6 +1,7 @@
 
 base_term(a::Any) = false
 base_term(a::Term{Real, Base.ImmutableDict{DataType,Any}}) = true
+base_term(a::Term{Real, Nothing}) = true
 base_term(a::Sym) = true
 base_term(a::Real) = true
 

src/transform/factor.jl

@@ -161,7 +161,6 @@ function apply_transform(transform::T, prob::ODESystem) where T<:AbstractTransfo
         else
             index = findall(x -> isequal(x.rhs, eqn.rhs), equations)
             push!(equations, Equation(eqn.lhs, equations[index[1]].lhs))
-
         end
     end
 
@@ -215,3 +214,71 @@ function apply_transform(transform::T, prob::ODESystem) where T<:AbstractTransfo
 
     return new_sys
 end
+
+# Separate case for applying transform to only a set of equations
+function apply_transform(transform::T, eqn_vector::Vector{Equation}) where T<:AbstractTransform
+
+    # Factorize all model equations to generate a new set of equations
+    equations = Equation[]
+    for eqn in eqn_vector
+        current = length(equations)
+        factor!(eqn.rhs, eqs=equations)
+        if length(equations) > current
+            push!(equations, Equation(eqn.lhs, equations[end].rhs))
+            deleteat!(equations, length(equations)-1)
+        else
+            index = findall(x -> isequal(x.rhs, eqn.rhs), equations)
+            push!(equations, Equation(eqn.lhs, equations[index[1]].lhs))
+        end
+    end
+
+    # Apply transform rules to the factored equations to make the final equation set
+    new_equations = Equation[]
+    for a in equations
+        zn = var_names(transform, zstr(a))
+        xn = var_names(transform, xstr(a))
+        if isone(arity(a)) 
+            targs = (transform, op(a), zn..., xn...)
+         else
+            targs = (transform, op(a), zn..., xn..., var_names(transform, ystr(a))...)
+        end
+        new = transform_rule(targs...)
+        for i in new
+            push!(new_equations, i)
+        end
+    end
+
+    return new_equations
+end
+function apply_transform(transform::T, eqn::Equation) where T<:AbstractTransform
+
+    # Factorize the equations to generate a new set of equations
+    equations = Equation[]
+    current = 0
+    factor!(eqn.rhs, eqs=equations)
+    if length(equations) > current
+        push!(equations, Equation(eqn.lhs, equations[end].rhs))
+        deleteat!(equations, length(equations)-1)
+    else
+        index = findall(x -> isequal(x.rhs, eqn.rhs), equations)
+        push!(equations, Equation(eqn.lhs, equations[index[1]].lhs))
+    end
+
+    # Apply transform rules to the factored equations to make the final equation set
+    new_equations = Equation[]
+    for a in equations
+        zn = var_names(transform, zstr(a))
+        xn = var_names(transform, xstr(a))
+        if isone(arity(a)) 
+            targs = (transform, op(a), zn..., xn...)
+         else
+            targs = (transform, op(a), zn..., xn..., var_names(transform, ystr(a))...)
+        end
+        new = transform_rule(targs...)
+        for i in new
+            push!(new_equations, i)
+        end
+    end
+
+    return new_equations
+end