Merge pull request #103 from adolgert/feature/sampler-spec

Feature/sampler spec

Author: adolgert

Project.toml

@@ -11,6 +11,7 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -19,6 +20,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Combinatorics = "^1.0"
 Distributions = "^0.25"
 Documenter = "^1.15"
+InteractiveUtils = "^1.10"
 Logging = "^1.0"
 Random = "^1.0"
 SpecialFunctions = "2.6.1"

docs/src/choosing_sampler.md

@@ -54,11 +54,11 @@ A sampler group is
  * A Symbol name for the sampler.
  * An inclusion function from `(ClockKey, Distribution)` to `Bool` that decides
    whether a given event belongs to this sampler.
- * An optional `sampler_spec` to say what kind of sampler this group should use.
+ * An optional `method` to say what kind of sampler this group should use.
 
 ```julia
 builder = SamplerBuilder(KeyType, Float64)
-add_group!(builder, :sparky => (x,d) -> x[1] == :recover, sampler_spec=(:nextreaction,))
+add_group!(builder, :sparky => (x,d) -> x[1] == :recover, method=NextReaction())
 add_group!(builder, :forthright=>(x,d) -> x[1] == :infect)
 sampler = SamplingContext(builder, rng)
 ```

docs/src/commonrandom.jl

@@ -31,7 +31,7 @@ using CompetingClocks
 example_clock = (3, 7)  # We will use clock IDs that are a tuple of 2 integers.
 model = MakeModel()
 (Key, Time) = (typeof(example_clock), Float64)
-builder = SamplerBuilder(Key, Time; sampler_spec=:firsttofire, common_random=true)
+builder = SamplerBuilder(Key, Time; common_random=true)
 rng = Xoshiro(9469922)
 sampler = SamplingContext(builder, rng)
 for trial_idx in 1:100

docs/src/gene_expression.jl

@@ -221,8 +221,8 @@ function run_epochs(epoch_cnt, use_importance, rng)
     model = GeneExpression(params)
     builder = SamplerBuilder(
         Tuple{Symbol,Int}, Float64;
-        sampler_spec=:firsttofire,
-        trajectory_likelihood=true,
+        method=FirstToFireMethod(),
+        path_likelihood=true,
         likelihood_cnt=2,
     )
     sampler = SamplingContext(builder, rng)

docs/src/integration-guide.md

