Merge pull request #103 from JuliaSymbolics/s/mpoly

Convert to and from MPoly

Author: shashi

Project.toml

@@ -4,7 +4,9 @@ authors = ["Shashi Gowda"]
 version = "0.3.3"
 
 [deps]
+AbstractAlgebra = "c3fe647b-3220-5bb0-a1ea-a7954cac585d"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"

benchmark/benchmarks.jl

@@ -52,5 +52,5 @@ let r = @rule(~x => ~x), rs = RuleSet([r]),
     ex = random_term(1000, atoms=[a, b, c, d, a^(-1), b^(-1), 1, 2.0], funs=[+, *])
 
     overhead["simplify_no_fixp"]["randterm:serial"] = @benchmarkable simplify($ex, threaded=false, fixpoint=false)
-    overhead["simplify_no_fixp"]["randterm:thread"] = @benchmarkable simplify($ex, threaded=true,  fixpoint=false)    
+    overhead["simplify_no_fixp"]["randterm:thread"] = @benchmarkable simplify($ex, threaded=true,  fixpoint=false)
 end

src/SymbolicUtils.jl

@@ -44,8 +44,12 @@ using Combinatorics: permutations
 export @rule, @acrule, RuleSet
 include("rule_dsl.jl")
 
-export simplify, substitute
+import AbstractAlgebra.Generic: MPoly, PolynomialRing, ZZ, exponent_vector
+using AbstractAlgebra: ismonomial
+using DataStructures
+include("abstractalgebra.jl")
 
+export simplify, substitute
 include("simplify.jl")
 
 include("rulesets.jl")

src/abstractalgebra.jl

@@ -0,0 +1,127 @@
+# Polynomial Normal Form
+
+"""
+    labels!(dict, t)
+
+Find all terms that are not + and * and replace them
+with a symbol, store the symbol => term mapping in `dict`.
+"""
+function labels! end
+
+# Turn a Term into a multivariate polynomial
+function labels!(dicts, t::Sym)
+    sym2term, term2sym = dicts
+    if !haskey(term2sym, t)
+        sym2term[t] = t
+        term2sym[t] = t
+    end
+    return t
+end
+
+function labels!(dicts, t)
+    if t isa Integer
+        return t
+    elseif t isa Term && (operation(t) == (*) || operation(t) == (+) || operation(t) == (-))
+        tt = arguments(t)
+        return Term{symtype(t)}(operation(t), map(x->labels!(dicts, x), arguments(t)))
+    elseif t isa Term && operation(t) == (^) && length(arguments(t)) > 1 && isnonnegint(arguments(t)[2])
+        return Term{symtype(t)}(operation(t), map(x->labels!(dicts, x), arguments(t)))
+    else
+        sym2term, term2sym = dicts
+        if haskey(term2sym, t)
+            return term2sym[t]
+        end
+        if t isa Term
+            tt = arguments(t)
+            sym = Sym{symtype(t)}(gensym(nameof(operation(t))))
+            sym2term[sym] = Term{symtype(t)}(operation(t),
+                                             map(x->to_mpoly(x, dicts)[1], arguments(t)))
+        else
+            sym = Sym{symtype(t)}(gensym("literal"))
+            sym2term[sym] = t
+        end
+
+        term2sym[t] = sym
+
+        return sym
+    end
+end
+
+ismpoly(x) = x isa MPoly || x isa Integer
+isnonnegint(x) = x isa Integer && x >= 0
+
+const mpoly_rules = RuleSet([@rule(~x::ismpoly - ~y::ismpoly => ~x + -1 * (~y))
+                             @acrule(~x::ismpoly + ~y::ismpoly => ~x + ~y)
+                             @rule(+(~x) => ~x)
+                             @acrule(~x::ismpoly * ~y::ismpoly => ~x * ~y)
+                             @rule(*(~x) => ~x)
+                             @rule((~x::ismpoly)^(~a::isnonnegint) => (~x)^(~a))])
+function to_mpoly(t, dicts=(OrderedDict{Sym, Any}(), OrderedDict{Any, Sym}()))
+    # term2sym is only used to assign the same
+    # symbol for the same term -- in other words,
+    # it does common subexpression elimination
+
+    sym2term, term2sym = dicts
+    labeled = labels!((sym2term, term2sym), t)
+
+    if isempty(sym2term)
+        return labeled, []
+    end
+
+    ks = sort(collect(keys(sym2term)), lt=<ₑ)
+    R, vars = PolynomialRing(ZZ, String.(nameof.(ks)))
+
+    replace_with_poly = Dict{Sym,MPoly}(zip(ks, vars))
+    t_poly = substitute(labeled, replace_with_poly, fold=false)
+    simplify(t_poly, EmptyCtx(), rules=mpoly_rules),
+        sym2term,
+        reverse(ks)
+end
+
+function to_term(x, dict, syms)
+    dict = copy(dict)
+    for (k, v) in dict
+        dict[k] = _to_term(v, dict, syms)
+    end
+    _to_term(x, dict, syms)
+end
+
+function _to_term(x::MPoly, dict, syms)
+
+    function mul_coeffs(exps)
+        monics = [e == 1 ? syms[i] : syms[i]^e for (i, e) in enumerate(reverse(exps)) if !iszero(e)]
+        if length(monics) == 1
+            return monics[1]
+        elseif length(monics) == 0
+            return 1
+        else
+            return Term(*, monics)
+        end
+    end
+
+    monoms = [mul_coeffs(exponent_vector(x, i)) for i in 1:x.length]
+    if length(monoms) == 0
+        return 0
+    end
+    if length(monoms) == 1
+        t = !isone(x.coeffs[1]) ?  monoms[1] * x.coeffs[1] : monoms[1]
+    else
+        t = Term(+, map((x,y)->isone(y) ? x : y*x, monoms, x.coeffs[1:length(monoms)]))
+    end
+
+    substitute(t, dict, fold=false)
+end
+
+function _to_term(x, dict, vars)
+    if haskey(dict, x)
+        return dict[x]
+    else
+        return x
+    end
+end
+
+function _to_term(x::Term, dict, vars)
+    t=Term{symtype(x)}(operation(x), _to_term.(arguments(x), (dict,), (vars,)))
+end
+
+<ₑ(a::MPoly, b::MPoly) = false

