Merge pull request #348 from JuliaSymbolics/s/div

Some more canonicalization in Div

Author: shashi

src/polyform.jl

@@ -1,4 +1,4 @@
-export PolyForm, simplify_fractions
+export PolyForm, simplify_fractions, quick_cancel
 using Bijections
 using DynamicPolynomials: PolyVar
 
@@ -230,10 +230,10 @@ expand(expr) = PolyForm(expr, Fs=Union{typeof(+), typeof(*), typeof(^)}, recurse
 
 function polyform_factors(d::Div, pvar2sym, sym2term)
     make(xs) = map(xs) do x
-        if x isa Pow && arguments(x)[2] isa Integer && arguments(x)[2] > 0
+        if x isa Pow && x.base isa Integer && x.exp > 0
             # here we do want to recurse one level, that's why it's wrong to just
             # use Fs = Union{typeof(+), typeof(*)} here.
-            Pow(PolyForm(arguments(x)[1], pvar2sym, sym2term), arguments(x)[2])
+            Pow(PolyForm(x.base, pvar2sym, sym2term), x.exp)
         else
             PolyForm(x, pvar2sym, sym2term)
         end
@@ -244,14 +244,14 @@ end
 
 _mul(xs...) = all(isempty, xs) ? 1 : *(Iterators.flatten(xs)...)
 
-function simplify_fractions(d::Div)
+function simplify_div(d::Div)
     d.simplified && return d
     ns, ds = polyform_factors(d, get_pvar2sym(), get_sym2term())
     ns, ds = rm_gcds(ns, ds)
     if all(_isone, ds)
         return isempty(ns) ? 1 : simplify_fractions(_mul(ns))
     else
-        return Div(simplify_fractions(_mul(ns)), simplify_fractions(_mul(ds)), true)
+        Div(simplify_fractions(_mul(ns)), simplify_fractions(_mul(ds)), true)
     end
 end
 
@@ -269,12 +269,26 @@ Find `Div` nodes and simplify them by cancelling a set of factors of numerators
 and denominators. It may leave some expressions in `PolyForm` format.
 """
 function simplify_fractions(x)
+    x = Postwalk(quick_cancel)(x)
+
+    !needs_div_rules(x) && return x
+
     isdiv(x) = x isa Div
 
-    rules = [@acrule ~a::isdiv + ~b::isdiv => add_divs(~a,~b)
-             @rule ~x::isdiv => simplify_fractions(~x)]
+    rules = [@rule ~x::isdiv => simplify_div(~x)
+             @acrule ~a::isdiv + ~b::isdiv => add_divs(~a,~b)]
+
+    Fixpoint(Postwalk(Chain(rules)))(x)
+end
+
+function needs_div_rules(x)
+    (x isa Div && !(x.num isa Number) && !(x.den isa Number)) ||
+    (istree(x) && operation(x) === (+) && count(has_div, unsorted_arguments(x)) > 1) ||
+    (istree(x) && any(needs_div_rules, unsorted_arguments(x)))
+end
 
-    Prewalk(RestartedChain(rules))(x)
+function has_div(x)
+    return x isa Div || (istree(x) && any(has_div, unsorted_arguments(x)))
 end
 
 flatten_pows(xs) = map(xs) do x
@@ -289,6 +303,110 @@ const MaybeGcd = Union{PolyForm, MP.AbstractPolynomialLike, Integer}
 _gcd(x::MaybeGcd, y::MaybeGcd) = (coefftype(x) <: Complex || coefftype(y) <: Complex) ? 1 : gcd(x, y)
 _gcd(x, y) = 1
 
+
+"""
+    quick_cancel(d::Div)
+
+Cancel out matching factors from numerator and denominator.
+This is not as effective as `simplify_fractions`, for example,
+it wouldn't simplify `(x^2 + 15 -  8x)  / (x - 5)` to `(x - 3)`.
+But it will simplify `(x - 5)^2*(x - 3) / (x - 5)` to `(x - 5)*(x - 3)`.
+Has optimized processes for `Mul` and `Pow` terms.
+"""
+quick_cancel(d::Div) = Div{symtype(d)}(quick_cancel(d.num, d.den)...)
+
+quick_cancel(x) = x
+
+quick_cancel(x, y) = isequal(x, y) ? (1,1) : (x, y)
+
+function quick_cancel(x::Pow, y)
+    x.exp isa Number || return (x, y)
+    isequal(x.base, y) && x.exp >= 1 ? (Pow{symtype(x)}(x.base, x.exp - 1),1) : (x, y)
+end
+
+quick_cancel(y, x::Pow) = reverse(quick_cancel(x,y))
+
+function quick_cancel(x::Pow, y::Pow)
+    if isequal(x.base, y.base)
+        !(x.exp isa Number && y.exp isa Number) && return (x, y)
+        if x.exp > y.exp
+            return Pow{symtype(x)}(x.base, x.exp-y.exp), 1
+        elseif x.exp == y.exp
+            return 1, 1
+        else # x.exp < y.exp
+            return 1, Pow{symtype(y)}(y.base, y.exp-x.exp)
+        end
+    end
+    return x, y
+end
+
+function quick_cancel(x::Mul, y)
+    if haskey(x.dict, y) && x.dict[y] >= 1
+        d = copy(x.dict)
+        if d[y] > 1
+            d[y] -= 1
+        elseif d[y] == 1
+            delete!(d, y)
+        else
+            error("Can't reach")
+        end
+
+        return Mul(symtype(x), x.coeff, d), 1
+    else
+        return x, y
+    end
+end
+
+function quick_cancel(x::Mul, y::Pow)
+    y.exp isa Number || return (x, y)
+    if haskey(x.dict, y.base)
+        d = copy(x.dict)
+        if x.dict[y.base] > y.exp
+            d[y.base] -= y.exp
+            den = 1
+        elseif x.dict[y.base] == y.exp
+            delete!(d, y.base)
+            den = 1
+        else
+            den = Pow{symtype(y)}(y.base, y.exp-d[y.base])
+            delete!(d, y.base)
+        end
+        return Mul(symtype(x), x.coeff, d), den
+    else
+        return x, y
+    end
+end
+
+quick_cancel(x::Pow, y::Mul) = reverse(quick_cancel(y,x))
+
+quick_cancel(y, x::Mul) = reverse(quick_cancel(x,y))
+
+function quick_cancel(x::Mul, y::Mul)
+    num_dict, den_dict = _merge_div(x.dict, y.dict)
+    Mul(symtype(x), x.coeff, num_dict), Mul(symtype(y), y.coeff, den_dict)
+end
+
+function _merge_div(ndict, ddict)
+    num = copy(ndict)
+    den = copy(ddict)
+    for (k, v) in den
+        if haskey(num, k)
+            nk = num[k]
+            if nk > v
+                num[k] -= v
+                delete!(den, k)
+            elseif nk == v
+                delete!(num, k)
+                delete!(den, k)
+            else
+                den[k] -= nk
+                delete!(num, k)
+            end
+        end
+    end
+    num, den
+end
+
 function rm_gcds(ns, ds)
     ns = flatten_pows(ns)
     ds = flatten_pows(ds)

src/types.jl

@@ -841,9 +841,37 @@ end
 Base.hash(x::Div, u::UInt64) = hash(x.num, hash(x.den, u))
 Base.isequal(x::Div, y::Div) = isequal(x.num, y.num) && isequal(x.den, y.den)
 
+const Rat = Union{Rational, Integer}
+
+ratcoeff(x) = false, NaN
+ratcoeff(x::Rat) = true, x
+ratcoeff(x::Mul) = ratcoeff(x.coeff)
+ratio(x::Integer,y::Integer) = iszero(rem(x,y)) ? div(x,y) : x//y
+ratio(x::Rat,y::Rat) = x//y
+function maybe_intcoeff(x::Mul)
+    x.coeff isa Rational && isone(x.coeff.den) ? Setfield.@set!(x.coeff = x.coeff.num) : x
+end
+maybe_intcoeff(x::Rational) = isone(x.den) ? x.num : x
+maybe_intcoeff(x) = x
+
 function (::Type{Div{T}})(n, d, simplified=false; metadata=nothing) where {T}
-    @assert !(n isa AbstractArray)
-    @assert !(d isa AbstractArray)
+    _iszero(n) && return zero(typeof(n))
+    _isone(d) && return n
+    d isa Number && _isone(-d) && return -1 * n
+    n isa Rat && d isa Rat && return n // d # maybe called by oblivious code in simplify
+
+    # GCD coefficient upon construction
+    rat, nc = ratcoeff(n)
+    if rat
+        rat, dc = ratcoeff(d)
+        if rat
+            g = gcd(nc, dc) * sign(dc) # make denominator positive
+            invdc = ratio(1, g)
+            n = maybe_intcoeff(invdc * n)
+            d = maybe_intcoeff(invdc * d)
+        end
+    end
+
     Div{T, typeof(n), typeof(d), typeof(metadata)}(n, d, simplified, metadata)
 end
 

test/basics.jl

@@ -162,7 +162,7 @@ end
     @test repr(-(a + b)) == "-a - b"
     @test repr((2a)^(-2a)) == "(2a)^(-2a)"
     @test repr(1/2a) == "1 / (2a)"
-    @test repr(2/(2*a)) == "2 / (2a)"
+    @test repr(2/(2*a)) == "1 / a"
     @test repr(Term(*, [1, 1])) == "1"
     @test repr(Term(*, [2, 1])) == "2*1"
     @test repr((a + b) - (b + c)) == "a - c"
@@ -226,3 +226,15 @@ end
     @syms a b c::Int
     @test isequal(arguments(s(a, b, c)), [a, b, c])
 end
+
+@testset "div" begin
+    @syms x y
+    @test (2x/2y).num isa Sym
+    @test (2x/3y).num.coeff == 2
+    @test (2x/3y).den.coeff == 3
+    @test (2x/-3x).num.coeff == -2
+    @test (2x/-3x).den.coeff == 3
+    @test (2.5x/3x).num.coeff == 2.5
+    @test (2.5x/3x).den.coeff == 3
+    @test (x/3x).den.coeff == 3
+end

test/polyform.jl

@@ -1,13 +1,15 @@
 using SymbolicUtils: PolyForm, Term, symtype
 using Test, SymbolicUtils
 
+include("utils.jl")
+
 @testset "div and polyform" begin
     @syms x y z
     @test repr(PolyForm(x-y)) == "x - y"
     @test repr(x/y*x/z) == "(x^2) / (y*z)"
     @test repr(simplify_fractions(((x-y+z)*(x+4z+1)) /
                                   (y*(2x - 3y + 3z) +
-                                   x*(x + z)))) == "(1 + x + 4z) / (x + 3.0y)"
+                                   x*(x + z)))) == repr(simplify_fractions((1 + x + 4z) / (x + 3.0y)))
     @test simplify_fractions(x/(x+3) + 3/(x+3)) == 1
     @test repr(simplify(simplify_fractions(cos(x)/sin(x) + sin(x)/cos(x)))) == "1 / (cos(x)*sin(x))"
 end
@@ -31,3 +33,23 @@ end
    #@test expand(identity(a * b) - b * a) == 0
     @test expand(a * b - b * a) == 0
 end
+
+@testset "simplify_fractions with quick-cancel" begin
+    @syms x y
+    @test simplify_fractions(x/2x) == 1//2
+    @test simplify_fractions(2x//x) == 2
+    @eqtest simplify_fractions((x+y) * (x-y) / (x+y)) == (x-y)
+    @eqtest simplify_fractions(x^3 * y / x) == y*x^2
+    @eqtest simplify_fractions(2x^3 * y / x) == 2y*x^2
+    @eqtest simplify_fractions(x / (3(x^3)*y)) == simplify_fractions(1/(3*(y*x^2)))
+    @eqtest simplify_fractions(2x / (3(x^3)*y)) == simplify_fractions(2/(3*(y*x^2)))
+    @eqtest simplify_fractions(x^2 / (3(x^3)*y)) == simplify_fractions(1/(3*(y*x)))
+    @eqtest simplify_fractions((3(x^3)*y) / x^2) == simplify_fractions(3*(y*x))
+    @eqtest simplify_fractions(x^2/x^4) == (1/x^2)
+    @eqtest simplify_fractions(x^2/x^3) == 1/x
+    @eqtest simplify_fractions(x^3/x^2) == x
+    @eqtest simplify_fractions(x^2/x^2) == 1
+    @eqtest simplify_fractions(3x^2/x^3) == 3/x
+    @eqtest simplify_fractions(3*(x^2)*(y^3)/(3*(x^3)*(y^2))) == y/x
+    @eqtest simplify_fractions(3*(x^x)/x*y) == 3*(x^x)/x*y
+end

test/rulesets.jl

@@ -42,7 +42,7 @@ end
     @eqtest simplify(a + b + 0*c + d) == simplify(a + b + d)
     @eqtest simplify(a * b * c^0 * d) == simplify(a * b * d)
     @eqtest simplify(a * b * 1*c * d) == simplify(a * b * c * d)
-    @eqtest repr(simplify_fractions(x^2.0/(x*y)^2.0)) == "1 / (y^2.0)"
+    @eqtest simplify_fractions(x^2.0/(x*y)^2.0) == simplify_fractions(1 / (y^2.0))
 
     @test simplify(Term(one, [a])) == 1
     @test simplify(Term(one, [b+1])) == 1

test/runtests.jl

@@ -17,18 +17,10 @@ if v"1.6" ≤ VERSION < v"1.7-beta3.0"
 else
     @warn "Skipping doctests"
 end
-
-# == / != syntax is nice, let's use it in tests
-macro eqtest(expr)
-    @assert expr.head == :call && expr.args[1] in [:(==), :(!=)]
-    if expr.args[1] == :(==)
-        :(@test isequal($(expr.args[2]), $(expr.args[3])))
-    else
-        :(@test !isequal($(expr.args[2]), $(expr.args[3])))
-    end |> esc
-end
 SymbolicUtils.show_simplified[] = false
 
+include("utils.jl")
+
 if haskey(ENV, "SU_BENCHMARK_ONLY")
     include("benchmark.jl")
 else

test/utils.jl

@@ -0,0 +1,10 @@
+
+# == / != syntax is nice, let's use it in tests
+macro eqtest(expr)
+    @assert expr.head == :call && expr.args[1] in [:(==), :(!=)]
+    if expr.args[1] == :(==)
+        :(@test isequal($(expr.args[2]), $(expr.args[3])))
+    else
+        :(@test !isequal($(expr.args[2]), $(expr.args[3])))
+    end |> esc
+end