diff --git a/packages/mrc-ide-covidsim/src/imperial.ts b/packages/mrc-ide-covidsim/src/imperial.ts index 2e7d24c..6c5f0c2 100644 --- a/packages/mrc-ide-covidsim/src/imperial.ts +++ b/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) diff --git a/packages/mrc-ide-covidsim/test/integration/imperial-integration-test.ts b/packages/mrc-ide-covidsim/test/integration/imperial-integration-test.ts index be1a868..819661f 100644 --- a/packages/mrc-ide-covidsim/test/integration/imperial-integration-test.ts +++ b/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()