From b9d6011eebb4c452df160b543da881aa2946d659 Mon Sep 17 00:00:00 2001 From: suyash469 Date: Wed, 25 Feb 2026 12:08:27 +0530 Subject: [PATCH] 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])