Fix capital gains basis read backfill

Author: MaxGhenis

policyengine_us_data/datasets/puf/puf.py

@@ -1340,6 +1340,88 @@ def _ensure_sstb_split_inputs(self) -> dict[str, np.ndarray]:
 
         return overrides
 
+    def _capital_gains_basis_overrides(
+        self,
+        existing_overrides: dict[str, np.ndarray] | None = None,
+    ) -> dict[str, np.ndarray]:
+        if not has_policyengine_us_variables(*CAPITAL_GAINS_BASIS_VARIABLES):
+            return {}
+        if not self.file_path.exists():
+            return {}
+
+        existing_overrides = existing_overrides or {}
+        with h5py.File(self.file_path, "r") as file_handle:
+            keys = set(file_handle.keys()) | set(existing_overrides)
+            if all(variable in keys for variable in CAPITAL_GAINS_BASIS_VARIABLES):
+                return {}
+            if (
+                "long_term_capital_gains" not in keys
+                or "person_tax_unit_id" not in keys
+            ):
+                return {}
+
+            gains = self._values_from_file_or_overrides(
+                file_handle,
+                "long_term_capital_gains",
+                existing_overrides,
+                0,
+            )
+            length = len(gains)
+            arrays = {
+                "long_term_capital_gains": gains,
+                "person_tax_unit_id": self._values_from_file_or_overrides(
+                    file_handle,
+                    "person_tax_unit_id",
+                    existing_overrides,
+                    length,
+                ),
+            }
+            for variable in (
+                "person_id",
+                "household_weight",
+                "person_household_id",
+                "household_id",
+                *CAPITAL_GAINS_BASIS_VARIABLES,
+            ):
+                if variable in keys:
+                    arrays[variable] = self._values_from_file_or_overrides(
+                        file_handle,
+                        variable,
+                        existing_overrides,
+                        length,
+                    )
+
+        arrays = _with_capital_gains_basis_inputs(arrays, self.time_period)
+        return {
+            variable: np.asarray(arrays[variable])
+            for variable in CAPITAL_GAINS_BASIS_VARIABLES
+            if variable not in keys and variable in arrays
+        }
+
+    def _ensure_capital_gains_basis_inputs(
+        self,
+        existing_overrides: dict[str, np.ndarray] | None = None,
+    ) -> dict[str, np.ndarray]:
+        overrides = self._capital_gains_basis_overrides(existing_overrides)
+        if not overrides:
+            return {}
+
+        try:
+            with h5py.File(self.file_path, "r+") as file_handle:
+                for key, values in overrides.items():
+                    self._replace_array(file_handle, key, values)
+        except OSError:
+            pass
+
+        return overrides
+
+    def _ensure_read_overrides(self) -> dict[str, np.ndarray]:
+        sstb_overrides = self._ensure_sstb_split_inputs()
+        capital_gains_overrides = self._ensure_capital_gains_basis_inputs(
+            sstb_overrides
+        )
+        return {**sstb_overrides, **capital_gains_overrides}
+
     class _OverrideView:
         def __init__(self, backing, overrides: dict[str, np.ndarray]):
             self._backing = backing
@@ -1393,15 +1475,15 @@ def __getattr__(self, name):
 
     def load(self, key=None, mode="r"):
         if mode == "r":
-            overrides = self._ensure_sstb_split_inputs()
+            overrides = self._ensure_read_overrides()
             if key in overrides:
                 return overrides[key]
             if key is None and overrides:
                 return self._OverrideView(super().load(key=key, mode=mode), overrides)
         return super().load(key=key, mode=mode)
 
     def load_dataset(self):
-        overrides = self._ensure_sstb_split_inputs()
+        overrides = self._ensure_read_overrides()
         arrays = super().load_dataset()
         arrays.update(overrides)
         return arrays

tests/unit/calibration/test_calibration_puf_impute.py

@@ -216,6 +216,59 @@ def test_capital_gains_basis_fields_are_stage_one_outputs(self):
         assert expected <= set(IMPUTED_VARIABLES)
         assert expected <= set(DETERMINISTIC_IMPUTED_VARIABLES)
 
