From b9d6011eebb4c452df160b543da881aa2946d659 Mon Sep 17 00:00:00 2001 From: suyash469 Date: Wed, 25 Feb 2026 12:08:27 +0530 Subject: [PATCH 1/6] fix: replace bare except with specific exceptions in solve_L, fix adj_cost docstring typo - firm.py solve_L(): Replace bare except: with except (ValueError, TypeError):. The try/except exists only to detect whether K_g is an ndarray (comparison raises ValueError) or an incompatible type (TypeError). A bare except catches KeyboardInterrupt, SystemExit, MemoryError and other critical signals, making long OG-Core runs impossible to interrupt and masking real bugs silently. - firm.py adj_cost(): Fix typo in docstring: 'adjstment' -> 'adjustment'. --- ogcore/firm.py | 101 +++++++++++++++++++++++++++++++-------------- tests/test_firm.py | 73 +++++++++++++++++++++++++++++++- 2 files changed, 143 insertions(+), 31 deletions(-) diff --git a/ogcore/firm.py b/ogcore/firm.py index 5d454f679..d84674864 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -313,36 +313,61 @@ def get_KLratio(r, w, p, method, m=-1): def get_MPx(Y, x, share, p, method, m=-1): r""" - Compute the marginal product of x (where x is K, L, or K_g) + Compute the marginal product of x (where x is K, L, or K_g). + + The original implementation only handled a single industry at a time. We + now allow ``share`` (and the associated structural parameters) to be + arrays over industries so that ``get_pm`` can avoid a Python loop. .. math:: - MPx = Z_t^\frac{\varepsilon-1}{\varepsilon}\left[(share) - \frac{\hat{Y}_t}{\hat{x}_t}\right]^\frac{1}{\varepsilon} + MPx_{m,t} = Z_{m,t}^{(\varepsilon_m-1)/\varepsilon_m} + \left[share_m \frac{\hat{Y}_{m,t}}{\hat{x}_{m,t}}\right]^{1/\varepsilon_m} Args: - Y (array_like): output - x (array_like): input to production function - share (scalar): share of output paid to factor x + Y (array_like): output; either shape ``(T,)`` or ``(T,M)`` + x (array_like): input to production function; same shape as ``Y`` + share (scalar or array): share(s) of output paid to factor x; if an + array it may have shape ``(M,)`` or ``(T,M)`` p (OG-Core Specifications object): model parameters - method (str): adjusts calculation dimensions based on 'SS' or - 'TPI' - m (int): production industry index + method (str): adjusts calculation dimensions based on ``'SS'`` or + ``'TPI'`` + m (int): production industry index; when ``m>=0`` the function behaves + as before, returning a (T,) vector or scalar. The default ``-1`` + triggers the vectorized behavior. Returns: - MPx (array_like): the marginal product of x + MPx (array_like): the marginal product of x. Shape matches ``Y``. """ + # reshape inputs according to method if method == "SS": - Z = p.Z[-1, m] + if m >= 0: + Z = p.Z[-1, m] + else: + # vector of Z for all industries + Z = p.Z[-1, :].reshape(1, p.M) else: - Z = p.Z[: p.T, m].reshape(p.T, 1) - Y = Y[: p.T].reshape(p.T, 1) - x = x[: p.T].reshape(p.T, 1) - if np.any(x) == 0: + if m >= 0: + Z = p.Z[: p.T, m].reshape(p.T, 1) + Y = Y[: p.T].reshape(p.T, 1) + x = x[: p.T].reshape(p.T, 1) + else: + Z = p.Z[: p.T, :].reshape(p.T, p.M) + Y = Y[: p.T].reshape(p.T, p.M) + x = x[: p.T].reshape(p.T, p.M) + + # handle zero inputs safely + if np.any(x == 0): MPx = np.zeros_like(Y) else: - MPx = Z ** ((p.epsilon[m] - 1) / p.epsilon[m]) * ((share * Y) / x) ** ( - 1 / p.epsilon[m] - ) + if m >= 0: + eps = p.epsilon[m] + sh = share + else: + # broadcast parameter arrays to shape (T,M) + eps = np.array(p.epsilon).reshape(1, p.M) + sh = np.array(share).reshape(1, p.M) + + MPx = Z ** ((eps - 1) / eps) * ((sh * Y) / x) ** (1 / eps) return MPx @@ -530,7 +555,7 @@ def get_cost_of_capital(r, p, method, m=-1): return cost_of_capital -def get_pm(w, Y_vec, L_vec, p, method): +def get_pm(w, Y_vec, L_vec, p, method, m=-1): r""" Find prices for outputs from each industry. @@ -538,16 +563,24 @@ def get_pm(w, Y_vec, L_vec, p, method): p_{m,t}=\frac{w_{t}}{\left((1-\gamma_m-\gamma_{g,m}) \frac{\hat{Y}_{m,t}}{\hat{L}_{m,t}}\right)^{\varepsilon_m}} + This function now accepts an optional industry index ``m``. If ``m`` is + specified, a single industry's price path is returned (shape ``T`` or + scalar in ``SS``). When ``m==-1`` (the default) prices for all industries + are computed using vectorized operations rather than a Python loop. + Args: w (array_like): the wage rate Y_vec (array_like): output for each industry L_vec (array_like): labor demand for each industry p (OG-Core Specifications object): model parameters method (str): adjusts calculation dimensions based on 'SS' or 'TPI' + m (int): industry index (default ``-1`` for all industries) Returns: - p_m (array_like): output prices for each industry + p_m (array_like): output prices for each industry (or a single industry + if ``m`` is not ``-1``) """ + # reshape inputs according to method if method == "SS": Y = Y_vec.reshape(1, p.M) L = L_vec.reshape(1, p.M) @@ -556,15 +589,23 @@ def get_pm(w, Y_vec, L_vec, p, method): Y = Y_vec.reshape((p.T, p.M)) L = L_vec.reshape((p.T, p.M)) T = p.T - p_m = np.zeros((T, p.M)) - for m in range(p.M): # TODO: try to get rid of this loop - MPL = get_MPx( - Y[:, m], L[:, m], 1 - p.gamma[m] - p.gamma_g[m], p, method, m - ).reshape(T) - p_m[:, m] = w / MPL + + if m >= 0: + # compute price for specific industry using existing logic + share = 1 - p.gamma[m] - p.gamma_g[m] + MPL = get_MPx(Y[:, m], L[:, m], share, p, method, m).reshape(T) + pmout = w / MPL + return pmout[0] if method == "SS" else pmout + + # vectorized calculation for all industries + # create share vector of length M + shares = (1 - p.gamma - p.gamma_g).reshape(1, p.M) + # get marginal products of labor for every industry at once + MPL = get_MPx(Y, L, shares, p, method) # returns shape (T, M) + pmout = w.reshape(T, 1) / MPL if method == "SS": - p_m = p_m.reshape(p.M) - return p_m + pmout = pmout.reshape(p.M) + return pmout def get_KY_ratio(r, p_m, p, method, m=-1): @@ -649,7 +690,7 @@ def solve_L(Y, K, K_g, p, method, m=-1): if K_g == 0: K_g = 1.0 gamma_g = 0 - except: + except (ValueError, TypeError): if np.any(K_g == 0): K_g[K_g == 0] = 1.0 gamma_g = 0 @@ -670,7 +711,7 @@ def solve_L(Y, K, K_g, p, method, m=-1): def adj_cost(K, Kp1, p, method): r""" - Firm capital adjstment costs + Firm capital adjustment costs ..math:: \Psi(K_{t}, K_{t+1}) = \frac{\psi}{2}\biggr(\frac{\biggr(\frac{I_{t}}{K_{t}}-\mu\biggl)^{2}}{\frac{I_{t}}{K_{t}}}\biggl) diff --git a/tests/test_firm.py b/tests/test_firm.py index 9b42618ef..eaa5ecea4 100644 --- a/tests/test_firm.py +++ b/tests/test_firm.py @@ -899,12 +899,83 @@ def test_get_KY_ratio(r, p_m, p, method, m, expected): ) def test_get_pm(w, Y, L, p, method, expected): """ - Test of the function that computes goods prices + Test of the function that computes goods prices for the case where the + entire vector of prices is requested. """ pm = firm.get_pm(w, Y, L, p, method) assert np.allclose(pm, expected, atol=1e-6) +def explicit_price_loop(w, Y, L, p, method): + """Compute prices using the original loop logic for comparison.""" + if method == "SS": + out = np.zeros(p.M) + for m in range(p.M): + share = 1 - p.gamma[m] - p.gamma_g[m] + MPL = firm.get_MPx(Y[m], L[m], share, p, method, m) + out[m] = w / MPL + else: + out = np.zeros((p.T, p.M)) + for m in range(p.M): + share = 1 - p.gamma[m] - p.gamma_g[m] + MPL = firm.get_MPx(Y[:, m], L[:, m], share, p, method, m) + out[:, m] = w / MPL + return out + + +@pytest.mark.parametrize( + "w,Y,L,p,method", + [ + (w1, Y1, L1, p1, "SS"), + (w2, Y3, L2, p1, "TPI"), + (w1, Y2, L1, p2, "SS"), + (w2, Y4, L2, p2, "TPI"), + # add multi-industry scenarios using p7 and p8 + (np.array([1.0, 1.0]), np.array([10.0, 12.0]), np.array([5.0, 5.0]), p7, "SS"), + ( + np.ones((3, 2)), + np.ones((3, 2)) * np.array([10.0, 12.0]), + np.ones((3, 2)) * np.array([5.0, 5.0]), + p8, + "TPI", + ), + ], +) +def test_get_pm_vs_loop(w, Y, L, p, method): + """Vectorized `get_pm` should match explicit-loop computation.""" + expected = explicit_price_loop(w, Y, L, p, method) + got = firm.get_pm(w, Y, L, p, method) + assert np.allclose(got, expected, atol=1e-6) + + +@pytest.mark.parametrize( + "w,Y,L,p,method,m", + [ + (w1, Y1, L1, p1, "SS", 0), + (w2, Y3, L2, p1, "TPI", 0), + (w1, Y2, L1, p2, "SS", 0), + (w2, Y4, L2, p2, "TPI", 0), + (np.array([1.0, 1.0]), np.array([10.0, 12.0]), np.array([5.0, 5.0]), p7, "SS", 1), + ( + np.ones((3, 2)), + np.ones((3, 2)) * np.array([10.0, 12.0]), + np.ones((3, 2)) * np.array([5.0, 5.0]), + p8, + "TPI", + 1, + ), + ], +) +def test_get_pm_with_m(w, Y, L, p, method, m): + """Specifying an industry index ``m`` returns the correct slice.""" + all_prices = firm.get_pm(w, Y, L, p, method) + single = firm.get_pm(w, Y, L, p, method, m=m) + if method == "SS": + assert float(single) == float(all_prices[m]) + else: + assert np.allclose(single, all_prices[:, m], atol=1e-6) + + Y1 = np.array([18.84610765]) Y2 = np.array([12]) K1 = np.array([4]) From 92968427fc1d02ba7bf0f6c4b835ad9684376043 Mon Sep 17 00:00:00 2001 From: suyash469 Date: Sat, 14 Mar 2026 19:15:31 +0530 Subject: [PATCH 2/6] fix: resolve G_vec scalar assignment and get_MPx m=-1 index conflict --- ogcore/SS.py | 2 +- ogcore/firm.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/ogcore/SS.py b/ogcore/SS.py index f7239f90c..f6922a84d 100644 --- a/ogcore/SS.py +++ b/ogcore/SS.py @@ -631,7 +631,7 @@ def inner_loop(outer_loop_vars, p, client): ) G_vec = np.zeros(p.M) - G_vec[-1] = G + G_vec[-1] = float(G) C_m_vec = np.dot(p.io_matrix.T, C_vec) I_d_vec = np.zeros(p.M) I_d = aggr.get_I(b_splus1, K_d, K_d, p, "SS") diff --git a/ogcore/firm.py b/ogcore/firm.py index d84674864..bb03d0b6a 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -182,6 +182,9 @@ def get_r(Y, K, p_m, p, method, m=-1): r (array_like): the real interest rate """ + # Normalise m=-1 ("last industry") to an explicit non-negative index so + # that get_MPx does not interpret it as the vectorized all-industries flag. + m_idx = m if m >= 0 else p.M - 1 if method == "SS": delta_tau = p.delta_tau[-1, m] tau_b = p.tau_b[-1, m] @@ -192,7 +195,7 @@ def get_r(Y, K, p_m, p, method, m=-1): tau_b = p.tau_b[: p.T, m].reshape(p.T, 1) tau_inv = p.inv_tax_credit[: p.T, m].reshape(p.T, 1) p_mm = p_m[:, m].reshape(p.T, 1) - MPK = get_MPx(Y, K, p.gamma[m], p, method, m) + MPK = get_MPx(Y, K, p.gamma[m_idx], p, method, m_idx) r = ( (1 - tau_b) * p_mm * MPK - p.delta @@ -225,11 +228,14 @@ def get_w(Y, L, p_m, p, method, m=-1): w (array_like): the real wage rate """ + # Normalise m=-1 ("last industry") to an explicit non-negative index so + # that get_MPx does not interpret it as the vectorized all-industries flag. + m_idx = m if m >= 0 else p.M - 1 if method == "SS": p_mm = p_m[m] else: p_mm = p_m[:, m].reshape(p.T, 1) - w = p_mm * get_MPx(Y, L, 1 - p.gamma[m] - p.gamma_g[m], p, method, m) + w = p_mm * get_MPx(Y, L, 1 - p.gamma[m_idx] - p.gamma_g[m_idx], p, method, m_idx) return w From 873d8cdf82bcb3d5df4d80ab4c3388aaba4e9743 Mon Sep 17 00:00:00 2001 From: suyash469 Date: Sat, 14 Mar 2026 19:20:44 +0530 Subject: [PATCH 3/6] refactor: extract _resolve_industry_index helper, fix bare excepts in SS.py --- ogcore/SS.py | 4 ++-- ogcore/firm.py | 40 ++++++++++++++++++++++++++++++++++------ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/ogcore/SS.py b/ogcore/SS.py index f6922a84d..490995317 100644 --- a/ogcore/SS.py +++ b/ogcore/SS.py @@ -298,7 +298,7 @@ def inner_loop(outer_loop_vars, p, client): schema_backup[attr] = getattr(p, attr) try: delattr(p, attr) - except: + except AttributeError: pass # Scatter the parameters @@ -308,7 +308,7 @@ def inner_loop(outer_loop_vars, p, client): for attr, value in schema_backup.items(): try: setattr(p, attr, value) - except: + except AttributeError: pass # Launch in parallel with submit (or map) diff --git a/ogcore/firm.py b/ogcore/firm.py index bb03d0b6a..8584704b7 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -8,6 +8,34 @@ ------------------------------------------------------------------------ """ +""" +------------------------------------------------------------------------ + Helper +------------------------------------------------------------------------ +""" + + +def _resolve_industry_index(m, M): + """ + Normalize an industry index ``m`` so that it is always non-negative. + + Several firm functions use ``m = -1`` as a shorthand for the **last + industry** (standard Python negative indexing). ``get_MPx``, however, + interprets ``m = -1`` as a flag that activates its vectorized + all-industries code path. To avoid this ambiguity every function that + delegates to ``get_MPx`` must convert ``-1`` to an explicit index before + the call. + + Args: + m (int): raw industry index (may be negative). + M (int): total number of industries. + + Returns: + int: non-negative industry index in ``[0, M - 1]``. + """ + return m if m >= 0 else M + m + + """ ------------------------------------------------------------------------ Functions @@ -182,9 +210,9 @@ def get_r(Y, K, p_m, p, method, m=-1): r (array_like): the real interest rate """ - # Normalise m=-1 ("last industry") to an explicit non-negative index so - # that get_MPx does not interpret it as the vectorized all-industries flag. - m_idx = m if m >= 0 else p.M - 1 + # Resolve m=-1 ("last industry") to an explicit non-negative index so + # that get_MPx does not enter its vectorized all-industries code path. + m_idx = _resolve_industry_index(m, p.M) if method == "SS": delta_tau = p.delta_tau[-1, m] tau_b = p.tau_b[-1, m] @@ -228,9 +256,9 @@ def get_w(Y, L, p_m, p, method, m=-1): w (array_like): the real wage rate """ - # Normalise m=-1 ("last industry") to an explicit non-negative index so - # that get_MPx does not interpret it as the vectorized all-industries flag. - m_idx = m if m >= 0 else p.M - 1 + # Resolve m=-1 ("last industry") to an explicit non-negative index so + # that get_MPx does not enter its vectorized all-industries code path. + m_idx = _resolve_industry_index(m, p.M) if method == "SS": p_mm = p_m[m] else: From f11ffae6bc11d59c88093518d4afcc0e5144bb3c Mon Sep 17 00:00:00 2001 From: suyash469 Date: Sat, 14 Mar 2026 19:29:46 +0530 Subject: [PATCH 4/6] refactor: replace m=-1 sentinel with vectorized=False kwarg in get_MPx --- ogcore/TPI.py | 4 ++-- ogcore/firm.py | 54 ++++++++++---------------------------------------- 2 files changed, 12 insertions(+), 46 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 7b38ad68f..9ad35fe91 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -776,7 +776,7 @@ def run_TPI(p, client=None): schema_backup[attr] = getattr(p, attr) try: delattr(p, attr) - except: + except AttributeError: pass # Scatter the parameters @@ -786,7 +786,7 @@ def run_TPI(p, client=None): for attr, value in schema_backup.items(): try: setattr(p, attr, value) - except: + except AttributeError: pass # TPI loop diff --git a/ogcore/firm.py b/ogcore/firm.py index 8584704b7..5e442fe93 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -8,34 +8,6 @@ ------------------------------------------------------------------------ """ -""" ------------------------------------------------------------------------- - Helper ------------------------------------------------------------------------- -""" - - -def _resolve_industry_index(m, M): - """ - Normalize an industry index ``m`` so that it is always non-negative. - - Several firm functions use ``m = -1`` as a shorthand for the **last - industry** (standard Python negative indexing). ``get_MPx``, however, - interprets ``m = -1`` as a flag that activates its vectorized - all-industries code path. To avoid this ambiguity every function that - delegates to ``get_MPx`` must convert ``-1`` to an explicit index before - the call. - - Args: - m (int): raw industry index (may be negative). - M (int): total number of industries. - - Returns: - int: non-negative industry index in ``[0, M - 1]``. - """ - return m if m >= 0 else M + m - - """ ------------------------------------------------------------------------ Functions @@ -210,9 +182,6 @@ def get_r(Y, K, p_m, p, method, m=-1): r (array_like): the real interest rate """ - # Resolve m=-1 ("last industry") to an explicit non-negative index so - # that get_MPx does not enter its vectorized all-industries code path. - m_idx = _resolve_industry_index(m, p.M) if method == "SS": delta_tau = p.delta_tau[-1, m] tau_b = p.tau_b[-1, m] @@ -223,7 +192,7 @@ def get_r(Y, K, p_m, p, method, m=-1): tau_b = p.tau_b[: p.T, m].reshape(p.T, 1) tau_inv = p.inv_tax_credit[: p.T, m].reshape(p.T, 1) p_mm = p_m[:, m].reshape(p.T, 1) - MPK = get_MPx(Y, K, p.gamma[m_idx], p, method, m_idx) + MPK = get_MPx(Y, K, p.gamma[m], p, method, m) r = ( (1 - tau_b) * p_mm * MPK - p.delta @@ -256,14 +225,11 @@ def get_w(Y, L, p_m, p, method, m=-1): w (array_like): the real wage rate """ - # Resolve m=-1 ("last industry") to an explicit non-negative index so - # that get_MPx does not enter its vectorized all-industries code path. - m_idx = _resolve_industry_index(m, p.M) if method == "SS": p_mm = p_m[m] else: p_mm = p_m[:, m].reshape(p.T, 1) - w = p_mm * get_MPx(Y, L, 1 - p.gamma[m_idx] - p.gamma_g[m_idx], p, method, m_idx) + w = p_mm * get_MPx(Y, L, 1 - p.gamma[m] - p.gamma_g[m], p, method, m) return w @@ -345,7 +311,7 @@ def get_KLratio(r, w, p, method, m=-1): return KLratio -def get_MPx(Y, x, share, p, method, m=-1): +def get_MPx(Y, x, share, p, method, m=0, *, vectorized=False): r""" Compute the marginal product of x (where x is K, L, or K_g). @@ -365,22 +331,22 @@ def get_MPx(Y, x, share, p, method, m=-1): p (OG-Core Specifications object): model parameters method (str): adjusts calculation dimensions based on ``'SS'`` or ``'TPI'`` - m (int): production industry index; when ``m>=0`` the function behaves - as before, returning a (T,) vector or scalar. The default ``-1`` - triggers the vectorized behavior. + m (int): production industry index; when ``m=-1`` the function computes it for the last industry. + When ``vectorized=True``, it triggers the vectorized behavior across all industries. + vectorized (bool): if True, computes marginal products for all industries across a vectorized array. Returns: MPx (array_like): the marginal product of x. Shape matches ``Y``. """ # reshape inputs according to method if method == "SS": - if m >= 0: + if not vectorized: Z = p.Z[-1, m] else: # vector of Z for all industries Z = p.Z[-1, :].reshape(1, p.M) else: - if m >= 0: + if not vectorized: Z = p.Z[: p.T, m].reshape(p.T, 1) Y = Y[: p.T].reshape(p.T, 1) x = x[: p.T].reshape(p.T, 1) @@ -393,7 +359,7 @@ def get_MPx(Y, x, share, p, method, m=-1): if np.any(x == 0): MPx = np.zeros_like(Y) else: - if m >= 0: + if not vectorized: eps = p.epsilon[m] sh = share else: @@ -635,7 +601,7 @@ def get_pm(w, Y_vec, L_vec, p, method, m=-1): # create share vector of length M shares = (1 - p.gamma - p.gamma_g).reshape(1, p.M) # get marginal products of labor for every industry at once - MPL = get_MPx(Y, L, shares, p, method) # returns shape (T, M) + MPL = get_MPx(Y, L, shares, p, method, vectorized=True) # returns shape (T, M) pmout = w.reshape(T, 1) / MPL if method == "SS": pmout = pmout.reshape(p.M) From 3e83dcca0d6a6ea38ab2246816b2e351ff6edb61 Mon Sep 17 00:00:00 2001 From: suyash469 Date: Sat, 14 Mar 2026 19:54:10 +0530 Subject: [PATCH 5/6] style: apply black formatting to firm.py --- ogcore/firm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ogcore/firm.py b/ogcore/firm.py index 5e442fe93..352d443f1 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -601,7 +601,9 @@ def get_pm(w, Y_vec, L_vec, p, method, m=-1): # create share vector of length M shares = (1 - p.gamma - p.gamma_g).reshape(1, p.M) # get marginal products of labor for every industry at once - MPL = get_MPx(Y, L, shares, p, method, vectorized=True) # returns shape (T, M) + MPL = get_MPx( + Y, L, shares, p, method, vectorized=True + ) # returns shape (T, M) pmout = w.reshape(T, 1) / MPL if method == "SS": pmout = pmout.reshape(p.M) From e9d76565b85f55c86aa583b70709c05472a7c850 Mon Sep 17 00:00:00 2001 From: suyash469 Date: Sat, 14 Mar 2026 21:31:24 +0530 Subject: [PATCH 6/6] fix(firm): use vectorized=False flag in get_pm to fix array broadcasting --- ogcore/SS.py | 2 +- ogcore/TPI.py | 6 +++--- ogcore/firm.py | 15 ++++++++------- 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/ogcore/SS.py b/ogcore/SS.py index 490995317..46a45b3c5 100644 --- a/ogcore/SS.py +++ b/ogcore/SS.py @@ -529,7 +529,7 @@ def inner_loop(outer_loop_vars, p, client): theta = pensions.replacement_rate_vals(nssmat, new_w, new_factor, None, p) # Find updated goods prices - new_p_m = firm.get_pm(new_w, Y_vec, L_vec, p, "SS") + new_p_m = firm.get_pm(new_w, Y_vec, L_vec, p, "SS", vectorized=True) new_p_m = new_p_m / new_p_m[-1] # normalize prices by industry M new_p_i = np.dot(p.io_matrix, new_p_m) new_p_tilde = aggr.get_ptilde(new_p_i, p.tau_c[-1, :], p.alpha_c) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 9ad35fe91..0499e71fd 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -636,7 +636,7 @@ def run_TPI(p, client=None): # compute goods prices p_m = np.ones((p.T + p.S, p.M)) * ss_vars["p_m"].reshape(1, p.M) p_m[: p.T, :] = firm.get_pm( - w[: p.T], Y_vec_init[: p.T, :], L_vec_init[: p.T, :], p, "TPI" + w[: p.T], Y_vec_init[: p.T, :], L_vec_init[: p.T, :], p, "TPI", vectorized=True ) p_m = p_m / p_m[:, -1].reshape( p.T + p.S, 1 @@ -652,7 +652,7 @@ def run_TPI(p, client=None): ) # repeat with updated w p_m[: p.T, :] = firm.get_pm( - w[: p.T], Y_vec_init[: p.T, :], L_vec_init[: p.T, :], p, "TPI" + w[: p.T], Y_vec_init[: p.T, :], L_vec_init[: p.T, :], p, "TPI", vectorized=True ) p_m = p_m / p_m[:, -1].reshape( p.T + p.S, 1 @@ -1150,7 +1150,7 @@ def run_TPI(p, client=None): ) # compute new prices - new_p_m = firm.get_pm(wnew, Y_vec, L_vec, p, "TPI") + new_p_m = firm.get_pm(wnew, Y_vec, L_vec, p, "TPI", vectorized=True) new_p_m = new_p_m / new_p_m[:, -1].reshape( p.T, 1 ) # normalize prices by industry M diff --git a/ogcore/firm.py b/ogcore/firm.py index 352d443f1..9960a20a2 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -555,7 +555,7 @@ def get_cost_of_capital(r, p, method, m=-1): return cost_of_capital -def get_pm(w, Y_vec, L_vec, p, method, m=-1): +def get_pm(w, Y_vec, L_vec, p, method, m=0, *, vectorized=False): r""" Find prices for outputs from each industry. @@ -565,7 +565,7 @@ def get_pm(w, Y_vec, L_vec, p, method, m=-1): This function now accepts an optional industry index ``m``. If ``m`` is specified, a single industry's price path is returned (shape ``T`` or - scalar in ``SS``). When ``m==-1`` (the default) prices for all industries + scalar in ``SS``). When ``vectorized=True`` prices for all industries are computed using vectorized operations rather than a Python loop. Args: @@ -574,11 +574,12 @@ def get_pm(w, Y_vec, L_vec, p, method, m=-1): L_vec (array_like): labor demand for each industry p (OG-Core Specifications object): model parameters method (str): adjusts calculation dimensions based on 'SS' or 'TPI' - m (int): industry index (default ``-1`` for all industries) + m (int): industry index (default 0). + vectorized (bool): if True, computes prices for all industries at once. Returns: p_m (array_like): output prices for each industry (or a single industry - if ``m`` is not ``-1``) + if not vectorized) """ # reshape inputs according to method if method == "SS": @@ -590,12 +591,12 @@ def get_pm(w, Y_vec, L_vec, p, method, m=-1): L = L_vec.reshape((p.T, p.M)) T = p.T - if m >= 0: + if not vectorized: # compute price for specific industry using existing logic share = 1 - p.gamma[m] - p.gamma_g[m] - MPL = get_MPx(Y[:, m], L[:, m], share, p, method, m).reshape(T) + MPL = get_MPx(Y[:, m], L[:, m], share, p, method, m=m, vectorized=False).reshape(T) pmout = w / MPL - return pmout[0] if method == "SS" else pmout + return float(pmout[0]) if method == "SS" else pmout # vectorized calculation for all industries # create share vector of length M