feat(exponential): remove NumberParameter logic

Author: domosedy

src/pysatl_core/distributions/support.py

@@ -32,7 +32,6 @@
     Interval1D,
     IntervalND,
     Number,
-    NumberParameter,
     NumericArray,
 )
 
@@ -63,7 +62,8 @@ class ContinuousSupport(Interval1D, Support):
     """
 
 
-class ContinuousNDSupport(IntervalND, Support):
+# Support want to have Number as a parameter of contains, but we decided that we should avoid this
+class ContinuousNDSupport(IntervalND, Support):  # type: ignore[misc]
     """
     Support for continuous distributions represented as an array of intervals.
 
@@ -455,18 +455,18 @@ def is_right_bounded(self) -> bool:
 
 @dataclass(slots=True)
 class SupportByPredicate(Support):
-    predicate: Callable[[NumberParameter], bool]
+    predicate: Callable[[NumericArray], bool]
 
     @overload
     def contains(self, x: Number) -> bool: ...
     @overload
     def contains(self, x: NumericArray) -> BoolArray: ...
 
-    def contains(self, x: NumberParameter) -> bool | BoolArray:
+    def contains(self, x: NumericArray) -> bool | BoolArray:  # type: ignore[misc]
         return self.predicate(x)
 
     def __contains__(self, item: object) -> bool | BoolArray:
-        return self.contains(cast(NumberParameter, item))
+        return self.contains(cast(NumericArray, item))
 
 
 __all__ = [

src/pysatl_core/families/exponential_family.py

@@ -36,7 +36,7 @@
 
 if TYPE_CHECKING:
     from pysatl_core.distributions.support import Support
-    from pysatl_core.types import Number, NumberParameter, NumericArray
+    from pysatl_core.types import Number, NumericArray
 
     type ParametrizedFunction = Callable[[Parametrization, Any], Any]
     type SupportArg = Callable[[Parametrization], Support | None] | None
@@ -52,10 +52,10 @@ class ExponentialFamilyParametrization(Parametrization):
         f(x|θ) = h(x) * exp(θᵀ T(x) - A(θ))
 
     Attributes:
-        theta (NumberParameter): Natural parameter vector (can be a scalar or array)
+        theta (NumericArray): Natural parameter vector (can be a scalar or array)
     """
 
-    theta: NumberParameter
+    theta: NumericArray
 
     def transform_to_base_parametrization(self) -> ExponentialFamilyParametrization:
         """Return the base parametrization (identity transform for canonical form)."""
@@ -120,9 +120,9 @@ class ContinuousExponentialClassFamily(ParametricFamily):
     def __init__(
         self,
         *,
-        log_partition: Callable[[NumberParameter], NumberParameter],
-        sufficient_statistics: Callable[[NumberParameter], NumberParameter],
-        normalization_constant: Callable[[NumberParameter], NumberParameter],
+        log_partition: Callable[[NumericArray], NumericArray],
+        sufficient_statistics: Callable[[NumericArray], NumericArray],
+        normalization_constant: Callable[[NumericArray], Number],
         support: SupportByPredicate,
         parameter_space: SupportByPredicate,
         sufficient_statistics_values: SupportByPredicate,
@@ -185,13 +185,13 @@ def log_density(self) -> ParametrizedFunction:
         and a point `x`, and returns log f(x|θ). Returns -inf for x outside the support.
 
         Returns:
-            Callable[[Parametrization, NumberParameter], Number]
+            Callable[[Parametrization, NumericArray], Number]
         """
 
-        def log_density_func(parametrization: Parametrization, x: NumberParameter) -> Number:
+        def log_density_func(parametrization: Parametrization, x: NumericArray) -> Number:
             parametrization = cast(ExponentialFamilyParametrization, parametrization)
             parametrization = parametrization.transform_to_base_parametrization()
-            if x not in self._support:
+            if np.array([x]) not in self._support:
                 return -np.inf
 
             theta = parametrization.theta
@@ -211,7 +211,7 @@ def density(self) -> ParametrizedFunction:
         Density function (exponentiated log‑density).
 
         Returns:
-            Callable[[Parametrization, NumberParameter], Number]
+            Callable[[Parametrization, NumericArray], Number]
         """
         return lambda parametrization, x: np.exp(self.log_density(parametrization, x))
 
@@ -230,8 +230,8 @@ def conjugate_prior_family(self) -> ContinuousExponentialClassFamily:
         """
 
         def conjugate_sufficient(
-            theta: NumberParameter,
-        ) -> NumberParameter:
+            theta: NumericArray,
+        ) -> NumericArray:
             if not hasattr(theta, "__len__"):
                 theta = np.array([theta])
 
