diff --git a/examples/Epidemiology/Scotland_run.jl b/examples/Epidemiology/Scotland_run.jl
index 8464c2a9..ced1d994 100644
--- a/examples/Epidemiology/Scotland_run.jl
+++ b/examples/Epidemiology/Scotland_run.jl
@@ -26,14 +26,51 @@ area = (AxisArrays.axes(scotpop, 1)[end] + AxisArrays.axes(scotpop, 1)[2] -
     (AxisArrays.axes(scotpop, 2)[end] + AxisArrays.axes(scotpop, 2)[2] -
     2 * AxisArrays.axes(scotpop, 2)[1]) * 1.0
 
-# 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
+# Set population to initially have no individuals
+abun_h = (
+    Susceptible = fill(0, age_categories),
+    Exposed = fill(0, age_categories),
+    Asymptomatic = fill(0, age_categories),
+    Presymptomatic = fill(0, age_categories),
+    Symptomatic = fill(0, age_categories),
+    Hospitalised = fill(0, age_categories),
+    Recovered = fill(0, age_categories),
+    Dead = fill(0, age_categories)
+)
+disease_classes = (
+    susceptible = ["Susceptible"],
+    infectious = ["Asymptomatic", "Presymptomatic", "Symptomatic"]
+)
+abun_v = (Virus = 0,)
+
+disease_axis = Axis{:class}(collect(keys(abun_h)))
+compartment_axes = (AxisArrays.axes(scotpop)..., disease_axis)
+# Populate susceptibles according to actual population spread
+initial_pop = cat(
+    scotpop,
+    zeros((size(scotpop)..., length(disease_axis)-1)),
+    dims=ndims(scotpop)+1
+)
+initial_pop = AxisArray(initial_pop, compartment_axes);
+# To be able to easily reshape to match abundances.matrix, the spatial axes have to be at
+# the end
+initial_pop = permutedims(initial_pop, [3, 4, 1, 2])
+
+susc = @view initial_pop[class=:Susceptible]
+exposed = @view initial_pop[class=:Exposed]
+# Weight samples by number of susceptibles
+samples = sample(CartesianIndices(susc), weights(1.0 .* vec(susc)), 100)
+for samp in samples
+    susc[samp] > 0 || continue
+    # Add to exposed
+    exposed[samp] += 1
+    # Remove from susceptible
+    susc[samp] -= 1
+end
 
 # Set simulation parameters
-numclasses = 8
-numvirus = 1
+numclasses = length(abun_h)
+numvirus = length(abun_v)
 birth_rates = fill(0.0/day, numclasses, age_categories)
 death_rates = fill(0.0/day, numclasses, age_categories)
 birth_rates[:, 2:4] .= uconvert(day^-1, 1/20years)
@@ -69,26 +106,7 @@ param = SEI3HRDGrowth(birth_rates, death_rates, ageing,
                       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)
-
-# Set population to initially have no individuals
-abun_h = (
-    Susceptible = fill(0, age_categories),
-    Exposed = fill(0, age_categories),
-    Asymptomatic = fill(0, age_categories),
-    Presymptomatic = fill(0, age_categories),
-    Symptomatic = fill(0, age_categories),
-    Hospitalised = fill(0, age_categories),
-    Recovered = fill(0, age_categories),
-    Dead = fill(0, age_categories)
-)
-
-disease_classes = (
-    susceptible = ["Susceptible"],
-    infectious = ["Asymptomatic", "Presymptomatic", "Symptomatic"]
-)
-
-abun_v = (Virus = 0,)
+epienv = simplehabitatAE(298.0K, area, NoControl(), initial_pop)
 
 # Dispersal kernels for virus and disease classes
 dispersal_dists = fill(1.0km, numclasses * age_categories)
@@ -106,52 +124,15 @@ rel = Gauss{eltype(epienv.habitat)}()
 # Create epi system with all information
 epi = EpiSystem(epilist, epienv, rel)
 
-# Populate susceptibles according to actual population spread
-reshaped_pop =
-    reshape(scotpop[1:size(epienv.active, 1), 1:size(epienv.active, 2), :],
-            size(epienv.active, 1) * size(epienv.active, 2), size(scotpop, 3))'
-epi.abundances.matrix[cat_idx[:, 1], :] = reshaped_pop
-
-# Add in initial infections randomly (samples weighted by population size)
-# Define generator for all pair age x cell
-age_and_cells = Iterators.product(1:age_categories,
-                                  1:size(epi.abundances.matrix, 2))
-# Take all susceptibles of each age per cell
-pop_weights = epi.abundances.matrix[vcat(cat_idx[:, 1]...), :]
-# It would be nice if it wasn't necessary to call collect here
-N_cells = size(epi.abundances.matrix, 2)
-samp = sample(collect(age_and_cells), weights(1.0 .* vec(pop_weights)), 100)
-age_ids = getfield.(samp, 1)
-cell_ids = getfield.(samp, 2)
-
-for i in eachindex(age_ids)
-    if (epi.abundances.matrix[cat_idx[age_ids[i], 1], cell_ids[i]] > 0)
-        # Add to exposed
-        epi.abundances.matrix[cat_idx[age_ids[i], 2], cell_ids[i]] += 1
-        # Remove from susceptible
-        epi.abundances.matrix[cat_idx[age_ids[i], 1], cell_ids[i]] -= 1
-    end
-end
-
 # Run simulation
 times = 2months; interval = 1day; timestep = 1day
-abuns = zeros(Int64, size(epi.abundances.matrix, 1), N_cells,
-              floor(Int, times/timestep) + 1)
+Nsteps = floor(Int, times/timestep) + 1
+abuns = zeros(Int64, size(epi.abundances.matrix)..., Nsteps)
 @time simulate_record!(abuns, epi, times, interval, timestep)
 
 if do_plot
     using Plots
     # View summed SIR dynamics for whole area
-    category_map = (
-        "Susceptible" => cat_idx[:, 1],
-        "Exposed" => cat_idx[:, 2],
-        "Asymptomatic" => cat_idx[:, 3],
-        "Presymptomatic" => cat_idx[:, 4],
-        "Symptomatic" => cat_idx[:, 5],
-        "Hospital" => cat_idx[:, 6],
-        "Recovered" => cat_idx[:, 7],
-        "Deaths" => cat_idx[:, 8],
-    )
-    display(plot_epidynamics(epi, abuns, category_map = category_map))
+    display(plot_epidynamics(epi, abuns, classes=keys(abun_h)))
     display(plot_epiheatmaps(epi, abuns, steps = [21]))
 end
diff --git a/src/Epidemiology/EpiEnv.jl b/src/Epidemiology/EpiEnv.jl
index 0a2f4290..3792a9d6 100644
--- a/src/Epidemiology/EpiEnv.jl
+++ b/src/Epidemiology/EpiEnv.jl
@@ -25,7 +25,7 @@ mutable struct GridEpiEnv{H, C} <: AbstractEpiEnv{H, C}
     habitat::H
     active::AbstractMatrix{Bool}
     control::C
-    initial_population::AbstractMatrix{Int}
+    initial_population::AbstractArray{Int}
     names::Vector{String}
     function (::Type{GridEpiEnv{H, C}})(
         habitat::H,
@@ -42,9 +42,9 @@ mutable struct GridEpiEnv{H, C} <: AbstractEpiEnv{H, C}
                 "size(active)=$(size(active))"
             ))
         end
