Merge pull request #2231 from anth-volk/fix/2230-pooled-reproduce

Display pooled 3-year CPS as dataset for state-level runs
PolicyEngine · Dec 16, 2024 · 114b62f · 114b62f
2 parents 153ea2e + aba9ccb
commit 114b62f
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 31 deletions.
diff --git a/src/__tests__/data/reformDefinitionCode.test.js b/src/__tests__/data/reformDefinitionCode.test.js
@@ -211,21 +211,31 @@ describe("Test getSituationCode", () => {
 });
 describe("Test getImplementationCode", () => {
   test("If not a policy type, return empty array", () => {
-    const output = getImplementationCode("household", "us", 2024);
+    const output = getImplementationCode("household", "us", metadataUS, 2024);
     expect(output).toBeInstanceOf(Array);
     expect(output.length).toBe(0);
   });
   test("If policy, return lines", () => {
     const output = getImplementationCode(
       "policy",
       "uk",
+      metadataUK,
       2024,
       baselinePolicyUK,
     );
     expect(output).toBeInstanceOf(Array);
+    // Ensure we don't code for reversed baseline and reform
     expect(output).not.toContain(
       "baseline = Microsimulation(reform=baseline_reform)",
     );
+    // Ensure we don't code for state-level runs
+    expect(output).not.toContain(
+      'baseline = Microsimulation(dataset="pooled_3_year_cps_2023")',
+    );
+    // And ensure we don't code for both simultaneously
+    expect(output).not.toContain(
+      'baseline = Microsimulation(reform=baseline_reform, dataset="pooled_3_year_cps_2023")',
+    );
   });
   test("If set baseline, return lines with baseline", () => {
     let testPolicy = JSON.parse(JSON.stringify(baselinePolicyUK));
@@ -242,21 +252,72 @@ describe("Test getImplementationCode", () => {
         id: 1,
       },
     };
-    const output = getImplementationCode("policy", "uk", 2024, testPolicy);
+    const output = getImplementationCode(
+      "policy",
+      "uk",
+      metadataUK,
+      2024,
+      testPolicy,
+    );
     expect(output).toBeInstanceOf(Array);
     expect(output).toContain("baseline = Microsimulation(reform=baseline)");
   });
+  test("If US state, return lines with pooled 3-year CPS dataset", () => {
+    let testPolicy = JSON.parse(JSON.stringify(baselinePolicyUS));
+    testPolicy = {
+      ...testPolicy,
+      reform: {
+        data: {
+          "sample.reform.item": {
+            "2020.01.01": true,
+            "2022.01.01": true,
+          },
+        },
+      },
+    };
+
+    const output = getImplementationCode(
+      "policy",
+      "ca",
+      metadataUS,
+      2024,
+      testPolicy,
+      null,
+    );
+    expect(output).toBeInstanceOf(Array);
+    expect(output).toContain(
+      'baseline = Microsimulation(dataset="pooled_3_year_cps_2023")',
+    );
+    expect(output).toContain(
+      'reformed = Microsimulation(reform=reform, dataset="pooled_3_year_cps_2023")',
+    );
+  });
   test("If dataset provided, return lines with dataset", () => {
     const output = getImplementationCode(
       "policy",
       "us",
+      metadataUS,
+      2024,
+      baselinePolicyUS,
+      "enhanced_cps",
+    );
+    expect(output).toBeInstanceOf(Array);
+    expect(output).toContain(
+      'baseline = Microsimulation(dataset="enhanced_cps_2024")',
+    );
+  });
+  test("If dataset provided alongside US state, return lines with dataset, not Pooled 3-Year CPS", () => {
+    const output = getImplementationCode(
+      "policy",
+      "ks",
+      metadataUS,
       2024,
       baselinePolicyUS,
       "enhanced_cps",
     );
     expect(output).toBeInstanceOf(Array);
     expect(output).toContain(
-      "baseline = Microsimulation(dataset='enhanced_cps_2024')",
+      'baseline = Microsimulation(dataset="enhanced_cps_2024")',
     );
   });
 });

diff --git a/src/data/reformDefinitionCode.js b/src/data/reformDefinitionCode.js
@@ -28,7 +28,7 @@ export function getReproducibilityCodeBlock(
       householdInput,
       earningVariation,
     ),
-    ...getImplementationCode(type, region, year, policy, dataset),
+    ...getImplementationCode(type, region, metadata, year, policy, dataset),
   ];
 }
 
@@ -160,50 +160,53 @@ export function getSituationCode(
 export function getImplementationCode(
   type,
   region,
+  metadata,
   timePeriod,
   policy,
   dataset,
 ) {
+  const countryId = metadata.economy_options.region[0].name;
+
   if (type !== "policy") {
     return [];
   }
 
   const hasBaseline = Object.keys(policy?.baseline?.data).length > 0;
   const hasReform = Object.keys(policy?.reform?.data).length > 0;
 
-  // Check if the region has a dataset specified; enhanced_us is legacy implemntation
+  // Check if the region has a dataset specified
+  const hasDatasetSpecified = Object.keys(DEFAULT_DATASETS).includes(dataset);
+
+  const COUNTRY_LEVEL_US_REGIONS = ["us", "enhanced_us"];
+
+  const isState =
+    countryId === "us" && !COUNTRY_LEVEL_US_REGIONS.includes(region);
+
+  let datasetText = "";
+
+  // enhanced_us is legacy implemntation
   // whereby enhanced_us region correlated with enhanced_cps dataset
-  const hasDatasetSpecified =
-    Object.keys(DEFAULT_DATASETS).includes(dataset) || region === "enhanced_us";
-
-  let formattedDataset = null;
-  if (region === "enhanced_us") {
-    formattedDataset = "enhanced_cps_2024";
-  } else if (hasDatasetSpecified) {
-    formattedDataset = DEFAULT_DATASETS[dataset];
+  if (hasDatasetSpecified) {
+    datasetText = DEFAULT_DATASETS[dataset];
+  } else if (isState) {
+    datasetText = "pooled_3_year_cps_2023";
+  } else if (region === "enhanced_us") {
+    datasetText = "enhanced_cps_2024";
   }
 
+  const datasetSpecifier = datasetText ? `dataset="${datasetText}"` : "";
+
+  const baselineSpecifier = hasBaseline ? "reform=baseline" : "";
+  const baselineComma = hasBaseline && datasetText ? ", " : "";
+
+  const reformSpecifier = hasReform ? "reform=reform" : "";
+  const reformComma = hasReform && datasetText ? ", " : "";
+
   return [
     "",
     "",
-    `baseline = Microsimulation(${
-      hasDatasetSpecified && hasBaseline
-        ? `reform=baseline, dataset='${formattedDataset}'`
-        : hasBaseline
-          ? `reform=baseline`
-          : hasDatasetSpecified
-            ? `dataset='${formattedDataset}'`
-            : ""
-    })`,
-    `reformed = Microsimulation(${
-      hasDatasetSpecified && hasReform
-        ? `reform=reform, dataset='${formattedDataset}'`
-        : hasReform
-          ? `reform=reform`
-          : hasDatasetSpecified
-            ? `dataset='${formattedDataset}'`
-            : ""
-    })`,
+    `baseline = Microsimulation(${baselineSpecifier}${baselineComma}${datasetSpecifier})`,
+    `reformed = Microsimulation(${reformSpecifier}${reformComma}${datasetSpecifier})`,
     `baseline_income = baseline.calculate("household_net_income", period=${timePeriod || defaultYear})`,
     `reformed_income = reformed.calculate("household_net_income", period=${timePeriod || defaultYear})`,
     "difference_income = reformed_income - baseline_income",