Use utils/trace.py capability in datasets/tmd.py module

tax_microdata_benchmarking/datasets/tmd.py

@@ -6,6 +6,7 @@
 from tax_microdata_benchmarking.datasets.taxcalc_dataset import (
     create_tc_dataset,
 )
+from tax_microdata_benchmarking.utils.trace import trace1
 from tax_microdata_benchmarking.utils.taxcalc_utils import add_taxcalc_outputs
 from tax_microdata_benchmarking.utils.reweight import reweight
 from tax_microdata_benchmarking.storage import STORAGE_FOLDER
@@ -33,13 +34,20 @@ def create_tmd_2021():
 
     combined = pd.concat([tc_puf_21, tc_cps_21], ignore_index=True)
 
+    trace1("A", combined)
+
     # Add Tax-Calculator outputs
     print("Adding Tax-Calculator outputs...")
     combined = add_taxcalc_outputs(combined, 2021)
     combined["s006_original"] = combined.s006.values
+
+    trace1("B", combined)
+
     print("Reweighting...")
     combined = reweight(combined, 2021, weight_deviation_penalty=0)
 
+    trace1("C", combined)
+
     return combined