+    def test_qrf_excludes_deterministic_capital_gains_basis_outputs(
+        self,
+        monkeypatch,
+    ):
+        import policyengine_us
+
+        data = _make_mock_data(n_persons=4, n_households=2)
+
+        class FakeCalculation:
+            values = np.array([100.0, 200.0, 300.0, 400.0], dtype=np.float32)
+
+        class FakeMicrosimulation:
+            def __init__(self, dataset):
+                self.dataset = dataset
+
+            def calculate(self, variable, map_to=None):
+                return FakeCalculation()
+
+            def calculate_dataframe(self, variables):
+                return pd.DataFrame(
+                    {variable: np.arange(4, dtype=np.float32) for variable in variables}
+                )
+
+        captured_output_vars = []
+
+        def fake_sequential_qrf(X_train, X_test, predictors, output_vars):
+            captured_output_vars.append(tuple(output_vars))
+            return {
+                variable: np.zeros(len(X_test), dtype=np.float32)
+                for variable in output_vars
+            }
+
+        monkeypatch.setattr(policyengine_us, "Microsimulation", FakeMicrosimulation)
+        monkeypatch.setattr(
+            puf_impute_module,
+            "_sequential_qrf",
+            fake_sequential_qrf,
+        )
+
+        puf_impute_module._run_qrf_imputation(
+            data=data,
+            time_period=2024,
+            puf_dataset=object(),
+        )
+
+        deterministic_outputs = set(DETERMINISTIC_IMPUTED_VARIABLES)
+        assert captured_output_vars
+        for output_vars in captured_output_vars:
+            assert deterministic_outputs.isdisjoint(output_vars)
+        assert set(captured_output_vars[0]) == (
+            set(IMPUTED_VARIABLES) - deterministic_outputs
+        )
+
     def test_overridden_subset_of_imputed(self):
         for var in OVERRIDDEN_IMPUTED_VARIABLES:
             assert var in IMPUTED_VARIABLES

tests/unit/datasets/test_irs_puf.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 
+from policyengine_us_data.datasets.puf import puf as puf_module
 from policyengine_us_data.datasets.puf.puf import (
     PUF,
     QBI_SIMULATION_VERSION,
@@ -14,6 +15,19 @@ def _mark_current_qbi_simulation(file_handle):
     file_handle.attrs[QBI_SIMULATION_VERSION_ATTR] = QBI_SIMULATION_VERSION
 
 
+def _write_capital_gains_basis_source_file(path):
+    with h5py.File(path, "w") as file_handle:
+        file_handle.create_dataset("person_id", data=np.array([1, 2, 3, 4]))
+        file_handle.create_dataset("person_tax_unit_id", data=np.array([1, 1, 2, 2]))
+        file_handle.create_dataset("person_household_id", data=np.array([1, 1, 2, 2]))
+        file_handle.create_dataset("household_id", data=np.array([1, 2]))
+        file_handle.create_dataset("household_weight", data=np.array([100.0, 200.0]))
+        file_handle.create_dataset(
+            "long_term_capital_gains",
+            data=np.array([100.0, -40.0, 0.0, 200.0]),
+        )
+
+
 @pytest.mark.skip(reason="This test requires private data.")
 @pytest.mark.parametrize("year", [2015])
 def test_irs_puf_generates(year: int):
@@ -50,6 +64,78 @@ def test_puf_person_split_keeps_capital_gains_holding_period_collapsed():
     )
 
 
+def test_puf_load_dataset_backfills_capital_gains_basis_inputs(
+    tmp_path,
+    monkeypatch,
+):
+    monkeypatch.setattr(
+        puf_module,
+        "has_policyengine_us_variables",
+        lambda *variables: True,
+    )
+
+    class DummyPUF(PUF):
+        label = "Dummy PUF"
+        name = "dummy_puf"
+        time_period = 2024
+        file_path = tmp_path / "dummy_puf.h5"
+
+    _write_capital_gains_basis_source_file(DummyPUF.file_path)
+
+    arrays = DummyPUF().load_dataset()
+
+    basis = arrays["long_term_capital_gains_basis"]
+    years = arrays["long_term_capital_gains_years_held"]
+    gains = arrays["long_term_capital_gains"]
+
+    assert np.all(basis[gains != 0] > 0)
+    assert np.all(years[gains != 0] > 0)
+    assert np.all(basis[gains == 0] == 0)
+    assert np.all(years[gains == 0] == 0)
+
+    with h5py.File(DummyPUF.file_path, "r") as file_handle:
+        assert "long_term_capital_gains_basis" in file_handle
+        assert "long_term_capital_gains_years_held" in file_handle
+
+
+def test_puf_load_key_backfills_read_only_capital_gains_basis_inputs(
+    tmp_path,
+    monkeypatch,
+):
+    monkeypatch.setattr(
+        puf_module,
+        "has_policyengine_us_variables",
+        lambda *variables: True,
+    )
+
+    class DummyPUF(PUF):
+        label = "Dummy PUF"
+        name = "dummy_puf"
+        time_period = 2024
+        file_path = tmp_path / "dummy_puf.h5"
+
+    _write_capital_gains_basis_source_file(DummyPUF.file_path)
+    DummyPUF.file_path.chmod(0o444)
+
+    dataset = DummyPUF()
+    try:
+        basis = dataset.load("long_term_capital_gains_basis")
+        years = dataset.load("long_term_capital_gains_years_held")
+        reader = dataset.load()
+        np.testing.assert_array_equal(
+            reader["long_term_capital_gains_basis"],
+            basis,
+        )
+        reader.close()
+    finally:
+        DummyPUF.file_path.chmod(0o644)
+
+    assert np.all(basis[[0, 1, 3]] > 0)
+    assert basis[2] == 0
+    assert np.all(years[[0, 1, 3]] > 0)
+    assert years[2] == 0
+
+
 def test_puf_load_dataset_backfills_sstb_split_inputs(tmp_path):
     class DummyPUF(PUF):
         label = "Dummy PUF"