-        if size(initial_population) != size(active)
+        if size(initial_population)[end-1:end] != size(active)
             throw(DimensionMismatch(
-                "size(initial_population)=$(size(initial_population)) != " *
+                "size(initial_population)[end-1:end]=$(size(initial_population)[end-1:end]) != " *
                 "size(active)=$(size(active))"
             ))
         end
@@ -112,7 +112,7 @@ end
         area::Unitful.Area{Float64},
         active::AbstractMatrix{Bool},
         control::C,
-        initial_population::AbstractMatrix{<:Integer}=zeros(Int, dimension),
+        initial_population::AbstractArray{<:Integer}=zeros(Int, dimension),
     )
 
 Function to create a simple `ContinuousHab` type epi environment. It creates a
@@ -129,7 +129,7 @@ function simplehabitatAE(
     area::Unitful.Area{Float64},
     active::AbstractMatrix{Bool},
     control::C,
-    initial_population::AbstractMatrix{<:Integer}=zeros(Int, dimension),
+    initial_population::AbstractArray{<:Integer}=zeros(Int, dimension),
 ) where C <: AbstractControl
     if typeof(val) <: Unitful.Temperature
         val = uconvert(K, val)
@@ -139,9 +139,10 @@ 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)
-    dimension = size(active)
+    # TODO fix this
+    #initial_population = _shrink_to_active(initial_population, 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)
@@ -162,7 +163,7 @@ end
         val::Union{Float64, Unitful.Quantity{Float64}},
         area::Unitful.Area{Float64},
         control::C,
-        initial_population::AbstractMatrix{<:Real},
+        initial_population::AbstractArray{<:Real},
     )
 
 Create a simple `ContinuousHab` type epi environment from a specified `initial_population`
