[refactor] refactoring for mypy

Author: domosedy

src/pysatl_core/families/exponential_family.py

@@ -2,21 +2,21 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Iterable, Sized, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
+from scipy.differentiate import jacobian
 from scipy.integrate import nquad
 from scipy.linalg import det
-from scipy.differentiate import jacobian
 
 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,
+    GenericCharacteristicName,
     ParametrizationName,
 )
 
@@ -137,9 +137,7 @@ def _transform_to_natural_parametrization(
     def log_density(self) -> ParametrizedFunction:
         def log_density_func(parametrization: Parametrization, x: Any) -> Any:
             parametrization = cast(ExponentialFamilyParametrization, parametrization)
-            parametrization = self._transform_to_natural_parametrization(
-                parametrization
-            )
+            parametrization = self._transform_to_natural_parametrization(parametrization)
             if not self._support.accepts(x):
                 return float("-inf")
 
@@ -150,9 +148,7 @@ def log_density_func(parametrization: Parametrization, x: Any) -> Any:
                 dot = dot[0]
 
             result = float(
-                np.log(self._normalization(x))
-                + dot
-                + self._log_partition(parametrization)
+                np.log(self._normalization(x)) + dot + self._log_partition(parametrization)
             )
             return result
 
@@ -188,31 +184,24 @@ def pdf(theta: Any) -> Any:
                     theta = [theta]
                 parametrization = ExponentialFamilyParametrization(theta=theta)
                 # parametrization.theta = theta
-                return np.exp(
-                    np.dot(theta, alpha) + beta * self._log_partition(parametrization)
-                )[0]
+                return np.exp(np.dot(theta, alpha) + beta * self._log_partition(parametrization))[0]
 
             all_value = nquad(
-                lambda x: pdf(x) if self._parameter_space.accepts(x) else 0,
+                lambda x: pdf(x) if self._parameter_space.accepts(x) else 0,  # type: ignore[arg-type]
                 [(float("-inf"), float("+inf"))],
             )[0]
             return -np.log(all_value)
 
-        # TODO: remove hardcoding - Done, all hardcoding is only on user's hands
-        # 1. pr with prototype/draft - in progress
-        # 2. write instruction about to add distributions as member of exponential family - not started
-        # 3. parametrization's spaces (передавать в конструктор) - maybe impossible, discuss this with desiment on meeting
-
         def conjugate_sufficient_accepts(
             parametrization: ExponentialFamilyParametrization,
         ) -> bool:
             theta = parametrization.theta
             xi = theta[:-1]
             nu = theta[-1]
 
-            return self._sufficient_statistics_values.accepts(
-                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,
@@ -269,10 +258,8 @@ 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 self._support.accepts(x)
-                    else 0
+                lambda x: (  # type: ignore[arg-type]
+                    np.dot(x, self.density(parametrization, x)) if self._support.accepts(x) else 0
                 ),
                 [(float("-inf"), float("inf"))] * dimension_size,
             )[0]
@@ -287,10 +274,8 @@ 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 self._support.accepts(x)
-                    else 0
+                lambda x: (  # type: ignore[arg-type]
+                    x**2 * self.density(parametrization, x) if self._support.accepts(x) else 0
                 ),
                 [(float("-inf"), float("inf"))] * dimension_size,
             )[0]
@@ -301,10 +286,7 @@ def func(parametrization: Parametrization, x: Any) -> 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
-            )
+            return self._second_moment(parametrization, x) - self._mean(parametrization, x) ** 2
 
         return func
 
