Refactor get_degrees() to make it faster

src/ordering.jl

@@ -17,49 +17,127 @@
 """
 $(SIGNATURES)
 
-Internal function used for printing symbolic expressions. This function determines
-the degrees of symbols within a given expression, implementing a variation on 
-degree lexicographic order.
+Get the degrees of symbols within a given expression.
+
+This internal function is used to define the order of terms in a symbolic expression,
+which is a variation on degree lexicographic order. It is used for printing and
+by [`sorted_arguments`](@ref).
+
+Returns a tuple of degree and lexicographically sorted *multiplier* ⇒ *power* pairs,
+where the *multiplier* is a tuple of the symbol optionally followed by its indices.
+For a sum expression, returns the `get_degrees()` result for term with the highest degree.
+
+See also `monomial_lt` and `lexlt`.
 """
 function get_degrees(expr)
+    degs_cache = Dict()
+    res = _get_degrees(expr, degs_cache)
+    if res isa AbstractVector
+        return Tuple(res)
+    else
+        return res
+    end
+end
+
+function _get_degrees(expr, degs_cache::AbstractDict)
     if issym(expr)
-        ((Symbol(expr),) => 1,)
+        return get!(() -> ((Symbol(expr),) => 1,), degs_cache, expr)
     elseif iscall(expr)
-        op = operation(expr)
-        args = sorted_arguments(expr)
-        if op == (^) && args[2] isa Number
-            return map(get_degrees(args[1])) do (base, pow)
-                (base => pow * args[2])
-            end
-        elseif op == (*)
-            return mapreduce(get_degrees,
-                             (x,y)->(x...,y...,), args)
-        elseif op == (+)
-            ds = map(get_degrees, args)
-            _, idx = findmax(x->sum(last.(x), init=0), ds)
-            return ds[idx]
-        elseif op == (getindex)
-            return ((Symbol.(args)...,) => 1,)
-        else
-            return ((Symbol("zzzzzzz", hash(expr)),) => 1,)
+        # operation-specific degree handling
+        return _get_degrees(operation(expr), expr, degs_cache)
+    else
+        return () # skip numbers and unsupported expressions
+    end
+end
+
+# fallback for unsupported operation
+_get_degrees(::Any, expr, degs_cache) =
+    ((Symbol("zzzzzzz", hash(expr)),) => 1,)
+
+_getindex_symbol(arr, i) = Symbol(arr[i])
+
+function _get_degrees(::typeof(getindex), expr, degs_cache)
+    args = arguments(expr)
+    @inbounds return get!(() -> (ntuple(Base.Fix1(_getindex_symbol, args), length(args)) => 1,),
+                degs_cache, expr)
+end
+
+function _get_degrees(::typeof(*), expr, degs_cache)
+    args = arguments(expr)
+    ds = sizehint!(Vector{Any}(), length(args))
+    for arg in args
+        degs = _get_degrees(arg, degs_cache)
+        append!(ds, degs)
+    end
+    return sort!(ds)
+end
+
+function _get_degrees(::typeof(+), expr, degs_cache)
+    # among the terms find the best in terms of monomial_lt
+    sel_degs = ()
+    sel_degsum = 0
+    for arg in arguments(expr)
+        degs = _get_degrees(arg, degs_cache)
+        degsum = sum(last, degs, init=0)
+        if (sel_degs == ()) || (degsum > sel_degsum) ||
+           (degsum == sel_degsum && lexlt(degs, sel_degs))
+            sel_degs, sel_degsum = degs, degsum
+        end
+    end
+    return sel_degs
+end
+
+function _get_degrees(::typeof(^), expr, degs_cache)
+    base_expr, pow_expr = arguments(expr)
+    if pow_expr isa Number
+        @inbounds degs = map(_get_degrees(base_expr, degs_cache)) do (base, pow)
+            (base => pow * pow_expr)
+        end
+        if pow_expr < 0 && length(degs) > 1
+            # fix the order after the powers were negated
+            isa(degs, AbstractVector) || (degs = collect(degs))
+            sort!(degs)
+        end
+        return degs
+    else
+        # expression in the power argument is not supported
+        return _get_degrees(nothing, expr, degs_cache)
+    end
+end
+
+function _get_degrees(::typeof(/), expr, degs_cache)
+    nom_expr, denom_expr = arguments(expr)
+    if denom_expr isa Number # constant denominator
+        return _get_degrees(nom_expr, degs_cache)
+    elseif nom_expr isa Number # constant nominator
+        @inbounds degs = map(_get_degrees(denom_expr, degs_cache)) do (base, pow)
+            (base => -pow)
         end
+        isa(degs, AbstractVector) || (degs = collect(degs))
+        return sort!(degs)
     else
-        return ()
+        # TODO expressions in both nom and denom are not yet supported
+        return _get_degrees(nothing, expr, degs_cache)
     end
 end
 
 function monomial_lt(degs1, degs2)
     d1 = sum(last, degs1, init=0)
     d2 = sum(last, degs2, init=0)
-    d1 != d2 ? d1 < d2 : lexlt(degs1, degs2)
+    d1 != d2 ?
+        # lower absolute degree first, or if equal, positive degree first
+        (abs(d1) < abs(d2) || abs(d1) == abs(d2) && d1 > d2) :
+        lexlt(degs1, degs2)
 end
 
 function lexlt(degs1, degs2)
-    for (a, b) in zip(degs1, degs2)
-        if a[1] == b[1] && a[2] != b[2]
-            return a[2] > b[2]
-        elseif a[1] != b[1]
-            return a < b
+    for ((a_base, a_deg), (b_base, b_deg)) in zip(degs1, degs2)
+        if a_base == b_base && a_deg != b_deg
+            # same base, higher absolute degree first, positive degree first
+            return abs(a_deg) > abs(b_deg) || abs(a_deg) == abs(b_deg) && a_deg > b_deg
+        elseif a_base != b_base
+            # lexicographic order for the base
+            return a_base < b_base
         end
     end
     return false # they are equal

test/basics.jl

@@ -221,6 +221,12 @@ end
     b1, b3, d1, d2 = get(b,1),get(b,3), get(d,1), get(d,2)
     @test repr(a + b3 + b1 + d2 + c) == "a + b[1] + b[3] + c + d[2]"
     @test repr(expand((c + b3 - d1)^3)) == "b[3]^3 + 3(b[3]^2)*c - 3(b[3]^2)*d[1] + 3b[3]*(c^2) - 6b[3]*c*d[1] + 3b[3]*(d[1]^2) + c^3 - 3(c^2)*d[1] + 3c*(d[1]^2) - (d[1]^3)"
+    # test negative powers sorting
+    @test repr((b3^2)^(-2) + a^(-3) + (c*d1)^(-2)) == "1 / (a^3) + 1 / (b[3]^4) + 1 / ((c^2)*(d[1]^2))"
+
+    # test that the "x^2 + y^-1 + sin(a)^3.5 + 2t + 1//1" expression from Symbolics.jl/build_targets.jl is properly sorted
+    @syms x1 y1 a1 t1
+    @test repr(x1^2 + y1^-1 + sin(a1)^3.5 + 2t1 + 1//1) == "(1//1) + 2t1 + 1 / y1 + x1^2 + sin(a1)^3.5"
 end
 
 @testset "inspect" begin