@@ -184,15 +185,20 @@ function simplehabitatAE(
     val::Union{Float64, Unitful.Quantity{Float64}},
     area::Unitful.Area{Float64},
     control::C,
-    initial_population::AbstractMatrix{<:Real},
+    initial_population::AbstractArray{<:Real},
 ) where C <: AbstractControl
-    inactive(x) = isnan(x) || ismissing(x)
+    inactive(x) = isnan(x) || ismissing(x) || x == 0
     if all(inactive.(initial_population))
-        throw(ArgumentError("initial_population is all NaN / missing"))
+        throw(ArgumentError("initial_population is all NaN / missing / 0"))
     end
-    dimension = size(initial_population)
-    active = Matrix{Bool}(.!inactive.(initial_population))
+    active = Array{Bool}(.!inactive.(initial_population))
     initial_population = _convert_population(initial_population, active)
+    # reduce to 2D matrix
+    # active cells have at least one active compartment
+    reducedims = Tuple(1:ndims(active)-2)
+    active = reduce(|, active, dims=reducedims)
+    active = dropdims(active, dims=reducedims)
+    dimension = size(active)
     return simplehabitatAE(val, dimension, area, active, control, initial_population)
 end
 
@@ -204,7 +210,7 @@ and rounding the active area.
 """
 function _convert_population(
     initial_population::Matrix{<:Real},
-    active::AbstractMatrix{Bool}
+    active::AbstractArray{Bool}
 )::Matrix{<:Int}
     initial_population[.!active] .= 0
     initial_population = Int.(round.(initial_population))
@@ -212,9 +218,9 @@ function _convert_population(
 end
 
 function _convert_population(
-    initial_population::AxisArray{<:Real, 2},
-    active::AbstractMatrix{Bool}
-)::AxisArray{<:Int, 2}
+    initial_population::AxisArray{<:Real},
+    active::AbstractArray{Bool}
+)::AxisArray{<:Int}
     # 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
diff --git a/src/Epidemiology/EpiHelper.jl b/src/Epidemiology/EpiHelper.jl
index 0babc488..e3d0596b 100644
--- a/src/Epidemiology/EpiHelper.jl
+++ b/src/Epidemiology/EpiHelper.jl
@@ -93,7 +93,7 @@ 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)
+  grid_id = map(Iterators.product(axisvalues(epi.epienv.initial_population)[end-1:end]...)) do (x,y)
       return string.(x, "-", y)
   end
   # TODO: confirm converting `grid_id` from matrix to vector in the way below gives the
diff --git a/src/Epidemiology/EpiLandscape.jl b/src/Epidemiology/EpiLandscape.jl
index d30fece1..25b1015c 100644
--- a/src/Epidemiology/EpiLandscape.jl
+++ b/src/Epidemiology/EpiLandscape.jl
@@ -7,13 +7,17 @@ Disease class abundances housed in the landscape. These are represented in both
 mutable struct EpiLandscape
   matrix::Matrix{Int64}
   matrix_v::Matrix{Int64}
-  grid::Array{Int64, 3}
+  grid::AxisArray{Int64}
   seed::Vector{MersenneTwister}
-  function EpiLandscape(human_abun::Matrix{Int64}, virus_abun::Matrix{Int64}, d1::Tuple)
-    return new(human_abun, virus_abun, reshape(human_abun, d1), [MersenneTwister(rand(UInt)) for _ in 1:Threads.nthreads()])
+  function EpiLandscape(human_abun::Matrix{Int64}, virus_abun::Matrix{Int64}, ax)
+    d = length.(ax)
+    grid = AxisArray(reshape(human_abun, d), ax)
+    return new(human_abun, virus_abun, grid, [MersenneTwister(rand(UInt)) for _ in 1:Threads.nthreads()])
   end
-  function EpiLandscape(human_abun::Matrix{Int64}, virus_abun::Matrix{Int64}, d1::Tuple, d2::Tuple, seed::Vector{MersenneTwister})
-    return new(human_abun, virus_abun, reshape(human_abun, d1), seed)
+  function EpiLandscape(human_abun::Matrix{Int64}, virus_abun::Matrix{Int64}, ax, d2::Tuple, seed::Vector{MersenneTwister})
+    d = length.(ax)
+    grid = AxisArray(reshape(human_abun, d), ax)
+    return new(human_abun, virus_abun, grid, seed)
   end
 end
 
@@ -38,6 +42,6 @@ EpiList.
 function emptyepilandscape(epienv::GridEpiEnv, epilist::EpiList)
   mat_human = zeros(Int64, counttypes(epilist.human, true), countsubcommunities(epienv))
   mat_virus = zeros(Int64, counttypes(epilist.virus, true), countsubcommunities(epienv))
-  dimension = (counttypes(epilist.human, true), _getdimension(epienv.habitat)...)
-  return EpiLandscape(mat_human, mat_virus, dimension)
+  ax = AxisArrays.axes(epienv.initial_population)
+  return EpiLandscape(mat_human, mat_virus, ax)
 end
diff --git a/src/Epidemiology/EpiPlots.jl b/src/Epidemiology/EpiPlots.jl
index f580ca6c..a14a555c 100644
--- a/src/Epidemiology/EpiPlots.jl
+++ b/src/Epidemiology/EpiPlots.jl
@@ -1,9 +1,3 @@
-function _compartment_idx(compartment, names)
-    idx = findfirst(names .== compartment)
-    isnothing(idx) && throw(ArgumentError("Compartment $compartment not in $names"))
-    return idx
-end
-
 function _default_steps(abuns)
     N_steps = size(abuns, 3)
     return Int.(floor.(range(1, N_steps; length=4)))
@@ -47,12 +41,11 @@ function _check_args(h)
 end
 @recipe function f(
     h::Plot_EpiHeatmaps;
-    compartment="Exposed",
+    class=:Exposed,
     steps=[],
 )
     _check_args(h)
     epi, abuns = h.args
-    idx = _compartment_idx(compartment, epi.epilist.human.names)
     if isempty(steps)
         steps = _default_steps(abuns)
     end
@@ -63,10 +56,18 @@ end
     match_dimensions --> true
     seriescolor --> :heat
 
+    Nsteps = size(abuns)[end]
+    ax = (AxisArrays.axes(epi.epienv.initial_population)..., Axis{:step}(1:Nsteps))
+    abuns_grid = AxisArray(reshape(abuns, length.(ax)), ax)
+
     subplot = 1
     gridsize = size(epi.epienv.habitat.matrix)
     for step in steps
-        data = Float64.(reshape(abuns[idx, :, step], gridsize...))
+        data = abuns_grid[class=class, step=step].data
+        reducedims = Tuple(1:ndims(data)-2)
+        data = sum(data, dims=reducedims)
+        data = dropdims(data, dims=reducedims)
+        data = Matrix{Any}(data)
         data[.!epi.epienv.active] .= NaN
         x = 1:size(data, 2)
         y = 1:size(data, 1)
@@ -75,7 +76,7 @@ end
         end
         @series begin
             seriestype := :heatmap
-            title := "Step $step ($compartment)"
+            title := "Step $step ($class)"
             subplot := subplot
             background_color --> :lightblue
             background_color_outside --> :white
@@ -116,22 +117,23 @@ Plot the dynamics of `abuns` summed over space, as a function of time.
 """
 plot_epidynamics
 @userplot Plot_EpiDynamics