src/rule_dsl.jl

@@ -16,7 +16,7 @@ getdepth(r::Rule) = r.depth
 
 function rule_depth(rule, d=0, maxdepth=0)
     if rule isa Term
-        maxdepth = maximum(rule_depth(r, d+1, maxdepth) for r in arguments(rule))
+        maxdepth = reduce(max, (rule_depth(r, d+1, maxdepth) for r in arguments(rule)), init=1)
     elseif rule isa Slot || rule isa Segment
         maxdepth = max(d, maxdepth)
     end
@@ -28,9 +28,10 @@ function Base.show(io::IO, r::Rule)
 end
 
 const EMPTY_DICT = ImmutableDict{Symbol, Any}(:____, nothing)
+struct DefaultCtx end
 struct EmptyCtx end
 
-function (r::Rule)(term, ctx=EmptyCtx())
+function (r::Rule)(term, ctx=DefaultCtx())
     rhs = r.rhs
 
     r.matcher((term,), EMPTY_DICT, ctx) do bindings, n
@@ -192,7 +193,7 @@ end
 
 Base.show(io::IO, acr::ACRule) = print(io, "ACRule(", acr.rule, ")")
 
-function (acr::ACRule)(term, ctx=EmptyCtx())
+function (acr::ACRule)(term, ctx=DefaultCtx())
     r = Rule(acr)
     if !(term isa Term)
         r(term)
@@ -220,7 +221,7 @@ end
 #### Rulesets
 
 """
-    RuleSet(rules::Vector{AbstractRules}, context=EmptyCtx())(expr; depth=typemax(Int), applyall=false, recurse=true)
+    RuleSet(rules::Vector{AbstractRules}, context=DefaultCtx())(expr; depth=typemax(Int), applyall=false, recurse=true)
 
 `RuleSet` is an `AbstractRule` which applies the given `rules` throughout an `expr` with the
 context `context`.
@@ -267,7 +268,7 @@ end
 
 const rule_repr = IdDict()
 
-function (r::RuleSet)(term, context=EmptyCtx();
+function (r::RuleSet)(term, context=DefaultCtx();
                       depth=typemax(Int),
                       applyall::Bool=false,
                       recurse::Bool=true,

src/rulesets.jl

@@ -23,8 +23,8 @@ const PLUS_RULES = RuleSet([
     @acrule(~a::isnumber + ~b::isnumber => ~a + ~b)
 
     @acrule(*(~~x) + *(~β, ~~x) => *(1 + ~β, (~~x)...))
-    @acrule(*(~α, ~~x) + *(~β, ~~x) => *(~α +  ~β, (~~x)...))
-    @acrule(*(~~x, ~α) + *(~~x, ~β) => *(~α +  ~β, (~~x)...))
+    @acrule(*(~α, ~~x) + *(~β, ~~x) => *(~α + ~β, (~~x)...))
+    @acrule(*(~~x, ~α) + *(~~x, ~β) => *(~α + ~β, (~~x)...))
 
     @acrule(~x + *(~β, ~x) => *(1 + ~β, ~x))
     @acrule(*(~α::isnumber, ~x) + ~x => *(~α + 1, ~x))

src/simplify.jl

@@ -11,14 +11,18 @@ of symtype Number.
 """
 default_rules(x, ctx) = SIMPLIFY_RULES
 
