expand tests, update top level source file

jangevaare · Apr 23, 2020 · bb08ea0 · bb08ea0 · jangevaare · Apr 23, 2020
1 parent 66dcd1a
commit bb08ea0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/src/Pathogen.jl b/src/Pathogen.jl
@@ -13,6 +13,7 @@ module Pathogen
         OnlineStats
 
   import StatsBase.mean,
+         StatsBase.mode,
          Base.summary,
          ProgressMeter.next!
 
@@ -66,6 +67,7 @@ module Pathogen
   include("visualization/population.jl")
   include("visualization/transmission_network.jl")
   include("visualization/transmission_network_distribution.jl")
+  include("visualization/degree_distribution.jl")
   include("visualization/trace.jl")
 
   export
@@ -174,13 +176,14 @@ module Pathogen
     ZeroMeanDiagNormalCanon,
     ZeroMeanFullNormal,
     ZeroMeanFullNormalCanon,
-    # Pathogen.jl types
+    # Pathogen.jl types/functions
     EpidemicModel,
     SEIR, SEI, SIR, SI,
     DiseaseState,
     State_S, State_E, State_I, State_R,
     RiskFunctions, RiskParameters, RiskPriors,
     Population,
+    AbstractTransmissionNetwork,
     TransmissionNetwork,
     TransmissionNetworkDistribution, TNDistribution,
     TransmissionNetworkPrior, TNPrior,
@@ -191,5 +194,5 @@ module Pathogen
     Events, EventObservations, EventExtents,
     observe,
     MCMC, start!, iterate!,
-    summary
+    summary, mean, mode
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -86,6 +86,8 @@ pop = Population(risks,
   @test all([mcmc.markov_chains[i].iterations for i = 1:3] .== 100)
   tnd = TNDistribution(mcmc, burnin=2)
   @test sum(tnd) ≈ sum(obs.infection .!== NaN)
+  @test sum(mode(tnd)) ≈ sum(obs.infection .!== NaN)
+  @test all(Pathogen._outdegree(tnd) .>= 0.0)
   mcmc = MCMC(obs, ee, pop, rf, rpriors, tnprior=tnd)
   start!(mcmc, attempts=100)
   iterate!(mcmc, 100, 1.0, condition_on_network=false)
@@ -159,6 +161,8 @@ end
   @test all([mcmc.markov_chains[i].iterations for i = 1:3] .== 100)
   tnd = TNDistribution(mcmc, burnin=2)
   @test sum(tnd) ≈ sum(obs.infection .!== NaN)
+  @test sum(mode(tnd)) ≈ sum(obs.infection .!== NaN)
+  @test all(Pathogen._outdegree(tnd) .>= 0.0)
   mcmc = MCMC(obs, ee, pop, rf, rpriors, tnprior=tnd)
   start!(mcmc, attempts=50)
   iterate!(mcmc, 100, 1.0, condition_on_network=false)
@@ -233,6 +237,8 @@ end
   @test all([mcmc.markov_chains[i].iterations for i = 1:3] .== 100)
   tnd = TNDistribution(mcmc, burnin=2)
   @test sum(tnd) ≈ sum(obs.infection .!== NaN)
+  @test sum(mode(tnd)) ≈ sum(obs.infection .!== NaN)
+  @test all(Pathogen._outdegree(tnd) .>= 0.0)
   mcmc = MCMC(obs, ee, pop, rf, rpriors, tnprior=tnd)
   start!(mcmc, attempts=50)
   iterate!(mcmc, 100, 1.0, condition_on_network=false)
@@ -297,6 +303,8 @@ end
   @test all([mcmc.markov_chains[i].iterations for i = 1:3] .== 100)
   tnd = TNDistribution(mcmc, burnin=2)
   @test sum(tnd) ≈ sum(obs.infection .!== NaN)
+  @test sum(mode(tnd)) ≈ sum(obs.infection .!== NaN)
+  @test all(Pathogen._outdegree(tnd) .>= 0.0)
   mcmc = MCMC(obs, ee, pop, rf, rpriors, tnprior=tnd)
   start!(mcmc, attempts=50)
   iterate!(mcmc, 100, 1.0, condition_on_network=false)