Implement Base.Iterator interface for branch and prune routine (#60)

* implement Iterator interface for branch and prune routine

* add example application for RootSearch iterator

* fix implementations of optional RootSearch methods

* add test cases for optional RootSearch methods

* remove stray letter from iterator eltype method

Author: gwater

examples/roots.jl

@@ -22,6 +22,15 @@ rts = roots(f, Xc, Bisection)
 rts = roots(f, rts, Newton)
 rts = roots(f, Xc)
 
+# track the number of working intervals during the iteration:
+f(x) = sin(x)
+contractor = Newton(f, x -> ForwardDiff.derivative(f, x))
+search = RootSearch(-10..10, contractor, 1e-3)
+for state in search
+    print(length(state.working), " ")
+end
+println("done.")
+
 # From R docs:
 
 # https://www.rdocumentation.org/packages/pracma/versions/1.9.9/topics/broyden

src/roots.jl

@@ -1,15 +1,78 @@
 
 import IntervalArithmetic: diam, isinterior
+import Base: start, next, done, copy, eltype, iteratorsize
 
-export branch_and_prune, Bisection, Newton
+export branch_and_prune, Bisection, Newton, RootSearch
+export start, next, done, copy, step!, eltype, iteratorsize
 
 diam(x::Root) = diam(x.interval)
 
 Base.size(x::Interval) = (1,)
 
 isinterior{N}(X::IntervalBox{N}, Y::IntervalBox{N}) = all(isinterior.(X, Y))
 
+struct RootSearchState{T <: Union{Interval,IntervalBox}}
+    working::Vector{T}
+    outputs::Vector{Root{T}}
+end
+RootSearchState{T<:Union{Interval,IntervalBox}}(region::T) =
+    RootSearchState([region], Root{T}[])
+
+copy(state::RootSearchState) =
+    RootSearchState(copy(state.working), copy(state.outputs))
 
+"""
+    RootSearch{R <: Union{Interval,IntervalBox}, S <: Contractor, T <: Real}
+
+Type implementing the `Base.Iterator` interface to the branch and prune routine.
+Returns the `RootSearchState` at each iteration. Note: Each iteration mutates
+the `RootSearchState`. Use `copy(state::RootSearchState)` to create an
+independent instance if necessary.
+"""
+struct RootSearch{R <: Union{Interval,IntervalBox}, S <: Contractor, T <: Real}
+    region::R
+    contractor::S
+    tolerance::T
+end
+
+eltype{R, T <: RootSearch{R}}(::Type{T}) = RootSearchState{R}
+iteratorsize{T <: RootSearch}(::Type{T}) = Base.SizeUnknown()
+
+function start(iter::RootSearch)
+    state = RootSearchState(iter.region)
+    sizehint!(state.outputs, 100)
+    sizehint!(state.working, 1000)
+    return state
+end
+
+"""
+    step!(state::RootSearchState, contractor, tolerance)
+
+Progress `state` by treating one of its `working` regions. Note: `state.working`
+is always modified. If a root is found, it is added to `state.outputs`.
+"""
+function step!(state::RootSearchState, contractor, tolerance)
+    X = pop!(state.working)
+    status, output = contractor(X, tolerance)
+    if status == :empty
+        return nothing
+    elseif status == :unique
+        push!(state.outputs, Root(output, :unique))
+    elseif diam(output) < tolerance
+        push!(state.outputs, Root(output, :unknown))
+    else # branch
+        X1, X2 = bisect(X)
+        push!(state.working, X1, X2)
+    end
+    return nothing
+end
+
+function next(iter::RootSearch, state::RootSearchState)
+    step!(state, iter.contractor, iter.tolerance)
+    return state, state
+end
+
+done(iter::RootSearch, state::RootSearchState) = isempty(state.working)
 
 """
     branch_and_prune(X, contract, tol=1e-3)
@@ -25,38 +88,13 @@ Inputs:
 
 """
 function branch_and_prune(X, contractor, tol=1e-3)
-    working = [X]
-    outputs = Root{typeof(X)}[]
-
-    sizehint!(outputs, 100)
-    sizehint!(working, 1000)
-
-    while !isempty(working)
-        # @show working
-        X = pop!(working)
-
-        status, output = contractor(X, tol)
-
-        if status == :empty
-            continue
-
-        elseif status == :unique
-            push!(outputs, Root(output, :unique))
-
-        elseif diam(output) < tol
-            push!(outputs, Root(output, :unknown))
-
-        else  # branch
-            X1, X2 = bisect(X)
-
-            push!(working, X1, X2)
-        end
-    end
-
-    return outputs
+    iter = RootSearch(X, contractor, tol)
+    local state
+    # complete iteration
+    for state in iter end
+    return state.outputs
 end
 
-
 export recursively_branch_and_prune
 
 function recursively_branch_and_prune(h, X, contractor=BisectionContractor, final_tol=1e-14)

test/roots.jl

@@ -95,3 +95,18 @@ end
     end
     @test all(d .< 1e-15)
 end
+
+@testset "RootSearch interface" begin
+    contractor = Newton(sin, cos)
+    search = RootSearch(-10..10, contractor, 1e-3)
+    state = start(search)
+
+    # check that original and copy are independent
+    state_copy = copy(state)
+    pop!(state_copy.working) # mutate copy
+    @test length(state.working) != length(state_copy.working)
+
+    # cover optional iterator methods
+    @test eltype(search) != Any
+    @test_nowarn iteratorsize(search)
+end