From bec4ff77163ae50881b1315770257b0769ce97d0 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 17:57:29 +0100
Subject: [PATCH 01/19] Remove initial_pop from EpiEnv

---
 src/Epidemiology/EpiEnv.jl | 134 ++-----------------------------------
 src/Epidemiology/shrink.jl |  49 ++++++++++++++
 src/Simulation.jl          |   3 +
 3 files changed, 58 insertions(+), 128 deletions(-)
 create mode 100644 src/Epidemiology/shrink.jl

diff --git a/src/Epidemiology/EpiEnv.jl b/src/Epidemiology/EpiEnv.jl
index 6dcb7196..55b1c2d6 100644
--- a/src/Epidemiology/EpiEnv.jl
+++ b/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,50 +53,6 @@ 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}},
@@ -112,7 +60,6 @@ end
         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
diff --git a/src/Epidemiology/shrink.jl b/src/Epidemiology/shrink.jl
new file mode 100644
index 00000000..c5943337
--- /dev/null
+++ b/src/Epidemiology/shrink.jl
@@ -0,0 +1,49 @@
+
+"""
+    _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 shrink_to_boundingbox(M::AM) where {AM <: AbstractMatrix}
+    inactive(x) = isnan(x) || ismissing(x) || x==0
+    return shrink_to_active(M, .!inactive.(M))
+end
diff --git a/src/Simulation.jl b/src/Simulation.jl
index aa252dad..30d2aaf7 100644
--- a/src/Simulation.jl
+++ b/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, shrink_to_boundingbox
+
 # Path into package
 path(paths...) = joinpath(@__DIR__, "..", paths...)
 

From a5b4d100443c7742ac9c99370112b46cac186a26 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 17:58:21 +0100
Subject: [PATCH 02/19] Pass initial_population to EpiSystem constructor

---
 src/Epidemiology/EpiSystem.jl | 39 +++++++++++++++++++++++++++++++++--
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/Epidemiology/EpiSystem.jl b/src/Epidemiology/EpiSystem.jl
index ed99726b..918230de 100644
--- a/src/Epidemiology/EpiSystem.jl
+++ b/src/Epidemiology/EpiSystem.jl
@@ -66,17 +66,52 @@ 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
+    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)
+    epi.abundances.grid[idx, :, :] .+= _convert_population(initial_population, epienv.active, intnum)
+    # Modify active cells based on new population
+    # Set any cells with zero population to inactive
+    epi.epienv.active .&= (epi.abundances.grid[idx, :, :] .!= 0)
     return epi
 end
 
+"""
+    _convert_population
+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{<:Real},
+    active::AbstractMatrix{Bool},
+    intnum::U = Int64(1)
+) where U <: Integer
+    initial_population[.!active] .= 0
+    initial_population = U.(round.(initial_population))
+    return initial_population
+end
+
+function _convert_population(
+    initial_population::AxisArray{<:Real, 2},
+    active::AbstractMatrix{Bool},
+    intnum::U = Int64(1)
+) where U <: Integer
+    # 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
 
 """
     isapprox(epi_1::AbstractEpiSystem, epi_2::AbstractEpiSystem; kwargs...)

From e50839b8b0d6b601c37ecfc8efbec88668e6b9d7 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 17:58:47 +0100
Subject: [PATCH 03/19] Remove reference to epienv.initial_population

---
 src/Epidemiology/EpiHelper.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/Epidemiology/EpiHelper.jl b/src/Epidemiology/EpiHelper.jl
index 95de1ef3..dbbbc980 100644
--- a/src/Epidemiology/EpiHelper.jl
+++ b/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

From 5fa327e71df039d8658e632a93101075c72b8f27 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 19:09:58 +0100
Subject: [PATCH 04/19] Add shrink_and_convert

---
 src/Epidemiology/EpiSystem.jl | 31 +---------------------------
 src/Epidemiology/shrink.jl    | 38 ++++++++++++++++++++++++++++++++---
 src/Simulation.jl             |  2 +-
 3 files changed, 37 insertions(+), 34 deletions(-)

diff --git a/src/Epidemiology/EpiSystem.jl b/src/Epidemiology/EpiSystem.jl
index 918230de..5936cdf3 100644
--- a/src/Epidemiology/EpiSystem.jl
+++ b/src/Epidemiology/EpiSystem.jl
@@ -77,42 +77,13 @@ function EpiSystem(epilist::EpiList, epienv::GridEpiEnv, rel::AbstractTraitRelat
         msg = "epilist has no Susceptible category. epilist.names = $(epilist.human.names)"
         throw(ArgumentError(msg))
     end
-    epi.abundances.grid[idx, :, :] .+= _convert_population(initial_population, epienv.active, intnum)
+    epi.abundances.grid[idx, :, :] .+= convert_population(initial_population, epienv.active, intnum)
     # Modify active cells based on new population
     # Set any cells with zero population to inactive
     epi.epienv.active .&= (epi.abundances.grid[idx, :, :] .!= 0)
     return epi
 end
 
-"""
-    _convert_population
-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{<:Real},
-    active::AbstractMatrix{Bool},
-    intnum::U = Int64(1)
-) where U <: Integer
-    initial_population[.!active] .= 0
-    initial_population = U.(round.(initial_population))
-    return initial_population
-end
-
-function _convert_population(
-    initial_population::AxisArray{<:Real, 2},
-    active::AbstractMatrix{Bool},
-    intnum::U = Int64(1)
-) where U <: Integer
-    # 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
-
 """
     isapprox(epi_1::AbstractEpiSystem, epi_2::AbstractEpiSystem; kwargs...)
 
diff --git a/src/Epidemiology/shrink.jl b/src/Epidemiology/shrink.jl
index c5943337..333f41c2 100644
--- a/src/Epidemiology/shrink.jl
+++ b/src/Epidemiology/shrink.jl
@@ -31,7 +31,7 @@ Return an AxisArray{T, 2}. The axes will be the selected subset of the original
 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
+function _construct_shrunk_matrix(M::AbstractMatrix, row_idxs, col_idxs)::AxisArray
     return AxisArray(
         M[row_idxs, col_idxs];
         row_idxs=row_idxs,
@@ -43,7 +43,39 @@ function _construct_shrunk_matrix(M::AxisArray, row_idxs, col_idxs)::AxisArray
     return M[row_idxs, col_idxs]
 end
 
-function shrink_to_boundingbox(M::AM) where {AM <: AbstractMatrix}
+function shrink_and_convert(M::AM, intnum::U=Int64(1)) where {AM <: AbstractMatrix, U <: Integer}
     inactive(x) = isnan(x) || ismissing(x) || x==0
-    return shrink_to_active(M, .!inactive.(M))
+    active = .!inactive.(M)
+    M_out = shrink_to_active(M, active)
+    active = shrink_to_active(active, active)
+    return convert_population(M_out, active, intnum)
+end
+
+"""
+    convert_population
+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{<:Real},
+    active::AbstractMatrix{Bool},
+    intnum::U = Int64(1)
+) where U <: Integer
+    initial_population[.!active] .= 0
+    initial_population = U.(round.(initial_population))
+    return initial_population
+end
+
+function convert_population(
+    initial_population::AxisArray{<:Real, 2},
+    active::AbstractMatrix{Bool},
+    intnum::U = Int64(1)
+) where U <: Integer
+    # 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
diff --git a/src/Simulation.jl b/src/Simulation.jl
index 30d2aaf7..8608681d 100644
--- a/src/Simulation.jl
+++ b/src/Simulation.jl
@@ -157,7 +157,7 @@ include("Epidemiology/Inference.jl")
 export SIR_wrapper, SIR_wrapper!
 
 include("Epidemiology/shrink.jl")
