Add plain SciPy wrapper factory

pykelihood/distributions/core.py

@@ -4,32 +4,19 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterator, Mapping
-from dataclasses import dataclass
-from types import MappingProxyType
 from typing import Union
 
 import numpy as np
 import numpy.typing as npt
-from scipy import stats
-from scipy.stats import rv_continuous
 
-from pykelihood.expr import Constant, Expr, Node, PathElem
+from pykelihood.expr import Expr, Node, PathElem
 from pykelihood.parameters import Parameter
-from pykelihood.state import PositiveTransform, Transform
 
 ParameterInput = Union[Expr, npt.ArrayLike, None]
 ParameterState = Mapping[Parameter, npt.NDArray[np.float64]]
 RandomState = Union[int, np.random.Generator, np.random.RandomState, None]
 
 
-@dataclass(frozen=True)
-class ParameterDefault:
-    """Initial value and optional transform for an omitted parameter."""
-
-    value: npt.ArrayLike
-    transform: Transform | None = None
-
-
 class Distribution(Node, ABC):
     """A probability law whose parameter expressions form a graph node."""
 
@@ -74,128 +61,3 @@ def ppf(
         self, q: npt.ArrayLike, *, state: ParameterState | None = None
     ) -> npt.NDArray[np.float64]:
         raise NotImplementedError
-
-
-class ScipyDistribution(Distribution):
-    """Continuous SciPy distribution evaluated from expression parameters."""
-
-    def __init__(
-        self,
-        scipy_distribution: rv_continuous,
-        parameters: Mapping[str, ParameterInput],
-        *,
-        defaults: Mapping[str, ParameterDefault] | None = None,
-    ) -> None:
-        shape_names = (
-            ()
-            if scipy_distribution.shapes is None
-            else tuple(name.strip() for name in scipy_distribution.shapes.split(","))
-        )
-        unknown = set(parameters) - set(shape_names) - {"loc", "scale"}
-        if unknown:
-            raise TypeError(
-                f"Unknown distribution parameters: {', '.join(sorted(unknown))}"
-            )
-        for name in shape_names:
-            if parameters.get(name) is None:
-                raise TypeError(f"Missing required distribution parameter: {name}")
-
-        parameter_defaults = {} if defaults is None else defaults
-        resolved: dict[str, Expr] = {}
-        for name, value in parameters.items():
-            if value is None:
-                if name not in parameter_defaults:
-                    raise TypeError(f"Missing required distribution parameter: {name}")
-                default = parameter_defaults[name]
-                resolved[name] = Parameter(
-                    init=default.value, transform=default.transform, name=name
-                )
-            elif isinstance(value, Expr):
-                resolved[name] = value
-            else:
-                resolved[name] = Constant(value)
-
-        self._scipy_distribution = scipy_distribution
-        self._parameters = MappingProxyType(resolved)
-
-    @property
-    def parameters(self) -> Mapping[str, Expr]:
-        return self._parameters
-
-    def _evaluated_parameters(
-        self, state: ParameterState | None
-    ) -> dict[str, npt.NDArray[np.float64]]:
-        parameter_state = {} if state is None else state
-        return {
-            name: np.asarray(parameter.eval(parameter_state), dtype=np.float64)
-            for name, parameter in self.parameters.items()
-        }
-
-    def rvs(
-        self,
-        size: int | tuple[int, ...] | None = None,
-        *,
-        state: ParameterState | None = None,
-        random_state: RandomState = None,
-    ) -> npt.NDArray[np.float64]:
-        parameters = self._evaluated_parameters(state)
-        batch_shape = np.broadcast_shapes(
-            *(value.shape for value in parameters.values())
-        )
-        sample_shape = (
-            () if size is None else (size,) if isinstance(size, int) else size
-        )
-        return np.asarray(
-            self._scipy_distribution.rvs(
-                **parameters, size=sample_shape + batch_shape, random_state=random_state
-            ),
-            dtype=np.float64,
-        )
-
-    def pdf(
-        self, x: npt.ArrayLike, *, state: ParameterState | None = None
-    ) -> npt.NDArray[np.float64]:
-        return np.asarray(
-            self._scipy_distribution.pdf(x, **self._evaluated_parameters(state)),
-            dtype=np.float64,
-        )
-
-    def logpdf(
-        self, x: npt.ArrayLike, *, state: ParameterState | None = None
-    ) -> npt.NDArray[np.float64]:
-        return np.asarray(
-            self._scipy_distribution.logpdf(x, **self._evaluated_parameters(state)),
-            dtype=np.float64,
-        )
-
-    def cdf(
-        self, x: npt.ArrayLike, *, state: ParameterState | None = None
-    ) -> npt.NDArray[np.float64]:
-        return np.asarray(
-            self._scipy_distribution.cdf(x, **self._evaluated_parameters(state)),
-            dtype=np.float64,
-        )
-
-    def ppf(
-        self, q: npt.ArrayLike, *, state: ParameterState | None = None
-    ) -> npt.NDArray[np.float64]:
-        return np.asarray(
-            self._scipy_distribution.ppf(q, **self._evaluated_parameters(state)),
-            dtype=np.float64,
-        )
-
-
-class Normal(ScipyDistribution):
-    """Normal distribution with free default location and scale parameters."""
-
-    def __init__(
-        self, loc: ParameterInput = None, scale: ParameterInput = None
-    ) -> None:
-        super().__init__(
-            stats.norm,
-            {"loc": loc, "scale": scale},
-            defaults={
-                "loc": ParameterDefault(0.0),
-                "scale": ParameterDefault(1.0, PositiveTransform()),
-            },
-        )

