Merge pull request #20 from covid-modeling/adityasharad/covidsim-regions

CovidSim connector: Refactor regional parameter file lookup

packages/mrc-ide-covidsim/src/imperial.ts

@@ -45,63 +45,91 @@ export class ImperialModel implements Model {
     this.threadCount = threadCount
   }
 
+  /** Gets the path to the administrative units parameter file for the given region. */
+  private getAdminPath(region: string): string {
+    if (region === 'US') {
+      return path.join(this.dataDir, 'admin_units', 'United_States_admin.txt')
+    } else if (COUNTRY_PARAMS_BY_ISO_CODE[region]) {
+      const { adminFileName } = COUNTRY_PARAMS_BY_ISO_CODE[region]
+      return path.join(this.dataDir, 'admin_units', adminFileName)
+    } else {
+      throw new Error(`Could not find admin file for region ${region}`)
+    }
+  }
+
+  /**
+   * Gets the path to the population density parameter file for the given region.
+   * Europe is used as the default.
+   */
+  private getPopulationDensityPath(region: string, subregion?: string): string {
+    let populationDensityFileName: string
+    if (
+      ['AS', 'GU', 'PR', 'VI'].includes(region) ||
+      (region === 'US' &&
+        ['US-AK', 'US-HI', 'US-AS', 'US-GU', 'US-PR', 'US-VI'].includes(
+          subregion
+        ))
+    ) {
+      populationDensityFileName = 'wpop_us_terr.txt'
+    } else if (['US', 'CA'].includes(region)) {
+      populationDensityFileName = 'wpop_usacan.txt'
+    } else {
+      populationDensityFileName = 'wpop_eur.txt'
+    }
+    return path.join(this.dataDir, 'populations', populationDensityFileName)
+  }
+
+  /**
+   * Gets the path to the pre-parameters template file for the given region.
+   * The UK is used as the default for known regions.
+   */
+  private getPreParametersTemplatePath(region: string): string {
+    let preParamsFileName: string
+    if (region === 'US') {
+      preParamsFileName = 'preUS_R0=2.0.txt'
+    } else if (COUNTRY_PARAMS_BY_ISO_CODE[region]) {
+      preParamsFileName =
+        COUNTRY_PARAMS_BY_ISO_CODE[region].preParamsFileName ??
+        'preUK_R0=2.0.txt'
+    } else {
+      throw new Error(
+        `Could not find pre-parameters template file for region ${region}`
+      )
+    }
+    return path.join(this.dataDir, 'param_files', preParamsFileName)
+  }
+
+  private getSubregionName(region: string, subregion: string): string {
+    if (region === 'US') {
+      return US_SUBREGIONS[subregion]
+    } else if (COUNTRY_PARAMS_BY_ISO_CODE[region]) {
+      const { subregions } = COUNTRY_PARAMS_BY_ISO_CODE[region]
+      return subregions[subregion]
+    } else {
+      throw new Error(`Could not find subregions for region ${region}`)
+    }
+  }
+
   inputs(input: input.ModelInput): ImperialRunnerModelInput {
     const inputFiles = []
-    // Select the correct executable and static inputs based on the region.
-    let adminPath,
-      populationDensityPath,
-      preParametersTemplatePath,
-      subregionName
 
     const modelPath = path.join(this.binDir, 'CovidSim')
 
+    // Select the correct static inputs based on the region.
+    let adminPath = this.getAdminPath(input.region)
+    const populationDensityPath = this.getPopulationDensityPath(
+      input.region,
+      input.subregion
+    )
     const parametersTemplatePath = path.join(
       this.dataDir,
       'param_files',
       'p_NoInt.txt'
     )
-
-    // The US has its own executable, population density file, and pre-params file.
-    if (input.region === 'US') {
-      adminPath = path.join(
-        this.dataDir,
-        'admin_units',
-        'United_States_admin.txt'
-      )
-      populationDensityPath = path.join(
-        this.dataDir,
-        'populations',
-        'wpop_usacan.txt'
-      )
-      preParametersTemplatePath = path.join(
-        this.dataDir,
-        'param_files',
-        'preUS_R0=2.0.txt'
-      )
-      subregionName = US_SUBREGIONS[input.subregion]
-    } else if (COUNTRY_PARAMS_BY_ISO_CODE[input.region]) {
-      const { adminFileName, subregions } = COUNTRY_PARAMS_BY_ISO_CODE[
-        input.region
-      ]
-
-      adminPath = path.join(this.dataDir, 'admin_units', adminFileName)
-      populationDensityPath = path.join(
-        this.dataDir,
-        'populations',
-        'wpop_eur.txt'
-      )
-
-      const preParamsFileName =
-        COUNTRY_PARAMS_BY_ISO_CODE[input.region].preParamsFileName ??
-        'preUK_R0=2.0.txt'
-
-      preParametersTemplatePath = path.join(
-        this.dataDir,
-        'param_files',
-        preParamsFileName
-      )
-      subregionName = subregions[input.subregion]
-    }
+    const preParametersTemplatePath = this.getPreParametersTemplatePath(
+      input.region
+    )
+    const subregionName = this.getSubregionName(input.region, input.subregion)
 
     // Generate the intervention-related pre-parameters based on the input.
     inputFiles.push(preParametersTemplatePath)

packages/mrc-ide-covidsim/test/integration/imperial-integration-test.ts

@@ -1,11 +1,100 @@
-import * as temp from 'temp'
 import * as fs from 'fs'
+import * as path from 'path'
+import * as temp from 'temp'
 import { assert } from 'chai'
 import { ImperialModel } from '../../src/imperial'
 import { BIN_DIR, MODEL_DATA_DIR } from '../../src/config'
 import { input } from '@covid-modeling/api'
 
 suite('imperial integration', () => {
+  const testRegions = [
+    {
+      region: 'CA',
+      subregion: undefined,
+      expectedAdminPath: 'Canada_admin.txt',
+      expectedPopulationDensityPath: 'wpop_usacan.txt',
+      expectedPreParameterPath: 'preUK_R0=2.0.txt',
+    },
+    {
+      region: 'NG',
+      subregion: undefined,
+      expectedAdminPath: 'Nigeria_admin.txt',
+      expectedPopulationDensityPath: 'wpop_eur.txt',
+      expectedPreParameterPath: 'preNG_R0=2.0.txt',
+    },
+    {
+      region: 'RU',
+      subregion: undefined,
+      expectedAdminPath: 'Russia_admin.txt',
+      expectedPopulationDensityPath: 'wpop_eur.txt',
+      expectedPreParameterPath: 'preUK_R0=2.0.txt',
+    },
+    {
+      region: 'GB',
+      subregion: undefined,
+      expectedAdminPath: 'United_Kingdom_admin.txt',
+      expectedPopulationDensityPath: 'wpop_eur.txt',
+      expectedPreParameterPath: 'preUK_R0=2.0.txt',
+    },
+    {
+      region: 'US',
+      subregion: undefined,
+      expectedAdminPath: 'United_States_admin.txt',
+      expectedPopulationDensityPath: 'wpop_usacan.txt',
+      expectedPreParameterPath: 'preUS_R0=2.0.txt',
+    },
+    {
+      region: 'US',
+      subregion: 'US-NY',
+      expectedAdminPath: 'admin-params.txt',
+      expectedPopulationDensityPath: 'wpop_usacan.txt',
+      expectedPreParameterPath: 'preUS_R0=2.0.txt',
+    },
+  ]
+  testRegions.forEach(testRegion => {
+    test(`finds parameter files for region ${testRegion.region} and subregion ${testRegion.subregion}`, () => {
+      const modelInput: input.ModelInput = {
+        region: testRegion.region,
+        subregion: testRegion.subregion,
+        parameters: {
+          calibrationDate: '2020-03-20',
+          calibrationCaseCount: 500,
+          calibrationDeathCount: 120,
+          r0: null,
+          interventionPeriods: [],
+        },
+      }
+
+      const logDir = temp.mkdirSync()
+      const inputDir = temp.mkdirSync()
+      const outputDir = temp.mkdirSync()
+      const binDir = temp.mkdirSync()
+
+      const model = new ImperialModel(
+        1,
+        binDir,
+        logDir,
+        MODEL_DATA_DIR,
+        inputDir,
+        outputDir
+      )
+
+      const runnerModelInput = model.inputs(modelInput)
+      assert.equal(
+        runnerModelInput.adminFilePath,
+        path.join(inputDir, testRegion.expectedAdminPath)
+      )
+      assert.equal(
+        runnerModelInput.populationDensityFilePath,
+        path.join(
+          MODEL_DATA_DIR,
+          'populations',
+          testRegion.expectedPopulationDensityPath
+        )
+      )
+    })
+  })
+
   test('run imperial model for Wyoming', async () => {
     const logDir = temp.mkdirSync()
     const inputDir = temp.mkdirSync()