unbreak some stuff

Co-authored-by: "Yingbo Ma" <[email protected]>

Author: shashi

src/build_function.jl

@@ -88,7 +88,7 @@ function add_integrator_header(ex, fargs, iip; X=gensym(:MTIIPVar))
 end
 
 # Scalar output
-function _build_function(target::JuliaTarget, op::Operation, args...;
+function _build_function(target::JuliaTarget, op, args...;
                          conv = toexpr, expression = Val{true},
                          checkbounds = false,
                          linenumbers = true, headerfun=addheader)
@@ -215,7 +215,7 @@ Special Keyword Argumnets:
 - `fillzeros`: Whether to perform `fill(out,0)` before the calculations to ensure
   safety with `skipzeros`.
 """
-function _build_function(target::JuliaTarget, rhss, args...;
+function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
                          conv = toexpr, expression = Val{true},
                          checkbounds = false,
                          linenumbers = false, multithread=nothing,
@@ -447,17 +447,24 @@ end
 
 vars_to_pairs(args) = vars_to_pairs(args[1],args[2])
 function vars_to_pairs(name,vs::AbstractArray)
-	_vs = convert.(Variable,vs)
-	names = [Symbol(u) for u ∈ _vs]
-	exs = [:($name[$i]) for (i, u) ∈ enumerate(_vs)]
-	names,exs
+    vs_names = [term_to_symbol(value(u)) for u ∈ vs]
+	exs = [:($name[$i]) for (i, u) ∈ enumerate(vs)]
+	vs_names,exs
+end
+
+function term_to_symbol(t::Term)
+    if operation(t) isa Sym
+        s = nameof(operation(t))
+        @show s
+    else
+        error("really?")
+    end
 end
 
+term_to_symbol(s::Sym) = nameof(s)
+
 function vars_to_pairs(name,vs)
-	_vs = convert(Variable,vs)
-	names = [Symbol(_vs)]
-	exs = [name]
-	names,exs
+    [term_to_symbol(value(vs))], [name]
 end
 
 get_varnumber(varop::Operation,vars::Vector{Operation}) =  findfirst(x->isequal(x,varop),vars)
@@ -480,7 +487,7 @@ end
 
 function numbered_expr(de::ModelingToolkit.Equation,args...;varordering = args[1],
                        lhsname=gensym("du"),rhsnames=[gensym("MTK") for i in 1:length(args)],offset=0)
-    i = findfirst(x->isequal(x isa Variable ? x.name : x.op.name,var_from_nested_derivative(de.lhs)[1].name),varordering)
+    i = findfirst(x->isequal(x isa Variable ? term_to_symbol(x) : term_to_symbol(x.op),term_to_symbol(var_from_nested_derivative(de.lhs)[1])),varordering)
     :($lhsname[$(i+offset)] = $(numbered_expr(de.rhs,args...;offset=offset,
 											  varordering = varordering,
 											  lhsname = lhsname,

src/systems/abstractsystem.jl

@@ -245,10 +245,7 @@ end
 function islinear(sys::AbstractSystem)
     rhs = [eq.rhs for eq ∈ equations(sys)]
 
-    iv = sys.iv
-    dvs = [dv(iv) for dv ∈ states(sys)]
-
-    all(islinear(r, dvs) for r in rhs)
+    all(islinear(r, states(sys)) for r in rhs)
 end
 
 function pins(sys::AbstractSystem,args...)

src/systems/diffeqs/abstractodesystem.jl

@@ -2,14 +2,18 @@ function calculate_tgrad(sys::AbstractODESystem;
                          simplify=true)
   isempty(sys.tgrad[]) || return sys.tgrad[]  # use cached tgrad, if possible
   rhs = [detime_dvs(eq.rhs) for eq ∈ equations(sys)]
-  iv = sys.iv()
+  iv = sys.iv
+  for r in rhs
+      @show r
+      @show expand_derivatives(Differential(iv)(r))
+  end
   notime_tgrad = [expand_derivatives(ModelingToolkit.Differential(iv)(r)) for r in rhs]
   tgrad = retime_dvs.(notime_tgrad,(states(sys),),iv)
   if simplify
       tgrad = ModelingToolkit.simplify.(tgrad)
   end
   sys.tgrad[] = tgrad
-  return @show(tgrad)
+  return tgrad
 end
 
 function calculate_jacobian(sys::AbstractODESystem;
@@ -18,7 +22,7 @@ function calculate_jacobian(sys::AbstractODESystem;
     rhs = [eq.rhs for eq ∈ equations(sys)]
 
     iv = sys.iv
-    dvs = [dv(iv) for dv ∈ states(sys)]
+    dvs = states(sys)
 
     if sparse
         jac = sparsejacobian(rhs, dvs, simplify=simplify)
@@ -90,7 +94,7 @@ function calculate_massmatrix(sys::AbstractODESystem; simplify=true)
         if eq.lhs isa Constant
             @assert eq.lhs.value == 0
         elseif eq.lhs.op isa Differential
-            j = findfirst(x->isequal(x.name,var_from_nested_derivative(eq.lhs)[1].name),dvs)
+            j = findfirst(x->isequal(term_to_symbol(x),term_to_symbol(var_from_nested_derivative(eq.lhs)[1])),dvs)
             M[i,j] = 1
         else
             error("Only semi-explicit constant mass matrices are currently supported. Faulty equation: $eq.")

src/systems/diffeqs/odesystem.jl

@@ -71,10 +71,10 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     iv′ = value(iv)
     dvs′ = value.(dvs)
     ps′ = value.(ps)
-    tgrad = RefValue(Vector{Expression}(undef, 0))
-    jac = RefValue{Any}(Matrix{Expression}(undef, 0, 0))
-    Wfact   = RefValue(Matrix{Expression}(undef, 0, 0))
-    Wfact_t = RefValue(Matrix{Expression}(undef, 0, 0))
+    tgrad = RefValue(Vector{Num}(undef, 0))
+    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
+    Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
+    Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
     ODESystem(deqs, iv′, dvs′, ps′, pins, observed, tgrad, jac, Wfact, Wfact_t, name, systems)
 end
 

src/utils.jl

@@ -31,14 +31,14 @@ function flatten_expr!(x)
     x
 end
 
-function detime_dvs(op::Operation)
-  if op.op isa Variable
-    Operation(Variable{vartype(op.op)}(op.op.name),Expression[])
+function detime_dvs(op::Term)
+  if op.op isa Sym
+      op.op
   else
-    Operation(op.op,detime_dvs.(op.args))
+    Term(op.op,detime_dvs.(op.args))
   end
 end
-detime_dvs(op::Constant) = op
+detime_dvs(op) = op
 
 function retime_dvs(op::Operation,dvs,iv)
   if op.op isa Variable && op.op ∈ dvs

test/direct.jl

@@ -98,7 +98,7 @@ Jiip(J2,[1.0,2.0,3.0],[1.0,2.0,3.0],1.0)
 
 # Function building
 
-@parameters σ() ρ() β()
+@parameters σ ρ β
 @variables x y z
 eqs = [σ*(y-x),
        x*(ρ-z)-y,
@@ -112,7 +112,7 @@ f(out,[1.0,2,3],[1.0,2,3])
 @test all(o1 .== out)
 
 function test_worldage()
-   @parameters σ() ρ() β()
+   @parameters σ ρ β
    @variables x y z
    eqs = [σ*(y-x),
           x*(ρ-z)-y,

test/odesystem.jl

@@ -24,8 +24,6 @@ generate_function(de)
 function test_diffeq_inference(name, sys, iv, dvs, ps)
     @testset "ODESystem construction: $name" begin
         @test independent_variable(sys) ==  value(iv)
-        @show sys.states
-        @show @which states(sys)
         @test Set(states(sys))          == Set(value.(dvs))
         @test Set(parameters(sys))      == Set(value.(ps))
     end
@@ -52,10 +50,11 @@ u  = SVector(1:3...)
 p  = SVector(4:6...)
 @test f(u, p, 0.1) === @SArray [4, 0, -16]
 
+@show y
 eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y*t,
        D(z) ~ x*y - β*z]
-de = ODESystem(eqs)
+de = ODESystem(eqs) # This is broken
 ModelingToolkit.calculate_tgrad(de)
 
 tgrad_oop, tgrad_iip = eval.(ModelingToolkit.generate_tgrad(de))