Continued updates to apply_transform

Author: RXGottlieb

src/interval/interval.jl

@@ -5,17 +5,76 @@ Rules for constructing interval bounding expressions
 # Structure used to indicate an overload with intervals is preferable
 struct IntervalTransform <: AbstractTransform end
 
-# Creates names for interval state variables
-function var_names(::IntervalTransform, s::String)
-    sL = Symbol(s*"_lo")
-    sU = Symbol(s*"_hi")
-    sL, sU
+# Creates names for interval state variables [DEPRECATED]
+# function var_names(::IntervalTransform, s::Num)
+#     if s.val.metadata.value[1]==:variables
+#         arg_list = Symbol[]
+#         for i in s.val.arguments
+#             push!(arg_list, get_name(i))
+#         end
+#         sL = genvar(Symbol(string(s.val.f)*"_lo"), arg_list)
+#         sU = genvar(Symbol(string(s.val.f)*"_hi"), arg_list)
+#     elseif s.val.metadata.parent.value[1]==:parameters
+#         sL = genparam(Symbol(string(s.val.name)*"_lo"))
+#         sU = genparam(Symbol(string(s.val.name)*"_hi"))
+#     else
+#         error("Not a variable or a parameter, check type of s")
+#     end
+#     return sL, sU
+# end
+
+function var_names(::IntervalTransform, s::Term{Real, Base.ImmutableDict{DataType, Any}}) #The variables
+    arg_list = Symbol[]
+    for i in s.arguments
+        push!(arg_list, get_name(i))
+    end
+    sL = genvar(Symbol(string(s.f)*"_lo"), arg_list)
+    sU = genvar(Symbol(string(s.f)*"_hi"), arg_list)
+    return sL, sU
+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
+        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"))
+    else
+        error("Type of argument invalid")
+    end
+
+    sL = s.f(var_lo)
+    sU = s.f(var_hi)
+    return sL, sU
+end
+function var_names(::IntervalTransform, s::Sym) #The parameters
+    sL = genparam(Symbol(string(s.name)*"_lo"))
+    sU = genparam(Symbol(string(s.name)*"_hi"))
+    return sL, sU
 end
 
-function var_names(::IntervalTransform, s::Number)
-    sL = s
-    sU = s
-    sL, sU
+
+
+# function var_names(::IntervalTransform, s::Number)
+#     sL = s
+#     sU = s
+#     sL, sU
+# end
+
+# Helper functions for navigating SymbolicUtils structures
+function get_name(s::Term)
+    return s.f.name
+end
+function get_name(s::Sym)
+    return s.name
 end
 
 include(joinpath(@__DIR__, "rules.jl"))

src/transform/factor.jl

@@ -26,7 +26,7 @@ function isfactor(ex::SymbolicUtils.Add)
     (iszero(ex.coeff)) && (length(ex.dict)>2) && return false
     for (key, val) in ex.dict
         ~(isone(val)) && return false
-        ~(typeof(key)<:Term) && return false
+        ~(typeof(key)<:Term) && ~(typeof(key)<:Sym) && return false
     end
     return true
 end
@@ -35,18 +35,18 @@ function isfactor(ex::SymbolicUtils.Mul)
     (isone(ex.coeff)) && (length(ex.dict)>2) && return false
     for (key, val) in ex.dict
         ~(isone(val)) && return false
-        ~(typeof(key)<:Term) && return false
+        ~(typeof(key)<:Term) && ~(typeof(key)<:Sym) && return false
     end
     return true
 end
 function isfactor(ex::SymbolicUtils.Div)
-    ~(typeof(ex.num)<:Term) && ~(typeof(ex.num)<:Real) && return false
-    ~(typeof(ex.den)<:Term) && ~(typeof(ex.num)<:Real) && return false
+    ~(typeof(ex.num)<:Term) && ~(typeof(key)<:Sym) && ~(typeof(ex.num)<:Real) && return false
+    ~(typeof(ex.den)<:Term) && ~(typeof(key)<:Sym) && ~(typeof(ex.num)<:Real) && return false
     return true
 end
 function isfactor(ex::SymbolicUtils.Pow)
-    ~(typeof(ex.base)<:Term) && ~(typeof(ex.base)<:Real) && return false
-    ~(typeof(ex.exp)<:Term) && ~(typeof(ex.exp)<:Real) && return false
+    ~(typeof(ex.base)<:Term) && ~(typeof(key)<:Sym) && ~(typeof(ex.base)<:Real) && return false
+    ~(typeof(ex.exp)<:Term) && ~(typeof(key)<:Sym) && ~(typeof(ex.exp)<:Real) && return false
     return true
 end
 
@@ -186,9 +186,9 @@ function factor!(ex::SymbolicUtils.Mul; assignments = Assignment[])
 
     new_terms = Dict{Any, Number}()
     for (key, val) in ex.dict
-        if (typeof(key)<:Term) && isone(val)
+        if ((typeof(key)<:Term) && isone(val)) || ((typeof(key)<:Sym) && isone(val))
             new_terms[key] = val
-        elseif (typeof(key)<:Term)
+        elseif (typeof(key)<:Term) || (typeof(key)<:Sym)
             index = findall(x -> isequal(x.rhs,key^val), assignments)
             if isempty(index)
                 newsym = gensym(:aux)

src/transform/transform.jl