-export shrink_to_active, shrink_to_boundingbox
+export shrink_to_active, convert_population, shrink_and_convert
 
 # Path into package
 path(paths...) = joinpath(@__DIR__, "..", paths...)

From b1e2f299df3fcff469c760a2550870acb6e66a99 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 19:13:11 +0100
Subject: [PATCH 05/19] Update UK example

---
 examples/Epidemiology/Scotland_run.jl | 7 ++++---
 examples/Epidemiology/UK_run.jl       | 5 +++--
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/examples/Epidemiology/Scotland_run.jl b/examples/Epidemiology/Scotland_run.jl
index 43bdecbb..ddaf857b 100644
--- a/examples/Epidemiology/Scotland_run.jl
+++ b/examples/Epidemiology/Scotland_run.jl
@@ -34,7 +34,8 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
     # Sum up age categories and turn into simple matrix
     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
+    total_pop = shrink_and_convert(total_pop);
 
     # Set simulation parameters
     numclasses = 8
@@ -74,7 +75,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 +112,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 =
diff --git a/examples/Epidemiology/UK_run.jl b/examples/Epidemiology/UK_run.jl
index 1f8a541e..017871e8 100644
--- a/examples/Epidemiology/UK_run.jl
+++ b/examples/Epidemiology/UK_run.jl
@@ -52,10 +52,11 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
                                        0.0, 7e5, 0, 1.25e6))
     # Coarsen grid to 10km
     ukpop = [sum(ukpop[i:i+9, j:j+9]) for i in 1:10:size(ukpop, 1), j in 1:10:size(ukpop, 2)]
+    ukpop = shrink_and_convert(ukpop)
 
     # 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 +88,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))

From 1cc944445fb46faac3d3625016626bcfbd93bc8e Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 19:16:56 +0100
Subject: [PATCH 06/19] Update SEI3HRD example