@@ -323,9 +305,7 @@ def posterior_hyperparameters(
             alpha_post = self._sufficient(sample)
             beta_post = 1
 
-        return ExponentialConjugateHyperparameters(
-            alpha=alpha + alpha_post, beta=beta + beta_post
-        )
+        return ExponentialConjugateHyperparameters(alpha=alpha + alpha_post, beta=beta + beta_post)
 
 
 class ExponentialFamily(NaturalExponentialFamily):
@@ -356,9 +336,7 @@ def natural_log_partition(
             return log_partition(ExponentialFamilyParametrization(theta=[theta]))
 
         natural_sufficient_statistics_values = SpacePredicate(
-            lambda eta: sufficient_statistics_values.accepts(
-                parameter_from_natural_parameter(eta)
-            )
+            lambda eta: sufficient_statistics_values.accepts(parameter_from_natural_parameter(eta))
         )
 
         self._natural_parameter = natural_parameter

tests/unit/families/test_exponential_family.py

@@ -16,14 +16,14 @@
 
 
 def gamma_pdf(alpha: float, beta: float, x: float) -> float:
-    return scipy.stats.gamma(a=alpha, scale=1 / beta).pdf(x).item()
+    return scipy.stats.gamma(a=alpha, scale=1 / beta).pdf(x).item()  # type: ignore[attr-defined]
 
 
 @pytest.fixture(scope="function")
 def conjugate_for_exponential() -> ExponentialFamily:
     def get_parameter_from_natural_parameter(
-        eta_parametrization: ExponentialFamilyParametrization,
-    ):
+        eta_parametrization: Any,
+    ) -> Any:
         if hasattr(eta_parametrization, "__len__"):
             if len(eta_parametrization) > 1:
                 return list(-1 * np.array(eta_parametrization))
@@ -34,18 +34,15 @@ def natural_parameter(
         theta_parametrization: Any,
     ) -> Any:
         if type(theta_parametrization) is ExponentialFamilyParametrization:
-            theta_parametrization = cast(
-                ExponentialFamilyParametrization, theta_parametrization
-            )
-            eta = -theta_parametrization.theta
+            eta = list(-np.array(theta_parametrization.theta))
             return ExponentialFamilyParametrization(theta=eta)
 
         return -1 * theta_parametrization
 
-    def transform_function(x: list[Any]) -> list[Any]:
-        if type(x) is not list:
-            return -x
-        return [-x[0]]
+    def transform_function(x: list[float] | float) -> list[float] | float:
+        if type(x) is list:
+            return [-x[0]]
+        return -x  # type: ignore[operator]
 
     fam = ExponentialFamily(
         log_partition=lambda parametrization: np.log(parametrization.theta[0]),
@@ -77,16 +74,12 @@ def test_exponential_pdf(theta1, theta2, conjugate_for_exponential):
     alpha = theta2 + 1
     beta = theta1
 
-    exponential = gamma_family(
-        theta=np.array([theta1, theta2]), parametrization_name="theta"
-    )
+    exponential = gamma_family(theta=np.array([theta1, theta2]), parametrization_name="theta")
     pdf = exponential.computation_strategy.query_method("pdf", distr=exponential)
 
     x = [i / 10 for i in range(100)]
 
-    assert_allclose(
-        [pdf(xx) for xx in x], [gamma_pdf(alpha, beta, xx) for xx in x], rtol=1e-6
-    )
+    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))
@@ -97,9 +90,7 @@ def test_exponential_mean(theta1, theta2, conjugate_for_exponential):
     alpha = theta2 + 1
     beta = theta1
 
-    exponential = gamma_family(
-        theta=np.array([theta1, theta2]), parametrization_name="theta"
-    )
+    exponential = gamma_family(theta=np.array([theta1, theta2]), parametrization_name="theta")
     mean = exponential.computation_strategy.query_method("mean", distr=exponential)
     assert np.isclose(mean(12), alpha / beta, rtol=1e-6)
 
@@ -112,8 +103,6 @@ def test_exponential_var(theta1, theta2, conjugate_for_exponential):
     alpha = theta2 + 1
     beta = theta1
 
-    exponential = gamma_family(
-        theta=np.array([theta1, theta2]), parametrization_name="theta"
-    )
+    exponential = gamma_family(theta=np.array([theta1, theta2]), parametrization_name="theta")
     var = exponential.computation_strategy.query_method("var", distr=exponential)
     assert np.isclose(var(12), alpha / beta**2, rtol=1e-6)