pykelihood/distributions/scipy_adapter.py

@@ -0,0 +1,138 @@
+"""SciPy implementation of explicit-state continuous distributions."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+
+import numpy as np
+import numpy.typing as npt
+from scipy.stats import rv_continuous
+
+from pykelihood.distributions.core import (
+    Distribution,
+    ParameterInput,
+    ParameterState,
+    RandomState,
+)
+from pykelihood.expr import Constant, Expr
+from pykelihood.parameters import Parameter
+from pykelihood.state import Transform
+
+
+@dataclass(frozen=True)
+class ParameterDefault:
+    """Initial value and optional transform for an omitted parameter."""
+
+    value: npt.ArrayLike
+    transform: Transform | None = None
+
+
+class ScipyDistribution(Distribution):
+    """Continuous SciPy distribution evaluated from expression parameters."""
+
+    def __init__(
+        self,
+        scipy_distribution: rv_continuous,
+        parameters: Mapping[str, ParameterInput],
+        *,
+        defaults: Mapping[str, ParameterDefault] | None = None,
+    ) -> None:
+        shape_names = (
+            ()
+            if scipy_distribution.shapes is None
+            else tuple(name.strip() for name in scipy_distribution.shapes.split(","))
+        )
+        unknown = set(parameters) - set(shape_names) - {"loc", "scale"}
+        if unknown:
+            raise TypeError(
+                f"Unknown distribution parameters: {', '.join(sorted(unknown))}"
+            )
+        for name in shape_names:
+            if parameters.get(name) is None:
+                raise TypeError(f"Missing required distribution parameter: {name}")
+
+        parameter_defaults = {} if defaults is None else defaults
+        resolved: dict[str, Expr] = {}
+        for name, value in parameters.items():
+            if value is None:
+                if name not in parameter_defaults:
+                    raise TypeError(f"Missing required distribution parameter: {name}")
+                default = parameter_defaults[name]
+                resolved[name] = Parameter(
+                    init=default.value, transform=default.transform, name=name
+                )
+            elif isinstance(value, Expr):
+                resolved[name] = value
+            else:
+                resolved[name] = Constant(value)
+
+        self._scipy_distribution = scipy_distribution
+        self._parameters = MappingProxyType(resolved)
+
+    @property
+    def parameters(self) -> Mapping[str, Expr]:
+        return self._parameters
+
+    def _evaluated_parameters(
+        self, state: ParameterState | None
+    ) -> dict[str, npt.NDArray[np.float64]]:
+        parameter_state = {} if state is None else state
+        return {
+            name: np.asarray(parameter.eval(parameter_state), dtype=np.float64)
+            for name, parameter in self.parameters.items()
+        }
+
+    def rvs(
+        self,
+        size: int | tuple[int, ...] | None = None,
+        *,
+        state: ParameterState | None = None,
+        random_state: RandomState = None,
+    ) -> npt.NDArray[np.float64]:
+        parameters = self._evaluated_parameters(state)
+        batch_shape = np.broadcast_shapes(
+            *(value.shape for value in parameters.values())
+        )
+        sample_shape = (
+            () if size is None else (size,) if isinstance(size, int) else size
+        )
+        return np.asarray(
+            self._scipy_distribution.rvs(
+                **parameters, size=sample_shape + batch_shape, random_state=random_state
+            ),
+            dtype=np.float64,
+        )
+
+    def pdf(
+        self, x: npt.ArrayLike, *, state: ParameterState | None = None
+    ) -> npt.NDArray[np.float64]:
+        return np.asarray(
+            self._scipy_distribution.pdf(x, **self._evaluated_parameters(state)),
+            dtype=np.float64,
+        )
+
+    def logpdf(
+        self, x: npt.ArrayLike, *, state: ParameterState | None = None
+    ) -> npt.NDArray[np.float64]:
+        return np.asarray(
+            self._scipy_distribution.logpdf(x, **self._evaluated_parameters(state)),
+            dtype=np.float64,
+        )
+
+    def cdf(
+        self, x: npt.ArrayLike, *, state: ParameterState | None = None
+    ) -> npt.NDArray[np.float64]:
+        return np.asarray(
+            self._scipy_distribution.cdf(x, **self._evaluated_parameters(state)),
+            dtype=np.float64,
+        )
+
+    def ppf(
+        self, q: npt.ArrayLike, *, state: ParameterState | None = None
+    ) -> npt.NDArray[np.float64]:
+        return np.asarray(
+            self._scipy_distribution.ppf(q, **self._evaluated_parameters(state)),
+            dtype=np.float64,
+        )

pykelihood/distributions/scipy_wrappers.py

@@ -0,0 +1,80 @@
+"""Opt-in structural wrappers for selected SciPy continuous distributions."""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any, ClassVar
+
+from scipy import stats
+from scipy.stats import rv_continuous
+
+from pykelihood.distributions.core import ParameterInput
+from pykelihood.distributions.scipy_adapter import ParameterDefault, ScipyDistribution
+from pykelihood.state import PositiveTransform
+
+
+def _name_from_scipy_dist(scipy_dist: rv_continuous) -> str:
+    return "".join(part.capitalize() for part in scipy_dist.name.split("_"))
+
+
+class _WrappedScipyDistribution(ScipyDistribution):
+    """Shared constructor for generated plain SciPy distributions."""
+
+    _base_module: ClassVar[rv_continuous]
+    __signature__: ClassVar[inspect.Signature]
+
+    def __init__(self, *args: ParameterInput, **kwargs: ParameterInput) -> None:
+        bound = self.__signature__.bind(*args, **kwargs)
+        bound.apply_defaults()
+        super().__init__(
+            self._base_module,
+            bound.arguments,
+            defaults={
+                "loc": ParameterDefault(0.0),
+                "scale": ParameterDefault(1.0, PositiveTransform()),
+            },
+        )
+
+
+def wrap_scipy_distribution(
+    scipy_dist: rv_continuous,
+) -> type[_WrappedScipyDistribution]:
+    """Create a structural distribution class for one SciPy continuous law.
+
+    Shape parameters follow SciPy's native names and are required. ``loc`` and
+    ``scale`` are optional free parameters initialized to 0 and 1 respectively.
+    Literal arguments become constants through :class:`ScipyDistribution`.
+    The generated constructor accepts arguments in SciPy's order: shape
+    parameters, then ``loc`` and ``scale``.
+    """
+    shape_names = (
+        ()
+        if scipy_dist.shapes is None
+        else tuple(name.strip() for name in scipy_dist.shapes.split(","))
+    )
+    parameter_names = (*shape_names, "loc", "scale")
+    signature_parameters = [
+        inspect.Parameter(
+            name,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=inspect.Parameter.empty if name in shape_names else None,
+        )
+        for name in parameter_names
+    ]
+    signature = inspect.Signature(signature_parameters)
+    wrapper_name = _name_from_scipy_dist(scipy_dist)
+    namespace: dict[str, Any] = {
+        "_base_module": scipy_dist,
+        "__doc__": f"Structural wrapper for ``scipy.stats.{scipy_dist.name}``.",
+        "__module__": __name__,
+        "__signature__": signature,
+    }
+    return type(wrapper_name, (_WrappedScipyDistribution,), namespace)
+
+
+Norm = wrap_scipy_distribution(stats.norm)
+Normal = Norm
+Gamma = wrap_scipy_distribution(stats.gamma)
+Genextreme = wrap_scipy_distribution(stats.genextreme)
+
+__all__ = ["Gamma", "Genextreme", "Norm", "Normal", "wrap_scipy_distribution"]

tests/test_distribution_core.py

@@ -3,7 +3,8 @@
 from numpy.testing import assert_allclose
 from scipy import stats
 
-from pykelihood.distributions.core import Normal, ParameterDefault, ScipyDistribution
+from pykelihood.distributions.scipy_adapter import ParameterDefault, ScipyDistribution
+from pykelihood.distributions.scipy_wrappers import Normal
 from pykelihood.effects import linear
 from pykelihood.expr import Constant
 from pykelihood.likelihood import negative_log_likelihood