Merge pull request #96 from tkoolen/tk/prune-zero

Add prune_zero, prune_zero!
tkoolen · Feb 12, 2019 · 086e766 · 086e766
2 parents 96f03a2 + 34095b3
commit 086e766
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 2 deletions.
diff --git a/docs/src/api/functions.md b/docs/src/api/functions.md
@@ -28,6 +28,11 @@ canonicalize(::QuadraticFunction)
 canonicalize!
 ```
 
+```@docs
+prune_zero
+prune_zero!
+```
+
 ## In-place math functions (unexported)
 
 ```@docs

diff --git a/src/functions.jl b/src/functions.jl
@@ -45,7 +45,9 @@ export
 
 export
     canonicalize,
-    canonicalize!
+    canonicalize!,
+    prune_zero,
+    prune_zero!
 
 using LinearAlgebra
 using DocStringExtensions
@@ -69,6 +71,21 @@ In-place version of [`canonicalize`](@ref).
 """
 function canonicalize! end
 
+"""
+$(METHODLIST)
+
+Prune terms with (approximately) zero coefficients.
+"""
+function prune_zero end
+
+"""
+$(METHODLIST)
+
+In-place version of [`prune_zero`](@ref).
+"""
+function prune_zero! end
+
+
 # Variable
 
 """
@@ -274,6 +291,13 @@ julia> canonicalize(f)
 """
 canonicalize(f::AffineFunction) = canonicalize!(AffineFunction(f))
 
+function prune_zero!(f::AffineFunction{T}; atol=zero(T)) where T
+    filter!(term -> abs(getcoeff(term)) > atol, f.linear)
+    f
+end
+
+prune_zero(f::AffineFunction; kwargs...) = prune_zero!(AffineFunction(f); kwargs...)
+
 # QuadraticFunction
 
 """
@@ -382,6 +406,14 @@ julia> canonicalize(f)
 """
 canonicalize(f::QuadraticFunction) = canonicalize!(QuadraticFunction(f))
 
+function prune_zero!(f::QuadraticFunction{T}; atol=zero(T)) where T
+    prune_zero!(f.affine)
+    filter!(term -> abs(getcoeff(term)) > atol, f.quadratic)
+    f
+end
+
+prune_zero(f::QuadraticFunction; kwargs...) = prune_zero!(QuadraticFunction(f); kwargs...)
+
 
 # copyto!
 Base.copyto!(f::AffineFunction, x::Number) = (zero!(f); f.constant[] = x; f)

diff --git a/test/functions.jl b/test/functions.jl
@@ -6,13 +6,17 @@ using StaticArrays: @SVector, SVector
 using Parametron
 using Parametron.Functions
 
-@testset "canonicalize" begin
+@testset "canonicalize, prune_zero" begin
     x = Variable(1)
     y = Variable(2)
 
     @test canonicalize(3 * x * y) === canonicalize(3 * y * x) == QuadraticTerm(3, x, y)
     @test canonicalize(y + x - 2 * y + 3) == x - y + 3
     @test canonicalize(x * y + y * x + y + y + x - y + 4) == 2 * x * y + x + y + 4
+    @test canonicalize(0 * y + x + 1) == x + 0 * y + 1
+    @test prune_zero(canonicalize(0 * y + x + 1)) == x + 1
+    @test canonicalize(0 * y^2 + x^2 + 0 * x + y + 1) == x^2 + 0 * y^2 + 0 * x + y + 1
+    @test prune_zero(canonicalize(0 * y^2 + x^2 + 0 * x + y + 1)) == x^2 + y + 1
 
     f = x * y + y * x + y + y + x - y + 4
     global allocs