---
 examples/Epidemiology/SEI3HRD_example.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/examples/Epidemiology/SEI3HRD_example.jl b/examples/Epidemiology/SEI3HRD_example.jl
index 719ec572..4dd4a721 100644
--- a/examples/Epidemiology/SEI3HRD_example.jl
+++ b/examples/Epidemiology/SEI3HRD_example.jl
@@ -44,10 +44,11 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
 
     # Read in population sizes for Scotland
     scotpop = Array{Float64, 2}(readfile(Simulation.path("test", "examples", "ScotlandDensity2011.tif"), 0.0, 7e5, 5e5, 1.25e6))
+    scotpop = shrink_and_convert(scotpop)
 
     # 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 +78,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)

From be0d466cf522675dc74b9b5f1c84e022a2f8a422 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 19:20:22 +0100
Subject: [PATCH 07/19] Add docstring

---
 src/Epidemiology/shrink.jl | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/Epidemiology/shrink.jl b/src/Epidemiology/shrink.jl
index 333f41c2..9fcf245b 100644
--- a/src/Epidemiology/shrink.jl
+++ b/src/Epidemiology/shrink.jl
@@ -43,6 +43,13 @@ function _construct_shrunk_matrix(M::AxisArray, row_idxs, col_idxs)::AxisArray
     return M[row_idxs, col_idxs]
 end
 
+"""
+    shrink_and_convert(M::AM, intnum::U=Int64(1)) where {AM <: AbstractMatrix, U <: Integer}
+
+Shrink the matrix M to its active bounding box, and then convert the entries to type U.
+Active entries are defined as being non-nan, non-inf, and non-zero.
+In the output, any inactive entries are set to zero.
+"""
 function shrink_and_convert(M::AM, intnum::U=Int64(1)) where {AM <: AbstractMatrix, U <: Integer}
     inactive(x) = isnan(x) || ismissing(x) || x==0
     active = .!inactive.(M)
