[WIP] implement fixed constraints right away.

Author: hmgaudecker

src/skillmodels/constraints.py

@@ -7,7 +7,9 @@
 import skillmodels.transition_functions as tf
 
 
-def get_constraints(dimensions, labels, anchoring_info, update_info, normalizations):
+def get_constraint_tuples(
+    dimensions, labels, anchoring_info, update_info, normalizations
+):
     """Generate constraints implied by the model specification.
 
     The result can easily be converted to optimagic-style constraints.
@@ -71,6 +73,9 @@ def add_bounds(params_df, bounds_distance=0.0):
         )
         if "lower_bound" not in df.columns:
             df["lower_bound"] = -np.inf
+        if "upper_bound" not in df.columns:
+            df["upper_bound"] = np.inf
+
         df.loc["meas_sds", "lower_bound"] = bounds_distance
         df.loc["shock_sds", "lower_bound"] = bounds_distance
 

src/skillmodels/maximization_inputs.py

@@ -1,13 +1,15 @@
 import functools
+from typing import Any
 
 import jax
 import jax.numpy as jnp
 import numpy as np
+import optimagic as om
 import pandas as pd
 
 import skillmodels.likelihood_function as lf
 import skillmodels.likelihood_function_debug as lfd
-from skillmodels.constraints import add_bounds, get_constraints
+from skillmodels.constraints import add_bounds, get_constraint_tuples
 from skillmodels.kalman_filters import calculate_sigma_scaling_factor_and_weights
 from skillmodels.params_index import get_params_index
 from skillmodels.parse_params import create_parsing_info
@@ -41,7 +43,7 @@ def get_maximization_inputs(model_dict, data):
         constraints (list): List of optimagic constraints that are implied by the
             model specification.
         params_template (pd.DataFrame): Parameter DataFrame with correct index and
-            bounds but with empty value column.
+            bounds. The value column is empty except for the fixed constraints.
         data_aug (pd.DataFrame): DataFrame with augmented data. If model contains
             investment factors, we double up the number of periods in order to add
 
@@ -123,26 +125,32 @@ def debug_loglike(params):
         tmp["value"] = float(tmp["value"])
         return process_debug_data(debug_data=tmp, model=model)
 
-    constr = get_constraints(
+    _constraints_tuples = get_constraint_tuples(
         dimensions=model["dimensions"],
         labels=model["labels"],
         anchoring_info=model["anchoring"],
         update_info=model["update_info"],
         normalizations=model["normalizations"],
     )
 
+    constraints = convert_old_style_constraints(_constraints_tuples)
+
     params_template = pd.DataFrame(columns=["value"], index=p_index)
     params_template = add_bounds(
         params_template,
         model["estimation_options"]["bounds_distance"],
     )
+    params_template = _fill_fixed_constraints(
+        params_template=params_template,
+        constraints_tuples=_constraints_tuples,
+    )
 
     out = {
         "loglike": loglike,
         "loglikeobs": loglikeobs,
         "debug_loglike": debug_loglike,
         "loglike_and_gradient": loglike_and_gradient,
-        "constraints": constr,
+        "constraints": constraints,
         "params_template": params_template,
     }
 
@@ -218,3 +226,49 @@ def _get_jnp_params_vec(params, target_index):
 
     vec = jnp.array(params.reindex(target_index)["value"].to_numpy())
     return vec
+
+
+def _sel(params, loc):
+    return params.loc[loc]
+
+
+def convert_old_style_constraints(old_style):
+    # Need this in many cases, anyhow -- so just impose to simplify code below!
+    new_style = []
+    for oc in old_style:
+        if oc["type"] == "pairwise_equality":
+            new_style.append(
+                om.PairwiseEqualityConstraint(
+                    selectors=[functools.partial(_sel, loc=loc) for loc in oc["locs"]]
+                )
+            )
+        else:
+            sel = functools.partial(_sel, loc=oc["loc"])
+            if oc["type"] == "fixed":
+                new_style.append(om.FixedConstraint(selector=sel))
+            elif oc["type"] == "equality":
+                new_style.append(om.EqualityConstraint(selector=sel))
+            elif oc["type"] == "probability":
+                new_style.append(om.ProbabilityConstraint(selector=sel))
+            elif oc["type"] == "increasing":
+                new_style.append(om.IncreasingConstraint(selector=sel))
+            else:
+                raise TypeError(oc["type"])
+    return new_style
+
+
+def _fill_fixed_constraints(
+    params_template: pd.DataFrame,
+    constraints_tuples: list[dict[str, Any]],
+) -> pd.DataFrame:
+    params = params_template.copy()
+    for constr in constraints_tuples:
+        if constr["type"] == "fixed":
+            params.loc[constr["loc"], "value"] = constr["value"]
+
+    # Check that fixed constraints are valid
+    fixed = params[params["value"].notna()]
+    invalid = fixed.query("value < lower_bound or value > upper_bound")
+    if len(invalid) > 0:
+        raise ValueError(f"Invalid fixed constraints:\n\n{invalid}")
+    return params

tests/test_constraints.py

@@ -28,6 +28,7 @@ def test_add_bounds():
     )
     expected = df.copy(deep=True)
     expected["lower_bound"] = [0.1] * 5 + [0.1, -np.inf, 0.1, -np.inf, 0.1]
+    expected["upper_bound"] = np.inf
 
     calculated = add_bounds(df, 0.1)
     assert_frame_equal(calculated, expected)