Function to optimize prior under constraints

RELEASE-NOTES.md

@@ -97,7 +97,8 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
   - Added partial dependence plots and individual conditional expectation plots [5091](https://github.com/pymc-devs/pymc3/pull/5091).
   - Modify how particle weights are computed. This improves accuracy of the modeled function (see [5177](https://github.com/pymc-devs/pymc3/pull/5177)).
 - `pm.Data` now passes additional kwargs to `aesara.shared`. [#5098](https://github.com/pymc-devs/pymc/pull/5098)
-- ...
+- The new `pm.find_optim_prior` function can be used to find the optimal 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

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

pymc/find_optim_prior.py

@@ -0,0 +1,96 @@
+#   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.
+
+from typing import List
+
+import aesara
+import aesara.tensor as aet
+import numpy as np
+
+from scipy import optimize
+
+import pymc as pm
+
+__all__ = ["find_optim_prior"]
+
+
+def find_optim_prior(
+    pm_dist: pm.Distribution,
+    lower: float,
+    upper: float,
+    init_guess: List[float],
+    mass: float = 0.95,
+) -> np.ndarray:
+    """
+    Find optimal parameters to get `mass` % of probability
+    of `pm_dist` between `lower` and `upper`.
+    Note: only works for two-parameter distributions, as there
+    are exactly two constraints.
+
+    Parameters
+    ----------
+    pm_dist : 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: List[float]
+        Initial guess for ``scipy.optimize.least_squares`` to find the
+        optimal parameters of `pm_dist` fitting the interval constraint.
+    mass: float, default to 0.95
+        Share of the probability mass we want between ``lower`` and ``upper``.
+        Defaults to 95%.
+
+    Returns
+    -------
+    The optimized distribution parameters as a numpy array.
+    """
+    if len(pm_dist.rv_op.ndims_params) != 2:
+        raise NotImplementedError(
+            "This function only works for two-parameter distributions, as there are exactly two "
+            "constraints. "
+        )
+
+    dist_params = aet.vector("dist_params")
+    lower_, upper_ = aet.scalars("lower", "upper")
+
+    dist_ = pm_dist.dist(*[dist_params[i] for i in range(len(init_guess))])
+    try:
+        logcdf_lower = pm.logcdf(dist_, lower_)
+        logcdf_upper = pm.logcdf(dist_, upper_)
+    except AttributeError:
+        raise AttributeError(
+            f"You cannot use `find_optim_params` with {pm_dist} -- it doesn't have a logcdf "
+            f"method yet. Open an issue or, even better, a PR on PyMC repo if you really need it."
+        )
+
+    alpha = 1 - mass
+    out = [logcdf_lower - np.log(alpha / 2), logcdf_upper - np.log(1 - alpha / 2)]
+    logcdf = aesara.function([dist_params, lower_, upper_], out)
+
+    try:
+        symb_grad = aet.as_tensor_variable([pm.gradient(o, [dist_params]) for o in out])
+        jac = aesara.function([dist_params, lower_, upper_], symb_grad)
+
+    # when PyMC cannot compute the gradient
+    except Exception:
+        jac = "2-point"
+
+    opt = optimize.least_squares(logcdf, init_guess, jac=jac, args=(lower, upper))
+    if not opt.success:
+        raise ValueError("Optimization of parameters failed.")
+
+    return opt.x

pymc/tests/test_util.py

@@ -142,3 +142,17 @@ def fn(a=UNSET):
     help(fn)
     captured = capsys.readouterr()
     assert "a=UNSET" in captured.out
+
+
+def test_find_optim_prior():
+    MASS = 0.95
+    opt_params = pm.find_optim_prior(pm.Gamma, lower=0.1, upper=0.4, mass=MASS, init_guess=[1, 10])
+    np.testing.assert_allclose(opt_params, np.array([8.47481597, 37.65435601]))
+
+    opt_params = pm.find_optim_prior(
+        pm.Normal, lower=155, upper=180, mass=MASS, init_guess=[170, 3]
+    )
+    np.testing.assert_allclose(opt_params, np.array([167.5000001, 6.37766828]))
+
+    with pytest.raises(NotImplementedError, match="only works for two-parameter distributions"):
+        pm.find_optim_prior(pm.Exponential, lower=0.1, upper=0.4, mass=MASS, init_guess=[1])