@@ -240,9 +240,9 @@ def conjugate_sufficient(
             return np.append(theta, self._log_partition(theta))
 
         def conjugate_log_partition(
-            parametrization: NumberParameter,
-        ) -> NumberParameter:
-            def pdf(theta: NumberParameter) -> NumberParameter:
+            parametrization: NumericArray,
+        ) -> NumericArray:
+            def pdf(theta: NumericArray) -> Number:
                 if not hasattr(theta, "__len__"):
                     theta = np.array([theta])
                 return cast(
@@ -259,23 +259,25 @@ def pdf(theta: NumberParameter) -> NumberParameter:
                 lambda x: pdf(x) if x in self._parameter_space else 0,  # type: ignore[arg-type]
                 [(float("-inf"), float("+inf"))],
             )[0]
-            return cast(np.float64, -np.log(all_value))
+            return np.array([cast(np.float64, -np.log(all_value))])
 
         def conjugate_sufficient_accepts(
             theta: NumericArray,
         ) -> bool:
             xi = theta[:-1]
             nu = theta[-1]
 
-            return xi in self._sufficient_statistics_values and nu in ContinuousSupport(0, np.inf)
+            return xi in self._sufficient_statistics_values and np.array([nu]) in ContinuousSupport(
+                0, np.inf
+            )
 
         return ContinuousExponentialClassFamily(
             log_partition=conjugate_log_partition,
             sufficient_statistics=conjugate_sufficient,
             normalization_constant=lambda _: 1,
             support=self._parameter_space,
-            sufficient_statistics_values=self._parameter_space,  # TODO: write convex hull for this
-            parameter_space=SupportByPredicate(predicate=conjugate_sufficient_accepts),  # type: ignore[arg-type]
+            sufficient_statistics_values=self._parameter_space,
+            parameter_space=SupportByPredicate(predicate=conjugate_sufficient_accepts),
             name=self.name,
             distr_type=self._distr_type,
             distr_parametrizations=self.parametrization_names,
@@ -284,7 +286,7 @@ def conjugate_sufficient_accepts(
 
     def transform(
         self,
-        transform_function: Callable[[NumberParameter], NumberParameter],
+        transform_function: Callable[[NumericArray], NumericArray],
     ) -> ContinuousExponentialClassFamily:
         """
         Transform the random variable by a monotonic, differentiable function.
@@ -301,20 +303,20 @@ def transform(
             ContinuousExponentialClassFamily: A new family for the transformed variable.
         """
 
-        def calculate_jacobian(x: NumberParameter) -> NumberParameter:
+        def calculate_jacobian(x: NumericArray) -> NumericArray:
             if not isinstance(x, Iterable):
                 x = np.array([x])
 
             return np.abs(det(jacobian(transform_function, x).df))
 
-        def new_support(x: NumberParameter) -> bool:
+        def new_support(x: NumericArray) -> bool:
             return transform_function(x) in self._support
 
-        def new_sufficient(x: NumberParameter) -> NumberParameter:
+        def new_sufficient(x: NumericArray) -> NumericArray:
             return self._sufficient(transform_function(x))
 
-        def new_normalization(x: NumberParameter) -> NumberParameter:
-            return self._normalization(x) * calculate_jacobian(x)
+        def new_normalization(x: NumericArray) -> Number:
+            return cast(np.float64, self._normalization(x) * calculate_jacobian(x))
 
         return ContinuousExponentialClassFamily(
             log_partition=self._log_partition,
@@ -339,7 +341,11 @@ def mean_func(parametrization: Parametrization, x: Any) -> Any:
             if hasattr(x, "__len__"):
                 dimension_size = len(x)
             return nquad(
-                lambda x: np.dot(x, self.density(parametrization, x)) if x in self._support else 0,
+                lambda x: (
+                    np.dot(x, self.density(parametrization, x))
+                    if np.array([x]) in self._support
+                    else 0
+                ),
                 [(float("-inf"), float("inf"))] * dimension_size,
             )[0]
 
@@ -355,7 +361,9 @@ def func(parametrization: Parametrization, x: Any) -> Any:
             if hasattr(x, "__len__"):
                 dimension_size = len(x)
             return nquad(
-                lambda x: x**2 * self.density(parametrization, x) if x in self._support else 0,
+                lambda x: (
+                    x**2 * self.density(parametrization, x) if np.array([x]) in self._support else 0
+                ),
                 [(float("-inf"), float("inf"))] * dimension_size,
             )[0]
 
@@ -394,7 +402,7 @@ def posterior_hyperparameters(
         posterior_effective_sample_size = parametrizaiton.effective_sample_size
         if hasattr(sample, "__iter__") and not isinstance(sample, str):
             posterior_effective_suff_stat_value += np.sum(
-                [self._sufficient(x) for x in sample],  # type: ignore[arg-type]
+                [self._sufficient(x) for x in sample],
                 axis=0,
             )
             posterior_effective_sample_size += len(sample)
@@ -424,13 +432,13 @@ def posterior_predictive(self) -> ParametricFamily:
 
         def conjugate_log_partition(
             parametrization: ExponentialConjugateHyperparameters,
-        ) -> NumberParameter:
+        ) -> NumericArray:
             conjugate_value = self.conjugate_prior_family._log_partition(
                 parametrization.transform_to_base_parametrization().theta
             )
             return np.exp(conjugate_value)
 
-        def posterior_density(parametrization: Parametrization, x: NumberParameter) -> Number:
+        def posterior_density(parametrization: Parametrization, x: NumericArray) -> Number:
             parametrization = cast(ExponentialConjugateHyperparameters, parametrization)
             return cast(
                 np.float32,

src/pysatl_core/types.py

@@ -121,9 +121,6 @@ class EuclideanDistributionType(DistributionType):
 type BoolArray = NDArray[np.bool_]
 """Type alias for boolean arrays."""
 
-type NumberParameter = Number | NumericArray
-"""Type alias for numeric or list parameter"""
-
 
 class ContinuousSupportShape1D(Enum):
     """
@@ -258,27 +255,22 @@ def shape(self) -> ContinuousSupportShape1D:
 class IntervalND:
     intervals: list[Interval1D]
 
-    @overload
-    def contains(self, x: Number) -> bool: ...
+    def contains(self, x: NumericArray) -> bool | BoolArray:
+        def contains_for_point(point: NumericArray) -> bool:
+            assert len(point) == len(self.intervals)
+            return all(
+                x_coordinate in interval
+                for interval, x_coordinate in zip(self.intervals, point, strict=True)
+            )
 
-    @overload
-    def contains(self, x: NumericArray) -> BoolArray: ...
+        if len(x.shape) == 1:
+            return contains_for_point(x)
 
-    def contains(self, x: Number | NumericArray) -> bool | BoolArray:
-        if not hasattr(x, "__iter__"):
-            x = np.array([x])
-
-        x = np.array(x)
-        assert len(x) == len(self.intervals)
-
-        return all(
-            x_coordinate in interval
-            for interval, x_coordinate in zip(self.intervals, x, strict=True)
-        )
+        return np.array([contains_for_point(point) for point in x])
 
     def __contains__(self, x: object) -> bool:
         """Check if a single point is in the interval."""
-        return bool(self.contains(cast(Number, x)))
+        return bool(self.contains(cast(NumericArray, x)))
 
 
 type GenericCharacteristicName = str

tests/unit/families/test_exponential_family.py

@@ -3,7 +3,6 @@
 __license__ = "SPDX-License-Identifier: MIT"
 
 import itertools
-from collections.abc import Iterable
 from typing import cast
 
 import numpy as np
@@ -16,7 +15,7 @@
     ContinuousExponentialClassFamily,
 )
 from pysatl_core.families.registry import ParametricFamilyRegister
-from pysatl_core.types import CharacteristicName, Interval1D, NumberParameter, UnivariateContinuous
+from pysatl_core.types import CharacteristicName, Interval1D, NumericArray, UnivariateContinuous
 
 
 def gamma_pdf(alpha: float, beta: float, x: float) -> float:
@@ -25,16 +24,14 @@ def gamma_pdf(alpha: float, beta: float, x: float) -> float:
 
 @pytest.fixture(scope="function")
 def conjugate_for_exponential() -> ContinuousExponentialClassFamily:
-    def transform_function(x: NumberParameter) -> NumberParameter:
-        if isinstance(x, Iterable):
-            return np.array([-x[0]])
+    def transform_function(x: NumericArray) -> NumericArray:
         return -x
 
     support_neg = SupportByPredicate(
-        predicate=lambda x: x in ContinuousNDSupport(intervals=[Interval1D(-np.inf, 0)])
+        predicate=lambda x: np.array([x]) in ContinuousNDSupport(intervals=[Interval1D(-np.inf, 0)])
     )
     support_pos = SupportByPredicate(
-        predicate=lambda x: x in ContinuousNDSupport(intervals=[Interval1D(0, np.inf)])
+        predicate=lambda x: np.array([x]) in ContinuousNDSupport(intervals=[Interval1D(0, np.inf)])
     )
     fam = ContinuousExponentialClassFamily(
         log_partition=lambda parametrization: np.log(-parametrization),