@@ -118,7 +118,7 @@ at any point during a simulation. They are always calulated relative to the
 last event to `fire!()`, even if `next()` has been called.
 ```julia
 builder = SamplerBuilder(KeyType, Float64;
-    trajectory_likelihood=true)
+    path_likelihood=true)
 sampler = SamplingContext(builder, rng)
 # After simulation, the log_prob is a Float64.
 log_prob = pathloglikelihood(sampler, end_time)

docs/src/simple_board.jl

@@ -48,8 +48,7 @@ end
 # * `enable!(sampler, event ID, distribution)`
 # * `disable!(sampler, event ID)`
 #
-# There are a lot of samplers in CompetingClocks to choose from. This example uses `CombinedNextReaction`
-# algorithm, which has good performance for a variety of distributions. Samplers in CompetingClocks
+# There are a lot of samplers in CompetingClocks to choose from. Samplers in CompetingClocks
 # require two type parameters, a key type for clocks and the type used to represent time.
 # In this case, the clock key type fully represents an event, giving the ID of the individual,
 # where they start, and which direction they may move.
@@ -62,7 +61,7 @@ end
 
 function run(event_count)
     rng = Xoshiro(2947223)
-    builder = SamplerBuilder(ClockKey, Float64; sampler_spec=(:nextreaction,))
+    builder = SamplerBuilder(ClockKey, Float64)
     sampler = SamplingContext(builder, rng)
     physical = PhysicalState(zeros(Int, 10, 10))
     sim = SimulationFSM(

src/CompetingClocks.jl

@@ -24,6 +24,7 @@ include("sample/multiple_direct.jl")
 include("sample/combinednr.jl")
 include("variance/with_common_random.jl")
 include("sample/petri.jl")
+include("samplerspec.jl")
 include("sampler_builder.jl")
 include("context.jl")
 

src/context.jl

@@ -32,7 +32,7 @@ function SamplingContext(builder::SamplerBuilder, rng::R) where {R<:AbstractRNG}
     sampler = build_sampler(builder)
     if builder.likelihood_cnt > 1
         likelihood = PathLikelihoods{K,T}(builder.likelihood_cnt)
-    elseif builder.trajectory_likelihood
+    elseif builder.path_likelihood
         likelihood = TrajectoryWatcher{K,T}()
     else
         likelihood = nothing
@@ -56,6 +56,11 @@ function SamplingContext(builder::SamplerBuilder, rng::R) where {R<:AbstractRNG}
 end
 
 
+function SamplingContext(::Type{K}, ::Type{T}, rng::R; kwargs...) where {K,T,R<:AbstractRNG}
+    return SamplingContext(SamplerBuilder(K, T; kwargs...), rng)
+end
+
+
 """
     clone(sampling, rng)
 
@@ -72,7 +77,7 @@ were making parallel RNGs with the `Random123` package, it might look like:
 master_seed = (0xd897a239, 0x77ff9238)
 rng = Philox4x((0, 0, 0, 0), master_seed)
 Key = Int64
-sampler = SamplingContext(SamplerBuilder(Key,Float64; trajectory_likelihood=true), rng)
+sampler = SamplingContext(SamplerBuilder(Key,Float64; path_likelihood=true), rng)
 observation_weight = zeros(Float64, particle_cnt)
 total_weight = zeros(Float64, particle_cnt)
 samplers = Vector{typeof(sampler)}(undef, particle_cnt)

src/sampler_builder.jl

@@ -1,4 +1,4 @@
-export SamplerBuilder, available_samplers, add_group!, build_sampler
+export SamplerBuilder, add_group!, build_sampler
 
 has_steploglikelihood(::Type) = false
 has_steploglikelihood(::Type{<:CombinedNextReaction}) = true
@@ -13,35 +13,34 @@ has_pathloglikelihood(::Type{MultipleDirect}) = true
 mutable struct SamplerBuilderGroup
     name::Symbol
     selector::Union{Function,Nothing}
-    sampler_spec::Tuple{Symbol}
+    method::Union{SamplerSpec,Nothing}
     instance::SSA
-    # Constructor sets the instance to undefined.
-    SamplerBuilderGroup(name::Symbol, selector, sampler_spec) = new(name, selector, sampler_spec)
+    # Constructor sets `instance` member to undefined.
+    SamplerBuilderGroup(name::Symbol, selector, method) = new(name, selector, method)
 end
 
 
 struct SamplerBuilder{K,T}
     clock_type::Type{K}
     time_type::Type{T}
     step_likelihood::Bool
-    trajectory_likelihood::Bool
+    path_likelihood::Bool
     debug::Bool
     recording::Bool
     common_random::Bool
     group::Vector{SamplerBuilderGroup}
-    samplers::Dict{Tuple{Symbol,Vararg{Symbol}},Function}
     start_time::T
     likelihood_cnt::Int
 end
 
 """
     SamplerBuilder(::Type{K}, ::Type{T};
         step_likelihood=false,
-        trajectory_likelihood=false,
+        path_likelihood=false,
         debug=false,
         recording=false,
         common_random=false,
-        sampler_spec=:none,
+        method=nothing,
         start_time::T,
         likelihood_cnt::Int
         )
@@ -51,50 +50,48 @@ an initial sampler.
 
  * `K` and `T` are the clock type and time type.
  * `step_likelihood` - whether you will call `steploglikelihood` before each `fire!`
- * `trajectory_likelihood` - whether you will call `pathloglikelihood`
+ * `path_likelihood` - whether you will call `pathloglikelihood`
     at the end of a simulation run.
  * `debug` - Print log messages at the debug level.
  * `recording` - Store every enable and disable for later examination.
  * `common_random` - Use common random numbers during sampling.
- * `sampler_spec` - If you want a single, particular sampler, put its Symbol name here.
+ * `method` - If you want a single, particular sampler, put its `SamplerSpec` here.
+   It will create a group called `:all` that has this sampling method.
  * `start_time` - Sometimes a simulation shouldn't start at zero.
  * `likelihood_cnt` - The number of likelihoods to compute, corresponds to number of
-   distributions in `enable!` calls. This turns on `trajectory_likelihood`.
+   distributions in `enable!` calls. This turns on `path_likelihood`.
 
 # Example
 
 ```julia
 builder = SamplerBuilder(Tuple,Float64)
