Function to optimize prior under constraints

RELEASE-NOTES.md

@@ -114,7 +114,8 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
   - Modify how particle weights are computed. This improves accuracy of the modeled function (see [5177](https://github.com/pymc-devs/pymc3/pull/5177)).
   - Improve sampling, increase default number of particles [5229](https://github.com/pymc-devs/pymc3/pull/5229).
 - `pm.Data` now passes additional kwargs to `aesara.shared`. [#5098](https://github.com/pymc-devs/pymc/pull/5098)
-- ...
+- The new `pm.find_constrained_prior` function can be used to find optimized prior parameters of a distribution under some
+  constraints (e.g lower and upper bound). See [#5231](https://github.com/pymc-devs/pymc/pull/5231).
 
 
 ### Internal changes

pymc/__init__.py

@@ -81,6 +81,7 @@ def __set_compiler_flags():
 from pymc.distributions import *
 from pymc.distributions import transforms
 from pymc.exceptions import *
+from pymc.func_utils import find_constrained_prior
 from pymc.math import (
     expand_packed_triangular,
     invlogit,

pymc/func_utils.py

@@ -0,0 +1,138 @@
+#   Copyright 2021 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   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 Dict, Optional
+
+import aesara.tensor as aet
+import numpy as np
+
+from scipy import optimize
+
+import pymc as pm
+
+__all__ = ["find_constrained_prior"]
+
+
+def find_constrained_prior(
+    distribution: pm.Distribution,
+    lower: float,
+    upper: float,
+    init_guess: Dict[str, float],
+    mass: float = 0.95,
+    fixed_params: Optional[Dict[str, float]] = None,
+) -> Dict[str, float]:
+    """
+    Find optimal parameters to get `mass` % of probability
+    of `pm_dist` between `lower` and `upper`.
+    Note: only works for one- and two-parameter distributions, as there
+    are exactly two constraints. Fix some combination of parameters
+    if you want to use it on >=3-parameter distributions.
+
+    Parameters
+    ----------
+    distribution : pm.Distribution
+        PyMC distribution you want to set a prior on.
+        Needs to have a ``logcdf`` method implemented in PyMC.
+    lower : float
+        Lower bound to get `mass` % of probability of `pm_dist`.
+    upper : float
+        Upper bound to get `mass` % of probability of `pm_dist`.
+    init_guess: Dict[str, float]
+        Initial guess for ``scipy.optimize.least_squares`` to find the
+        optimal parameters of `pm_dist` fitting the interval constraint.
+        Must be a dictionary with the name of the PyMC distribution's
+        parameter as keys and the initial guess as values.
+    mass: float, default to 0.95
+        Share of the probability mass we want between ``lower`` and ``upper``.
+        Defaults to 95%.
+    fixed_params: Dict[str, float], Optional, default None
+        Only used when `pm_dist` has at least three parameters.
+        Dictionary of fixed parameters, so that there are only 2 to optimize.
+        For instance, for a StudenT, you fix nu to a constant and get the optimized
+        mu and sigma.
+
+    Returns
+    -------
+    The optimized distribution parameters as a dictionary with the parameters'
+    name as key and the optimized value as value.
+    """
+    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}."
+    )
+
+    # exit when any parameter is not scalar:
+    if np.any(np.asarray(distribution.rv_op.ndims_params) != 0):
+        raise NotImplementedError(
+            "`pm.find_constrained_prior` does not work with non-scalar parameters yet.\n"
+            "Feel free to open a pull request on PyMC repo if you really need this feature."
+        )
+
+    dist_params = aet.vector("dist_params")
+    params_to_optim = {
+        arg_name: dist_params[i] for arg_name, i in zip(init_guess.keys(), range(len(init_guess)))
+    }
+
+    if fixed_params is not None:
+        params_to_optim.update(fixed_params)
+
+    dist = distribution.dist(**params_to_optim)
+
+    try:
+        logcdf_lower = pm.logcdf(dist, pm.floatX(lower))
+        logcdf_upper = pm.logcdf(dist, pm.floatX(upper))
+    except AttributeError:
+        raise AttributeError(
+            f"You cannot use `find_constrained_prior` with {distribution} -- it doesn't have a logcdf "
+            "method yet.\nOpen an issue or, even better, a pull request on PyMC repo if you really "
+            "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)
+
+    try:
+        aesara_jac = pm.gradient(cdf_error, [dist_params])
+        jac = pm.aesaraf.compile_pymc([dist_params], aesara_jac, allow_input_downcast=True)
+    # when PyMC cannot compute the gradient
+    # TODO: use specific gradient, not implemented exception
+    except Exception:
+        jac = "2-point"
+
+    opt = optimize.least_squares(cdf_error_fn, x0=list(init_guess.values()), jac=jac)
+    if not opt.success:
+        raise ValueError("Optimization of parameters failed.")
+
+    # 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_util.py

@@ -142,3 +142,130 @@ def fn(a=UNSET):
     help(fn)
     captured = capsys.readouterr()
     assert "a=UNSET" in captured.out
+
+
+@pytest.mark.parametrize(
+    "distribution, lower, upper, init_guess, fixed_params",
+    [
+        (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}, {}), PyMC Exponential gradient is failing miserably, need to figure out why
+        (pm.HalfNormal, 0, 1, {"sigma": 1}, {}),
+        (pm.Binomial, 0, 8, {"p": 0.5}, {"n": 10}),
+    ],
+)
+@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
+
+    opt_distribution = distribution.dist(**opt_params)
+    mass_in_interval = (
+        pm.math.exp(pm.logcdf(opt_distribution, upper))
+        - pm.math.exp(pm.logcdf(opt_distribution, lower))
+    ).eval()
+    assert np.abs(mass_in_interval - mass) <= 1e-5
+
+
+# test Poisson separately -- hard to optimize precisely when float32
+@pytest.mark.parametrize(
+    "lower, upper, init_guess",
+    [
+        (1, 15, {"mu": 10}),
+        (19, 41, {"mu": 30}),
+    ],
+)
+def test_constrained_prior_poisson(lower, upper, init_guess):
+    distribution = pm.Poisson
+    mass = 0.95
+    with pytest.warns(None) as record:
+        opt_params = pm.find_constrained_prior(
+            distribution,
+            lower=lower,
+            upper=upper,
+            init_guess=init_guess,
+        )
+    assert len(record) == 0
+
+    opt_distribution = distribution.dist(**opt_params)
+    mass_in_interval = (
+        pm.math.exp(pm.logcdf(opt_distribution, upper))
+        - pm.math.exp(pm.logcdf(opt_distribution, lower))
+    ).eval()
+    assert np.abs(mass_in_interval - mass) <= 1e-2  # reduce to 1% tolerance for float32
+
+
+@pytest.mark.parametrize(
+    "distribution, lower, upper, init_guess, fixed_params",
+    [
+        (pm.Gamma, 0.1, 0.4, {"alpha": 1}, {"beta": 10}),
+        (pm.Exponential, 0.1, 1, {"lam": 1}, {}),
+        (pm.Binomial, 0, 2, {"p": 0.8}, {"n": 10}),
+    ],
+)
+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%"):
+        pm.find_constrained_prior(
+            distribution,
+            lower=lower,
+            upper=upper,
+            mass=0.95,
+            init_guess=init_guess,
+            fixed_params=fixed_params,
+        )
+
+
+def test_find_constrained_prior_input_errors():
+    # missing param
+    with pytest.raises(TypeError, match="required positional argument"):
+        pm.find_constrained_prior(
+            pm.StudentT,
+            lower=0.1,
+            upper=0.4,
+            mass=0.95,
+            init_guess={"mu": 170, "sigma": 3},
+        )
+
+    # mass too high
+    with pytest.raises(AssertionError, match="has to be between 0.01 and 0.99"):
+        pm.find_constrained_prior(
+            pm.StudentT,
+            lower=0.1,
+            upper=0.4,
+            mass=0.995,
+            init_guess={"mu": 170, "sigma": 3},
+            fixed_params={"nu": 7},
+        )
+
+    # mass too low
+    with pytest.raises(AssertionError, match="has to be between 0.01 and 0.99"):
+        pm.find_constrained_prior(
+            pm.StudentT,
+            lower=0.1,
+            upper=0.4,
+            mass=0.005,
+            init_guess={"mu": 170, "sigma": 3},
+            fixed_params={"nu": 7},
+        )
+
+    # non-scalar params
+    with pytest.raises(NotImplementedError, match="does not work with non-scalar parameters yet"):
+        pm.find_constrained_prior(
+            pm.MvNormal,
+            lower=0,
+            upper=1,
+            mass=0.95,
+            init_guess={"mu": 5, "cov": np.asarray([[1, 0.2], [0.2, 1]])},
+        )