@@ -52,7 +59,12 @@ function shrink_and_convert(M::AM, intnum::U=Int64(1)) where {AM <: AbstractMatr
 end
 
 """
-    convert_population
+    function convert_population(
+        initial_population,
+        active::AbstractMatrix{Bool},
+        intnum::U = Int64(1)
+    )
+
 Convert population matrix to Int matrix by filling in the inactive area with 0 population
 and rounding the active area.
 """

From aef55ad3a5865fcfd661b30ead69cbd38dc13006 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Tue, 23 Jun 2020 19:24:49 +0100
Subject: [PATCH 08/19] Fix EpiSystem tests

---
 test/TestCases.jl      | 4 ++--
 test/test_EpiSystem.jl | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/test/TestCases.jl b/test/TestCases.jl
index f57fd84f..5d50ea0b 100644
--- a/test/TestCases.jl
+++ b/test/TestCases.jl
@@ -92,7 +92,7 @@ function TestEpiSystemFromPopulation(initial_pop::Matrix{<:Real})
     param = transition(param)
 
     area = 10.0km^2
-    epienv = simplehabitatAE(298.0K, area, NoControl(), initial_pop)
+    epienv = simplehabitatAE(298.0K, size(initial_pop), area, NoControl())
 
     # Zero susceptible so we can test the specified initial_pop
     abun_h = (
@@ -115,7 +115,7 @@ function TestEpiSystemFromPopulation(initial_pop::Matrix{<:Real})
     epilist = EpiList(traits, abun_v, abun_h, disease_classes, movement, param)
 
     rel = Gauss{eltype(epienv.habitat)}()
-    epi = EpiSystem(epilist, epienv, rel)
+    epi = EpiSystem(epilist, epienv, rel, initial_pop)
 
     return epi
 end
diff --git a/test/test_EpiSystem.jl b/test/test_EpiSystem.jl
index 4cea35d9..3bdf211f 100644
--- a/test/test_EpiSystem.jl
+++ b/test/test_EpiSystem.jl
@@ -27,7 +27,6 @@ epi = TestEpiSystem()
 @testset "EpiSystem from initial population" begin
     initial_pop = rand(10, 10) * 10
     epi = TestEpiSystemFromPopulation(initial_pop)
-    @test epi.epienv.initial_population == round.(initial_pop)
     @test epi.abundances.grid[1, :, :] == round.(initial_pop)
 end
 

From deba8a83b5ca557fbb57e0ebf866ea0684315e69 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 11:01:37 +0100
Subject: [PATCH 09/19] Remove shrink_and_convert, add shrink_to_active method

---
 src/Epidemiology/shrink.jl | 34 ++++++++++++++--------------------
 1 file changed, 14 insertions(+), 20 deletions(-)

diff --git a/src/Epidemiology/shrink.jl b/src/Epidemiology/shrink.jl
index 9fcf245b..5f96c31f 100644
--- a/src/Epidemiology/shrink.jl
+++ b/src/Epidemiology/shrink.jl
@@ -4,6 +4,9 @@
 
 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)
@@ -21,6 +24,13 @@ function shrink_to_active(M::AM, active::A) where {AM <: AbstractMatrix, A <: Ab
     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
 
@@ -43,25 +53,9 @@ function _construct_shrunk_matrix(M::AxisArray, row_idxs, col_idxs)::AxisArray
     return M[row_idxs, col_idxs]
 end
 
-"""
-    shrink_and_convert(M::AM, intnum::U=Int64(1)) where {AM <: AbstractMatrix, U <: Integer}
-
-Shrink the matrix M to its active bounding box, and then convert the entries to type U.
-Active entries are defined as being non-nan, non-inf, and non-zero.
-In the output, any inactive entries are set to zero.
-"""
-function shrink_and_convert(M::AM, intnum::U=Int64(1)) where {AM <: AbstractMatrix, U <: Integer}
-    inactive(x) = isnan(x) || ismissing(x) || x==0
-    active = .!inactive.(M)
-    M_out = shrink_to_active(M, active)
-    active = shrink_to_active(active, active)
-    return convert_population(M_out, active, intnum)
-end
-
 """
     function convert_population(
         initial_population,
-        active::AbstractMatrix{Bool},
         intnum::U = Int64(1)
     )
 
@@ -69,20 +63,20 @@ Convert population matrix to Int matrix by filling in the inactive area with 0 p
 and rounding the active area.
 """
 function convert_population(
-    initial_population::Matrix{<:Real},
-    active::AbstractMatrix{Bool},
+    initial_population::Matrix,
     intnum::U = Int64(1)
 ) where U <: Integer
+    active = .!_inactive.(initial_population)
     initial_population[.!active] .= 0
     initial_population = U.(round.(initial_population))
     return initial_population
 end
 
 function convert_population(
-    initial_population::AxisArray{<:Real, 2},
-    active::AbstractMatrix{Bool},
+    initial_population::AxisArray,
     intnum::U = Int64(1)
 ) where U <: Integer
+    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

From 0b86f234c6e9f976fd7a27c87028166ad7f44021 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 11:02:05 +0100
Subject: [PATCH 10/19] Fix setting initial population

---
 src/Epidemiology/EpiSystem.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/Epidemiology/EpiSystem.jl b/src/Epidemiology/EpiSystem.jl
index 5936cdf3..607a1a1b 100644
--- a/src/Epidemiology/EpiSystem.jl
+++ b/src/Epidemiology/EpiSystem.jl
@@ -77,10 +77,10 @@ function EpiSystem(epilist::EpiList, epienv::GridEpiEnv, rel::AbstractTraitRelat
         msg = "epilist has no Susceptible category. epilist.names = $(epilist.human.names)"
         throw(ArgumentError(msg))
     end
-    epi.abundances.grid[idx, :, :] .+= convert_population(initial_population, epienv.active, intnum)
     # Modify active cells based on new population
-    # Set any cells with zero population to inactive
-    epi.epienv.active .&= (epi.abundances.grid[idx, :, :] .!= 0)
+    epi.epienv.active .&= .!_inactive.(initial_population)
+    initial_population = convert_population(initial_population, intnum)
+    epi.abundances.grid[idx, :, :] .+= initial_population
     return epi
 end
 

From cac0dc554df1695c1cdd6555ad403ad179c67c73 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 11:02:16 +0100
Subject: [PATCH 11/19] Fix exports

---
 src/Simulation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Simulation.jl b/src/Simulation.jl
index 8608681d..879ada09 100644
--- a/src/Simulation.jl
+++ b/src/Simulation.jl
@@ -157,7 +157,7 @@ include("Epidemiology/Inference.jl")
 export SIR_wrapper, SIR_wrapper!
 
 include("Epidemiology/shrink.jl")
-export shrink_to_active, convert_population, shrink_and_convert
+export shrink_to_active, convert_population
 
 # Path into package
 path(paths...) = joinpath(@__DIR__, "..", paths...)

From eb56130d557e066ae987890a786930288ee67d92 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 11:38:04 +0100
Subject: [PATCH 12/19] Move tests to right places and add some more

---
 test/test_EpiEnv.jl    | 83 ------------------------------------------
 test/test_EpiSystem.jl | 76 ++++++++++++++++++++++++++++++++++++--
 test/test_shrink.jl    | 75 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 148 insertions(+), 86 deletions(-)
 create mode 100644 test/test_shrink.jl

diff --git a/test/test_EpiEnv.jl b/test/test_EpiEnv.jl
index 1ef0f751..4920b50b 100644
--- a/test/test_EpiEnv.jl
+++ b/test/test_EpiEnv.jl
@@ -21,26 +21,6 @@ abenv = simplehabitatAE(fillval, grid, area, control)
 @test all(abenv.active)
 @test abenv.habitat.size^2 * grid[1] * grid[2] == area
 
-@testset "Constructing from initial population" begin
-    fillval = 1.0
-    area = 25.0km^2
-    control = NoControl()
-    initial_pop = rand(10, 15) * 10
-    # Fill some NaNs, these should indicate inactive areas
-    missing_idx = CartesianIndex.([1, 4, 8], [5, 8, 2])
-    initial_pop[missing_idx] .= NaN
-
-    expected_active = Matrix{Bool}(.!isnan.(initial_pop))
-    expected_size = size(initial_pop)
-    expected_pop = round.(replace(initial_pop, NaN => 0))
-
-    epienv = simplehabitatAE(fillval, area, control, initial_pop)
-    @test all(epienv.habitat.matrix .== fillval)
-    @test size(epienv.habitat.matrix) == expected_size
-    @test epienv.active == expected_active
-    @test epienv.initial_population == expected_pop
-end
-
 @testset "Shrink grid" begin
     grid = (10, 10)
     area = 100.0km^2
@@ -71,30 +51,6 @@ end
     control = NoControl()
     expected_matrix = fill(fillval, expected_grid)
 
-    @testset "_shrink_to_active" begin
-        @testset "shrink a normal matrix" begin
-            M = rand(grid...)
-            M_shrunk = Simulation._shrink_to_active(M, active)
-            @test M_shrunk isa AxisArray
-            @test axisvalues(M_shrunk) == (4:9, 2:9)
-            @test size(M_shrunk) == expected_grid
-            @test M_shrunk == M[4:9, 2:9]
-        end
-        @testset "shrink an AxisArray matrix" begin
-            M = AxisArray(
-                rand(grid...);
-                x=Symbol.("x_", 1:grid[1]),
-                y=Symbol.("y_", 1:grid[2]),
-            )
-            M_shrunk = Simulation._shrink_to_active(M, active)
-            @test M_shrunk isa AxisArray
-            @test axisvalues(M_shrunk) == (Symbol.("x_", 4:9), Symbol.("y_", 2:9))
-            @test size(M_shrunk) == expected_grid
-            @test M_shrunk == M[4:9, 2:9]
-        end
-
-    end
-
     @testset "Construct directly" begin
         epienv = simplehabitatAE(fillval, grid, area, active, control)
 
@@ -104,43 +60,4 @@ end
         @test epienv.habitat.matrix == expected_matrix
         @test epienv.initial_population isa AxisArray
     end
-
-    @testset "Construct from initial population" begin
-        # Set an initial population of zeros
-        # The zeros should not be masked out (but the NaNs should)
-        @testset "Initial population is provided as a normal Matrix" begin
-            initial_pop = zeros(grid...)
-            initial_pop[.!active] .= NaN
-
-            epienv = simplehabitatAE(fillval, area, control, initial_pop)
-            expected_initial_pop = zeros(expected_grid)
-
-            @test epienv.active == expected_active
-            @test size(epienv.habitat.matrix) == expected_grid
-            @test epienv.habitat.size == expected_gridlength
-            @test epienv.habitat.matrix == expected_matrix
-            @test epienv.initial_population == expected_initial_pop
-            @test epienv.initial_population isa AxisArray
-            @test axisvalues(epienv.initial_population) == (4:9, 2:9)
-        end
-        @testset "Initial population is provided as a AxisArray Matrix" begin
-            initial_pop = AxisArray(
-                zeros(grid...);
-                x=Symbol.("x_", 1:grid[1]),
-                y=Symbol.("y_", 1:grid[2]),
-            )
-            initial_pop.data[.!active] .= NaN
-
-            epienv = simplehabitatAE(fillval, area, control, initial_pop)
-            expected_initial_pop = zeros(expected_grid)
-            @test epienv.active == expected_active
-            @test size(epienv.habitat.matrix) == expected_grid
-            @test epienv.habitat.size == expected_gridlength
-            @test epienv.habitat.matrix == expected_matrix
-            @test epienv.initial_population == expected_initial_pop
-            @test axisvalues(epienv.initial_population) ==
-                (Symbol.("x_", 4:9), Symbol.("y_", 2:9))
-        end
-
-    end
 end
diff --git a/test/test_EpiSystem.jl b/test/test_EpiSystem.jl
index 3bdf211f..ef8a2821 100644
--- a/test/test_EpiSystem.jl
+++ b/test/test_EpiSystem.jl
@@ -25,9 +25,79 @@ epi = TestEpiSystem()
 @test_nowarn getdispersalvar(epi, "Susceptible")
 
 @testset "EpiSystem from initial population" begin
-    initial_pop = rand(10, 10) * 10
-    epi = TestEpiSystemFromPopulation(initial_pop)
-    @test epi.abundances.grid[1, :, :] == round.(initial_pop)
+    @testset "All active" begin
+        initial_pop = rand(10, 10) * 10
+        epi = TestEpiSystemFromPopulation(initial_pop)
+        @test epi.abundances.grid[1, :, :] == round.(initial_pop)
+    end
+
+    @testset "Some inactive" begin
+        initial_pop = Matrix{Union{Float64, Missing}}(rand(10, 15) * 10)
+        # Fill some NaNs and missing, these should indicate inactive areas
+        missing_idx = CartesianIndex.([1, 4, 8], [5, 8, 2])
+        initial_pop[missing_idx[1]] .= NaN
+        initial_pop[missing_idx[2]] .= missing
+        initial_pop[missing_idx[3]] .= missing
+        initial_pop_ref = copy(initial_pop)
+
+        expected_active = fill(true, size(initial_pop))
+        expected_active[missing_idx] .= false
+        expected_size = size(initial_pop)
+        expected_pop = round.(replace(initial_pop, NaN => 0, missing => 0))
+
+        epi = TestEpiSystemFromPopulation(initial_pop)
+        @test size(epi.epienv.habitat.matrix) == expected_size
+        @test epi.epienv.active == expected_active
+        @test epi.abundances.grid[1, :, :] == expected_pop
+        # Should not have modified initial_pop
+        @test initial_pop == initial_pop_ref
+
+        # If we specify inactive regions when constructing the EpiEnv, these should be
+        # preserved
+        epienv_active = fill(true, size(initial_pop))
+        epienv_active[1, 1] = false
+        epienv_active[2, 2] = false
+        epienv_active[3, 3] = false
+        expected_active = copy(epienv_active)
+        expected_active[missing_idx] .= false
+        @test expected_active != epienv_active
+
+        epi = TestEpiSystemFromPopulation(initial_pop; epienv_active=epienv_active)
+        @test size(epi.epienv.habitat.matrix) == expected_size
+        @test epi.epienv.active == expected_active
+        @test epi.abundances.grid[1, :, :] == expected_pop
+        # Should not have modified initial_pop
+        @test initial_pop == initial_pop_ref
+    end
+
+    @testset "Initial population is provided as a AxisArray Matrix" begin
+        grid = (10, 10)
+        active = fill(true, grid)
+        # Set inactive cells on the outer layers
+        # These should not be trimmed off by the EpiSystem constructor
+        for row in (1, 2, 3, 10)
+            active[row, :] .= false
+        end
+        for col in (1, 2, 10)
+            active[:, col] .= false
+        end
+        initial_pop = AxisArray(
+            zeros(grid...);
+            x=Symbol.("x_", 1:grid[1]),
+            y=Symbol.("y_", 1:grid[2]),
+        )
+        initial_pop.data[.!active] .= NaN
+
+        expected_active = active
+        expected_initial_pop = Int.(zeros(grid))
+        expected_initial_pop[.!active] .= 0
+
+        epi = TestEpiSystemFromPopulation(initial_pop)
+
+        @test epi.epienv.active == expected_active
+        @test size(epi.epienv.habitat.matrix) == grid
+        @test epi.abundances.grid[1, :, :] == expected_initial_pop
+    end
 end
 
 @testset "Approximate equality" begin
diff --git a/test/test_shrink.jl b/test/test_shrink.jl
new file mode 100644
index 00000000..7874d78c
--- /dev/null
+++ b/test/test_shrink.jl
@@ -0,0 +1,75 @@
+using Simulation
+using Unitful.DefaultSymbols
+using Test
+using Simulation.Units
+using Simulation.ClimatePref
+using AxisArrays
+
+grid = (10, 10)
+area = 100.0km^2
+fillval = 1.0
+
+active = fill(true, grid)
+# Set inactive cells on the outer layers
+# Make this asymmetric to be more general
+# These should be trimmed off
+for row in (1, 2, 3, 10)
+    active[row, :] .= false
+end
+for col in (1, 2, 10)
+    active[:, col] .= false
+end
+# Set some more inactive cells, these shouldn't result in any further trimming
+active[4, 4] = false
+active[5, 4] = false
+active[5, 5] = false
+active[8, 8] = false
+# This should stop col 2 being trimmed
+active[5, 2] = true
+# Should be trimmed from 10x10 to 6x8
+expected_grid= (6, 8)
+expected_active = active[4:9, 2:9]
+# Shouldn't change
+expected_gridlength = 1km
+control = NoControl()
+expected_matrix = fill(fillval, expected_grid)
+
+@testset "_shrink_to_active" begin
+    @testset "shrink a normal matrix" begin
+        M = rand(grid...)
+        M_ref = copy(M)
+        M_shrunk = shrink_to_active(M, active)
+        @test M_shrunk isa AxisArray
+        @test axisvalues(M_shrunk) == (4:9, 2:9)
+        @test size(M_shrunk) == expected_grid
+        @test M_shrunk == M[4:9, 2:9]
+        @test M == M_ref
+    end
+    @testset "automatically identify inactive regions" begin
+        M = Matrix{Union{Float64, Missing}}(rand(grid...))
+        M[.!active] .= NaN
+        M[1, 1] = missing
+        M_ref = copy(M)
+        M_shrunk = shrink_to_active(M, active)
+        @test M_shrunk isa AxisArray
+        @test axisvalues(M_shrunk) == (4:9, 2:9)
+        @test size(M_shrunk) == expected_grid
+        # Should have preserved the NaNs
+        @test any(isnan.(M_shrunk))
+        @test isequal(M_shrunk, M[4:9, 2:9])
+        @test isequal(M, M_ref)
+    end
+    @testset "shrink an AxisArray matrix" begin
+        M = AxisArray(
+            rand(grid...);
+            x=Symbol.("x_", 1:grid[1]),
+            y=Symbol.("y_", 1:grid[2]),
+        )
+        M_ref = copy(M)
+        M_shrunk = shrink_to_active(M, active)
+        @test M_shrunk isa AxisArray
+        @test axisvalues(M_shrunk) == (Symbol.("x_", 4:9), Symbol.("y_", 2:9))
+        @test size(M_shrunk) == expected_grid
+        @test M_shrunk == M[4:9, 2:9]
+    end
+end

From 99bd0045c47fe979fde36a82530912ece8d9f8aa Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 11:45:25 +0100
Subject: [PATCH 13/19] Fix test

---
 test/test_EpiEnv.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/test_EpiEnv.jl b/test/test_EpiEnv.jl
index 4920b50b..de909733 100644
--- a/test/test_EpiEnv.jl
+++ b/test/test_EpiEnv.jl
@@ -58,6 +58,5 @@ abenv = simplehabitatAE(fillval, grid, area, control)
         @test size(epienv.habitat.matrix) == expected_grid
         @test epienv.habitat.size == expected_gridlength
         @test epienv.habitat.matrix == expected_matrix
-        @test epienv.initial_population isa AxisArray
     end
 end

From e2017b028e3c9fd8ce874d904d79a43314640abd Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 12:01:28 +0100
Subject: [PATCH 14/19] Don't modify input

---
 src/Epidemiology/shrink.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/Epidemiology/shrink.jl b/src/Epidemiology/shrink.jl
index 5f96c31f..fbec722f 100644
--- a/src/Epidemiology/shrink.jl
+++ b/src/Epidemiology/shrink.jl
@@ -66,6 +66,8 @@ 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))
@@ -76,6 +78,8 @@ 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

From 6c058a40afad480e2eb0acd998f06efd17c682ca Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 12:01:42 +0100
Subject: [PATCH 15/19] Test convert_population

---
 test/test_shrink.jl | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/test/test_shrink.jl b/test/test_shrink.jl
index 7874d78c..b0dae6fd 100644
--- a/test/test_shrink.jl
+++ b/test/test_shrink.jl
@@ -73,3 +73,21 @@ expected_matrix = fill(fillval, expected_grid)
         @test M_shrunk == M[4:9, 2:9]
     end
 end
+
+@testset "convert_population" begin
+    M = Matrix{Union{Float64, Missing}}(rand(grid...))
+    M[.!active] .= NaN
+    M[1, 1] = missing
+    M_ref = copy(M)
+
+    @testset "Integer type $U" for U in (Int64, UInt64)
+        M_expected = U.(round.(replace(M, NaN => 0, missing => 0)))
+
+        M_converted = convert_population(M, U(1))
+
+        @test M_converted isa Matrix{U}
+        @test M_converted == M_expected
+        # Should not modify M
+        @test isequal(M, M_ref)
+    end
+end

From db9ec2af5da85d9ecb513e089d24965cb359ffe9 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 12:02:02 +0100
Subject: [PATCH 16/19] Update test case

---
 test/TestCases.jl | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/test/TestCases.jl b/test/TestCases.jl
index 5d50ea0b..8693309f 100644
--- a/test/TestCases.jl
+++ b/test/TestCases.jl
@@ -78,7 +78,10 @@ function TestEpiSystem()
 
     return epi
 end
-function TestEpiSystemFromPopulation(initial_pop::Matrix{<:Real})
+function TestEpiSystemFromPopulation(
+    initial_pop::AbstractMatrix;
+    epienv_active=fill(true, size(initial_pop))
+)
     numclasses = 4
     numvirus = 2
     birth = [fill(1e-5/day, numclasses - 1); 0.0/day]
@@ -92,7 +95,7 @@ function TestEpiSystemFromPopulation(initial_pop::Matrix{<:Real})
     param = transition(param)
 
     area = 10.0km^2
-    epienv = simplehabitatAE(298.0K, size(initial_pop), area, NoControl())
+    epienv = simplehabitatAE(298.0K, size(initial_pop), area, epienv_active, NoControl())
 
     # Zero susceptible so we can test the specified initial_pop
     abun_h = (

From 13e3ba881c7d2112f54959d1e5804f0dfb38e969 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 12:12:42 +0100
Subject: [PATCH 17/19] Fix test

---
 test/test_EpiSystem.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/test_EpiSystem.jl b/test/test_EpiSystem.jl
index ef8a2821..9669ed84 100644
--- a/test/test_EpiSystem.jl
+++ b/test/test_EpiSystem.jl
@@ -50,7 +50,7 @@ epi = TestEpiSystem()
         @test epi.epienv.active == expected_active
         @test epi.abundances.grid[1, :, :] == expected_pop
         # Should not have modified initial_pop
-        @test initial_pop == initial_pop_ref
+        @test isequal(initial_pop, initial_pop_ref)
 
         # If we specify inactive regions when constructing the EpiEnv, these should be
         # preserved
@@ -67,7 +67,7 @@ epi = TestEpiSystem()
         @test epi.epienv.active == expected_active
         @test epi.abundances.grid[1, :, :] == expected_pop
         # Should not have modified initial_pop
-        @test initial_pop == initial_pop_ref
+        @test isequal(initial_pop, initial_pop_ref)
     end
 
     @testset "Initial population is provided as a AxisArray Matrix" begin

From e1ccfd37ba3de802337976f26e8ef435ce435a64 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 12:23:15 +0100
Subject: [PATCH 18/19] Update examples

---
 examples/Epidemiology/SEI3HRD_example.jl | 1 -
 examples/Epidemiology/Scotland_run.jl    | 5 +++--
 examples/Epidemiology/UK_run.jl          | 1 -
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/examples/Epidemiology/SEI3HRD_example.jl b/examples/Epidemiology/SEI3HRD_example.jl
index 4dd4a721..d9f5af44 100644
--- a/examples/Epidemiology/SEI3HRD_example.jl
+++ b/examples/Epidemiology/SEI3HRD_example.jl
@@ -44,7 +44,6 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
 
     # Read in population sizes for Scotland
     scotpop = Array{Float64, 2}(readfile(Simulation.path("test", "examples", "ScotlandDensity2011.tif"), 0.0, 7e5, 5e5, 1.25e6))
-    scotpop = shrink_and_convert(scotpop)
 
     # Set up simple gridded environment
     area = 525_000.0km^2
diff --git a/examples/Epidemiology/Scotland_run.jl b/examples/Epidemiology/Scotland_run.jl
index ddaf857b..125754fe 100644
--- a/examples/Epidemiology/Scotland_run.jl
+++ b/examples/Epidemiology/Scotland_run.jl
@@ -34,8 +34,9 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
     # Sum up age categories and turn into simple matrix
     total_pop = dropdims(sum(Float64.(scotpop), dims=3), dims=3)
     total_pop = AxisArray(total_pop, AxisArrays.axes(scotpop)[1], AxisArrays.axes(scotpop)[2])
-    # Shrink to smallest bounding box
-    total_pop = shrink_and_convert(total_pop);
+    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
diff --git a/examples/Epidemiology/UK_run.jl b/examples/Epidemiology/UK_run.jl
index 017871e8..88380586 100644
--- a/examples/Epidemiology/UK_run.jl
+++ b/examples/Epidemiology/UK_run.jl
@@ -52,7 +52,6 @@ function run_model(times::Unitful.Time, interval::Unitful.Time, timestep::Unitfu
                                        0.0, 7e5, 0, 1.25e6))
     # Coarsen grid to 10km
     ukpop = [sum(ukpop[i:i+9, j:j+9]) for i in 1:10:size(ukpop, 1), j in 1:10:size(ukpop, 2)]
-    ukpop = shrink_and_convert(ukpop)
 
     # Set up simple gridded environment
     area = 875_000.0km^2

From bbbbc6b584ef62dfe1416224a24f334d12ed47e6 Mon Sep 17 00:00:00 2001
From: Sean Lovett <sean.lovett@invenialabs.co.uk>
Date: Wed, 24 Jun 2020 12:59:09 +0100
Subject: [PATCH 19/19] Check initial_population is the right size

---
 src/Epidemiology/EpiSystem.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/Epidemiology/EpiSystem.jl b/src/Epidemiology/EpiSystem.jl
index 607a1a1b..b0cb7ca0 100644
--- a/src/Epidemiology/EpiSystem.jl
+++ b/src/Epidemiology/EpiSystem.jl
@@ -70,6 +70,11 @@ function EpiSystem(epilist::EpiList, epienv::GridEpiEnv, rel::AbstractTraitRelat
 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")