Allow fixed entries in probability constraints

CHANGES.md

@@ -5,6 +5,14 @@ chronological order. We follow [semantic versioning](https://semver.org/) and al
 releases are available on [Anaconda.org](https://anaconda.org/optimagic-dev/optimagic).
 
 
+## Unreleased
+
+- Allows a `FixedConstraint` pinning selected elements to any value in `[0, 1)` (with
+  their sum strictly less than one) to coexist with a `ProbabilityConstraint` on the
+  same parameters. The fixed entries are held at their values and the remaining free
+  entries are optimised on a simplex summing to `1 - sum(fixed values)`.
+
+
 ## 0.5.3
 
 This release introduces **multi-backend plotting** with support for matplotlib, bokeh,

docs/source/how_to/how_to_constraints.md

@@ -254,6 +254,14 @@ full parameter vector and returns the subset of parameters that should be constr
         >>> res.params.round(2) # doctest: +SKIP
         array([0.53, 0.33, 0.13, 0.  , 0.2 , 0.  ])
 
+    You can combine a ``ProbabilityConstraint`` with a ``FixedConstraint`` that
+    pins some of the selected entries. The fixed values must each be in
+    ``[0, 1)``, sum to strictly less than one, and leave at least two free
+    entries. The remaining free entries are then optimised on the simplex
+    that sums to ``1 - sum(fixed values)``. This is useful when part of a
+    larger model does not need to contribute to the probability, or when
+    some weights are set externally.
+
 
 ```
 

src/optimagic/parameters/check_constraints.py

@@ -226,7 +226,9 @@ def check_fixes_and_bounds(constr_info, transformations, parnames):
             if subset["is_fixed_to_value"].any():
                 problematic = subset["index"][subset["is_fixed_to_value"]]
                 raise InvalidConstraintError(
-                    prob_msg.format(constr["type"], problematic)
+                    "Fixed values inside a probability constraint should have "
+                    "been folded into the selector before this check; the "
+                    f"following parameters still carry a fix:\n{problematic}"
                 )
             finite_bounds = np.isfinite(subset["lower_bounds"]) | np.isfinite(
                 subset["upper_bounds"]

src/optimagic/parameters/consolidate_constraints.py

@@ -68,6 +68,13 @@ def consolidate_constraints(
         if not constr_info["is_fixed_to_value"][c["index"]].all()
     ]
 
+    other_constraints = _fold_fixes_into_probability_constraints(
+        other_constraints,
+        fixed_values=constr_info["fixed_values"],
+        is_fixed_to_value=constr_info["is_fixed_to_value"],
+        param_names=param_names,
+    )
+
     (
         other_constraints,
         lower_bounds,
@@ -223,6 +230,88 @@ def _consolidate_fixes_with_equality_constraints(
     return fixed_value
 
 
+def _fold_fixes_into_probability_constraints(
+    constraints, fixed_values, is_fixed_to_value, param_names
+):
+    """Shrink probability selectors to exclude parameters with compatible fixes.
+
+    A probability constraint requires that the selected parameters are non-negative
+    and sum to 1. If some of the selected parameters are additionally pinned via a
+    fixed constraint to values in ``[0, 1)`` that sum to less than 1, the constraint
+    reduces to a probability constraint over the remaining free parameters, summing
+    to ``1 - sum(fixed values)``. This function rewrites such constraints by
+    removing fixed indices from the selector and attaching the implied simplex
+    target as ``sum_target`` on the returned transformation dict. When all fixed
+    values are 0.0, the ``sum_target`` is 1.0 and the key is omitted so that the
+    no-op path produces a transformation dict identical to the pre-existing one.
+
+    Args:
+        constraints (list): List of constraint dictionaries where selectors have
+            been processed into an ``"index"`` field.
+        fixed_values (np.ndarray): 1d array with fixed values (NaN where free).
+        is_fixed_to_value (np.ndarray): 1d boolean array that is True for
+            parameters fixed to a value.
+        param_names (list): Names of the flattened parameters. Used for error
+            messages only.
+
+    Returns:
+        list: Constraints with probability selectors reduced to their non-fixed
+        subsets, carrying ``sum_target`` when the fixed values are non-zero.
+        Non-probability constraints are returned unchanged.
+
+    """
+    out = []
+    for constr in constraints:
+        if constr["type"] != "probability":
+            out.append(constr)
+            continue
+
+        index = list(constr["index"])
+        fixed_mask = is_fixed_to_value[index]
+        if not fixed_mask.any():
+            out.append(constr)
+            continue
+
+        fixed_idx = [
+            i for i, is_fixed in zip(index, fixed_mask, strict=True) if is_fixed
+        ]
+        free_idx = [
+            i for i, is_fixed in zip(index, fixed_mask, strict=True) if not is_fixed
+        ]
+
+        negative = [i for i in fixed_idx if fixed_values[i] < 0.0]
+        if negative:
+            problematic = [param_names[i] for i in negative]
+            raise InvalidConstraintError(
+                "Parameters inside a probability constraint that are fixed to a "
+                "value must be fixed to a value in [0, 1). This is violated for:"
+                f"\n{problematic}"
+            )
+
+        fixed_sum = float(sum(fixed_values[i] for i in fixed_idx))
+        if fixed_sum >= 1.0:
+            problematic = [param_names[i] for i in fixed_idx]
+            raise InvalidConstraintError(
+                "The fixed values inside a probability constraint must sum to "
+                f"strictly less than 1; got sum={fixed_sum} for:\n{problematic}"
+            )
+
+        if len(free_idx) < 2:
+            problematic = [param_names[i] for i in index]
+            raise InvalidConstraintError(
+                "A probability constraint must have at least two non-fixed "
+                "selected parameters after folding in fixes. This is violated "
+                f"for:\n{problematic}"
+            )
+
+        new_constr = {**constr, "index": free_idx}
+        if fixed_sum > 0.0:
+            new_constr["sum_target"] = 1.0 - fixed_sum
+        out.append(new_constr)
+
+    return out
+
+
 def _consolidate_bounds_with_equality_constraints(
     equality_constraints, lower_bounds, upper_bounds
 ):

src/optimagic/parameters/kernel_transformations.py

@@ -311,22 +311,33 @@ def probability_to_internal_jacobian(external_values, constr):
 
 
 def probability_from_internal(internal_values, constr):
-    """Reparametrize probability constrained parameters from internal."""
-    return internal_values / internal_values.sum()
+    """Reparametrize probability constrained parameters from internal.
+
+    The result is rescaled to sum to ``constr["sum_target"]`` when present
+    (defaults to 1). This supports probability constraints where some of the
+    selected parameters are fixed to non-zero values and the remaining free
+    entries must sum to ``1 - sum(fixed values)``.
+    """
+    sum_target = 1.0 if constr is None else constr.get("sum_target", 1.0)
+    return sum_target * internal_values / internal_values.sum()
 
 
 def probability_from_internal_jacobian(internal_values, constr):
     r"""Jacobian of ``probability_from_internal``.
 
-    Let :math:`x := \text{internal}`. The function ``probability_from_internal``
-    has the following structure
+    Let :math:`x := \text{internal}` and :math:`S := \text{sum\_target}`. The
+    function ``probability_from_internal`` has the following structure
 
-    .. math::`f: \mathbb{R}^m \to \mathbb{R}^m, x \mapsto \frac{1}{x^\top 1} x`
+    .. math::  f: \mathbb{R}^m \to \mathbb{R}^m,
+               x \mapsto \frac{S}{x^\top 1} x
 
     where :math:`1` denotes a vector of all ones and :math:`I_m` the identity
     matrix. The jacobian can be computed as
 
-    .. math::  J(f)(x) = \frac{1}{\sigma} I_m - \frac{1}{\sigma^2} 1 x^\top
+    .. math::  J(f)(x) = \frac{S}{\sigma} I_m - \frac{S}{\sigma^2} 1 x^\top
+
+    When ``sum_target`` is absent from ``constr`` (or ``constr`` is ``None``),
+    :math:`S = 1` and the Jacobian reduces to the unscaled form.
 
     Args:
         internal_values (np.ndarray): Internal (positive) values.
@@ -341,7 +352,8 @@ def probability_from_internal_jacobian(internal_values, constr):
     left = np.eye(dim)
     right = (np.ones((dim, dim)) * (internal_values / sigma)).T
 
-    deriv = (left - right) / sigma
+    sum_target = 1.0 if constr is None else constr.get("sum_target", 1.0)
+    deriv = sum_target * (left - right) / sigma
     return deriv
 
 

tests/optimagic/optimization/test_with_constraints.py

@@ -345,6 +345,52 @@ def criterion(x):
         )
 
 
+def test_probability_constraint_with_zero_fix_on_selector_element():
+    """Some selected entries are fixed to zero; the remaining simplex is optimised."""
+
+    def criterion(params):
+        # Distances to a target simplex weight away from the fixed entries.
+        target = np.array([0.0, 0.2, 0.0, 0.5, 0.3])
+        return np.sum((params - target) ** 2)
+
+    res = minimize(
+        fun=criterion,
+        params=np.array([0.0, 0.3, 0.0, 0.3, 0.4]),
+        algorithm="scipy_lbfgsb",
+        constraints=[
+            om.ProbabilityConstraint(lambda x: x[[0, 1, 2, 3, 4]]),
+            om.FixedConstraint(lambda x: x[[0, 2]]),
+        ],
+    )
+
+    assert res.params[0] == 0.0
+    assert res.params[2] == 0.0
+    aaae(res.params[[1, 3, 4]].sum(), 1.0)
+    aaae(res.params, [0.0, 0.2, 0.0, 0.5, 0.3], decimal=4)
+
+
+def test_probability_constraint_with_non_zero_fix_on_selector_element():
+    """One selected entry is fixed to 0.2; the remaining free entries sum to 0.8."""
+
+    def criterion(params):
+        target = np.array([0.2, 0.2, 0.3, 0.3])
+        return np.sum((params - target) ** 2)
+
+    res = minimize(
+        fun=criterion,
+        params=np.array([0.2, 0.3, 0.2, 0.3]),
+        algorithm="scipy_lbfgsb",
+        constraints=[
+            om.ProbabilityConstraint(lambda x: x[[0, 1, 2, 3]]),
+            om.FixedConstraint(lambda x: x[[0]]),
+        ],
+    )
+
+    assert res.params[0] == 0.2
+    aaae(res.params[1:].sum(), 0.8)
+    aaae(res.params, [0.2, 0.2, 0.3, 0.3], decimal=4)
+
+
 def test_covariance_constraint_in_2_by_2_case():
     spector_data = sm.datasets.spector.load_pandas()
     spector_data.exog = sm.add_constant(spector_data.exog)

tests/optimagic/parameters/test_check_constraints.py

@@ -2,7 +2,7 @@
 import pytest
 
 import optimagic as om
-from optimagic.exceptions import InvalidParamsError
+from optimagic.exceptions import InvalidConstraintError, InvalidParamsError
 from optimagic.parameters.check_constraints import _iloc
 from optimagic.parameters.constraint_tools import check_constraints
 
@@ -70,6 +70,50 @@ def test_check_constraints_are_satisfied_type_probability():
         )
 
 
+def test_probability_constraint_accepts_zero_fix_on_selector_element():
+    check_constraints(
+        params=np.array([0.0, 0.3, 0.0, 0.3, 0.4]),
+        constraints=[
+            om.ProbabilityConstraint(lambda x: x[[0, 1, 2, 3, 4]]),
+            om.FixedConstraint(lambda x: x[[0, 2]]),
+        ],
+    )
+
+
+def test_probability_constraint_accepts_non_zero_fix_on_selector_element():
+    check_constraints(
+        params=np.array([0.1, 0.3, 0.2, 0.4]),
+        constraints=[
+            om.ProbabilityConstraint(lambda x: x[[0, 1, 2, 3]]),
+            om.FixedConstraint(lambda x: x[[0]]),
+        ],
+    )
+
+
+def test_probability_constraint_rejects_fixes_summing_to_one_or_more():
+    with pytest.raises(InvalidConstraintError, match="sum to strictly less than 1"):
+        check_constraints(
+            params=np.array([0.5, 0.5, 0.0, 0.0]),
+            constraints=[
+                om.ProbabilityConstraint(lambda x: x[[0, 1, 2, 3]]),
+                om.FixedConstraint(lambda x: x[[0, 1]]),
+            ],
+        )
+
+
+def test_probability_constraint_rejects_too_few_free_after_fold():
+    with pytest.raises(
+        InvalidConstraintError, match="at least two non-fixed selected parameters"
+    ):
+        check_constraints(
+            params=np.array([0.0, 0.0, 1.0]),
+            constraints=[
+                om.ProbabilityConstraint(lambda x: x[[0, 1, 2]]),
+                om.FixedConstraint(lambda x: x[[0, 1]]),
+            ],
+        )
+
+
 def test_check_constraints_are_satisfied_type_linear_lower_bound():
     with pytest.raises(InvalidParamsError):
         check_constraints(

tests/optimagic/parameters/test_kernel_transformations.py

@@ -97,6 +97,31 @@ def test_probability_to_internal_jacobian(dim, seed):  # noqa: ARG001
     aaae(deriv, numerical_deriv.derivative, decimal=3)
 
 
+@pytest.mark.parametrize("sum_target", [0.2, 0.5, 0.9])
+def test_probability_from_internal_with_sum_target(sum_target):
+    internal = get_internal_probability(dim=5)
+    constr = {"type": "probability", "index": [0, 1, 2, 3, 4], "sum_target": sum_target}
+
+    external = kt.probability_from_internal(internal, constr)
+
+    assert np.isclose(external.sum(), sum_target)
+    # Internal pivot stays at 1 regardless of sum_target — so the inverse map is
+    # a simple division by the last external entry.
+    assert np.allclose(external / external[-1], internal / internal[-1])
+
+
+@pytest.mark.parametrize("sum_target", [0.2, 0.5, 0.9])
+def test_probability_from_internal_jacobian_with_sum_target(sum_target):
+    internal = get_internal_probability(dim=10)
+    constr = {"type": "probability", "index": list(range(10)), "sum_target": sum_target}
+
+    func = partial(kt.probability_from_internal, constr=constr)
+    numerical_deriv = first_derivative(func, internal)
+    deriv = kt.probability_from_internal_jacobian(internal, constr)
+
+    aaae(deriv, numerical_deriv.derivative, decimal=3)
+
+
 @pytest.mark.parametrize("dim, seed", to_test)
 def test_sdcorr_from_internal_jacobian(dim, seed):  # noqa: ARG001
     internal = get_internal_cholesky(dim)