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

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,24 +555,32 @@ 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.
 
     .. 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 ``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)

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