Fix FixedConstraintWithValue.loc type and test expectations

Remove list from loc type union, convert callers to tuple().
Update anchoring test expectations from list to tuple.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: hmgaudecker

CLAUDE.md

@@ -93,7 +93,7 @@ simulate_dataset() → Simulate states (with optional policy effects)
   computation using Kalman filtering. The debug variant is not jitted and returns
   intermediate results (residuals, contributions, filtered states).
 - **constraints.py**: Generates parameter constraints (bounds, equalities from stagemap,
-  fixed values) for optimization. Exports `get_constraints()`, `get_constraints()`,
+  fixed values) for optimization. Exports `get_constraints()`,
   `enforce_fixed_constraints()`, `add_bounds()`, `FixedConstraintWithValue`.
 - **parse_params.py**: Converts flat parameter vectors to structured model parameters.
   Exports `create_parsing_info()` and `parse_params()`.
@@ -196,7 +196,7 @@ These are not in `__all__` but are imported directly by application projects:
 - `skillmodels.types.ProcessedModel`, `EndogenousFactorsInfo`
 - `skillmodels.decorators.register_params` — essential for custom transition functions
 - `skillmodels.constraints.get_constraints`, `enforce_fixed_constraints`,
-  `FixedConstraintWithValue`, `_sel`
+  `FixedConstraintWithValue`, `select_by_loc`
 - `skillmodels.utilities.extract_factors`, `update_parameter_values`
 - `skillmodels.process_data.pre_process_data`
 - `skillmodels.correlation_heatmap.get_measurements_corr`, `get_quasi_scores_corr`,

src/skillmodels/constraints.py

