fix pow on termcombination and use similarterm to avoid defining <_e

Author: shashi

src/ModelingToolkit.jl

@@ -25,6 +25,8 @@ using RecursiveArrayTools
 import SymbolicUtils
 import SymbolicUtils: Term, Add, Mul, Pow, Sym, to_symbolic, FnType, @rule, Rewriters, substitute, similarterm
 
+import SymbolicUtils.Rewriters: Chain, Postwalk, Prewalk, Fixpoint
+
 using LinearAlgebra: LU, BlasInt
 
 import LightGraphs: SimpleDiGraph, add_edge!
@@ -88,7 +90,7 @@ end
 <ₑ(s::Num, x::Num) = value(s) <ₑ value(x)
 
 for T in (Integer, Rational)
-    @eval Base.:(^)(n::Num, i::$T) = Num(Term{symtype(n)}(^, [value(n),i]))
+    @eval Base.:(^)(n::Num, i::$T) = Num(value(n)^i)
 end
 
 macro num_method(f, expr, Ts=nothing)

src/direct.jl

@@ -128,7 +128,8 @@ let
 
     _scalar = one(TermCombination)
 
-    linearity_propagator = [
+    simterm(t, f, args) = Term{Any}(f, args)
+    linearity_rules = [
           @rule +(~~xs) => reduce(+, filter(isidx, ~~xs), init=_scalar)
           @rule *(~~xs) => reduce(*, filter(isidx, ~~xs), init=_scalar)
           @rule (~f)(~x::(!isidx)) => _scalar
@@ -146,7 +147,8 @@ let
               else
                   error("Function of unknown linearity used: ", ~f)
               end
-          end] |> Rewriters.Chain |> Rewriters.Postwalk |> Rewriters.Fixpoint
+          end]
+    linearity_propagator = Fixpoint(Postwalk(Chain(linearity_rules); similarterm=simterm))
 
     global hessian_sparsity
 
@@ -164,7 +166,7 @@ let
         u = map(value, u)
         idx(i) = TermCombination(Set([Dict(i=>1)]))
         dict = Dict(u .=> idx.(1:length(u)))
-        f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x)(f)
+        f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x; similarterm=simterm)(f)
         lp = linearity_propagator(f)
         _sparse(lp, length(u))
     end

src/linearity.jl

@@ -78,11 +78,14 @@ function Base.:(==)(comb1::TermCombination, comb2::TermCombination)
 end
 =#
 
-# we don't care about the ordering in this case
-SymbolicUtils.:<ₑ(comb1::TermCombination, comb2::TermCombination) = true
 # to make Mul and Add work
 Base.:*(::Number, comb::TermCombination) = comb
-Base.:^(comb::TermCombination, ::Number) = comb
+function Base.:^(comb::TermCombination, ::Number)
+    isone(comb) && return comb
+    iszero(comb) && return _scalar
+    return comb *  comb
+end
+
 function Base.:+(comb1::TermCombination, comb2::TermCombination)
     if isone(comb1) && !iszero(comb2)
         return comb2

test/direct.jl

@@ -60,7 +60,7 @@ reference_hes = ModelingToolkit.hessian(rr, X)
 
 sp_hess = ModelingToolkit.sparsehessian(rr, X)
 @test findnz(sparse(reference_hes))[1:2] == findnz(sp_hess)[1:2]
-# @test isequal(map(spoly, findnz(sparse(reference_hes))[3]), map(spoly, findnz(sp_hess)[3]))
+@test isequal(map(spoly, findnz(sparse(reference_hes))[3]), map(spoly, findnz(sp_hess)[3]))
 
 Joop, Jiip = eval.(ModelingToolkit.build_function(∂,[x,y,z],[σ,ρ,β],t))
 J = Joop([1.0,2.0,3.0],[1.0,2.0,3.0],1.0)