Enable reparametrization of SciPy distributions (#56)

* Make fit an instance method and deprecate fit_instance

* Update documentation

* Make params_names an instance-level property

* Feedback from copilot

* Enable reparametrization of SciPy distributions

* Self is not available in 3.9

* Move reparametrization logic to ScipyDistribution

* Fix tests

* Fix module on wrapped distributions

* Revert unnecessary change

Author: jonathanberthias

pykelihood/distributions/base.py

@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Generic, TypeVar
+from typing import TYPE_CHECKING, Callable, Generic, Protocol, TypeVar
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -129,7 +129,7 @@ def fit(
         The result of the fit. A new instance is created with the fitted parameters.
         """
         init_parms = self._process_fit_params(**fixed_values)
-        init = type(self)(**init_parms)
+        init = self.with_params(**init_parms)
         data = init._apply_constraints(data)
 
         if x0 is None:
@@ -236,16 +236,91 @@ def __getattr__(self, item: str):
         return getattr(self.fitted, item)
 
 
+class Reparametrization(Protocol):
+    """The blueprint for reparametrization functions.
+
+    A reparametrization function turns the parameters of the reparametrized
+    distribution into the parameters of the base distribution.
+
+    The provided parameters have already been evaluated. Parameters may be
+    renamed, modified, removed, or added to the returned dictionary to match
+    the expected parameters of the base distribution.
+    """
+
+    def __call__(self, params: dict[str, ArrayLike]) -> dict[str, ArrayLike]: ...
+
+
+def no_reparametrization(params: dict[str, ArrayLike]) -> dict[str, ArrayLike]:
+    """
+    A no-op reparametrization function that returns the parameters unchanged.
+    This is useful when no reparametrization is needed.
+    """
+    return params
+
+
+def _extract_scipy_shape_params_names(
+    scipy_dist: stats.rv_continuous,
+) -> tuple[str, ...]:
+    return tuple(scipy_dist.shapes.split(", ") if scipy_dist.shapes else ())
+
+
 class ScipyDistribution(Distribution, ABC):
     """
     Base class for distributions based on SciPy.
     """
 
     _base_module: stats.rv_continuous
 
-    @abstractmethod
-    def _to_scipy_args(self, **params) -> dict[str, ArrayLike]:
-        raise NotImplementedError
+    def __init__(
+        self, *args, reparametrization: Reparametrization | None = None, **params
+    ):
+        if args and reparametrization:
+            raise ValueError(
+                "Cannot use both positional parameters and reparametrization."
+            )
+        self.reparametrization = reparametrization or no_reparametrization
+        if reparametrization is None:
+            self._params_names = ("loc", "scale") + _extract_scipy_shape_params_names(
+                self._base_module
+            )
+
+            # Insert the positional arguments into params
+            for arg, name in zip(args, self._params_names):
+                if name not in params:
+                    params[name] = arg
+                else:
+                    raise ValueError(
+                        f"Parameter `{name}` passed as positional and keyword argument when initializing {type(self).__name__} distribution."
+                    )
+
+            # Set default values for loc and scale if not provided
+            params["loc"] = params.get("loc", 0.0)
+            params["scale"] = params.get("scale", 1.0)
+
+            # Ensure all shape parameters are present
+            for arg in self._params_names[len(args) :]:
+                if arg not in params:
+                    raise ValueError(
+                        f"Missing shape parameter `{arg}` when initializing {type(self).__name__} distribution."
+                    )
+
+            # Reorder params to match the order of _params_names
+            params = {a: params[a] for a in self._params_names}
+        else:
+            self._params_names = tuple(params)
+        super().__init__(*params.values())
+
+    def _build_instance(self, **params) -> Self:
+        return type(self)(reparametrization=self.reparametrization, **params)
+
+    @property
+    def params_names(self) -> tuple[str, ...]:
+        """Return the names of the parameters."""
+        return self._params_names
+
+    def _to_scipy_args(self, **kwargs):
+        values = {k: kwargs.get(k, getattr(self, k)()) for k in self.params_names}
+        return self.reparametrization(values)
 
     def rvs(self, size=None, random_state=None, **kwargs):
         return self._base_module.rvs(

pykelihood/distributions/custom.py

@@ -231,8 +231,20 @@ class GEV(ScipyDistribution):
 
     _base_module = _stats.genextreme
 
+    def _reparametrization(self, params):
+        return {
+            "c": -params["shape"],
+            "loc": params["loc"],
+            "scale": params["scale"],
+        }
+
     def __init__(self, loc=0.0, scale=1.0, shape=0.0):
-        super().__init__(loc, scale, shape)
+        super().__init__(
+            loc=loc, scale=scale, shape=shape, reparametrization=self._reparametrization
+        )
+
+    def _build_instance(self, **params):
+        return type(self)(**params)
 
     @property
     def params_names(self):
@@ -284,32 +296,6 @@ def ub_shape(self, data):
         else:
             return self.scale / (x_max - self.loc())
 
-    def _to_scipy_args(self, loc=None, scale=None, shape=None):
-        """
-        Convert to scipy arguments.
-
-        Parameters
-        ----------
-        loc : float, optional
-            Location parameter, by default None.
-        scale : float, optional
-            Scale parameter, by default None.
-        shape : float, optional
-            Shape parameter, by default None.
-
-        Returns
-        -------
-        dict
-            Dictionary of scipy arguments.
-        """
-        if shape is not None:
-            shape = -shape
-        return {
-            "c": ifnone(shape, -self.shape()),
-            "loc": ifnone(loc, self.loc()),
-            "scale": ifnone(scale, self.scale()),
-        }
-
 
 class GPD(ScipyDistribution):
     """

pykelihood/distributions/scipy.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import scipy.special
+import scipy
 from packaging.version import Version
 from scipy import stats
 
@@ -13,9 +13,7 @@ def _name_from_scipy_dist(scipy_dist: stats.rv_continuous) -> str:
     return "".join(map(str.capitalize, scipy_dist_name.split("_")))
 
 
-def wrap_scipy_distribution(
-    scipy_dist: stats.rv_continuous,
-) -> type[ScipyDistribution]:
+def wrap_scipy_distribution(scipy_dist: stats.rv_continuous) -> type[ScipyDistribution]:
     """Wrap a scipy distribution class to create a ScipyDistribution subclass."""
     scipy_dist_name = type(scipy_dist).__name__.removesuffix("_gen")
     clean_dist_name = _name_from_scipy_dist(scipy_dist)
@@ -43,33 +41,15 @@ def format_param_docstring(param: str) -> str:
     for param in dist_params_names[2:]:
         docstring += format_param_docstring(param)
 
-    class Wrapper(ScipyDistribution):
-        _base_module = scipy_dist
-        __doc__ = docstring
-
-        def __init__(self, loc=0.0, scale=1.0, **kwargs):
-            self._params_names = dist_params_names
-            assert self._params_names[:2] == ("loc", "scale")
-            shape_args = self._params_names[2:]
-            for arg in shape_args:
-                if arg not in kwargs:
-                    raise ValueError(
-                        f"Missing shape parameter `{arg}` when initializing {type(self).__name__} distribution."
-                    )
-            args = [kwargs[a] for a in shape_args]
-            super().__init__(loc, scale, *args)
-
-        @property
-        def params_names(self) -> tuple[str, ...]:
-            """Return the names of the parameters."""
-            return self._params_names
-
-        def _to_scipy_args(self, **kwargs):
-            return {k: kwargs.get(k, getattr(self, k)()) for k in self.params_names}
-
-    Wrapper.__name__ = clean_dist_name
-    Wrapper.__qualname__ = f"{Wrapper.__module__}.{Wrapper.__name__}"
-    return Wrapper
+    return type(
+        clean_dist_name,
+        (ScipyDistribution,),
+        {
+            "_base_module": scipy_dist,
+            "__doc__": docstring,
+            "__module__": wrap_scipy_distribution.__module__,
+        },
+    )
 
 
 Alpha = wrap_scipy_distribution(stats.alpha)

tests/test_parametrization.py

@@ -0,0 +1,71 @@
+import numpy as np
+import numpy.testing
+from numpy.typing import ArrayLike
+
+from pykelihood import kernels
+from pykelihood.distributions import Uniform
+from pykelihood.metrics import log_likelihood
+
+
+def uniform_bounds(params: dict[str, ArrayLike]) -> dict[str, float]:
+    return {"loc": params["a"], "scale": params["b"] - params["a"]}
+
+
+def test_reparametrization_names():
+    u = Uniform(a=0, b=1, reparametrization=uniform_bounds)
+    assert u.params_names == ("a", "b")
+
+
+def test_reparametrization_rvs():
+    u = Uniform(a=3, b=5, reparametrization=uniform_bounds)
+    sample = u.rvs(10000)
+    numpy.testing.assert_array_less(sample, 5)
+    numpy.testing.assert_array_less(3, sample)
+
+
+def test_reparametrization_pdf():
+    u = Uniform(a=3, b=5, reparametrization=uniform_bounds)
+    u_standard = Uniform(loc=3, scale=2)
+    x = np.linspace(2, 6, 100)
+    numpy.testing.assert_array_almost_equal(u.pdf(x), u_standard.pdf(x))
+
+
+def test_reparametrization_kernel():
+    x = np.linspace(2, 6, 100)
+    u = Uniform(a=2, b=kernels.linear(x, a=2, b=1), reparametrization=uniform_bounds)
+    u.pdf(x)
+    u.rvs(x.size)
+
+
+def test_reparametrization_with_params():
+    u = Uniform(a=3, b=5, reparametrization=uniform_bounds)
+    v = u.with_params(a=2, b=4)
+    assert v.a() == 2
+    assert v.b() == 4
+
+
+def test_reparametrization_fit():
+    u = Uniform(a=0, b=5, reparametrization=uniform_bounds)
+    data = np.random.uniform(1, 3, 1000)
+    fit = u.fit(data)
+    assert 0 < fit.a() <= 1.1
+    assert 2.9 <= fit.b() < 5
+    assert log_likelihood(fit, data) > log_likelihood(u, data)
+
+
+def test_reparametrization_fit_fixed_param():
+    u = Uniform(a=0, b=5, reparametrization=uniform_bounds)
+    data = np.random.uniform(1, 3, 1000)
+    fit = u.fit(data, a=0.5)
+    assert fit.a() == 0.5
+    assert 2.9 <= fit.b() < 5
+    assert log_likelihood(fit, data) > log_likelihood(u, data)
+
+
+def test_reparametrization_fit_kernel():
+    u = Uniform(a=0, b=kernels.constant(value=5), reparametrization=uniform_bounds)
+    data = np.random.uniform(1, 3, 1000)
+    fit = u.fit(data, a=0.5)
+    assert fit.a() == 0.5
+    assert 2.9 <= fit.b() < 5
+    assert log_likelihood(fit, data) > log_likelihood(u, data)