Merge pull request #370 from JuliaSymbolics/s/fix-expand

fix expand for expressions with array indexing

Author: shashi

Project.toml

@@ -1,7 +1,7 @@
 name = "SymbolicUtils"
 uuid = "d1185830-fcd6-423d-90d6-eec64667417b"
 authors = ["Shashi Gowda"]
-version = "0.15.4"
+version = "0.15.5"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/polyform.jl

@@ -160,6 +160,10 @@ function PolyForm(x,
         recurse=false,
         metadata=metadata(x))
 
+    if !(symtype(x) <: Number)
+        return x
+    end
+
     # Polyize and return a PolyForm
     p = polyize(x, pvar2sym, sym2term, vtype, pow, Fs, recurse)
     PolyForm{symtype(x)}(p, pvar2sym, sym2term, metadata)
@@ -225,8 +229,12 @@ Expand expressions by distributing multiplication over addition, e.g.,
 multivariate polynomials implementation.
 `variable_type` can be any subtype of `MultivariatePolynomials.AbstractVariable`.
 """
-expand(expr) = Postwalk(identity)(PolyForm(expr, Fs=Union{typeof(+), typeof(*), typeof(^)}, recurse=true))
+expand(expr) = unpolyize(PolyForm(expr, Fs=Union{typeof(+), typeof(*), typeof(^)}, recurse=true))
 
+function unpolyize(x)
+    simterm(x, f, args; kw...) = similarterm(x, f, args, symtype(x); kw...)
+    Postwalk(identity, similarterm=simterm)(x)
+end
 
 ## Rational Polynomial form with Div
 
@@ -280,7 +288,7 @@ function simplify_fractions(x; polyform=false)
 
     expr = Postwalk(sdiv ∘ quick_cancel)(Postwalk(add_with_div)(x))
 
-    polyform ? expr : Postwalk(identity)(expr)
+    polyform ? expr : unpolyize(expr)
 end
 
 function add_with_div(x, flatten=true)

test/polyform.jl

@@ -32,6 +32,10 @@ end
    #@test expand(Term{Number}(zero, 0)) == 0
    #@test expand(identity(a * b) - b * a) == 0
     @test expand(a * b - b * a) == 0
+
+    @syms A::Vector{Real}
+    # test that the following works
+    expand(Term{Real}(getindex, [A, 3]) - 3)
 end
 
 @testset "simplify_fractions with quick-cancel" begin