implement coeff

src/Symbolics.jl

@@ -74,7 +74,6 @@ export Equation, ConstrainedEquation
 include("equations.jl")
 
 include("utils.jl")
-export degree
 
 using ConstructionBase
 include("arrays.jl")

src/utils.jl

@@ -225,45 +225,42 @@ function var_from_nested_derivative(x,i=0)
     end
 end
 
-function degree(p::Sym, sym=nothing)
-    if sym === nothing
-        return 1
-    else
-        return Int(isequal(p, sym))
-    end
-end
+degree(p::Union{Term,Sym}, sym=nothing) = sym === nothing ? 1 : Int(isequal(p, sym))
+degree(p::Add, sym=nothing) = maximum(degree.(arguments(p), sym))
+degree(p::Mul, sym=nothing) = sum(degree(k^v, sym) for (k, v) in p.dict)
+degree(p::Pow, sym=nothing) = p.exp * degree(p.base, sym)
 
-function degree(p::Pow, sym=nothing)
-    return p.exp * degree(p.base, sym)
-end
+"""
+    degree(p, sym=nothing)
 
-function degree(p::Add, sym=nothing)
-    return maximum(degree(key, sym) for key in keys(p.dict))
+Extract the degree of `p` with respect to `sym`.
+"""
+function degree(p, sym=nothing)
+    p, sym = value(p), value(sym)
+    p isa Number && return 0
+    isequal(p, sym) && return 1
+    p isa Symbolic && return degree(p, sym)
+    throw(DomainError(p, "Datatype $(typeof(p)) not accepted."))
 end
 
-function degree(p::Mul, sym=nothing)
-    return sum(degree(k^v, sym) for (k, v) in zip(keys(p.dict), values(p.dict)))
+coeff(p::Union{Term,Sym}, sym=nothing) = sym === nothing ? 0 : Int(isequal(p, sym))
+coeff(p::Pow, sym=nothing) = sym === nothing ? 0 : Int(isequal(p, sym))
+coeff(p::Add, sym=nothing) = sum(coeff.(arguments(p), sym))
+function coeff(p::Mul, sym=nothing)
+    args = arguments(p)
+    I = findall(a -> !isequal(a, sym), args)
+    length(I) == length(args) ? 0 : prod(args[I])
 end
 
-function degree(p::Term, sym=nothing)
-    if sym === nothing
-        return 1
-    else
-        return Int(isequal(p, sym))
-    end
-end
+"""
+    coeff(p, sym=nothing)
 
-function degree(p, sym=nothing)
-    p = value(p)
-    sym = value(sym)
-    if p isa Number
-        return 0
-    end
-    if isequal(p, sym)
-        return 1
-    end
-    if p isa Symbolic
-        return degree(p, sym)
-    end
+Extract the coefficient of `p` with respect to `sym`.
+Note that `p` might need to be expanded and/or simplified with `expand` and/or `simplify`.
+"""
+function coeff(p, sym=nothing)
+    p, sym = value(p), value(sym)
+    p isa Number && return sym === nothing ? p : 0
+    p isa Symbolic && return coeff(p, sym)
     throw(DomainError(p, "Datatype $(typeof(p)) not accepted."))
 end

test/coeff.jl

