From 79c141ed0459d9d84fcd0e6b8b95e11e25d89e4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 7 Dec 2023 15:51:21 +0100
Subject: [PATCH] Add filter_terms (#287)

* Add filter_terms

* Fix format
---
 docs/src/types.md  |  2 ++
 src/polynomial.jl  | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 test/polynomial.jl |  7 ++++++-
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/docs/src/types.md b/docs/src/types.md
index 08302dee..e7bfb725 100644
--- a/docs/src/types.md
+++ b/docs/src/types.md
@@ -81,6 +81,8 @@ leading_coefficient
 leading_monomial
 remove_leading_term
 remove_monomials
+filter_terms
+OfDegree
 monic
 map_coefficients
 map_coefficients!
diff --git a/src/polynomial.jl b/src/polynomial.jl
index e3830d0f..aef6aa92 100644
--- a/src/polynomial.jl
+++ b/src/polynomial.jl
@@ -474,6 +474,52 @@ function remove_monomials(
     return q
 end
 
+"""
+    function filter_terms(f::Function, p::AbstractPolynomialLike)
+
+Filter the polynomial `p` by only keep the terms `t` such that `f(p)` is
+`true`.
+
+See also [`OfDegree`](@ref).
+
+### Examples
+
+```julia
+julia> p = 1 - 2x + x * y - 3y^2 + x^2 * y
+1 - 2x - 3y² + xy + x²y
+
+julia> filter_terms(OfDegree(2), p)
+-3y² + xy
+
+julia> filter_terms(!OfDegree(2), p)
+1 - 2x + x²y
+
+julia> filter_terms(!OfDegree(0:2), p)
+x²y
+
+julia> filter_terms(iseven ∘ coefficient, p)
+-2x
+```
+"""
+function filter_terms(f::F, p::AbstractPolynomialLike) where {F<:Function}
+    return polynomial(filter(f, terms(p)), SortedUniqState())
+end
+
+"""
+    struct OfDegree{D} <: Function
+        degree::D
+    end
+
+A function `d::OfDegree` is such that `d(t)` returns
+`degree(t) == d.degree`. Note that `!d` creates the negation.
+See also [`filter_terms`](@ref).
+"""
+struct OfDegree{D} <: Function
+    degree::D
+end
+
+(d::OfDegree)(mono::AbstractTermLike) = in(degree(mono), d.degree)
+
 """
     monic(p::AbstractPolynomialLike)
 
diff --git a/test/polynomial.jl b/test/polynomial.jl
index 5fcab373..90434ba5 100644
--- a/test/polynomial.jl
+++ b/test/polynomial.jl
@@ -149,7 +149,12 @@ const MP = MultivariatePolynomials
     @test transpose(x + y) == x + y
     @test transpose([1 2; 3 4] * x) == [1 3; 2 4] * x
 
-    @test remove_monomials(4x^2 * y + x * y + 2x, [x * y]) == 4x^2 * y + 2x
+    p = 4x^2 * y + x * y + 2x
+    @test remove_monomials(p, [x * y]) == 4x^2 * y + 2x
+    @test filter_terms(OfDegree(2), 4x^2 * y + x * y + 2x) == x * y
+    @test filter_terms(!OfDegree(2), 4x^2 * y + x * y + 2x) == 4x^2 * y + 2x
+    @test filter_terms(iseven ∘ MP.coefficient, 4x^2 * y + x * y + 2x) ==
+          4x^2 * y + 2x
 
     @test_throws InexactError push!([1], x + 1)