Reff e2e and unit tests passing

.gitignore

@@ -1,4 +1,3 @@
 node_modules
 *.swp
-*.json
 .DS_Store

README.md

@@ -105,7 +105,7 @@ The model is parameterised with a country specific `beta` value.
 To translate Rt values into beta values, you can divide them by the country
 specific eigenvalue, e.g:
 
-```
+```js
 import nigeriaData from './data/NGA.json';
 const r0 = 3;
 const rt = [r0, r0/2];
@@ -114,8 +114,33 @@ const beta = rt.map(r => { return r / nigeriaData.eigenvalue });
 
 You can translate back to Rt by multiplying by the eigenvalue:
 
-```
+```js
 import nigeriaData from './data/NGA.json';
 const beta = [nigeriaData.beta, nigeriaData.beta/2];
 const rt = beta.map(r => { return r * nigeriaData.eigenvalue });
 ```
+
+#### Reff
+
+We provide function to estimate Reff from model outputs:
+
+```js
+function reff(
+  output,
+  rt,
+  beta,
+  population,
+  mixingMatrix
+)
+```
+
+Parameters:
+
+ * output - the `y` component of the model output returned from runModel
+ * rt - an array containing Rt for the first timesteps of the output
+ * beta - an array with corresponding beta values for rt
+ * population - an array with population counts for each age group (called
+   `population` in the country json files)
+ * mixingMatrix - a 2D array representing an age scaled mixing matrix (called `contactMatrixScaledAge` in the country json files)
+
+This returns an array of Reff values for the corresponding rt and beta arguments

e2e/test.js

@@ -26,7 +26,7 @@ async function load() {
 async function test() {
   let scenario = 0;
   let failed = false;
-  browser.evaluate("let p; let r;");
+  browser.evaluate("let p; let r; let output;");
   for (const country of [ 'LCA', 'NGA', 'IND' ]) {
     for (const bed of [ 100, 100000, 100000000 ]) {
       for (const R0 of [ 4, 3, 2, 1 ]) {
@@ -44,7 +44,7 @@ async function test() {
         ).prob_severe_death_treatment;
 
         let actual = browser.evaluate(
-          `runModel(
+          `output = runModel(
             ${country}.population,
             ${country}.contactMatrix,
             [0, 50],
@@ -90,6 +90,26 @@ async function test() {
         } else {
           console.log('passed');
         }
+
+
+        /*
+         * Basic test of Reff
+        */
+        const reff = browser.evaluate(
+          `reff(
+            output.y,
+            [${4}],
+            [${beta[0]}],
+            ${country}.population,
+            ${country}.contactMatrixScaledAge
+          )
+          `
+        )
+        if (!reff[0] == 4) {
+          failed = true;
+          console.log(`Expected ${reff[0]} == 4 for t = 0`)
+        }
+
         scenario++;
       }
     }

e2e/test_script.js

@@ -1,10 +1,11 @@
-import { runModel } from '../build/squire.js'
+import { runModel, reff } from '../build/squire.js'
 import LCA from '../data/LCA.json'
 import NGA from '../data/NGA.json'
 import IND from '../data/IND.json'
 
 
 window.runModel = runModel;
+window.reff = reff;
 window.LCA = LCA
 window.NGA = NGA;
 window.IND = IND;