Merge pull request #88 from adolgert/feature/sets-and-dicts

Feature/sets and dicts

Author: adolgert

docs/src/algorithms.md

@@ -15,4 +15,6 @@ choose
 setindex!
 rand
 set_multiple!
+SetOfSets
+PrefixEnabled
 ```

docs/src/reference.md

@@ -29,6 +29,7 @@ freeze
 haskey
 misscount
 misses
+enabled
 MultiSampler
 ChatReaction
 Petri

src/CompetingClocks.jl

@@ -3,6 +3,7 @@ using Documenter
 
 const ContinuousTime = AbstractFloat
 
+include("setofsets.jl")
 include("prefixsearch/binarytreeprefixsearch.jl")
 include("prefixsearch/cumsumprefixsearch.jl")
 include("prefixsearch/keyedprefixsearch.jl")

src/prefixsearch/cumsumprefixsearch.jl

@@ -13,7 +13,7 @@ and, each time the Direct method samples, this evaluates the cumulative
 sum of the array.
 """
 struct CumSumPrefixSearch{T<:Real}
-	array::Vector{T}
+    array::Vector{T}
     cumulant::Vector{T}
 end
 
@@ -49,7 +49,7 @@ function Base.setindex!(ps::CumSumPrefixSearch{T}, value::T, index) where {T}
     ps.array[index] = value
     value
 end
-
+Base.getindex(ps::CumSumPrefixSearch, index) = ps.array[index]
 
 function Base.sum!(ps::CumSumPrefixSearch{T})::T where {T}
     cumsum!(ps.cumulant, ps.array)

src/prefixsearch/keyedprefixsearch.jl

@@ -12,7 +12,7 @@ re-enables the same set of clocks, this is the faster choice.
 """
 struct KeyedKeepPrefixSearch{T,P} <: KeyedPrefixSearch
     # Map from clock name to index in propensity array.
-    index::Dict{T, Int}
+    index::Dict{T,Int}
     # Map from index in propensity array to clock name.
     key::Vector{T}
     prefix::P
@@ -37,6 +37,7 @@ end
 
 
 Base.length(kp::KeyedKeepPrefixSearch) = length(kp.index)
+key_type(kp::KeyedKeepPrefixSearch{T,P}) where {T,P} = T
 time_type(kp::KeyedKeepPrefixSearch{T,P}) where {T,P} = time_type(P)
 
 
@@ -54,9 +55,51 @@ end
 
 
 isenabled(kp::KeyedKeepPrefixSearch, clock) = (
-    haskey(kp.index, clock) && kp.prefix[clock] > zero(time_type(kp))
+    haskey(kp.index, clock) && kp.prefix[kp.index[clock]] > zero(time_type(kp))
 )
 
+"""
+Construct a set that checks which values are zeroed out because this
+prefix sum doesn't mark what has been deleted. That's faster for small
+sets of keys but makes getting the set that's enabled more difficult.
+A hazard rate that is set to zero at enabling will give a funny count
+because it's `enabled!()` by the user but set to never fire.
+"""
+struct PrefixEnabled{K,P,KK} <: AbstractSet{K}
+    prefix::P
+    keys::KK
+end
+
+# Implements the interface to return a set of enabled clock keys.
+function enabled(prefix::KeyedKeepPrefixSearch{T,P}) where {T,P}
+    kk = keys(prefix.index)
+    PrefixEnabled{T,typeof(prefix),typeof(kk)}(prefix, kk)
+end
+
+function Base.iterate(nre::PrefixEnabled)
+    res = iterate(nre.keys)
+    res === nothing && return res
+    while !isenabled(nre.prefix, res[1])
+        res = iterate(nre.keys, res[2])
+        res === nothing && return res
+    end
+    return res
+end
+
+
+function Base.iterate(nre::PrefixEnabled, state)
+    res = iterate(nre.keys, state)
+    res === nothing && return res
+    while !isenabled(nre.prefix, res[1])
+        res = iterate(nre.keys, res[2])
+        res === nothing && return res
+    end
+    return res
+end
+
+Base.length(nre::PrefixEnabled) = count(x -> isenabled(nre.prefix, x), nre.keys)
+Base.in(x, nre::PrefixEnabled) = isenabled(nre.prefix, x)
+Base.eltype(::Type{PrefixEnabled{C}}) where {C} = C
 
 Base.delete!(kp::KeyedKeepPrefixSearch, clock) = kp.prefix[kp.index[clock]] = zero(time_type(kp))
 function Base.sum!(kp::KeyedKeepPrefixSearch)
@@ -72,7 +115,7 @@ end
 
 function Random.rand(
     rng::AbstractRNG, d::Random.SamplerTrivial{KeyedKeepPrefixSearch{T,P}}
-    ) where {T,P}
+) where {T,P}
     total = sum!(d[])
     LocalTime = time_type(P)
     choose(d[], rand(rng, Uniform{LocalTime}(zero(LocalTime), total)))
@@ -87,7 +130,7 @@ a large key space, this will use less memory.
 """
 struct KeyedRemovalPrefixSearch{T,P} <: KeyedPrefixSearch
     # Map from clock name to index in propensity array.
-    index::Dict{T, Int}
+    index::Dict{T,Int}
     # Map from index in propensity array to clock name.
     key::Vector{T}
     free::Set{Int}
@@ -114,6 +157,7 @@ function Base.copy!(dst::KeyedRemovalPrefixSearch{T,P}, src::KeyedRemovalPrefixS
 end
 
 Base.length(kp::KeyedRemovalPrefixSearch) = length(kp.index)
+key_type(kp::KeyedRemovalPrefixSearch{T,P}) where {T,P} = T
 time_type(kp::KeyedRemovalPrefixSearch{T,P}) where {T,P} = time_type(P)
 
 function Base.setindex!(kp::KeyedRemovalPrefixSearch, val, clock)
@@ -134,10 +178,8 @@ function Base.setindex!(kp::KeyedRemovalPrefixSearch, val, clock)
 end
 
 
-isenabled(kp::KeyedRemovalPrefixSearch, clock) = (
-    haskey(kp.index, clock) && kp.prefix[clock] > zero(time_type(kp))
-)
-
+isenabled(kp::KeyedRemovalPrefixSearch, clock) = haskey(kp.index, clock)
+enabled(kp::KeyedRemovalPrefixSearch) = keys(kp.index)
 
 function Base.getindex(kp::KeyedRemovalPrefixSearch, clock)
     if haskey(kp.index, clock)
@@ -170,7 +212,7 @@ end
 
 function Random.rand(
     rng::AbstractRNG, d::Random.SamplerTrivial{KeyedRemovalPrefixSearch{T,P}}
-    ) where {T,P}
+) where {T,P}
     total = sum!(d[])
     LocalTime = time_type(P)
     choose(d[], rand(rng, Uniform{LocalTime}(zero(LocalTime), total)))

src/sample/combinednr.jl

@@ -2,7 +2,7 @@
 using DataStructures: MutableBinaryMinHeap, extract_all!, update!
 
 export sampling_space
-export CombinedNextReaction
+export CombinedNextReaction, enabled
 
 """
 This function decides whether a particular distribution can be sampled faster
@@ -384,6 +384,55 @@ function Base.length(nr::CombinedNextReaction)
     return length(nr.transition_entry)
 end
 
+
+function isenabled(nr::CombinedNextReaction, clock)
+    haskey(nr.transition_entry, clock) && nr.transition_entry[clock].heap_handle > 0
+end
+
+
+# A set of all enabled clock keys for a CombinedNextReaction method.
+# We make a custom Set implementation because the information is in the
+# CombinedNextReaction object, but it's spread across a Heap and a Dictionary.
+# This helper class should make it much more efficient to iterate the set.
+struct NextReactionEnabled{C,T,K} <: AbstractSet{C}
+    nr::T
+    keys::K
+end
+
+_has_handle(nre::NextReactionEnabled, key) = nre.nr.transition_entry[key].heap_handle > 0
+
+function Base.iterate(nre::NextReactionEnabled)
+    res = iterate(nre.keys)
+    res === nothing && return res
+    while !_has_handle(nre, res[1])
+        res = iterate(nre.keys, res[2])
+        res === nothing && return res
+    end
+    return res
+end
+
+
+function Base.iterate(nre::NextReactionEnabled, state)
+    res = iterate(nre.keys, state)
+    res === nothing && return res
+    while !_has_handle(nre, res[1])
+        res = iterate(nre.keys, res[2])
+        res === nothing && return res
+    end
+    return res
+end
+
+Base.length(nre::NextReactionEnabled) = length(nre.nr.firing_queue)
+Base.in(x, nre::NextReactionEnabled) = isenabled(nre.nr, x)
+Base.eltype(::Type{NextReactionEnabled{C}}) where {C} = C
+
+
+function enabled(nr::CombinedNextReaction{K,T}) where {K,T}
+    kks = keys(nr.transition_entry)
+    NextReactionEnabled{K,typeof(nr),typeof(kks)}(nr, kks)
+end
+
+
 function Base.haskey(nr::CombinedNextReaction{K,T}, clock::K) where {K,T}
     return haskey(nr.transition_entry, clock) && nr.transition_entry[clock].heap_handle > 0
 end

src/sample/direct.jl

@@ -1,7 +1,7 @@
 using Random: rand, AbstractRNG
 using Distributions: Uniform, Exponential, rate
 
-export DirectCall, enable!, disable!, next
+export DirectCall, enable!, disable!, next, enabled
 
 
 """
@@ -126,13 +126,18 @@ function Base.getindex(dc::DirectCall{K,T,P}, clock::K) where {K,T,P}
 end
 
 function Base.keys(dc::DirectCall)
-    return collect(keys(dc.prefix_tree.index))
+    return keys(dc.prefix_tree.index)
 end
 
 function Base.length(dc::DirectCall)
     return length(dc.prefix_tree)
 end
 
+# Implements the interface to return a set of enabled clock keys.
+enabled(dc::DirectCall{K,T,P}) where {K,T,P} = enabled(dc.prefix_tree)
+
+isenabled(dc::DirectCall{K,T,P}, clock::K) where {K,T,P} = isenabled(dc.prefix_tree, clock)
+
 
 function steploglikelihood(dc::DirectCall, now, when, which)
     total = sum!(dc.prefix_tree)

src/sample/firstreaction.jl

@@ -103,4 +103,4 @@ function next(fr::ChatReaction{K,T}, when::T, rng) where {K,T}
 end
 
 
-Base.haskey(fr::ChatReaction, clock) = isenabled(fr.enabled, clock)
+Base.haskey(fr::ChatReaction, clock) = haskey(fr.enabled, clock)

src/sample/firsttofire.jl

@@ -46,7 +46,7 @@ function next(propagator::FirstToFire{K,T}, when::T, rng::AbstractRNG) where {K,
         OrderedSample(nothing, typemax(T))
     end
     @debug("FirstToFire.next queue length ",
-            length(propagator.firing_queue), " least ", least)
+        length(propagator.firing_queue), " least ", least)
     (least.time, least.key)
 end
 
@@ -101,3 +101,5 @@ end
 
 Base.haskey(propagator::FirstToFire{K,T}, clock::K) where {K,T} = haskey(propagator.transition_entry, clock)
 Base.haskey(propagator::FirstToFire{K,T}, clock) where {K,T} = false
+
+enabled(propagator::FirstToFire) = keys(propagator.transition_entry)

src/sample/interface.jl

@@ -2,7 +2,7 @@ using Random: AbstractRNG
 using Distributions: UnivariateDistribution
 import Base: getindex, keys, length, keytype, haskey
 
-export SSA, enable!, disable!, next, 
+export SSA, enable!, disable!, next,
     getindex, keys, length, keytype
 
 """
@@ -39,7 +39,7 @@ function enable!(
     te::T, # enabling time
     when::T, # current simulation time
     rng::AbstractRNG
-    ) where {K,T}
+) where {K,T}
     @assert false
 end
 
@@ -139,4 +139,11 @@ timetype(::SSA{K,T}) where {K,T} = T
 
 Return a boolean.
 """
-Base.haskey(sampler::SSA{K,T}, key) where {K,T}
+Base.haskey(sampler::SSA{K,T}, key) where {K,T}
+
+"""
+    enabled(sampler)
+
+Returns a read-only set of currently-enabled clocks.
+"""
+enabled(sampler::SSA{K,T}) where {K,T}