@@ -0,0 +1,41 @@
+using Symbolics, Test
+import Symbolics: coeff
+
+@variables x y z a b
+
+@test isequal(coeff(2), 2)
+@test isequal(coeff(2, x), 0)
+
+@test isequal(coeff(2a, x), 0)
+@test isequal(coeff(a*x, x), a)
+@test isequal(coeff(2x*a, x), 2a)
+
+@test isequal(coeff(a + x, x), 1)
+@test isequal(coeff(2(a + x), x), 2)
+
+e = 4 + x + 3x^2 + 2x^4 + a*x^2 + b
+@test isequal(coeff(e), 4)
+@test isequal(coeff(e, x^1), 1)
+@test isequal(coeff(e, x^2), 3 + a)
+@test isequal(coeff(e, x^3), 0)
+@test isequal(coeff(e, x^4), 2)
+
+e = x*y^2 + 2x + y^3*x^3
+@test isequal(coeff(e, x), 2 + y^2)
+@test isequal(coeff(e, x^3), y^3)
+@test isequal(coeff(e, y^2), x)
+@test isequal(coeff(e, y^3), x^3)
+
+@test isequal(coeff(x^2 + y^2 + z^2, sin(z)), 0)
+@test isequal(coeff(sin(z), z), 0)
+@test isequal(coeff(sin(z)), 0)
+
+# issue #236
+@test isequal(coeff(3x + 2y, x), 3)
+@test isequal(coeff(x*y, x), y)
+@test isequal(coeff(x^2 + y, x^2), 1)
+
+# expand - simplify needed
+@test isequal(coeff(expand(((x + 1)^4 + x)^3), x^2), 93)
+@test isequal(coeff(simplify((x^2 - 1) / (x - 1)), x), 1)
+@test isequal(coeff(expand((x^(1//2) + y^0.5)^2), x), 1)

test/degree.jl

@@ -1,32 +1,40 @@
-using Symbolics
-using Test
+using Symbolics, Test
+import Symbolics: degree
 
 @variables x, y, z
 
+@test isequal(degree(x^0), 0)
+@test isequal(degree(x^1), 1)
+@test isequal(degree(x^2), 2)
+
 @test isequal(degree(1), 0)
 @test isequal(degree(x), 1)
 @test isequal(degree(x, x), 1)
+@test isequal(degree(2x, x), 1)
 @test isequal(degree(x, y), 0)
 
+@test isequal(degree(x/2, x), 1)
+@test_broken isequal(degree(x/y, x), 1)  # FIXME: `StackOverflowError`
+
 @test isequal(degree(x*y, y), 1)
 @test isequal(degree(x*y, x), 1)
 @test isequal(degree(x*y, x*y), 1)
 @test isequal(degree(x*y, z), 0)
 
-@test isequal(degree(x*y^2+2*x+y^3*x^3), 6)
-@test isequal(degree(x*y^2+2*x+y^3*x^3, y), 3)
-@test isequal(degree(x*y^2+2*x+y^3*x^6, x), 6)
+@test isequal(degree(x*y^2 + 2x + y^3*x^3), 6)
+@test isequal(degree(x*y^2 + 2x + y^3*x^3, y), 3)
+@test isequal(degree(x*y^2 + 2x + y^3*x^6, x), 6)
 
-@test isequal(degree(x*y^2+2*x+y^3*x^6, z), 0)
-@test isequal(degree(x*y^2+2*x+y^3*x^6+z, z), 1)
-@test isequal(degree(x*y^2+2*x+y^3*x^6*z, z), 1)
+@test isequal(degree(x*y^2 + 2x + y^3*x^6, z), 0)
+@test isequal(degree(x*y^2 + 2x + y^3*x^6 + z, z), 1)
+@test isequal(degree(x*y^2 + 2x + y^3*x^6*z, z), 1)
 
-@test isequal(degree(x^2 + y^2 +z^2, sin(z)), 0)
+@test isequal(degree(x^2 + y^2 + z^2, sin(z)), 0)
 @test isequal(degree(sin(z), z), 0)
 @test isequal(degree(sin(z)), 1)
 
 @test isequal(degree(x^2*sin(z), x), 2)
 @test isequal(degree(x+exp(z), x), 1)
 
 @test isequal(degree((x - y)^2*((y + x*y)^3)), 8)
-@test isequal(degree((x + z)*((y + x*y)^3), x), 4)
+@test isequal(degree((x + z)*((y + x*y)^3), x), 4)

test/runtests.jl

@@ -27,6 +27,7 @@ if GROUP == "All" || GROUP == "Core"
     @safetestset "Differentiation Test" begin include("diff.jl") end
     @safetestset "Difference Test" begin include("difference.jl") end
     @safetestset "Degree Test" begin include("degree.jl") end
+    @safetestset "Coeff Test" begin include("coeff.jl") end
     @safetestset "Is Linear or Affine Test" begin include("islinear_affine.jl") end
     @safetestset "Linear Solver Test" begin include("linear_solver.jl") end
     @safetestset "Groebner Bases Test" begin include("groebner_basis.jl") end