From 34e478a819331a99c91ac8c169214608c9bc30e9 Mon Sep 17 00:00:00 2001
From: Thomas Vandal <thomas.vandal@umontreal.ca>
Date: Wed, 6 Mar 2024 10:00:15 -0500
Subject: [PATCH] Add back estimate_inverse_gamma_parameters in utils (#46)

---
 src/pymc_ext/utils.py | 53 ++++++++++++++++++++++++++++++++++++++++++-
 tests/utils_test.py   | 25 +++++++++++++++++++-
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/src/pymc_ext/utils.py b/src/pymc_ext/utils.py
index 0d2146d..61e18ef 100644
--- a/src/pymc_ext/utils.py
+++ b/src/pymc_ext/utils.py
@@ -1,7 +1,14 @@
-__all__ = ["Evaluator", "eval_in_model", "sample_inference_data"]
+__all__ = [
+    "Evaluator",
+    "eval_in_model",
+    "sample_inference_data",
+    "estimate_inverse_gamma_parameters",
+]
 
 import numpy as np
 import pymc as pm
+from scipy.optimize import root
+from scipy.special import gammaincc
 
 
 class Evaluator:
@@ -75,3 +82,47 @@ def sample(**kwargs):
     kwargs["target_accept"] = kwargs.get("target_accept", 0.9)
     kwargs["init"] = kwargs.get("init", "adapt_full")
     return pm.sample(**kwargs)
+
+
+def estimate_inverse_gamma_parameters(
+    lower, upper, target=0.01, initial=None, **kwargs
+):
+    r"""Estimate an inverse Gamma with desired tail probabilities
+    This method numerically solves for the parameters of an inverse Gamma
+    distribution where the tails have a given probability. In other words
+    :math:`P(x < \mathrm{lower}) = \mathrm{target}` and similarly for the
+    upper bound. More information can be found in `part 4 of this blog post
+    <https://betanalpha.github.io/assets/case_studies/gp_part3/part3.html>`_.
+    Args:
+        lower (float): The location of the lower tail
+        upper (float): The location of the upper tail
+        target (float, optional): The desired tail probability
+        initial (ndarray, optional): An initial guess for the parameters
+            ``alpha`` and ``beta``
+    Raises:
+        RuntimeError: If the solver does not converge.
+    Returns:
+        dict: A dictionary with the keys ``alpha`` and ``beta`` for the
+        parameters of the distribution.
+    """
+    lower, upper = np.sort([lower, upper])
+    if initial is None:
+        initial = np.array([2.0, 0.5 * (lower + upper)])
+    if np.shape(initial) != (2,) or np.any(np.asarray(initial) <= 0.0):
+        raise ValueError("invalid initial guess")
+
+    def obj(x):
+        a, b = np.exp(x)
+        return np.array(
+            [
+                gammaincc(a, b / lower) - target,
+                1 - gammaincc(a, b / upper) - target,
+            ]
+        )
+
+    result = root(obj, np.log(initial), method="hybr", **kwargs)
+    if not result.success:
+        raise RuntimeError(
+            "failed to find parameter estimates: \n{0}".format(result.message)
+        )
+    return dict(zip(("alpha", "beta"), np.exp(result.x)))
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 099224e..afbc7c5 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -1,7 +1,9 @@
 import numpy as np
 import pymc as pm
+import pytest
+from scipy.stats import invgamma
 
-from pymc_ext.utils import eval_in_model
+from pymc_ext.utils import estimate_inverse_gamma_parameters, eval_in_model
 
 
 def test_eval_in_model(seed=123409):
@@ -39,3 +41,24 @@ def test_eval_in_model_list(seed=123409):
         x_eval, y_eval = eval_in_model([x, y])
         assert np.allclose(x_eval, x_val)
         assert np.allclose(y_eval, y_val)
+
+
+@pytest.mark.parametrize(
+    "lower, upper, target",
+    [(1.0, 2.0, 0.01), (0.01, 0.1, 0.1), (10.0, 25.0, 0.01)],
+)
+def test_estimate_inverse_gamma_parameters(lower, upper, target):
+    np.random.seed(20200409)
+
+    params = estimate_inverse_gamma_parameters(lower, upper, target=target)
+    dist = invgamma(params["alpha"], scale=params["beta"])
+    assert np.allclose(dist.cdf(lower), target)
+    assert np.allclose(1 - dist.cdf(upper), target)
+
+    samples = pm.draw(pm.InverseGamma.dist(**params), draws=10000)
+    assert np.allclose(
+        (samples < lower).sum() / len(samples), target, atol=1e-2
+    )
+    assert np.allclose(
+        (samples > upper).sum() / len(samples), target, atol=1e-2
+    )