Introduce Axis Arrays and/or NamedTuples to track quantities.

[AlexRobson]

There are several points in the code that refer to classes that have hard coded column indexes. Resolving this issue will mean that e.g. abundancy matricies can be referred to by name, rather than index.

Example: from test_SEI2HRD.jl:

idx_sus = 2
idx_rec = 7
idx_dead = 8

# Test susceptible population decreasing or constant only [Source]
@test all(diff(vec(sum(abuns[idx_sus, :, :], dims = 1))) .<= 0)
@test sum(abuns[idx_sus, :, 1]) == initial_pops.susceptible

idx_sus can be a Axis name e.g. idx_sus = :susceptible.

[github-actions[bot]]

Heads up @claireh93 @richardreeve @jwscook @sdl1 @wytbella - the "Simulation.jl" label was applied to this issue.

[richardreeve]

If we're doing a bigger fix, then we need to think more broadly. Another alternative would be to have Enums for each disease type. I don't know whether there is a lookup cost when AxisArrays use symbols to index arrays, but that isn't true of enums, because they happen at compile time I think.

[claireh93]

It would be great to have a way of doing this so we can easily select out all of a certain disease type, e.g. when each has multiple age categories, too. Would that work with AxisArrays/Enums? I'm not sure I know how enums work.

[richardreeve]

Enums are basically nicely typed integers with names:

ulia> @enum SIR SIRsusceptible SIRinfected SIRrecovered

julia> instances(SIR)
(SIRsusceptible, SIRinfected, SIRrecovered)

julia> SIRsusceptible
SIRsusceptible::SIR = 0

julia> Int.(instances(SIR))
(0, 1, 2)

We could use their associated integers in the small matrices to index within a demographic group, because I think we don't have a "big" matrix any longer after the discussion last week. AxisArrays would actually provide the index directly, but I just don't know how they work with named indices.

[eperim]

I was actually planning on using AxisArrays in the sketch of the abundance matrix with all risk factors for the formulae issue (while NamedTuples are already there as a proposal for the risk factors in the EpiList). The advantage of this is that it makes it really easy for humans to read abundances.

I don't know whether there is a lookup cost when AxisArrays

@iamed2 probably knows the answer for this.

[richardreeve]

It's obviously important to know about any potential speed costs, but depending on what progress you've made, we should probably anyway have a chat about this whole issue either at the meeting tomorrow or at a separate meeting if there's not sufficient time...

[iamed2]

Lookup by index values (e.g., having names for each column rather than names for the dimension) would be done by a package like NamedArrays. I'm not sure what the cost of those lookups are.

AxisArrays and NamedDims have named dimensions, and that lookup is compile-time.

[richardreeve]

Ah, that makes sense - and it may be an issue then... we need to work this out as we move to this formula-based scheme for labelled human compartments (or whatever it turns out to be).

[sdl1]

My initial plan with this is to make EpiLandscape.grid a multidimensional AxisArray view of the underlying 2D EpiLandscape.matrix. The axes (apart from the disease classes) are already present in scotpop so can just feed through. This should allow indexing operations like

abundances.grid[class=:susceptible, age=10year] # Get a 2D spatial matrix

Also should allow things like

sum(abundances.grid; dims=Axis{:age}) # -> 3D matrix of (class, grid_x, grid_y)

although unfortunately it appears this particular functionality is broken at the moment with Unitful axes, https://github.com/JuliaArrays/AxisArrays.jl/issues/113

If that sounds good, I'll implement this and see how it looks, and we can decide if we want to keep it / extend it

[richardreeve]

Gah. I wasted ages trying to work out what I was doing wrong in Scotland_run.jl at https://github.com/ScottishCovidResponse/Simulation.jl/blob/dev/examples/Epidemiology/Scotland_run.jl#L30 and eventually gave up and converted it to an Array and back to sum over a dimension. It never occurred to me that it was a problem in AxisArrays. I don't suppose there's any chance of getting someone to look at it?

[richardreeve]

I'm still uncertain whether this will have a cost associated with it (and if so, how big), but it would be excellent - a really clear way of seeing what's going on. I would go ahead and give it a go.