From a8595ff7985d58b118beb810b80d60a7338679ee Mon Sep 17 00:00:00 2001
From: arihantlodha-cmd <arihantlodha48@gmail.com>
Date: Fri, 24 Jul 2026 19:13:32 +0900
Subject: [PATCH] Make SS and TPI output dictionaries expose the same keys

Closes #1152. Adds labor_income, capital_income, and income_payroll_taxes
to the SS output dict and theta and factor to the TPI output dict, so both
dictionaries expose the same 72 keys.

Also fixes a latent bug: SS's income_tax_ss was computed via
tax.income_tax_liab but passed bqssmat (bequests) as the n (labor supply)
argument. The value was previously never used; it is now surfaced as
income_payroll_taxes, so the call is corrected to pass nssmat (matching the
TPI income_tax_mat calculation).
---
 ogcore/SS.py  | 9 ++++++++-
 ogcore/TPI.py | 2 ++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/ogcore/SS.py b/ogcore/SS.py
index a3a0a9394..12a050b27 100644
--- a/ogcore/SS.py
+++ b/ogcore/SS.py
@@ -982,7 +982,7 @@ def SS_solver(
         r_p_ss,
         wss,
         bssmat_s,
-        bqssmat,
+        nssmat,
         factor_ss,
         0,
         None,
@@ -1037,6 +1037,10 @@ def SS_solver(
     yss_before_tax_mat = household.get_y(
         r_p_ss, wss, bssmat_s, nssmat, p, "SS"
     )
+    labor_income_ss = (
+        wss * nssmat * np.squeeze(p.e[-1, :, :]).reshape((p.S, p.J))
+    )
+    capital_income_ss = r_p_ss * bssmat_s
     Css = aggr.get_C(cssmat, p, "SS")
     c_i_ss_mat = household.get_ci(
         cssmat, p_i_ss, p_tilde_ss, p.tau_c[-1, :], p.alpha_c, p.c_min
@@ -1224,7 +1228,10 @@ def SS_solver(
         "tr": trssmat,
         "ubi": ubissmat,
         "before_tax_income": yss_before_tax_mat,
+        "labor_income": labor_income_ss,
+        "capital_income": capital_income_ss,
         "hh_net_taxes": taxss,
+        "income_payroll_taxes": income_tax_ss,
         "sales_tax": sales_tax_ss[: p.T, ...],
         "wealth_tax": wealth_tax_ss[: p.T, ...],
         "bequest_tax": bq_tax_ss[: p.T, ...],
diff --git a/ogcore/TPI.py b/ogcore/TPI.py
index 5ecf7b936..68c8116b8 100644
--- a/ogcore/TPI.py
+++ b/ogcore/TPI.py
@@ -1768,6 +1768,8 @@ def run_TPI(p, client=None):
         "etr": etr_path[: p.T, ...],
         "mtrx": mtrx_path[: p.T, ...],
         "mtry": mtry_path[: p.T, ...],
+        "theta": theta,
+        "factor": factor,
         "euler_savings": eul_savings[: p.T, ...],
         "euler_labor_leisure": eul_laborleisure[: p.T, ...],
         "resource_constraint_error": RC_error[: p.T, ...],