Constrain priors with symmetric mass distribution (#5981)

* find_constrained_prior now assumes symmetric probability mass above and below upper and lower by default

* Fix typo in docstring

Co-authored-by: Ricardo Vieira <[email protected]>

Co-authored-by: Alexandre Andorra <[email protected]>
Co-authored-by: Ricardo Vieira <[email protected]>

Author: 3 people

pymc/func_utils.py

@@ -11,8 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import warnings
-
 from typing import Callable, Dict, Optional, Union
 
 import aesara.tensor as aet
@@ -33,6 +31,8 @@ def find_constrained_prior(
     init_guess: Dict[str, float],
     mass: float = 0.95,
     fixed_params: Optional[Dict[str, float]] = None,
+    mass_below_lower: Optional[float] = None,
+    **kwargs,
 ) -> Dict[str, float]:
     """
     Find optimal parameters to get `mass` % of probability
@@ -64,6 +64,11 @@ def find_constrained_prior(
         Dictionary of fixed parameters, so that there are only 2 to optimize.
         For instance, for a StudentT, you fix nu to a constant and get the optimized
         mu and sigma.
+    mass_below_lower : float, optional, default None
+        The probability mass below the ``lower`` bound. If ``None``,
+        defaults to ``(1 - mass) / 2``, which implies that the probability
+        mass below the ``lower`` value will be equal to the probability
+        mass above the ``upper`` value.
 
     Returns
     -------
@@ -72,6 +77,11 @@ def find_constrained_prior(
         Dictionary keys are the parameter names and
         dictionary values are the optimized parameter values.
 
+    Notes
+    -----
+    Optional keyword arguments can be passed to ``find_constrained_prior``. These will be
+    delivered to the underlying call to :external:py:func:`scipy.optimize.minimize`.
+
     Examples
     --------
     .. code-block:: python
@@ -96,11 +106,31 @@ def find_constrained_prior(
             init_guess={"mu": 5, "sigma": 2},
             fixed_params={"nu": 7},
         )
+
+    Under some circumstances, you might not want to have the same cumulative
+    probability below the ``lower`` threshold and above the ``upper`` threshold.
+    For example, you might want to constrain an Exponential distribution to
+    find the parameter that yields 90% of the mass below the ``upper`` bound,
+    and have zero mass below ``lower``. You can do that with the following call
+    to ``find_constrained_prior``
+
+    .. code-block:: python
+
+        opt_params = pm.find_constrained_prior(
+            pm.Exponential,
+            lower=0,
+            upper=3.,
+            mass=0.9,
+            init_guess={"lam": 1},
+            mass_below_lower=0,
+        )
     """
     assert 0.01 <= mass <= 0.99, (
         "This function optimizes the mass of the given distribution +/- "
         f"1%, so `mass` has to be between 0.01 and 0.99. You provided {mass}."
     )
+    if mass_below_lower is None:
+        mass_below_lower = (1 - mass) / 2
 
     # exit when any parameter is not scalar:
     if np.any(np.asarray(distribution.rv_op.ndims_params) != 0):
@@ -129,39 +159,39 @@ def find_constrained_prior(
             "need it."
         )
 
-    cdf_error = (pm.math.exp(logcdf_upper) - pm.math.exp(logcdf_lower)) - mass
-    cdf_error_fn = pm.aesaraf.compile_pymc([dist_params], cdf_error, allow_input_downcast=True)
+    target = (aet.exp(logcdf_lower) - mass_below_lower) ** 2
+    target_fn = pm.aesaraf.compile_pymc([dist_params], target, allow_input_downcast=True)
+
+    constraint = aet.exp(logcdf_upper) - aet.exp(logcdf_lower)
+    constraint_fn = pm.aesaraf.compile_pymc([dist_params], constraint, allow_input_downcast=True)
 
     jac: Union[str, Callable]
+    constraint_jac: Union[str, Callable]
     try:
-        aesara_jac = pm.gradient(cdf_error, [dist_params])
+        aesara_jac = pm.gradient(target, [dist_params])
         jac = pm.aesaraf.compile_pymc([dist_params], aesara_jac, allow_input_downcast=True)
+        aesara_constraint_jac = pm.gradient(constraint, [dist_params])
+        constraint_jac = pm.aesaraf.compile_pymc(
+            [dist_params], aesara_constraint_jac, allow_input_downcast=True
+        )
     # when PyMC cannot compute the gradient
     except (NotImplementedError, NullTypeGradError):
         jac = "2-point"
+        constraint_jac = "2-point"
+    cons = optimize.NonlinearConstraint(constraint_fn, lb=mass, ub=mass, jac=constraint_jac)
 
-    opt = optimize.least_squares(cdf_error_fn, x0=list(init_guess.values()), jac=jac)
+    opt = optimize.minimize(
+        target_fn, x0=list(init_guess.values()), jac=jac, constraints=cons, **kwargs
+    )
     if not opt.success:
-        raise ValueError("Optimization of parameters failed.")
+        raise ValueError(
+            f"Optimization of parameters failed.\nOptimization termination details:\n{opt}"
+        )
 
     # save optimal parameters
     opt_params = {
         param_name: param_value for param_name, param_value in zip(init_guess.keys(), opt.x)
     }
     if fixed_params is not None:
         opt_params.update(fixed_params)
-
-    # check mass in interval is not too far from `mass`
-    opt_dist = distribution.dist(**opt_params)
-    mass_in_interval = (
-        pm.math.exp(pm.logcdf(opt_dist, upper)) - pm.math.exp(pm.logcdf(opt_dist, lower))
-    ).eval()
-    if (np.abs(mass_in_interval - mass)) > 0.01:
-        warnings.warn(
-            f"Final optimization has {(mass_in_interval if mass_in_interval.ndim < 1 else mass_in_interval[0])* 100:.0f}% of probability mass between "
-            f"{lower} and {upper} instead of the requested {mass * 100:.0f}%.\n"
-            "You may need to use a more flexible distribution, change the fixed parameters in the "
-            "`fixed_params` dictionary, or provide better initial guesses."
-        )
-
     return opt_params

pymc/tests/test_func_utils.py

@@ -19,31 +19,32 @@
 
 
 @pytest.mark.parametrize(
-    "distribution, lower, upper, init_guess, fixed_params",
+    "distribution, lower, upper, init_guess, fixed_params, mass_below_lower",
     [
-        (pm.Gamma, 0.1, 0.4, {"alpha": 1, "beta": 10}, {}),
-        (pm.Normal, 155, 180, {"mu": 170, "sigma": 3}, {}),
-        (pm.StudentT, 0.1, 0.4, {"mu": 10, "sigma": 3}, {"nu": 7}),
-        (pm.StudentT, 0, 1, {"mu": 5, "sigma": 2, "nu": 7}, {}),
-        (pm.Exponential, 0, 1, {"lam": 1}, {}),
-        (pm.HalfNormal, 0, 1, {"sigma": 1}, {}),
-        (pm.Binomial, 0, 8, {"p": 0.5}, {"n": 10}),
-        (pm.Poisson, 1, 15, {"mu": 10}, {}),
-        (pm.Poisson, 19, 41, {"mu": 30}, {}),
+        (pm.Gamma, 0.1, 0.4, {"alpha": 1, "beta": 10}, {}, None),
+        (pm.Normal, 155, 180, {"mu": 170, "sigma": 3}, {}, None),
+        (pm.StudentT, 0.1, 0.4, {"mu": 10, "sigma": 3}, {"nu": 7}, None),
+        (pm.StudentT, 0, 1, {"mu": 5, "sigma": 2, "nu": 7}, {}, None),
+        (pm.Exponential, 0, 1, {"lam": 1}, {}, 0),
+        (pm.HalfNormal, 0, 1, {"sigma": 1}, {}, 0),
+        (pm.Binomial, 0, 8, {"p": 0.5}, {"n": 10}, None),
+        (pm.Poisson, 1, 15, {"mu": 10}, {}, None),
+        (pm.Poisson, 19, 41, {"mu": 30}, {}, None),
     ],
 )
 @pytest.mark.parametrize("mass", [0.5, 0.75, 0.95])
-def test_find_constrained_prior(distribution, lower, upper, init_guess, fixed_params, mass):
-    with pytest.warns(None) as record:
-        opt_params = pm.find_constrained_prior(
-            distribution,
-            lower=lower,
-            upper=upper,
-            mass=mass,
-            init_guess=init_guess,
-            fixed_params=fixed_params,
-        )
-    assert len(record) == 0
+def test_find_constrained_prior(
+    distribution, lower, upper, init_guess, fixed_params, mass, mass_below_lower
+):
+    opt_params = pm.find_constrained_prior(
+        distribution,
+        lower=lower,
+        upper=upper,
+        mass=mass,
+        init_guess=init_guess,
+        fixed_params=fixed_params,
+        mass_below_lower=mass_below_lower,
+    )
 
     opt_distribution = distribution.dist(**opt_params)
     mass_in_interval = (
@@ -64,7 +65,9 @@ def test_find_constrained_prior(distribution, lower, upper, init_guess, fixed_pa
 def test_find_constrained_prior_error_too_large(
     distribution, lower, upper, init_guess, fixed_params
 ):
-    with pytest.warns(UserWarning, match="instead of the requested 95%"):
+    with pytest.raises(
+        ValueError, match="Optimization of parameters failed.\nOptimization termination details:\n"
+    ):
         pm.find_constrained_prior(
             distribution,
             lower=lower,