Hard stop on simulation time, next! -> update!

Project.toml

@@ -1,7 +1,7 @@
 name = "Pathogen"
 uuid = "c8189ad9-0349-5cbb-98f2-579781f27e5c"
 authors = ["Justin Angevaare <justinangevaare@gmail.com>"]
-version = "0.4.3"
+version = "0.4.4"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"

examples/SIR.md

@@ -79,11 +79,11 @@ sim = Simulation(pop, starting_states, rf, rparams)
 simulate!(sim, tmax=200.0)
 ```
 
-    SIR epidemic simulation @ time = 204.88
+    SIR epidemic simulation @ time = 200.0
 
     S = 18
-    I = 7
-    R = 75
+    I = 8
+    R = 74
 
 <br><br>
 

src/Pathogen.jl

@@ -45,7 +45,6 @@ module Pathogen
   include("functions/initialize.jl")
   include("functions/observe.jl")
   include("functions/update!.jl")
-  include("functions/next!.jl")
   include("functions/simulate!.jl")
   include("functions/loglikelihood.jl")
   include("functions/logpriors.jl")

src/functions/iterate!.jl

@@ -19,9 +19,9 @@ function iterate!(mc::MarkovChain{T},
   end
   for i = 1:n
     if use_adapted_cov
-      next!(mc, mcmc, adapted_cov, σ, condition_on_network = condition_on_network, event_batches = event_batches)
+      update!(mc, mcmc, adapted_cov, σ, condition_on_network = condition_on_network, event_batches = event_batches)
     else
-      next!(mc, mcmc, Σ, σ, condition_on_network = condition_on_network, event_batches = event_batches)
+      update!(mc, mcmc, Σ, σ, condition_on_network = condition_on_network, event_batches = event_batches)
     end
     next!(pmeter)
     @logmsg LogLevel(-5000) "MCMC progress" progress = i/n
@@ -57,9 +57,9 @@ function iterate!(mc::MarkovChain{T},
   end
   for i = 1:n
     if use_adapted_cov
-      next!(mc, mcmc, adapted_cov, σ, condition_on_network = condition_on_network, event_batches = event_batches)
+      update!(mc, mcmc, adapted_cov, σ, condition_on_network = condition_on_network, event_batches = event_batches)
     else
-      next!(mc, mcmc, Σ, σ, condition_on_network = condition_on_network, event_batches = event_batches)
+      update!(mc, mcmc, Σ, σ, condition_on_network = condition_on_network, event_batches = event_batches)
     end
     put!(progress_channel, true)
     if adapt_cov != 0 && mod(i, adapt_cov) == 0

src/functions/simulate!.jl

@@ -10,11 +10,15 @@ function simulate!(sim::Simulation{T}; tmax::Float64=Inf, nmax::Int64=1000, pmax
   end
   stoptime = time() + pmax
   while true
-    next!(sim)
-    if sim.time >= tmax
+    event = generate(Event, sim.event_rates, sim.time)
+    if event.time > tmax
+      sim.time = tmax
       @debug "Simulation stopped: simulation time maximum reached"
       break
-    elseif sim.iterations >= nmax
+    else
+      update!(sim, event)
+    end
+    if sim.iterations >= nmax
       @debug "Simulation stopped: iteration maximum reached"
       break
     elseif pmax < Inf && time() >= stoptime

src/functions/update!.jl

@@ -115,3 +115,155 @@ function update!(tr::TransmissionRates,
   update!(rates, tr, event, states, pop, rf, rp)
   return tr, rates
 end
+
+function update!(s::Simulation{T}, event::Event) where T <: EpidemicModel
+  if event.time < s.time
+    @error "Time of a new event must be >= that of the previous event"
+  elseif event.time < Inf
+    update!(s.events, event)
+    update!(s.disease_states, event)
+    update!(s.transmission_network,
+            generate(Transmission,
+                     s.transmission_rates,
+                     event))
+    update!(s.transmission_rates,
+            s.event_rates,
+            event,
+            s.disease_states,
+            s.population,
+            s.risk_functions,
+            s.risk_parameters)
+    # Count the iteration as having occurred.
+    s.iterations += 1
+  end
+  # Update simulation time
+  s.time = event.time
+  # Return updated Simulation object
+  return s
+end
+
+function update!(mc::MarkovChain{T},
+                 mcmc::MCMC{T},
+                 Σ::Array{Float64, 2},
+                 σ::Float64;
+                 condition_on_network::Bool=false,
+                 event_batches::Int64=1) where T <: EpidemicModel
+  # Initialize
+  new_events = mc.events[end]
+  new_events_array = new_events[_state_progressions[T][2:end]]
+  new_params = mc.risk_parameters[end]
+  new_network = mc.transmission_network[end]
+  new_lposterior = mc.log_posterior[end]
+  # Randomize event time augmentation
+  event_indices = findall(new_events_array[:] .> -Inf)
+  aug_order = sample(event_indices, length(event_indices), replace=false)
+  @debug "$(length(aug_order)) events to augment"
+  if event_batches < 0
+    @error "Cannot have negative amount of event batches"
+  end
+  if event_batches > length(aug_order)
+    @warn "More event batches than there are events to augment ($(event_batches) > $(length(aug_order))), setting to maximum ($(length(aug_order)))"
+    event_batches = length(aug_order)
+  end
+  batch_size = fld(length(aug_order), event_batches)
+  @debug "Performing data augmentation in batches of $batch_size events at a time"
+  for i = 1:(event_batches + 1)
+    if i <= event_batches
+      for j = (batch_size*(i-1) + 1):minimum([(batch_size*i + 1); length(aug_order)])
+        id, state_index = Tuple(CartesianIndices((new_events.individuals,
+                                            length(_state_progressions[T][2:end])))[aug_order[j]])
+        new_state = _state_progressions[T][state_index+1]
+        time = new_events[new_state][id]
+        # Conditioning on network means that only event times valid under current network will be proposed. This is useful for models which may have additional contributions to the posterior based on network, and require more modest proposals (e.g. phylodynamic models).
+        if condition_on_network
+          proposed_event = generate(Event,
+                                    Event{T}(time, id, new_state),
+                                    σ,
+                                    mcmc.event_extents,
+                                    mcmc.event_observations,
+                                    new_events,
+                                    new_network)
+        else
+          proposed_event = generate(Event,
+                                    Event{T}(time, id, new_state),
+                                    σ,
+                                    mcmc.event_extents,
+                                    mcmc.event_observations,
+                                    new_events)
+        end
+        # For the first event of batch, we'll create an events array from current Markov chain position
+        # For other events in batch, we'll update proposal itself
+        if j == (batch_size*(i-1) + 1)
+          proposed_events_array = reshape([new_events_array[1:(aug_order[j]-1)]
+                                           proposed_event.time
+                                           new_events_array[(aug_order[j]+1):end]],
+                                          size(new_events_array))
+        else
+          proposed_events_array = reshape([proposed_events_array[1:(aug_order[j]-1)]
+                                           proposed_event.time
+                                           proposed_events_array[(aug_order[j]+1):end]],
+                                          size(new_events_array))
+        end
+        # Only need to generate new `Events` on last event of batch
+        if j == minimum([(batch_size*i + 1); length(aug_order)])
+          proposed_events = Events{T}(proposed_events_array)
+        end
+      end
+      proposed_params = new_params
+    else
+      # Propose new risk parameters
+      proposed_events = new_events
+      proposed_events_array = new_events_array
+      proposed_params = generate(RiskParameters{T}, new_params, Σ)
+    end
+    proposed_lprior = logpriors(proposed_params, mcmc.risk_priors)
+    # Based on the logprior and competiting MCMC iteration, this loglikelihood is required for acceptance
+    # Calculating this in advance allows us to cut loglikelihood calculation short if it goes below threshold
+    ll_acceptance_threshold = log(rand()) + new_lposterior - proposed_lprior
+    if ll_acceptance_threshold < Inf
+      proposed_llikelihood = loglikelihood(proposed_params,
+                                           mcmc.risk_functions,
+                                           proposed_events,
+                                           mcmc.population,
+                                           mcmc.starting_states,
+                                           transmission_network_output = false,
+                                           early_decision_value = ll_acceptance_threshold)
+      proposed_lposterior = proposed_lprior + proposed_llikelihood
+    else
+      proposed_llikelihood = -Inf
+      proposed_lposterior = -Inf
+    end
+    if proposed_llikelihood >= ll_acceptance_threshold
+      @debug "MCMC proposal accepted (acceptance probability = $(round(min(1.0, exp(proposed_lposterior - new_lposterior)), digits=3)))"
+      new_params = proposed_params
+      new_events = proposed_events
+      new_events_array = proposed_events_array
+      new_lposterior = proposed_lposterior
+    else
+      @debug "next!: MCMC proposal rejected (acceptance probability = $(round(min(1.0, exp(proposed_lposterior - new_lposterior)), digits=3)))"
+    end
+  end
+  mc.iterations += 1
+  new_network = loglikelihood(new_params,
+                              mcmc.risk_functions,
+                              new_events,
+                              mcmc.population,
+                              mcmc.starting_states,
+                              loglikelihood_output = false,
+                              transmission_network_output = true)
+  push!(mc.events, new_events)
+  push!(mc.transmission_network, new_network)
+  push!(mc.risk_parameters, new_params)
+  push!(mc.log_posterior, new_lposterior)
+  return mc
+end
+
+function update!(mcmc::MCMC{T},
+                 Σ::Array{Float64, 2},
+                 σ::Float64) where T <: EpidemicModel
+  @simd for mc in mcmc.markov_chains
+    next!(mc, mcmc, Σ, σ)
+  end
+  return mcmc
+end
+