-function default_rules(x, ctx::EmptyCtx)
+function default_rules(x, ctx::DefaultCtx)
     has_trig(x) ?
         SIMPLIFY_RULES_TRIG :
         SIMPLIFY_RULES
 end
 
+function default_rules(x, ctx::EmptyCtx)
+    identity
+end
+
 """
-    simplify(x, ctx=EmptyCtx();
+    simplify(x, ctx=DefaultCtx();
         rules=default_rules(x, ctx),
         fixpoint=true,
         applyall=true,
@@ -28,7 +32,7 @@ Simplify an expression by applying `rules` until there are no changes.
 The second argument, the context is passed to every [`Contextual`](#Contextual)
 predicate and can be accessed as `(@ctx)` in the right hand side of `@rule` expression.
 
-By default the context is an `EmptyCtx()` -- which means there is no contextual information.
+By default the context is an `DefaultCtx()` -- which means there is no contextual information.
 Any arbitrary type can be used as a context, and packages defining their own contexts
 should define `default_rules(ctx::TheContextType)` to return a `RuleSet` that will
 be used by default while simplifying under that context.
@@ -39,7 +43,15 @@ Applies them once if `fixpoint=false`.
 The `applyall` and `recurse` keywords are forwarded to the enclosed
 `RuleSet`, they are mainly used for internal optimization.
 """
-function simplify(x, ctx=EmptyCtx(); rules=default_rules(x, ctx), fixpoint=true, applyall=true, kwargs...)
+function simplify(x, ctx=DefaultCtx();
+                  rules=default_rules(x, ctx),
+                  fixpoint=true,
+                  applyall=true,
+                  mpoly=false,
+                  kwargs...)
+    if mpoly
+        x = to_term(to_mpoly(x)...)
+    end
     if fixpoint
         SymbolicUtils.fixpoint(rules, x, ctx; applyall=applyall)
     else
@@ -56,8 +68,13 @@ Base.@deprecate simplify(x, rules::RuleSet; kwargs...)  simplify(x, rules=rules;
 substitute any subexpression that matches a key in `dict` with
 the corresponding value.
 """
-function substitute(expr, dict)
-    RuleSet([@rule ~x::(x->haskey(dict, x)) => dict[~x]])(expr) |> fold
+function substitute(expr, dict; fold=true)
+    rs = RuleSet([@rule ~x::(x->haskey(dict, x)) => dict[~x]])
+    if fold
+        rs(expr) |> SymbolicUtils.fold
+    else
+        rs(expr)
+    end
 end
 
 fold(x) = x

test/basics.jl

@@ -1,4 +1,4 @@
-using SymbolicUtils: Sym, FnType, Term, symtype, Contextual, EmptyCtx
+using SymbolicUtils: Sym, FnType, Term, symtype, Contextual, DefaultCtx
 using SymbolicUtils
 using Test
 
@@ -87,7 +87,7 @@ end
 @testset "Contexts" begin
     @syms a b c
 
-    @test @rule(~x::Contextual((x, ctx) -> ctx==EmptyCtx()) => "yes")(1) == "yes"
+    @test @rule(~x::Contextual((x, ctx) -> ctx==DefaultCtx()) => "yes")(1) == "yes"
     @test @rule(~x::Contextual((x, ctx) -> haskey(ctx, x)) => true)(a, Dict(a=>1))
     @test @rule(~x::Contextual((x, ctx) -> haskey(ctx, x)) => true)(b, Dict(a=>1)) === nothing
     @test_throws UndefVarError @rule(~x => __CTX__)(a, "test")

test/fuzzlib.jl

@@ -7,7 +7,11 @@ function rand_input(T)
     if T == Bool
         return rand(Bool)
     elseif T <: Integer
-        return rand(-100:100)
+        x = rand(-100:100)
+        while iszero(x)
+            x = rand(-100:100)
+        end
+        return x
     elseif T == Rational
         return Rational(rand_input(Int), rand(1:50)) # no 0 denominator tests yet!
     elseif T == Real

test/nf.jl

@@ -0,0 +1,8 @@
+using SymbolicUtils, Test
+@testset "polyform" begin
+    @syms a b c d
+    @test simplify(a * (b + -1 * c) + -1 * (b * a + -1 * c * a), mpoly=true) == 0
+    @eqtest simplify(sin((a+b)^2)^2; mpoly=true) == sin(a^2+b^2+2*a*b)^2
+    # fixme: can this be made faster?
+    @test simplify(sin((a+b)^2)^2 + cos((a+b)^2)^2; mpoly=true) == 1
+end