some more genericness

Author: shashi

src/build_function.jl

@@ -476,7 +476,7 @@ end
 
 get_varnumber(varop, vars::Vector) =  findfirst(x->isequal(x,varop),vars)
 
-function numbered_expr(O::Union{Term,Sym},args...;varordering = args[1],offset = 0,
+function numbered_expr(O::Symbolic,args...;varordering = args[1],offset = 0,
                        lhsname=gensym("du"),rhsnames=[gensym("MTK") for i in 1:length(args)])
     O = value(O)
   if O isa Sym || isa(operation(O), Sym)

src/systems/abstractsystem.jl

@@ -117,8 +117,13 @@ Generate a function to evaluate the system's equations.
 """
 function generate_function end
 
-getname(x::Sym) = nameof(x)
-getname(t::Term) = operation(t) isa Sym ? getname(operation(t)) : error("Cannot get name of $t")
+function getname(t)
+    if istree(t)
+        operation(t) isa Sym ? getname(operation(t)) : error("Cannot get name of $t")
+    else
+        nameof(t)
+    end
+end
 
 function Base.getproperty(sys::AbstractSystem, name::Symbol)
 
@@ -188,14 +193,17 @@ function namespace_expr(O::Sym,name,ivname)
     O.name == ivname ? O : rename(O,renamespace(name,O.name))
 end
 
-function namespace_expr(O::Term{T},name,ivname) where {T}
-    if operation(O) isa Sym
-        Term{T}(rename(operation(O),renamespace(name,operation(O).name)),namespace_expr.(arguments(O),name,ivname))
+function namespace_expr(O,name,ivname) where {T}
+    if istree(O)
+        if operation(O) isa Sym
+            Term{T}(rename(operation(O),renamespace(name,operation(O).name)),namespace_expr.(arguments(O),name,ivname))
+        else
+            similarterm(O,operation(O),namespace_expr.(arguments(O),name,ivname))
+        end
     else
-        Term{T}(operation(O),namespace_expr.(arguments(O),name,ivname))
+        O
     end
 end
-namespace_expr(O,name,ivname) = O
 
 independent_variable(sys::AbstractSystem) = sys.iv
 function states(sys::AbstractSystem)
@@ -277,12 +285,12 @@ struct AbstractSysToExpr
     states::Vector
 end
 AbstractSysToExpr(sys) = AbstractSysToExpr(sys,states(sys))
-function (f::AbstractSysToExpr)(O::Term)
+function (f::AbstractSysToExpr)(O)
+    !istree(O) && return toexpr(O)
     any(isequal(O), f.states) && return operation(O).name  # variables
     if isa(operation(O), Sym)
         return build_expr(:call, Any[operation(O).name; f.(arguments(O))])
     end
     return build_expr(:call, Any[operation(O); f.(arguments(O))])
 end
-(f::AbstractSysToExpr)(x) = toexpr(x)
 

src/systems/control/controlsystem.jl

@@ -47,7 +47,7 @@ sys = ControlSystem(loss,eqs,t,[x,v],[u],[])
 """
 struct ControlSystem <: AbstractControlSystem
     """The Loss function"""
-    loss::Term
+    loss::Any
     """The ODEs defining the system."""
     eqs::Vector{Equation}
     """Independent variable."""
@@ -89,7 +89,8 @@ struct ControlToExpr
     controls::Vector
 end
 ControlToExpr(@nospecialize(sys)) = ControlToExpr(sys,states(sys),controls(sys))
-function (f::ControlToExpr)(O::Term)
+function (f::ControlToExpr)(O)
+    !istree(O) && return O
     res = if isa(operation(O), Sym)
         # normal variables and control variables
         (any(isequal(O), f.states) || any(isequal(O), f.controls)) && return tosymbol(O)
@@ -99,7 +100,6 @@ function (f::ControlToExpr)(O::Term)
     end
 end
 (f::ControlToExpr)(x::Sym) = x.name
-(f::ControlToExpr)(x) = x
 
 function constructRadauIIA5(T::Type = Float64)
   sq6 = sqrt(6)

src/systems/optimization/optimizationsystem.jl

@@ -64,7 +64,7 @@ end
 function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys);
                           sparse = false, kwargs...)
     if sparse
-        hess = sparsehessian(equations(sys),[dv() for dv in states(sys)])
+        hess = sparsehessian(equations(sys),states(sys))
     else
         hess = calculate_hessian(sys)
     end
@@ -77,11 +77,11 @@ function generate_function(sys::OptimizationSystem, vs = states(sys), ps = param
                           conv = AbstractSysToExpr(sys),kwargs...)
 end
 
-equations(sys::OptimizationSystem) = isempty(sys.systems) ? operation(sys) : operation(sys) + reduce(+,namespace_expr.(sys.systems))
-namespace_expr(sys::OptimizationSystem) = namespace_expr(operation(sys),sys.name,nothing)
+equations(sys::OptimizationSystem) = isempty(sys.systems) ? sys.op : sys.op + reduce(+,namespace_expr.(sys.systems))
+namespace_expr(sys::OptimizationSystem) = namespace_expr(sys.op,sys.name,nothing)
 
 hessian_sparsity(sys::OptimizationSystem) =
-    hessian_sparsity(operation(sys), states(sys))
+    hessian_sparsity(sys.op, states(sys))
 
 struct AutoModelingToolkit <: DiffEqBase.AbstractADType end