-@recipe function f(h::Plot_EpiDynamics; category_map=nothing)
+@recipe function f(h::Plot_EpiDynamics; classes=[:Susceptible])
     _check_args(h)
     epi, abuns = h.args
 
-    if isnothing(category_map)
-        # Make each compartment its own category
-        category_map = (name => [idx] for (idx, name) in enumerate(epi.epilist.human.names))
-    end
+    Nsteps = size(abuns)[end]
+    ax = (AxisArrays.axes(epi.epienv.initial_population)..., Axis{:step}(1:Nsteps))
+    abuns_grid = AxisArray(reshape(abuns, length.(ax)), ax)
 
-    for (name, idx) in category_map
-        data = vec(mapslices(sum, abuns[idx, :, :], dims = (1, 2)))
+    for class in classes
+        # Need .data here due to https://github.com/JuliaArrays/AxisArrays.jl/issues/113
+        data = abuns_grid[class=class].data
+        data = vec(mapslices(sum, data, dims = Tuple(1:ndims(data)-1)))
         title --> "Infection dynamics"
         xguide --> "Step"
         yguide --> "Totals"
         @series begin
-            label := name
+            label := class
             data
         end
     end
diff --git a/src/Epidemiology/EpiSystem.jl b/src/Epidemiology/EpiSystem.jl
index 71ab50d9..b0a80d4d 100644
--- a/src/Epidemiology/EpiSystem.jl
+++ b/src/Epidemiology/EpiSystem.jl
@@ -68,12 +68,7 @@ end
 function EpiSystem(epilist::EpiList, epienv::GridEpiEnv, rel::AbstractTraitRelationship)
     epi = EpiSystem(populate!, epilist, epienv, rel)
     # 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, :, :] .+= epienv.initial_population
+    epi.abundances.grid .+= epienv.initial_population
     return epi
 end