Fix fiitax to include Additional Medicare Tax

changelog.d/fix-addl-medicare-tax-in-fiitax.fixed.md

@@ -0,0 +1 @@
+Include Additional Medicare Tax (0.9% on wages above $200K/$250K) in fiitax output to match TAXSIM's treatment.

policyengine_taxsim/runners/policyengine_runner.py

@@ -1236,7 +1236,9 @@ def _extract_vectorized_results(
                                             )
                                             var_sum += arr
                                         except Exception as e:
-                                            if "does not exist" in str(e) or "was not found" in str(e):
+                                            if "does not exist" in str(
+                                                e
+                                            ) or "was not found" in str(e):
                                                 if self.logs:
                                                     print(
                                                         f"Variable {resolved} not implemented, setting to 0"
@@ -1258,7 +1260,9 @@ def _extract_vectorized_results(
                                         )
                                         result_array[state_mask] = arr[state_mask]
                                     except Exception as e:
-                                        if "does not exist" in str(e) or "was not found" in str(e):
+                                        if "does not exist" in str(
+                                            e
+                                        ) or "was not found" in str(e):
                                             if self.logs:
                                                 print(
                                                     f"Variable {resolved} not implemented, setting to 0"
@@ -1295,12 +1299,14 @@ def _extract_vectorized_results(
                         ) from e
 
             # Apply fiitax special calculation (income_tax + additional_medicare_tax)
+            # TAXSIM includes Additional Medicare Tax (0.9% on wages above
+            # $200K/$250K) in fiitax. PE's income_tax does not include it,
+            # so we add it here.
+            addl_med = self._calc_tax_unit(sim, "additional_medicare_tax", year_str)
             if "fiitax" in columns:
-                pass  # Already computed above
+                columns["fiitax"] = np.round(columns["fiitax"] + addl_med, 2)
             else:
-                fiitax_arr = self._calc_tax_unit(
-                    sim, "income_tax", year_str
-                ) + self._calc_tax_unit(sim, "additional_medicare_tax", year_str)
+                fiitax_arr = self._calc_tax_unit(sim, "income_tax", year_str) + addl_med
                 columns["fiitax"] = np.round(fiitax_arr, 2)
 
             # Compute marginal rates if any idtl level requests them

tests/test_performance.py

@@ -91,6 +91,64 @@ def test_deterministic_output(self, golden_output):
         pd.testing.assert_frame_equal(golden_output, output2)
 
 
+class TestFederalOutputAdjustments:
+    def test_fiitax_includes_additional_medicare_tax_when_precomputed(self):
+        """
+        fiitax is normally populated from the income_tax mapping before the
+        special post-processing step runs. That post-processing still needs to
+        add additional_medicare_tax instead of skipping fiitax entirely.
+        """
+        records = pd.DataFrame(
+            {
+                "taxsimid": [1],
+                "year": [2023],
+                "state": [5],  # CA
+                "mstat": [1],
+                "depx": [0],
+                "page": [45],
+                "sage": [0],
+                "pwages": [300000.0],
+                "swages": [0.0],
+                "idtl": [0],
+            }
+        )
+        runner = PolicyEngineRunner(records.copy(), logs=False, disable_salt=True)
+        runner.mappings = {
+            "policyengine_to_taxsim": {
+                "fiitax": {
+                    "implemented": True,
+                    "variable": "income_tax",
+                    "idtl": [{"standard_output": 0}, {"full_output": 2}],
+                }
+            }
+        }
+
+        fake_sim = types.SimpleNamespace(
+            tax_benefit_system=types.SimpleNamespace(
+                variables={
+                    "income_tax": object(),
+                    "additional_medicare_tax": object(),
+                }
+            )
+        )
+        calc_calls = []
+
+        def fake_calc_tu(self_runner, sim, var_name, period):
+            calc_calls.append(var_name)
+            if var_name == "income_tax":
+                return np.array([1000.0])
+            if var_name == "additional_medicare_tax":
+                return np.array([900.0])
+            raise AssertionError(f"Unexpected variable: {var_name}")
+
+        runner._calc_tax_unit = types.MethodType(fake_calc_tu, runner)
+
+        result = runner._extract_vectorized_results(fake_sim, runner.input_df)
+
+        assert result["fiitax"].iloc[0] == pytest.approx(1900.0)
+        assert calc_calls == ["income_tax", "additional_medicare_tax"]
+
+
 class TestGeneratePhaseEfficiency:
     """Verify that the generate phase scales sub-linearly (vectorized)."""