Pass initial population to EpiSystem, not EpiEnv

examples/Epidemiology/SEI3HRD_example.jl

@@ -47,7 +47,7 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
 
     # Set up simple gridded environment
     area = 525_000.0km^2
-    epienv = simplehabitatAE(298.0K, area, NoControl(), scotpop)
+    epienv = simplehabitatAE(298.0K, size(scotpop), area, NoControl())
 
     # Set population to initially have no individuals
     abun_h = (
@@ -77,7 +77,7 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
     rel = Gauss{eltype(epienv.habitat)}()
 
     # Create epi system with all information
-    epi = EpiSystem(epilist, epienv, rel)
+    epi = EpiSystem(epilist, epienv, rel, scotpop)
 
     # Add in initial infections randomly (samples weighted by population size)
     N_cells = size(epi.abundances.matrix, 2)

examples/Epidemiology/Scotland_run.jl

@@ -35,6 +35,8 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
     total_pop = dropdims(sum(Float64.(scotpop), dims=3), dims=3)
     total_pop = AxisArray(total_pop, AxisArrays.axes(scotpop)[1], AxisArrays.axes(scotpop)[2])
     total_pop.data[total_pop .≈ 0.0] .= NaN
+    # Shrink to smallest bounding box. The NaNs are inactive.
+    total_pop = shrink_to_active(total_pop);
 
     # Set simulation parameters
     numclasses = 8
@@ -74,7 +76,7 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
                           T_lat, T_asym, T_presym, T_sym, T_hosp, T_rec)
     param = transition(param, age_categories)
 
