[feat] added transform method to ExponentialFamily

Author: domosedy

src/pysatl_core/families/__init__.py

@@ -14,13 +14,6 @@
 from .builtins import __all__ as _builtins_all
 from .configuration import configure_families_register
 from .distribution import ParametricFamilyDistribution
-from .parametric_family import ParametricFamily
-from .parametrizations import (
-    Parametrization,
-    ParametrizationConstraint,
-    constraint,
-    parametrization,
-)
 from .exponential_family import (
     ExponentialConjugateHyperparameters,
     ExponentialFamily,
@@ -29,6 +22,13 @@
     SpacePredicate,
     SpacePredicateArray,
 )
+from .parametric_family import ParametricFamily
+from .parametrizations import (
+    Parametrization,
+    ParametrizationConstraint,
+    constraint,
+    parametrization,
+)
 from .registry import ParametricFamilyRegister
 
 __all__ = [

src/pysatl_core/families/exponential_family.py

@@ -1,19 +1,24 @@
 from __future__ import annotations
+
 from collections.abc import Callable
-from typing import Any, cast, TYPE_CHECKING
-from scipy.integrate import nquad, quad
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Iterable, Sized, cast
+
 import numpy as np
+from scipy.integrate import nquad
+from scipy.linalg import det
+from scipy.differentiate import jacobian
 
-from pysatl_core.distributions.fitters import _ppf_brentq_from_cdf
+from pysatl_core.distributions import (
+    SamplingStrategy,
+)
 from pysatl_core.families.parametric_family import ParametricFamily
 from pysatl_core.families.parametrizations import Parametrization, parametrization
 from pysatl_core.types import (
+    GenericCharacteristicName,
     DistributionType,
     ParametrizationName,
 )
-from pysatl_core.distributions import (
-    SamplingStrategy,
-)
 
 if TYPE_CHECKING:
     from pysatl_core.distributions.support import Support
@@ -32,34 +37,37 @@
 KURT = "kurtosis"
 
 
+@dataclass
 class ExponentialFamilyParametrization(Parametrization):
     """
     Standard parametrization of Exponential Family.
     """
 
-    theta: list[Callable[[float], float]]  # TODO: mb more clever
+    theta: list[float]  # TODO: mb more clever
 
 
 class ExponentialConjugateHyperparameters:
     def __init__(self, alpha: Any, beta: int):
         self.alpha = alpha
         self.beta = beta
 
-    def __str__(self):
+    def __str__(self) -> str:
         return f"alpha={self.alpha}, beta={self.beta}"
 
 
-def doesAccept(x, support):
+def doesAccept(x: list[float] | float, support: list[tuple[float, float]]) -> bool:
     if not hasattr(x, "__len__"):
         x = [x]
 
-    def accept_1D(x, borders):
+    x = cast(list[float], x)
+
+    def accept_1D(x: float, borders: tuple[float, float]) -> bool:
         left, right = borders
         if abs(x) == 0 and (abs(left) == 0 or abs(right) == 0):
             return False
         return left <= x <= right
 
-    return all(accept_1D(x_i, border) for x_i, border in zip(x, support))
+    return all(accept_1D(x_i, border) for x_i, border in zip(x, support, strict=False))
 
 
 class SpacePredicate:
@@ -100,7 +108,10 @@ def __init__(
         self._parameter_space = parameter_space
         self._sufficient_statistics_values = sufficient_statistics_values
 
-        distr_characteristics = {
+        distr_characteristics: dict[
+            GenericCharacteristicName,
+            dict[ParametrizationName, ParametrizedFunction] | ParametrizedFunction,
+        ] = {
             PDF: self.density,
             MEAN: self._mean,
             VAR: self._var,
@@ -115,20 +126,29 @@ def __init__(
             sampling_strategy=sampling_strategy,
             support_by_parametrization=support_by_parametrization,
         )
-        parametrization(family=self, name="theta")((ExponentialFamilyParametrization))
+        parametrization(family=self, name="theta")(ExponentialFamilyParametrization)
+
+    def _transform_to_natural_parametrization(
+        self, theta_parametrization: ExponentialFamilyParametrization
+    ) -> ExponentialFamilyParametrization:
+        return theta_parametrization
 
     @property
     def log_density(self) -> ParametrizedFunction:
-        def log_density_func(
-            parametrization: ExponentialFamilyParametrization, x: Any
-        ) -> Any:
+        def log_density_func(parametrization: Parametrization, x: Any) -> Any:
+            parametrization = cast(ExponentialFamilyParametrization, parametrization)
+            parametrization = self._transform_to_natural_parametrization(
+                parametrization
+            )
             if not self._support.accepts(x):
                 return float("-inf")
 
-            params = cast(ExponentialFamilyParametrization, parametrization)
-            theta = params.parameters.get("theta")
+            theta = parametrization.theta
             sufficient = self._sufficient(x)
             dot = np.dot(theta, sufficient)
+            if hasattr(dot, "__len__"):
+                dot = dot[0]
+
             result = float(
                 np.log(self._normalization(x))
                 + dot
@@ -143,26 +163,31 @@ def density(self) -> ParametrizedFunction:
         return lambda parametrization, x: np.exp(self.log_density(parametrization, x))
 
     @property
-    def conjugate_prior_family(self):
-        def conjugate_sufficient(theta: Any):
+    def conjugate_prior_family(self) -> NaturalExponentialFamily:
+        def conjugate_sufficient(
+            theta: float,
+        ) -> list[Any]:
             if not self._parameter_space.accepts(theta):
                 return [float("-inf"), float("-inf")]
 
+            parametrization = ExponentialFamilyParametrization([theta])
+            # parametrization.theta = [theta]
             return [
                 theta,
-                self._log_partition(ExponentialFamilyParametrization(theta=[theta])),
+                self._log_partition(parametrization),
             ]
 
-        def conjugate_log_partition(parametrization: ExponentialFamilyParametrization):
+        def conjugate_log_partition(
+            parametrization: ExponentialFamilyParametrization,
+        ) -> Any:
             alpha = parametrization.theta[0]
             beta = parametrization.theta[1]
 
-            def pdf(theta: Any):
+            def pdf(theta: Any) -> Any:
                 if not hasattr(theta, "__len__"):
                     theta = [theta]
-                parametrization = ExponentialFamilyParametrization(
-                    theta=theta,
-                )
+                parametrization = ExponentialFamilyParametrization(theta=theta)
+                # parametrization.theta = theta
                 return np.exp(
                     np.dot(theta, alpha) + beta * self._log_partition(parametrization)
                 )[0]
@@ -180,15 +205,14 @@ def pdf(theta: Any):
 
         def conjugate_sufficient_accepts(
             parametrization: ExponentialFamilyParametrization,
-        ):
-            parametrization = cast(parametrization, ExponentialFamilyParametrization)
-            theta = parametrization.parameters.get("theta")
+        ) -> bool:
+            theta = parametrization.theta
             xi = theta[:-1]
             nu = theta[-1]
 
-            return self._sufficient_statistics_values(xi) and SpacePredicateArray(
-                [(0, float("+inf"))]
-            ).accepts(nu)
+            return self._sufficient_statistics_values.accepts(
+                xi
+            ) and SpacePredicateArray([(0, float("+inf"))]).accepts(nu)
 
         return NaturalExponentialFamily(
             log_partition=conjugate_log_partition,
@@ -197,15 +221,50 @@ def conjugate_sufficient_accepts(
             support=self._parameter_space,
             sufficient_statistics_values=self._parameter_space,  # TODO: write convex hull for this
             parameter_space=SpacePredicate(conjugate_sufficient_accepts),
+            name=self.name,
             sampling_strategy=self.sampling_strategy,
             distr_type=self._distr_type,
             distr_parametrizations=self.parametrization_names,
             support_by_parametrization=self.support_resolver,
         )
 
+    def transform(
+        self,
+        transform_function: Callable[[Any], Any],
+    ) -> NaturalExponentialFamily:
+        def calculate_jacobian(x: Any) -> Any:
+            if type(x) is not list:
+                x = np.array([x])
+
+            return np.abs(det(jacobian(transform_function, x).df))
+
+        def new_support(x: Any) -> bool:
+            return self._support.accepts(transform_function(x))
+
+        def new_sufficient(x: Any) -> Any:
+            return self._sufficient(transform_function(x))
+
+        def new_normalization(x: Any) -> Any:
+            return self._normalization(x) * calculate_jacobian(x)
+
+        return NaturalExponentialFamily(
+            log_partition=self._log_partition,
+            sufficient_statistics=new_sufficient,
+            normalization_constant=new_normalization,
+            support=SpacePredicate(new_support),
+            parameter_space=self._parameter_space,
+            sufficient_statistics_values=self._sufficient_statistics_values,
+            name=f"Transformed{self._name}",
+            distr_type=self._distr_type,
+            distr_parametrizations=self.parametrization_names,
+            sampling_strategy=self.sampling_strategy,
+            support_by_parametrization=self.support_resolver,
+        )
+
     @property
     def _mean(self) -> ParametrizedFunction:
         def mean_func(parametrization: Parametrization, x: Any) -> Any:
+            parametrization = cast(ExponentialFamilyParametrization, parametrization)
             dimension_size = 1
             if hasattr(x, "__len__"):
                 dimension_size = len(x)
@@ -223,6 +282,7 @@ def mean_func(parametrization: Parametrization, x: Any) -> Any:
     @property
     def _second_moment(self) -> ParametrizedFunction:
         def func(parametrization: Parametrization, x: Any) -> Any:
+            parametrization = cast(ExponentialFamilyParametrization, parametrization)
             dimension_size = 1
             if hasattr(x, "__len__"):
                 dimension_size = len(x)
@@ -238,8 +298,9 @@ def func(parametrization: Parametrization, x: Any) -> Any:
         return func
 
     @property
-    def _var(self):
-        def func(parametrization, x: Any):
+    def _var(self) -> ParametrizedFunction:
+        def func(parametrization: Parametrization, x: Any) -> Any:
+            parametrization = cast(ExponentialFamilyParametrization, parametrization)
             return (
                 self._second_moment(parametrization, x)
                 - self._mean(parametrization, x) ** 2
@@ -248,8 +309,8 @@ def func(parametrization, x: Any):
         return func
 
     def posterior_hyperparameters(
-        self, prior_hyper: ExponentialConjugateHyperparameters, sample
-    ):
+        self, prior_hyper: ExponentialConjugateHyperparameters, sample: list[Any]
+    ) -> ExponentialConjugateHyperparameters:
         alpha = prior_hyper.alpha
         beta = prior_hyper.beta
 
@@ -275,6 +336,9 @@ def __init__(
         sufficient_statistics: Callable[[Any], Any],
         normalization_constant: Callable[[Any], Any],
         parameter_from_natural_parameter: Callable[[Any], Any],
+        natural_parameter: Callable[
+            [ExponentialFamilyParametrization], ExponentialFamilyParametrization
+        ],
         support: SpacePredicate,
         parameter_space: SpacePredicate,
         sufficient_statistics_values: SpacePredicate,
@@ -284,11 +348,10 @@ def __init__(
         name: str = "ExponentialFamily",
         support_by_parametrization: SupportArg = None,
     ):
-        def natural_log_partition(eta_parametrizaion: ExponentialFamilyParametrization):
-            eta_parametrizaion = cast(
-                ExponentialFamilyParametrization, eta_parametrizaion
-            )
-            eta = eta_parametrizaion.parameters.get("theta")
+        def natural_log_partition(
+            eta_parametrizaion: ExponentialFamilyParametrization,
+        ) -> Any:
+            eta = eta_parametrizaion.theta
             theta = parameter_from_natural_parameter(eta)
             return log_partition(ExponentialFamilyParametrization(theta=[theta]))
 
@@ -297,6 +360,8 @@ def natural_log_partition(eta_parametrizaion: ExponentialFamilyParametrization):
                 parameter_from_natural_parameter(eta)
             )
         )
+
+        self._natural_parameter = natural_parameter
         natural_parameter_space = SpacePredicate(
             lambda eta: parameter_space.accepts(parameter_from_natural_parameter(eta)),
         )
@@ -315,3 +380,8 @@ def natural_log_partition(eta_parametrizaion: ExponentialFamilyParametrization):
             sampling_strategy=sampling_strategy,
             support_by_parametrization=support_by_parametrization,
         )
+
+    def _transform_to_natural_parametrization(
+        self, theta_parametrization: ExponentialFamilyParametrization
+    ) -> ExponentialFamilyParametrization:
+        return self._natural_parameter(theta_parametrization)