diff --git a/ogcore/SS.py b/ogcore/SS.py index f7239f90c..46a45b3c5 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) @@ -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) @@ -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/TPI.py b/ogcore/TPI.py index 7b38ad68f..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 @@ -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 @@ -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 5d454f679..9960a20a2 100644 --- a/ogcore/firm.py +++ b/ogcore/firm.py @@ -311,38 +311,63 @@ 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) + 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=-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 + 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 not vectorized: + 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 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) + 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 not vectorized: + 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=0, *, vectorized=False): r""" Find prices for outputs from each industry. @@ -538,16 +563,25 @@ 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 ``vectorized=True`` 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 0). + vectorized (bool): if True, computes prices for all industries at once. Returns: - p_m (array_like): output prices for each industry + p_m (array_like): output prices for each industry (or a single industry + if not vectorized) """ + # reshape inputs according to method if method == "SS": Y = Y_vec.reshape(1, p.M) L = L_vec.reshape(1, p.M) @@ -556,15 +590,25 @@ 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 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=m, vectorized=False).reshape(T) + pmout = w / MPL + return float(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, vectorized=True + ) # 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 +693,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 +714,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])