-    epienv = simplehabitatAE(298.0K, area, NoControl(), total_pop)
+    epienv = simplehabitatAE(298.0K, size(total_pop), area, NoControl())
 
     # Set population to initially have no individuals
     abun_h = (
@@ -111,7 +113,7 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
     rel = Gauss{eltype(epienv.habitat)}()
 
     # Create epi system with all information
-    epi = EpiSystem(epilist, epienv, rel, UInt16(1))
+    epi = EpiSystem(epilist, epienv, rel, total_pop, UInt16(1))
 
     # Populate susceptibles according to actual population spread
     reshaped_pop =

examples/Epidemiology/UK_run.jl

@@ -55,7 +55,7 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
 
     # Set up simple gridded environment
     area = 875_000.0km^2
-    epienv = simplehabitatAE(298.0K, area, NoControl(), ukpop)
+    epienv = simplehabitatAE(298.0K, size(ukpop), area, NoControl())
 
     # Set population to initially have no individuals
     abun_h = (
@@ -87,7 +87,7 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
     rel = Gauss{eltype(epienv.habitat)}()
 
     # Create epi system with all information
-    epi = EpiSystem(epilist, epienv, rel)
+    epi = EpiSystem(epilist, epienv, rel, ukpop)
 
     # Spread susceptibles randomly over age categories
     split_pop = rand.(Multinomial.(Int.(epi.abundances.matrix[1, :]), 10))

src/Epidemiology/EpiEnv.jl

@@ -21,19 +21,17 @@ abstract type AbstractEpiEnv{H <: AbstractHabitat, C <: AbstractControl} <:
 This epi environment type holds a habitat and control strategy, as well as a string of
 subcommunity names, and initial susceptible population.
 """
-mutable struct GridEpiEnv{H, C, A, P} <: AbstractEpiEnv{H, C}
+mutable struct GridEpiEnv{H, C, A} <: AbstractEpiEnv{H, C}
     habitat::H
     active::A
     control::C
-    initial_population::P
     names::Vector{String}
     function (::Type{GridEpiEnv{H, C}})(
         habitat::H,
         active::A,
         control::C,
-        initial_population::P=zeros(Int, _getdimension(habitat)),
         names::Vector{String}=map(x -> "$x", 1:countsubcommunities(habitat))
-    ) where {H, C, A <: AbstractMatrix{Bool}, P <: AbstractMatrix{Int}}
+    ) where {H, C, A <: AbstractMatrix{Bool}}
         countsubcommunities(habitat) == length(names) ||
             error("Number of subcommunities must match subcommunity names")
         if size(habitat.matrix) != size(active)
@@ -42,13 +40,7 @@ mutable struct GridEpiEnv{H, C, A, P} <: AbstractEpiEnv{H, C}
                 "size(active)=$(size(active))"
             ))
         end
-        if size(initial_population) != size(active)
-            throw(DimensionMismatch(
-                "size(initial_population)=$(size(initial_population)) != " *
-                "size(active)=$(size(active))"
-            ))
-        end
-        return new{H, C, A, P}(habitat, active, control, initial_population, names)
+        return new{H, C, A}(habitat, active, control, names)
     end
 end
 
@@ -61,58 +53,13 @@ function _getsubcommunitynames(epienv::GridEpiEnv)
     return epienv.names
 end
 
-"""
-    _shrink_to_active(M::AbstractMatrix, active::AbstractMatrix{<:Bool})
-
-Shrink the matrix `M` to the minimum rectangular region which contains all active cells, as
-defined by `active`. Returns the shrunk matrix.
-"""
-function _shrink_to_active(M::AM, active::A) where {AM <: AbstractMatrix, A <: AbstractMatrix{<: Bool}}
-    if size(M) != size(active)
-        throw(DimensionMismatch("size(M)=$(size(M)) != size(active)=$(size(active))"))
-    end
-    # Find indices of non-missing values
-    idx = Tuple.(findall(active))
-    # Separate into row and column indices
-    row_idx = first.(idx)
-    col_idx = last.(idx)
-    # Return the shrunk region
-    shrunk_rows = minimum(row_idx):maximum(row_idx)
-    shrunk_cols = minimum(col_idx):maximum(col_idx)
-    #return M[shrunk_rows, shrunk_cols]
-    return _construct_shrunk_matrix(M, shrunk_rows, shrunk_cols)
-end
-
-"""
-    _construct_shrunk_matrix
-
-Construct a shrunk matrix by selecting certain rows and columns specified by `row_idxs` and
-`col_idxs` from AbstractMatrix `M`.
-
-Return an AxisArray{T, 2}. The axes will be the selected subset of the original axes if `M`
-is an AxisArray. If `M` is a normal matrix, the axes of the returned AxisArray are the
-selected coordinates.
-"""
-function _construct_shrunk_matrix(M::Matrix, row_idxs, col_idxs)::AxisArray
-    return AxisArray(
-        M[row_idxs, col_idxs];
-        row_idxs=row_idxs,
-        col_idxs=col_idxs,
-    )
-end
-
-function _construct_shrunk_matrix(M::AxisArray, row_idxs, col_idxs)::AxisArray
-    return M[row_idxs, col_idxs]
-end
-
 """
     function simplehabitatAE(
         val::Union{Float64, Unitful.Quantity{Float64}},
         dimension::Tuple{Int64, Int64},
         area::Unitful.Area{Float64},
         active::AbstractMatrix{Bool},
         control::C,
-        initial_population::AbstractMatrix{<:Integer}=zeros(Int, dimension),
     )
 
 Function to create a simple `ContinuousHab` type epi environment. It creates a
@@ -129,8 +76,7 @@ function simplehabitatAE(
     area::Unitful.Area{Float64},
     active::M,
     control::C,
-    initial_population::P=zeros(Int, dimension),
-) where {C <: AbstractControl, M <: AbstractMatrix{Bool}, P <: AbstractMatrix{<:Integer}}
+) where {C <: AbstractControl, M <: AbstractMatrix{Bool}}
     if typeof(val) <: Unitful.Temperature
         val = uconvert(K, val)
     end
@@ -139,12 +85,11 @@ function simplehabitatAE(
 
     # Shrink to active region
     # This doesn't change the gridsquaresize
-    initial_population = _shrink_to_active(initial_population, active)
-    active = _shrink_to_active(active, active)
+    active = shrink_to_active(active, active)
     dimension = size(active)
 
     hab = simplehabitat(val, gridsquaresize, dimension)
-    return GridEpiEnv{typeof(hab), typeof(control)}(hab, active, control, initial_population)
+    return GridEpiEnv{typeof(hab), typeof(control)}(hab, active, control)
 end
 
 function simplehabitatAE(
@@ -156,70 +101,3 @@ function simplehabitatAE(
     active = fill(true, dimension)
     return simplehabitatAE(val, dimension, area, active, control)
 end
-
-"""
-    simplehabitatAE(
-        val::Union{Float64, Unitful.Quantity{Float64}},
-        area::Unitful.Area{Float64},
-        control::C,
-        initial_population::AbstractMatrix{<:Real},
-    )
-
-Create a simple `ContinuousHab` type epi environment from a specified `initial_population`
-matrix.
-
-## Inputs
-- `val`: Fill the habitat with this value
-- `initial_population`: Used to derive the dimensions of the habitat, and the initial
-    susceptible population. Values in `initial_population` which are `NaN` or `Missing` are
-    used to mask off inactive areas. `initial_population` will be rounded to integers.
-- `area`: The area of the habitat
-- `control`: The control to apply
-
-!!! note
-    The simulation grid will be shrunk so that it tightly wraps the active values in
-    `initial_population`.
-"""
-function simplehabitatAE(
-    val::Union{Float64, Unitful.Quantity{Float64}},
-    area::Unitful.Area{Float64},
-    control::C,
-    initial_population::M,
-) where {C <: AbstractControl, M <: AbstractMatrix{<:Real}}
-    inactive(x) = isnan(x) || ismissing(x)
-    if all(inactive.(initial_population))
-        throw(ArgumentError("initial_population is all NaN / missing"))
-    end
-    dimension = size(initial_population)
-    active = Matrix{Bool}(.!inactive.(initial_population))
-    initial_population = _convert_population(initial_population, active)
-    return simplehabitatAE(val, dimension, area, active, control, initial_population)
-end
-
-"""
-    _convert_population
-
-Convert populatioin matrix to Int matrix by filling in the inactive area with 0 population
-and rounding the active area.
-"""
-function _convert_population(
-    initial_population::Matrix{<:Real},
-    active::AbstractMatrix{Bool}
-)::Matrix{<:Int}
-    initial_population[.!active] .= 0
-    initial_population = Int.(round.(initial_population))
-    return initial_population
-end
-
-function _convert_population(
-    initial_population::AxisArray{<:Real, 2},
-    active::AbstractMatrix{Bool}
-)::AxisArray{<:Int, 2}
-    # NOTE: this is a workaround as logical indexing directly on AxisArray leads to
-    #   stackoverflow. see issue: https://github.com/JuliaArrays/AxisArrays.jl/issues/179
-    initial_population.data[.!active] .= 0
-    return AxisArray(
-            Int.(round.(initial_population.data)),
-            initial_population.axes
-        )
-end

src/Epidemiology/EpiHelper.jl

@@ -93,7 +93,8 @@ function simulate_record!(
 
   # - initialise and save the first timestep abuns/storage to HDF5
   # construct axes for abuns/storage matrix
-  grid_id = map(Iterators.product(axisvalues(epi.epienv.initial_population)...)) do (x,y)
+  ax = AxisArrays.axes(epi.abundances.grid)[end-1:end]
+  grid_id = map(Iterators.product(ax...)) do (x,y)
       return string.(x, "-", y)
   end
   # TODO: confirm converting `grid_id` from matrix to vector in the way below gives the

src/Epidemiology/EpiSystem.jl

@@ -66,18 +66,29 @@ function EpiSystem(popfun::F, epilist::EpiList, epienv::GridEpiEnv,
 end
 
 function EpiSystem(epilist::EpiList, epienv::GridEpiEnv, rel::AbstractTraitRelationship, intnum::U = Int64(1)) where U <: Integer
-    epi = EpiSystem(populate!, epilist, epienv, rel, intnum)
+    return EpiSystem(populate!, epilist, epienv, rel, intnum)
+end
+
+function EpiSystem(epilist::EpiList, epienv::GridEpiEnv, rel::AbstractTraitRelationship, initial_population, intnum::U = Int64(1)) where U <: Integer
+    if size(initial_population) != size(epienv.active)
+        msg = "size(initial_population)==$(size(initial_population)) != " *
+            "size(epienv.active)==$(size(epienv.active))"
+        throw(DimensionMismatch(msg))
+    end
+    epi = EpiSystem(epilist, epienv, rel, intnum)
     # Add in the initial susceptible population
     idx = findfirst(epilist.human.names .== "Susceptible")
     if idx == nothing
         msg = "epilist has no Susceptible category. epilist.names = $(epilist.human.names)"
         throw(ArgumentError(msg))
     end
-    epi.abundances.grid[idx, :, :] .+= U.(epienv.initial_population)
+    # Modify active cells based on new population
+    epi.epienv.active .&= .!_inactive.(initial_population)
+    initial_population = convert_population(initial_population, intnum)
+    epi.abundances.grid[idx, :, :] .+= initial_population
     return epi
 end
 
-
 """
     isapprox(epi_1::AbstractEpiSystem, epi_2::AbstractEpiSystem; kwargs...)
 

src/Epidemiology/shrink.jl

@@ -0,0 +1,91 @@
+
+"""
+    _shrink_to_active(M::AbstractMatrix, active::AbstractMatrix{<:Bool})
+
+Shrink the matrix `M` to the minimum rectangular region which contains all active cells, as
+defined by `active`. Returns the shrunk matrix.
+
+If `active is not provided, automatically determines the active region by masking out
+entries which are `NaN` or `missing`.
+"""
+function shrink_to_active(M::AM, active::A) where {AM <: AbstractMatrix, A <: AbstractMatrix{<: Bool}}
+    if size(M) != size(active)
+        throw(DimensionMismatch("size(M)=$(size(M)) != size(active)=$(size(active))"))
+    end
+    # Find indices of non-missing values
+    idx = Tuple.(findall(active))
+    # Separate into row and column indices
+    row_idx = first.(idx)
+    col_idx = last.(idx)
+    # Return the shrunk region
+    shrunk_rows = minimum(row_idx):maximum(row_idx)
+    shrunk_cols = minimum(col_idx):maximum(col_idx)
+    #return M[shrunk_rows, shrunk_cols]
+    return _construct_shrunk_matrix(M, shrunk_rows, shrunk_cols)
+end
+
+function shrink_to_active(M::AM) where {AM <: AbstractMatrix}
+    active = .!_inactive.(M)
+    return shrink_to_active(M, active)
+end
+
+_inactive(x) = ismissing(x) || isnan(x)
+
+"""
+    _construct_shrunk_matrix
+
+Construct a shrunk matrix by selecting certain rows and columns specified by `row_idxs` and
+`col_idxs` from AbstractMatrix `M`.
+
+Return an AxisArray{T, 2}. The axes will be the selected subset of the original axes if `M`
+is an AxisArray. If `M` is a normal matrix, the axes of the returned AxisArray are the
+selected coordinates.
+"""
+function _construct_shrunk_matrix(M::AbstractMatrix, row_idxs, col_idxs)::AxisArray
+    return AxisArray(
+        M[row_idxs, col_idxs];
+        row_idxs=row_idxs,
+        col_idxs=col_idxs,
+    )
+end
+
+function _construct_shrunk_matrix(M::AxisArray, row_idxs, col_idxs)::AxisArray
+    return M[row_idxs, col_idxs]
+end
+
+"""
+    function convert_population(
+        initial_population,
+        intnum::U = Int64(1)
+    )
+
+Convert population matrix to Int matrix by filling in the inactive area with 0 population
+and rounding the active area.
+"""
+function convert_population(
+    initial_population::Matrix,
+    intnum::U = Int64(1)
+) where U <: Integer
+    # Don't modify the arg
+    initial_population = copy(initial_population)
+    active = .!_inactive.(initial_population)
+    initial_population[.!active] .= 0
+    initial_population = U.(round.(initial_population))
+    return initial_population
+end
+
+function convert_population(
+    initial_population::AxisArray,
+    intnum::U = Int64(1)
+) where U <: Integer
+    # Don't modify the arg
+    initial_population = copy(initial_population)
+    active = .!_inactive.(initial_population)
+    # NOTE: this is a workaround as logical indexing directly on AxisArray leads to
+    #   stackoverflow. see issue: https://github.com/JuliaArrays/AxisArrays.jl/issues/179
+    initial_population.data[.!active] .= 0
+    return AxisArray(
+        U.(round.(initial_population.data)),
+        initial_population.axes
+    )
+end

src/Simulation.jl

@@ -156,6 +156,9 @@ include("Epidemiology/EpiPlots.jl")
 include("Epidemiology/Inference.jl")
 export SIR_wrapper, SIR_wrapper!
 
+include("Epidemiology/shrink.jl")
+export shrink_to_active, convert_population
+
 # Path into package
 path(paths...) = joinpath(@__DIR__, "..", paths...)