-add_group!(builder, :sparky => (x,d) -> x[1] == :recover, sampler_spec=(:nextreaction,))
+add_group!(builder, :sparky => (x,d) -> x[1] == :recover, method=NextReaction())
 add_group!(builder, :forthright=>(x,d) -> x[1] == :infect)
 context = SamplingContext(builder, rng)
 ```
 """
 function SamplerBuilder(::Type{K}, ::Type{T};
     step_likelihood=false,
-    trajectory_likelihood=false,
+    path_likelihood=false,
     debug=false,
     recording=false,
     common_random=false,
-    sampler_spec::Union{Symbol,Tuple{Symbol}}=(:none,),    # Ask for specific sampler.
+    method::Union{SamplerSpec,Nothing}=nothing,    # Ask for specific sampler.
     start_time::T=zero(T),
     likelihood_cnt=1,
 ) where {K,T}
     group = SamplerBuilderGroup[]
-    avail = make_builder_dict()
-    trajectory_likelihood = trajectory_likelihood || likelihood_cnt > 1
+    path_likelihood = path_likelihood || likelihood_cnt > 1
     builder = SamplerBuilder(
-        K, T, step_likelihood, trajectory_likelihood, debug, recording,
-        common_random, group, avail, start_time, likelihood_cnt
+        K, T, step_likelihood, path_likelihood, debug, recording,
+        common_random, group, start_time, likelihood_cnt
     )
-    if sampler_spec != (:none,)
-        add_group!(builder, :all => (x, d) -> true; sampler_spec=sampler_spec)
+    if !isnothing(method)
+        add_group!(builder, :all => (x, d) -> true; method=method)
     end
     return builder
 end
 
-available_samplers(builder::SamplerBuilder) = keys(builder.samplers)
-
 
 """
 The `selector` defines the group of clocks that go to this sampler using
@@ -103,37 +100,17 @@ an inclusion rule, so it's a function from a clock key and distribution to a Boo
 function add_group!(
     builder::SamplerBuilder,
     selector::Union{Pair,Nothing}=nothing;      # Which clocks use this sampler.
-    sampler_spec::Union{Symbol,Tuple{Symbol}}=(:any,),    # Ask for specific sampler.
+    method::Union{SamplerSpec,Nothing}=nothing, # Ask for specific sampler.
 )
-    sampler_spec = sampler_spec isa Symbol ? (sampler_spec,) : sampler_spec
-    sampler_spec = sampler_spec == (:any,) ? (:firsttofire,) : sampler_spec
-    if sampler_spec ∉ keys(builder.samplers)
-        error("Looking for a sampler in this list: $(keys(builder.samplers))")
-    end
     if length(builder.group) >= 1 && (builder.group[1].selector === nothing || selector === nothing)
         error("Need a selector on all samplers if there is more than one sampler.")
     end
     name = selector isa Pair ? selector.first : :all
     selector_func = selector isa Pair ? selector.second : selector
-    push!(builder.group, SamplerBuilderGroup(name, selector_func, sampler_spec))
+    push!(builder.group, SamplerBuilderGroup(name, selector_func, method))
     return nothing
 end
 
-
-function make_builder_dict()
-    return Dict([
-        (:nextreaction,) => (K, T) -> CombinedNextReaction{K,T}(),
-        (:direct,) => (K, T) -> DirectCallExplicit(K, T, KeyedRemovalPrefixSearch, BinaryTreePrefixSearch),
-        (:direct, :remove, :tree) => (K, T) -> DirectCallExplicit(K, T, KeyedRemovalPrefixSearch, BinaryTreePrefixSearch),
-        (:direct, :keep, :tree) => (K, T) -> DirectCallExplicit(K, T, KeyedKeepPrefixSearch, BinaryTreePrefixSearch),
-        (:direct, :remove, :array) => (K, T) -> DirectCallExplicit(K, T, KeyedRemovalPrefixSearch, CumSumPrefixSearch),
-        (:direct, :keep, :array) => (K, T) -> DirectCallExplicit(K, T, KeyedKeepPrefixSearch, CumSumPrefixSearch),
-        (:firstreaction,) => (K, T) -> FirstReaction{K,T}(),
-        (:firsttofire,) => (K, T) -> FirstToFire{K,T}(),
-        (:petri,) => (K, T) -> Petri{K,T}(),
-    ])
-end
-
 """
 See sampler.jl for the MultiSampler to understand how we're making a chooser.
 We would like to implement this:
@@ -149,17 +126,32 @@ function CompetingClocks.choose_sampler(
     return chooser.matcher(clock, distribution)
 end
 
+function auto_select_method(builder::SamplerBuilder)
+    # Auto-select a sampler method based on builder requirements
+    if builder.path_likelihood
+        return DirectMethod()
+    elseif builder.step_likelihood
+        return NextReactionMethod()
+    else
+        return FirstToFireMethod()
+    end
+end
+
 function build_sampler(builder::SamplerBuilder)
-    isempty(builder.group) && error("Need to add_group! on the builder.")
     K = builder.clock_type
     T = builder.time_type
-    if length(builder.group) == 1
-        sampler = builder.samplers[builder.group[1].sampler_spec](K, T)
+    if length(builder.group) == 0
+        sampler = FirstToFireMethod()(K, T)
+        matcher = nothing
+    elseif length(builder.group) == 1
+        method = isnothing(builder.group[1].method) ? auto_select_method(builder) : builder.group[1].method
+        sampler = method(K, T)
         matcher = nothing
     else
         competes = builder.group
         for compete in competes
-            compete.instance = builder.samplers[compete.sampler_spec](K, T)
+            method = isnothing(compete.method) ? auto_select_method(builder) : compete.method
+            compete.instance = method(K, T)
         end
         # Any direct method gets added to the others for combination.
         inclusion = Dict(samp.name => samp.selector for samp in competes)