@@ -35,7 +35,7 @@ class FixedConstraintWithValue(om.FixedConstraint):
     DataFrame) and `value` (the value to set before optimization).
     """
 
-    loc: pd.MultiIndex | list | tuple | str | None = None
+    loc: pd.MultiIndex | tuple | str | None = None
     """Parameter location in the params DataFrame."""
     value: float | None = None
     """Value to enforce on the parameter."""
@@ -202,7 +202,9 @@ def _get_mixture_weights_constraints(
                 value=1.0,
             ),
         ]
-    return [om.ProbabilityConstraint(selector=functools.partial(select_by_loc, loc=loc))]
+    return [
+        om.ProbabilityConstraint(selector=functools.partial(select_by_loc, loc=loc))
+    ]
 
 
 def _get_stage_constraints(
@@ -231,12 +233,16 @@ def _get_stage_constraints(
             loc_q = [("shock_sds", p) for p in stage_periods]
             constraints.append(
                 om.PairwiseEqualityConstraint(
-                    selectors=[functools.partial(select_by_loc, loc=loc) for loc in loc_trans],
+                    selectors=[
+                        functools.partial(select_by_loc, loc=loc) for loc in loc_trans
+                    ],
                 ),
             )
             constraints.append(
                 om.PairwiseEqualityConstraint(
-                    selectors=[functools.partial(select_by_loc, loc=loc) for loc in loc_q],
+                    selectors=[
+                        functools.partial(select_by_loc, loc=loc) for loc in loc_q
+                    ],
                 ),
             )
 
@@ -292,7 +298,9 @@ def _get_initial_states_constraints(
             ("initial_states", 0, f"mixture_{emf}", factors[0])
             for emf in range(n_mixtures)
         ]
-        return [om.IncreasingConstraint(selector=functools.partial(select_by_loc, loc=locs))]
+        return [
+            om.IncreasingConstraint(selector=functools.partial(select_by_loc, loc=locs))
+        ]
     return []
 
 
@@ -352,10 +360,11 @@ def _get_anchoring_constraints(  # noqa: C901
         for period, meas in anchoring_updates:
             locs.append(("controls", period, meas, "constant"))
         if locs:
+            loc = tuple(locs)
             constraints.append(
                 FixedConstraintWithValue(
-                    selector=functools.partial(select_by_loc, loc=locs),
-                    loc=locs,
+                    selector=functools.partial(select_by_loc, loc=loc),
+                    loc=loc,
                     value=0,
                 ),
             )
@@ -366,10 +375,11 @@ def _get_anchoring_constraints(  # noqa: C901
             for cont in [c for c in controls if c != "constant"]:
                 ind_tups.append(("controls", period, meas, cont))
         if ind_tups:
+            loc = tuple(ind_tups)
             constraints.append(
                 FixedConstraintWithValue(
-                    selector=functools.partial(select_by_loc, loc=ind_tups),
-                    loc=ind_tups,
+                    selector=functools.partial(select_by_loc, loc=loc),
+                    loc=loc,
                     value=0,
                 ),
             )
@@ -383,10 +393,11 @@ def _get_anchoring_constraints(  # noqa: C901
                 ind_tups.append(("loadings", period, meas, factor))
 
         if ind_tups:
+            loc = tuple(ind_tups)
             constraints.append(
                 FixedConstraintWithValue(
-                    selector=functools.partial(select_by_loc, loc=ind_tups),
-                    loc=ind_tups,
+                    selector=functools.partial(select_by_loc, loc=loc),
+                    loc=loc,
                     value=1,
                 ),
             )

src/skillmodels/transition_functions.py

@@ -37,7 +37,7 @@
 from jax import Array
 
 
-def _sel(params: Any, loc: Any) -> Any:  # noqa: ANN401
+def select_by_loc(params: Any, loc: Any) -> Any:  # noqa: ANN401
     """Select parameters by location."""
     return params.loc[loc]
 
@@ -68,7 +68,7 @@ def identity_constraints_linear(
         loc = ("transition", aug_period, factor, regressor)
         constraints.append(
             FixedConstraintWithValue(
-                selector=functools.partial(_sel, loc=loc),
+                selector=functools.partial(select_by_loc, loc=loc),
                 loc=loc,
                 value=val,
             )
@@ -122,7 +122,7 @@ def identity_constraints_translog(
         loc = ("transition", aug_period, factor, regressor)
         constraints.append(
             FixedConstraintWithValue(
-                selector=functools.partial(_sel, loc=loc),
+                selector=functools.partial(select_by_loc, loc=loc),
                 loc=loc,
                 value=val,
             )
@@ -158,7 +158,7 @@ def constraints_log_ces(
     """Constraints for log_ces production function."""
     names = params_log_ces(factors)
     loc = [("transition", aug_period, factor, name) for name in names[:-1]]
-    return om.ProbabilityConstraint(selector=functools.partial(_sel, loc=loc))
+    return om.ProbabilityConstraint(selector=functools.partial(select_by_loc, loc=loc))
 
 
 def identity_constraints_log_ces(
@@ -244,7 +244,7 @@ def identity_constraints_linear_and_squares(
         loc = ("transition", aug_period, factor, regressor)
         constraints.append(
             FixedConstraintWithValue(
-                selector=functools.partial(_sel, loc=loc),
+                selector=functools.partial(select_by_loc, loc=loc),
                 loc=loc,
                 value=val,
             )

tests/test_constraints.py

@@ -334,12 +334,12 @@ def test_anchoring_constraints_for_constants(anch_uinfo) -> None:
 
     expected = [
         {
-            "loc": [
+            "loc": (
                 ("controls", 0, "outcome_f1", "constant"),
                 ("controls", 0, "outcome_f2", "constant"),
                 ("controls", 1, "outcome_f1", "constant"),
                 ("controls", 1, "outcome_f2", "constant"),
-            ],
+            ),
             "type": "fixed",
             "value": 0,
         },
@@ -368,7 +368,7 @@ def test_anchoring_constraints_for_controls(anch_uinfo) -> None:
 
     expected = [
         {
-            "loc": [
+            "loc": (
                 ("controls", 0, "outcome_f1", "c1"),
                 ("controls", 0, "outcome_f1", "c2"),
                 ("controls", 0, "outcome_f2", "c1"),
@@ -377,7 +377,7 @@ def test_anchoring_constraints_for_controls(anch_uinfo) -> None:
                 ("controls", 1, "outcome_f1", "c2"),
                 ("controls", 1, "outcome_f2", "c1"),
                 ("controls", 1, "outcome_f2", "c2"),
-            ],
+            ),
             "type": "fixed",
             "value": 0,
         },
@@ -401,12 +401,12 @@ def test_anchoring_constraints_for_loadings(anch_uinfo) -> None:
 
     expected = [
         {
-            "loc": [
+            "loc": (
                 ("loadings", 0, "outcome_f1", "f1"),
                 ("loadings", 0, "outcome_f2", "f2"),
                 ("loadings", 1, "outcome_f1", "f1"),
                 ("loadings", 1, "outcome_f2", "f2"),
-            ],
+            ),
             "type": "fixed",
             "value": 1,
         },