fix: replace bare except with specific exceptions in solve_L, fix adj_cost docstring typo

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")

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

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,24 +555,33 @@ 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.
 
     .. math::
          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)

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])