Merge pull request PSLmodels#1149 from SeaCelo/fix/J-group-vector-length

Validate beta_annual and chi_b length against the number of income groups

Author: jdebacker

ogcore/parameters.py

@@ -87,6 +87,28 @@ def compute_default_params(self):
         self.T = int(self.T)
         self.J = len(self.lambdas)
 
+        # beta_annual and chi_b carry one value per lifetime-income group, but
+        # paramtools validates values and nesting depth, not list length, so a
+        # wrong-length vector used to flow through silently: a short one
+        # crashed much later (IndexError deep in the SS solve, far from the
+        # cause) and a long one was silently truncated to the first J entries
+        # -- wrong economics with no symptom. Guard here, where J is first
+        # known: one value broadcasts to all J groups, J values pass through,
+        # a uniform vector of any other length is reshaped losslessly
+        # (smaller-J runs that keep the uniform defaults rely on this), and a
+        # non-uniform mismatch raises immediately.
+        for param in ("beta_annual", "chi_b"):
+            val = np.asarray(getattr(self, param)).ravel()
+            if val.shape[0] != self.J:
+                if val.shape[0] == 1 or np.all(val == val[0]):
+                    setattr(self, param, np.full(self.J, val[0]))
+                else:
+                    raise ValueError(
+                        f"{param} must have either 1 value (applied to all "
+                        f"J groups) or J={self.J} values; got "
+                        f"{val.shape[0]} non-uniform values."
+                    )
+
         # get parameters of elliptical utility function
         self.b_ellipse, self.upsilon = elliptical_u_est.estimation(
             self.frisch, self.ltilde

tests/test_parameters.py

@@ -194,3 +194,28 @@ def test_expand_taxfunc_params():
     assert len(specs.etr_params) == specs.T + specs.S
     assert len(specs.etr_params[0]) == specs.S
     assert specs.etr_params[0][0][0] == 0.35
+
+
+def test_J_dimensioned_length_guard():
+    """beta_annual / chi_b must have 1 or J values; see issue #1146.
+
+    Wrong-length vectors used to flow through silently: a short one crashed
+    much later (IndexError deep in the SS solve) and a long one was silently
+    truncated to the first J entries.
+    """
+    # one value broadcasts to all J groups
+    specs = Specifications()
+    specs.update_specifications({"beta_annual": [0.95], "chi_b": [50.0]})
+    assert specs.beta.shape == (specs.J,)
+    assert specs.chi_b.shape == (specs.J,)
+    assert np.allclose(specs.chi_b, 50.0)
+
+    # a uniform vector of the wrong length is reshaped losslessly
+    specs2 = Specifications()
+    specs2.update_specifications({"beta_annual": [0.96, 0.96, 0.96]})
+    assert specs2.beta.shape == (specs2.J,)
+
+    # a non-uniform vector of the wrong length raises immediately
+    specs3 = Specifications()
+    with pytest.raises(ValueError, match="beta_annual"):
+        specs3.update_specifications({"beta_annual": [0.94, 0.95, 0.96]})