refactor(exponential): fix pr issues

Author: domosedy

src/pysatl_core/distributions/support.py

@@ -27,7 +27,14 @@
 
 import numpy as np
 
-from pysatl_core.types import BoolArray, Interval1D, IntervalND, Number, NumericArray
+from pysatl_core.types import (
+    BoolArray,
+    Interval1D,
+    IntervalND,
+    Number,
+    NumberParameter,
+    NumericArray,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator
@@ -56,7 +63,7 @@ class ContinuousSupport(Interval1D, Support):
     """
 
 
-class ContinuousNDSupport(IntervalND, Support):  # type: ignore[misc]
+class ContinuousNDSupport(IntervalND, Support):
     """
     Support for continuous distributions represented as an array of intervals.
 
@@ -446,17 +453,20 @@ def is_right_bounded(self) -> bool:
     __iter__ = iter_points
 
 
-class SupportByPredicate:
-    def __init__(self, predicate: Callable[[NumericArray | Number], bool]):
-        self._predicate = predicate
+@dataclass(slots=True)
+class SupportByPredicate(Support):
+    predicate: Callable[[NumberParameter], bool]
 
-    def __contains__(self, item: NumericArray | Number) -> bool:
-        return self._predicate(item)
+    @overload
+    def contains(self, x: Number) -> bool: ...
+    @overload
+    def contains(self, x: NumericArray) -> BoolArray: ...
 
+    def contains(self, x: NumberParameter) -> bool | BoolArray:
+        return self.predicate(x)
 
-class SupportByIntervals(SupportByPredicate):
-    def __init__(self, support: ContinuousNDSupport):
-        SupportByPredicate.__init__(self, lambda x: x in support)
+    def __contains__(self, item: object) -> bool | BoolArray:
+        return self.contains(cast(NumberParameter, item))
 
 
 __all__ = [
@@ -465,7 +475,6 @@ def __init__(self, support: ContinuousNDSupport):
     "ContinuousSupport",
     "ContinuousNDSupport",
     "SupportByPredicate",
-    "SupportByIntervals",
     # Discrete support protocol and implementations
     "DiscreteSupport",
     "ExplicitTableDiscreteSupport",

src/pysatl_core/families/__init__.py

@@ -43,7 +43,6 @@
     "ContinuousExponentialClassFamily",
     "ExponentialFamilyParametrization",
     "ExponentialConjugateHyperparameters",
-    # "CanonicalContinuousExponentialClassFamily",
 ]
 
 del _builtins_all

src/pysatl_core/families/exponential_family.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Callable
+__author__ = "Vinogradov Ilya"
+__copyright__ = "Copyright (c) 2025 PySATL project"
+__license__ = "SPDX-License-Identifier: MIT"
+
+
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 
@@ -25,11 +30,10 @@
 
 if TYPE_CHECKING:
     from pysatl_core.distributions.support import Support
-    from pysatl_core.types import Number, NumericArray
+    from pysatl_core.types import Number, NumberParameter, NumericArray
 
     type ParametrizedFunction = Callable[[Parametrization, Any], Any]
     type SupportArg = Callable[[Parametrization], Support | None] | None
-    type NumberParameter = Number | NumericArray
 
 
 @dataclass
@@ -77,6 +81,7 @@ def __init__(
         distr_type: DistributionType | Callable[[Parametrization], DistributionType],
         distr_parametrizations: list[ParametrizationName],
         support_by_parametrization: SupportArg = None,
+        base_score: Callable[[Parametrization, NumericArray], NumericArray] | None = None,
     ):
         self._sufficient = sufficient_statistics
         self._log_partition = log_partition
@@ -91,8 +96,8 @@ def __init__(
             dict[ParametrizationName, ParametrizedFunction] | ParametrizedFunction,
         ] = {
             CharacteristicName.PDF: self.density,
-            CharacteristicName.MEAN: self._mean,
-            CharacteristicName.VAR: self._var,
+            CharacteristicName.MEAN_DEFAULT: self._mean,
+            CharacteristicName.VAR_DEFAULT: self._var,
         }
 
         ParametricFamily.__init__(
@@ -102,6 +107,7 @@ def __init__(
             distr_parametrizations=distr_parametrizations,
             distr_characteristics=distr_characteristics,
             support_by_parametrization=support_by_parametrization,
+            base_score=base_score,
         )
         parametrization(family=self, name="theta")(ExponentialFamilyParametrization)
 
@@ -176,7 +182,7 @@ def conjugate_sufficient_accepts(
             normalization_constant=lambda _: 1,
             support=self._parameter_space,
             sufficient_statistics_values=self._parameter_space,  # TODO: write convex hull for this
-            parameter_space=SupportByPredicate(conjugate_sufficient_accepts),  # type: ignore[arg-type]
+            parameter_space=SupportByPredicate(predicate=conjugate_sufficient_accepts),  # type: ignore[arg-type]
             name=self.name,
             distr_type=self._distr_type,
             distr_parametrizations=self.parametrization_names,
@@ -185,28 +191,28 @@ def conjugate_sufficient_accepts(
 
     def transform(
         self,
-        transform_function: Callable[[Any], Any],
+        transform_function: Callable[[NumberParameter], NumberParameter],
     ) -> ContinuousExponentialClassFamily:
-        def calculate_jacobian(x: Any) -> Any:
-            if type(x) is not list:
+        def calculate_jacobian(x: NumberParameter) -> NumberParameter:
+            if not isinstance(x, Iterable):
                 x = np.array([x])
 
             return np.abs(det(jacobian(transform_function, x).df))
 
-        def new_support(x: Any) -> bool:
+        def new_support(x: NumberParameter) -> bool:
             return transform_function(x) in self._support
 
-        def new_sufficient(x: Any) -> Any:
+        def new_sufficient(x: NumberParameter) -> NumberParameter:
             return self._sufficient(transform_function(x))
 
-        def new_normalization(x: Any) -> Any:
+        def new_normalization(x: NumberParameter) -> NumberParameter:
             return self._normalization(x) * calculate_jacobian(x)
 
         return ContinuousExponentialClassFamily(
             log_partition=self._log_partition,
             sufficient_statistics=new_sufficient,
             normalization_constant=new_normalization,
-            support=SupportByPredicate(new_support),
+            support=SupportByPredicate(predicate=new_support),
             parameter_space=self._parameter_space,
             sufficient_statistics_values=self._sufficient_statistics_values,
             name=f"Transformed{self._name}",

src/pysatl_core/types.py

@@ -112,6 +112,9 @@ 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):
     """
@@ -246,6 +249,12 @@ def shape(self) -> ContinuousSupportShape1D:
 class IntervalND:
     intervals: list[Interval1D]
 
+    @overload
+    def contains(self, x: Number) -> bool: ...
+
+    @overload
+    def contains(self, x: NumericArray) -> BoolArray: ...
+
     def contains(self, x: Number | NumericArray) -> bool | BoolArray:
         if not hasattr(x, "__iter__"):
             x = np.array([x])
@@ -307,6 +316,8 @@ class CharacteristicName(StrEnum):
     CDF = "cdf"
     PPF = "ppf"
     PMF = "pmf"
+    MEAN_DEFAULT = "MEAN_DEFAULT"  # defined in class implementation of mean
+    VAR_DEFAULT = "VAR_DEFAULT"  # defined in class implementation of var
     LPDF = "lpdf"  # unimplemented in graph yet
     CF = "cf"  # unimplemented in graph yet
     SF = "sf"  # unimplemented in graph yet

tests/unit/families/test_exponential_family.py

@@ -1,16 +1,22 @@
+__author__ = "Leonid Elkin"
+__copyright__ = "Copyright (c) 2025 PySATL project"
+__license__ = "SPDX-License-Identifier: MIT"
+
+import itertools
+from collections.abc import Iterable
 from typing import cast
 
 import numpy as np
 import pytest
 import scipy
 from numpy.testing import assert_allclose
 
-from pysatl_core.distributions.support import ContinuousNDSupport, SupportByIntervals
+from pysatl_core.distributions.support import ContinuousNDSupport, SupportByPredicate
 from pysatl_core.families import (
     ContinuousExponentialClassFamily,
 )
 from pysatl_core.families.registry import ParametricFamilyRegister
-from pysatl_core.types import Interval1D, UnivariateContinuous
+from pysatl_core.types import CharacteristicName, Interval1D, NumberParameter, UnivariateContinuous
 
 
 def gamma_pdf(alpha: float, beta: float, x: float) -> float:
@@ -19,13 +25,17 @@ def gamma_pdf(alpha: float, beta: float, x: float) -> float:
 
 @pytest.fixture(scope="function")
 def conjugate_for_exponential() -> ContinuousExponentialClassFamily:
-    def transform_function(x: list[float] | float) -> list[float] | float:
-        if type(x) is list:
-            return [-x[0]]
-        return -x  # type: ignore[operator]
+    def transform_function(x: NumberParameter) -> NumberParameter:
+        if isinstance(x, Iterable):
+            return np.array([-x[0]])
+        return -x
 
-    support_neg = SupportByIntervals(ContinuousNDSupport(intervals=[Interval1D(-np.inf, 0)]))
-    support_pos = SupportByIntervals(ContinuousNDSupport(intervals=[Interval1D(0, np.inf)]))
+    support_neg = SupportByPredicate(
+        predicate=lambda x: x in ContinuousNDSupport(intervals=[Interval1D(-np.inf, 0)])
+    )
+    support_pos = SupportByPredicate(
+        predicate=lambda x: x in ContinuousNDSupport(intervals=[Interval1D(0, np.inf)])
+    )
     fam = ContinuousExponentialClassFamily(
         log_partition=lambda parametrization: np.log(-parametrization),
         sufficient_statistics=lambda x: x,
@@ -45,8 +55,10 @@ def transform_function(x: list[float] | float) -> list[float] | float:
     )
 
 
-@pytest.mark.parametrize("theta1", range(2, 5))
-@pytest.mark.parametrize("theta2", range(2, 5))
+@pytest.mark.parametrize(
+    ("theta1", "theta2"),
+    itertools.product(range(2, 5), range(2, 5)),
+)
 def test_exponential_pdf(theta1, theta2, conjugate_for_exponential):
     gamma_family: ContinuousExponentialClassFamily = conjugate_for_exponential
 
@@ -61,27 +73,35 @@ def test_exponential_pdf(theta1, theta2, conjugate_for_exponential):
     assert_allclose([pdf(xx) for xx in x], [gamma_pdf(alpha, beta, xx) for xx in x], rtol=1e-6)
 
 
-@pytest.mark.parametrize("theta1", range(2, 5))
-@pytest.mark.parametrize("theta2", range(2, 5))
+@pytest.mark.parametrize(
+    ("theta1", "theta2"),
+    itertools.product(range(2, 5), range(2, 5)),
+)
 def test_exponential_mean(theta1, theta2, conjugate_for_exponential):
     gamma_family: ContinuousExponentialClassFamily = conjugate_for_exponential
 
     alpha = theta2 + 1
     beta = theta1
 
     exponential = gamma_family(theta=np.array([theta1, theta2]), parametrization_name="theta")
-    mean = exponential.computation_strategy.query_method("mean", distr=exponential)
+    mean = exponential.computation_strategy.query_method(
+        CharacteristicName.MEAN_DEFAULT, distr=exponential
+    )
     assert np.isclose(mean(12), alpha / beta, rtol=1e-6)
 
 
-@pytest.mark.parametrize("theta1", range(2, 5))
-@pytest.mark.parametrize("theta2", range(2, 5))
+@pytest.mark.parametrize(
+    ("theta1", "theta2"),
+    itertools.product(range(2, 5), range(2, 5)),
+)
 def test_exponential_var(theta1, theta2, conjugate_for_exponential):
     gamma_family: ContinuousExponentialClassFamily = conjugate_for_exponential
 
     alpha = theta2 + 1
     beta = theta1
 
     exponential = gamma_family(theta=np.array([theta1, theta2]), parametrization_name="theta")
-    var = exponential.computation_strategy.query_method("var", distr=exponential)
+    var = exponential.computation_strategy.query_method(
+        CharacteristicName.VAR_DEFAULT, distr=exponential
+    )
     assert np.isclose(var(12), alpha / beta**2, rtol=1e-6)