@@ -145,50 +145,61 @@ end
 ```
         
 =#
-function apply_transform(t::T, prob::ODESystem) where T<:AbstractTransform
+function apply_transform(transform::T, prob::ODESystem) where T<:AbstractTransform
 
+    # Factor out the equations
     assignments = Assignment[]
     for eqn in prob.eqs
         # Flesh out the original RHS
         current = length(assignments)
-        factor!(toexpr(eqn.rhs), assignments=assignments)
+        factor!(eqn.rhs, assignments=assignments)
 
         # If new equations were added, stick on the original LHS to the last point 
         # in assignments by stealing its RHS from the last item and taking its place
         if length(assignments) > current
-            push!(assignments, Assignment(toexpr(eqn.lhs), assignments[end].rhs))
+            push!(assignments, Assignment(eqn.lhs, assignments[end].rhs))
             deleteat!(assignments, length(assignments)-1)
         end
     end
 
-    # new_assignments = AssignmentPair[]
+    # Develop equations for the transforms
     new_assignments = Assignment[]
     for a in assignments
-        zn = var_names(t, zstr(a)) #LHS
-        xn = var_names(t, xstr(a)) #RHS(1)
+        # Get zn, first. Which is the LHS.
+        zn = var_names(transform, zstr(a))
+        xn = var_names(transform, xstr(a))
 
-        first = zn[1]
-        second = zn[2]
-
-        # Define the zn's as variables
-        @variables $first $second
         if isone(arity(a)) 
-            targs = (t, op(a), zn..., xn...)
+            targs = (transform, op(a), zn..., xn...)
          else
-            targs = (t, op(a), zn..., xn..., var_names(t, ystr(a))...)
+            targs = (transform, op(a), zn..., xn..., var_names(transform, ystr(a))...)
         end
-        println("targs: $targs")
 
         # push!(new_assignments, transform_rule(targs...))
         new = transform_rule(targs...)
         push!(new_assignments, new.l)
         push!(new_assignments, new.u)
     end
 
+    # Combine all transforms into a new set of equations and create a new ODE system
+    new_eqs = Equation[]
+    for i in new_assignments
+        push!(new_eqs, Equation(i.lhs, i.rhs))
+    end
+
     @named new_sys = ODESystem(new_eqs)
+    println("Completed.")
+
+    println(new_sys)
+    println(typeof(new_sys))
+    println(new_sys.eqs)
+
     # Form ODE system from new assignments
     # CSE - MTK.structural_simplify()
 
     # Figure out a way to give the new ODE system the proper parameters, variables, etc.
-    return the_new_ODE_System
+
+    # println(new_assignments)
+
+    return new_sys
 end

src/transform/utilities.jl

@@ -8,16 +8,53 @@ arity(a::SymbolicUtils.Mul) = length(a.dict) + (~isone(a.coeff))
 
 op(a::Expr) = a.args[1]
 op(a::Assignment) = op(a.rhs)
-op(a::Number) = "const"
-op(a::Symbol) = "const"
+op(::Number) = "const"
+op(::Symbol) = "const"
+op(::SymbolicUtils.Add) = +
+op(::SymbolicUtils.Mul) = *
+op(::SymbolicUtils.Pow) = ^
+op(::SymbolicUtils.Div) = /
+
+xstr(a::Assignment) = sub_1(a.rhs)
+ystr(a::Assignment) = sub_2(a.rhs)
+zstr(a::Assignment) = a.lhs
+
+function sub_1(a::SymbolicUtils.Add)
+    sorted_dict = sort(collect(a.dict), by=x->string(x[1]))
+    return sorted_dict[1].first
+end
+function sub_2(a::SymbolicUtils.Add)
+    ~(iszero(a.coeff)) && return a.coeff
+    sorted_dict = sort(collect(a.dict), by=x->string(x[1]))
+    return sorted_dict[2].first
+end
+
+function sub_1(a::SymbolicUtils.Mul)
+    sorted_dict = sort(collect(a.dict), by=x->string(x[1]))
+    return sorted_dict[1].first
+end
+function sub_2(a::SymbolicUtils.Mul)
+    ~(isone(a.coeff)) && return a.coeff
+    sorted_dict = sort(collect(a.dict), by=x->string(x[1]))
+    return sorted_dict[2].first
+end
+
+# xstr(a::Assignment) = string_or_num(a.rhs.args[2])
+# ystr(a::Assignment) = string_or_num(a.rhs.args[3])
+# zstr(a::Assignment) = string_or_num(a.lhs)
+
+# string_or_num(a::Num) = a
+# string_or_num(a::Expr) = string(a)
+# string_or_num(a::Symbol) = string(a)
+# string_or_num(a::Number) = a
+
+# # A function to identify whether a term... is.... variable or parameter, hm
+# function identify(x::Num)
+
+# end
+
 
-xstr(a::Assignment) = string_or_num(a.rhs.args[2])
-ystr(a::Assignment) = string_or_num(a.rhs.args[3])
-zstr(a::Assignment) = string_or_num(a.lhs)
 
-string_or_num(a::Expr) = string(a)
-string_or_num(a::Symbol) = string(a)
-string_or_num(a::Number) = a
 
 # Uses Symbolics functions to generate a variable as a function of the dependent variables of choice (default: t)
 function genvar(a::Symbol, b=:t)