diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 300729abbd..f9d46986ad 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -113,8 +113,10 @@ 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)). - Improve sampling, increase default number of particles [5229](https://github.com/pymc-devs/pymc3/pull/5229). +- 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). - New features for `pm.Data` containers: - - With `pm.Data(..., mutable=True/False)`, or by using `pm.MutableData` vs. `pm.ConstantData` one can now create `TensorConstant` data variables. They can be more performant and compatible in situtations where a variable doesn't need to be changed via `pm.set_data()`. See [#5295](https://github.com/pymc-devs/pymc/pull/5295). + - With `pm.Data(..., mutable=True/False)`, or by using `pm.MutableData` vs. `pm.ConstantData` one can now create `TensorConstant` data variables. They can be more performant and compatible in situations where a variable doesn't need to be changed via `pm.set_data()`. See [#5295](https://github.com/pymc-devs/pymc/pull/5295). - New named dimensions can be introduced to the model via `pm.Data(..., dims=...)`. For mutable data variables (see above) the lengths of these dimensions are symbolic, so they can be re-sized via `pm.set_data()`. - `pm.Data` now passes additional kwargs to `aesara.shared`/`at.as_tensor`. [#5098](https://github.com/pymc-devs/pymc/pull/5098). - ... diff --git a/pymc/__init__.py b/pymc/__init__.py index 1f7665d7ee..a55aaf2200 100644 --- a/pymc/__init__.py +++ b/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, diff --git a/pymc/func_utils.py b/pymc/func_utils.py new file mode 100644 index 0000000000..63af83d305 --- /dev/null +++ b/pymc/func_utils.py @@ -0,0 +1,163 @@ +# 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 aesara.gradient import NullTypeGradError +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. + + Examples + -------- + .. code-block:: python + + # get parameters obeying constraints + opt_params = pm.find_constrained_prior( + pm.Gamma, lower=0.1, upper=0.4, mass=0.75, init_guess={"alpha": 1, "beta": 10} + ) + + # use these parameters to draw random samples + samples = pm.Gamma.dist(**opt_params, size=100).eval() + + # use these parameters in a model + with pm.Model(): + x = pm.Gamma('x', **opt_params) + + # specify fixed values before optimization + opt_params = pm.find_constrained_prior( + pm.StudentT, + lower=0, + upper=1, + init_guess={"mu": 5, "sigma": 2}, + fixed_params={"nu": 7}, + ) + """ + 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 + except (NotImplementedError, NullTypeGradError): + 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 diff --git a/pymc/tests/test_func_utils.py b/pymc/tests/test_func_utils.py new file mode 100644 index 0000000000..d57e7212fa --- /dev/null +++ b/pymc/tests/test_func_utils.py @@ -0,0 +1,120 @@ +# Copyright 2020 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 numpy as np +import pytest + +import pymc as pm + + +@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}), + (pm.Poisson, 1, 15, {"mu": 10}, {}), + (pm.Poisson, 19, 41, {"mu": 30}, {}), + ], +) +@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 + + +@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]])}, + )