From 492a58a0d4be2e923bf724e976fbbb401dc112f9 Mon Sep 17 00:00:00 2001
From: arihantlodha-cmd <arihantlodha48@gmail.com>
Date: Fri, 25 Sep 2026 20:35:39 +0900
Subject: [PATCH] Report aggregate hours worked (N) in addition to effective
 labor (L)

Closes #709. Aggregate labor L weights each household's hours by its
effective-labor-units profile e, so it does not directly map to hours
worked when productivity varies across households. This adds an aggregate
hours measure N that sums hours without the e weighting:

    N_t = sum_s sum_j omega_{s,t} lambda_j n_{j,s,t}

- aggregates.py: add get_N, mirroring get_L but without the e factor.
- SS.py and TPI.py: compute N and add it to the steady-state and
  transition output dictionaries (both, so the two keep the same keys).

Adds test_get_N (SS and TPI) checking the result against an explicit
index-by-index loop. With the default normalized e profile N and L
coincide at baseline (L = 0.3330, N = 0.3333); they diverge once e is not
mean-one, which is the case the separate measure is for.
---
 ogcore/SS.py             |  2 ++
 ogcore/TPI.py            |  4 ++++
 ogcore/aggregates.py     | 33 +++++++++++++++++++++++++++++++++
 tests/test_aggregates.py | 24 ++++++++++++++++++++++++
 4 files changed, 63 insertions(+)

diff --git a/ogcore/SS.py b/ogcore/SS.py
index bbed8aea7..64975c15a 100644
--- a/ogcore/SS.py
+++ b/ogcore/SS.py
@@ -877,6 +877,7 @@ def SS_solver(
     I_g_ss = fiscal.get_I_g(Yss, Ig_baseline, p, "SS")
     K_g_ss = fiscal.get_K_g(0, I_g_ss, p, "SS")
     Lss = aggr.get_L(nssmat, p, "SS")
+    Nss = aggr.get_N(nssmat, p, "SS")
     Bss = aggr.get_B(bssmat_splus1, p, "SS", False)
     (
         Dss,
@@ -1207,6 +1208,7 @@ def SS_solver(
         "K_f": K_f_ss,
         "K_d": K_d_ss,
         "L": Lss,
+        "N": Nss,
         "C": Css,
         "I": Iss,
         "I_total": Iss_total,
diff --git a/ogcore/TPI.py b/ogcore/TPI.py
index a405431dd..b77ed3727 100644
--- a/ogcore/TPI.py
+++ b/ogcore/TPI.py
@@ -1745,6 +1745,9 @@ def run_TPI(p, client=None):
         f"Max Euler error labor supply: {np.abs(eul_laborleisure).max()}"
     )
 
+    # Aggregate hours worked (not weighted by effective-labor units).
+    N = aggr.get_N(n_mat[: p.T], p, "TPI")
+
     """
     ------------------------------------------------------------------------
     Save variables/values so they can be used in other modules
@@ -1758,6 +1761,7 @@ def run_TPI(p, client=None):
         "K_f": K_f[: p.T],
         "K_d": K_d[: p.T, ...],
         "L": L[: p.T, ...],
+        "N": N[: p.T, ...],
         "C": C[: p.T, ...],
         "I": I[: p.T, ...],
         "I_total": I_total[: p.T, ...],
diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py
index b7f015726..44c1154b7 100644
--- a/ogcore/aggregates.py
+++ b/ogcore/aggregates.py
@@ -45,6 +45,39 @@ def get_L(n, p, method):
     return L
 
 
+def get_N(n, p, method):
+    r"""
+    Calculate aggregate hours worked.
+
+    Unlike aggregate labor supply :math:`L`, this measure does not weight
+    hours by the effective-labor-units profile :math:`e_{j,s,t}`, so it
+    reports raw hours rather than productivity-adjusted labor:
+
+    .. math::
+        N_{t} = \sum_{s=E}^{E+S}\sum_{j=0}^{J}\omega_{s,t}\lambda_{j}n_{j,s,t}
+
+    Args:
+        n (Numpy array): labor supply of households
+        p (OG-Core Specifications object): model parameters
+        method (str): adjusts calculation dimensions based on 'SS' or
+            'TPI'
+
+    Returns:
+        N (array_like): aggregate hours worked
+
+    """
+    if method == "SS":
+        # p.omega_SS is the (S, J) population weight and already carries the
+        # lambda_j type shares (its columns sum to lambdas), so no separate
+        # lambda multiplication is needed, mirroring get_L.
+        N_presum = p.omega_SS * n
+        N = N_presum.sum()
+    elif method == "TPI":
+        N_presum = n * p.omega[: p.T, :, :]
+        N = N_presum.sum(1).sum(1)
+    return N
+
+
 def get_I(b_splus1, K_p1, K, p, method):
     r"""
     Calculate aggregate investment.
diff --git a/tests/test_aggregates.py b/tests/test_aggregates.py
index 82b5695ad..890dc77af 100644
--- a/tests/test_aggregates.py
+++ b/tests/test_aggregates.py
@@ -57,6 +57,30 @@ def test_get_L(n, p, method, expected):
     assert np.allclose(L, expected)
 
 
+# Aggregate hours worked is the same population-weighted sum as get_L but
+# without the effective-labor-units (p.e) weighting.
+N_loop = np.ones(p.T * p.S * p.J).reshape(p.T, p.S, p.J)
+for t in range(p.T):
+    for i in range(p.S):
+        for k in range(p.J):
+            N_loop[t, i, k] *= p.omega[t, i, k].item() * n[t, i, k].item()
+expected_N_SS = N_loop[-1, :, :].sum()
+expected_N_TPI = N_loop.sum(1).sum(1)
+test_data_N = [
+    (n[-1, :, :], p, "SS", expected_N_SS),
+    (n, p, "TPI", expected_N_TPI),
+]
+
+
+@pytest.mark.parametrize("n,p,method,expected", test_data_N, ids=["SS", "TPI"])
+def test_get_N(n, p, method, expected):
+    """
+    Test aggregate hours-worked function.
+    """
+    N = aggr.get_N(n, p, method)
+    assert np.allclose(N, expected)
+
+
 def test_get_L_J1_regression():
     """
     Regression test for issue #1143: aggregate labor for J=1 must use a