[WIP] start implementing contact matrix and age-specific interventions

src/algorithms/initialize.ts

@@ -9,6 +9,10 @@ const monthToDay = (m: number) => {
   return m * 30 + 15
 }
 
+export function sum(arr: number[]): number {
+  return arr.reduce((a, b) => a + b, 0)
+}
+
 const jan2020 = new Date('2020-01-01').valueOf() // time in ms
 
 export function withUncertainty(scenario: ScenarioFlat): boolean {
@@ -33,13 +37,15 @@ export const msPerDay = 1000 * 60 * 60 * 24
  */
 export function infectionRate(
   time: number,
-  avgInfectionRate: number,
+  infectedFractionByAge: number[],
+  contactMatrix: number[][],
   peakMonth: number,
   seasonalForcing: number,
-): number {
+): number[] {
   // this is super hacky
   const phase = ((time - jan2020) / msPerDay / 365 - monthToDay(peakMonth) / 365) * 2 * Math.PI
-  return avgInfectionRate * (1 + seasonalForcing * Math.cos(phase))
+  const baseValue = 1 + seasonalForcing * Math.cos(phase)
+  return contactMatrix.map((v) => baseValue * sum(v.map((d, i) => d * infectedFractionByAge[i])))
 }
 
 export function getPopulationParams(
@@ -88,6 +94,7 @@ export function getPopulationParams(
       isolated: [],
     },
     rate: {
+      contactMatrix: Object.keys(AgeGroup).map(() => Object.keys(AgeGroup).map(() => 1.0)),
       latency: 1 / latencyDays,
       infection: () => -Infinity, // Dummy infectionRate function. This is set below.
       recovery: [],
@@ -145,8 +152,8 @@ export function getPopulationParams(
     const avgInfectionRate = tmpR0 / infectiousPeriodDays
 
     const containment = containmentRealization.length > 1 ? containmentRealization[i] : containmentRealization[0]
-    elt.rate.infection = (time: number) =>
-      containment(time) * infectionRate(time, avgInfectionRate, peakMonth, seasonalForcing)
+    elt.rate.infection = (time: number, infectedFraction: number[]) =>
+      containment(time) * infectionRate(time, infectedFraction, elt.rate.contactMatrix, avgInfectionRate, peakMonth, seasonalForcing)
 
     return elt
   })

src/algorithms/model.ts

@@ -7,7 +7,7 @@ import {
   ExportedTimePoint,
 } from './types/Result.types'
 
-import { msPerDay } from './initialize'
+import { msPerDay, sum } from './initialize'
 
 const eulerStep = 0.5
 export const eulerStepsPerDay = Math.round(1 / eulerStep)
@@ -73,10 +73,6 @@ function stepODE(pop: SimulationTimePoint, P: ModelParams, dt: number): Simulati
   return state
 }
 
-export function sum(arr: number[]): number {
-  return arr.reduce((a, b) => a + b, 0)
-}
-
 export function gz(x: number): number {
   return x > 0 ? x : 0
 }
@@ -294,16 +290,17 @@ function fluxes(time: number, pop: SimulationTimePoint, P: ModelParams): StateFl
   }
 
   // Compute all fluxes (apart from overflow states) barring no hospital bed constraints
-  const fracInfected = sum(pop.current.infectious) / P.populationServed
+  const infectionPressure = P.rate.infection(
+    time,
+    pop.current.infectious.map((d) => d / P.populationServed),
+  )
 
   for (let age = 0; age < pop.current.infectious.length; age++) {
     // Initialize all multi-faceted states with internal arrays
     flux.exposed[age] = Array(pop.current.exposed[age].length)
 
     // Susceptible -> Exposed
-    flux.susceptible[age] =
-      P.importsPerDay[age] +
-      (1 - P.frac.isolated[age]) * P.rate.infection(time) * pop.current.susceptible[age] * fracInfected
+    flux.susceptible[age] = P.importsPerDay[age] + infectionPressure[age] * pop.current.susceptible[age]
 
     // Exposed -> Internal -> Infectious
     pop.current.exposed[age].forEach((exposed, i, exposedArray) => {

src/algorithms/types/Result.types.ts

@@ -56,8 +56,9 @@ export interface ModelFracs {
 }
 
 export interface ModelRates {
+  contactMatrix: number[][]
   latency: number
-  infection: (t: number) => number
+  infection: (t: number) => number[]
   recovery: number[]
   severe: number[]
   discharge: number[]