feat: Implement dynamic integrator selection, enabling the use of custom integrators

Author: Amellgo

src/algorithms/semiparam_algorithms/nvm_semi_param_algorithms/g_estimation_given_mu_rqmc_based.py

@@ -7,7 +7,8 @@
 from scipy.integrate import quad_vec
 from scipy.special import gamma
 
-from src.algorithms.support_algorithms.rqmc import RQMC
+from src.algorithms.support_algorithms.integrator import Integrator
+from src.algorithms.support_algorithms.rqmc import RQMCIntegrator
 from src.estimators.estimate_result import EstimateResult
 
 MU_DEFAULT_VALUE = 1.0
@@ -30,6 +31,7 @@ def v_sequence_default_value(n: float) -> float:
 GRID_SIZE_DEFAULT_VALUE: int = 200
 INTEGRATION_TOLERANCE_DEFAULT_VALUE: float = 1e-2
 INTEGRATION_LIMIT_DEFAULT_VALUE: int = 50
+INTEGRATOR_DEFAULT = RQMCIntegrator()
 
 
 class SemiParametricGEstimationGivenMuRQMCBased:
@@ -69,6 +71,7 @@ def __init__(self, sample: Optional[_typing.NDArray[np.float64]] = None, **kwarg
             self.grid_size,
             self.integration_tolerance,
             self.integration_limit,
+            self.integrator
         ) = self._validate_kwargs(self.n, **kwargs)
         self.denominator: float = 2 * math.pi * self.n
         self.precompute_gamma_grid()
@@ -78,7 +81,7 @@ def __init__(self, sample: Optional[_typing.NDArray[np.float64]] = None, **kwarg
     @staticmethod
     def _validate_kwargs(
         n: int, **kwargs: Unpack[ParamsAnnotation]
-    ) -> tuple[float, float, float, float, List[float], int, float, int]:
+    ) -> tuple[float, float, float, float, List[float], int, float, int, Integrator]:
         mu: float = kwargs.get("mu", MU_DEFAULT_VALUE)
         gmm: float = kwargs.get("gmm", GAMMA_DEFAULT_VALUE)
         u_value: float = kwargs.get("u_value", U_SEQUENCE_DEFAULT_VALUE(n))
@@ -87,7 +90,8 @@ def _validate_kwargs(
         grid_size: int = kwargs.get("grid_size", GRID_SIZE_DEFAULT_VALUE)
         integration_tolerance: float = kwargs.get("integration_tolerance", INTEGRATION_TOLERANCE_DEFAULT_VALUE)
         integration_limit: int = kwargs.get("integration_limit", INTEGRATION_LIMIT_DEFAULT_VALUE)
-        return mu, gmm, u_value, v_value, x_data, grid_size, integration_tolerance, integration_limit
+        integrator: Integrator = kwargs.get("integrator_type", INTEGRATOR_DEFAULT)
+        return mu, gmm, u_value, v_value, x_data, grid_size, integration_tolerance, integration_limit, integrator
 
     def conjugate_psi(self, u: float) -> complex:
         return complex((u**2) / 2, self.mu * u)
@@ -162,15 +166,16 @@ def second_v_integrand(self, v: float, x: float) -> np.ndarray:
         x_power = self.x_powers[x][idx]
         return (self.second_u_integrals[idx] * x_power) / gamma_val
 
-    def compute_integrals_for_x(self, x: float) -> float:
+    def compute_integrals_for_x(self, x: float, integrator: Integrator = None) -> float:
         """Compute integrals using RQMC for v-integration."""
-        first_integral = RQMC(lambda t: np.sum(self.first_v_integrand(t * self.v_value, x)) * self.v_value).rqmc()[0]
+        integrator = integrator or self.integrator
+        first_integral = integrator.compute_integral(func=lambda t: np.sum(self.first_v_integrand(t * self.v_value, x)) * self.v_value).value
 
-        second_integral = RQMC(lambda t: np.sum(self.second_v_integrand(-t * self.v_value, x)) * self.v_value).rqmc()[0]
+        second_integral = integrator.compute_integral(func=lambda t: np.sum(self.second_v_integrand(-t * self.v_value, x)) * self.v_value).value
 
         total = (first_integral + second_integral) / self.denominator
         return max(0.0, total.real)
 
     def algorithm(self, sample: np._typing.NDArray) -> EstimateResult:
-        y_data = [self.compute_integrals_for_x(x) for x in self.x_data]
+        y_data = [self.compute_integrals_for_x(x, self.integrator) for x in self.x_data]
         return EstimateResult(list_value=y_data, success=True)

src/algorithms/support_algorithms/integrator.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from typing import Any, Protocol, Callable, Optional
+
+
+@dataclass
+class IntegrationResult:
+    value: float
+    error: float
+    message: Optional[dict[str, Any]] | None = None
+
+
+class Integrator(Protocol):
+
+    """Base class for integral calculation"""
+
+    def compute_integral(self, func: Callable, params: dict) -> IntegrationResult:
+        ...

src/algorithms/support_algorithms/quad_integrator.py

@@ -0,0 +1,27 @@
+from typing import Callable
+
+from scipy.integrate import quad
+from src.algorithms.support_algorithms.integrator import IntegrationResult
+
+class QuadIntegrator:
+
+    def compute_integral(self, func: Callable, params: dict) -> IntegrationResult:
+
+        """
+        Compute integral via quad integrator
+
+        Args:
+            func: integrated function
+            params: Parameters of integration algorithm
+
+        Returns: moment approximation and error tolerance
+        """
+
+        full_output_requested = params.pop('full_output', False)
+        quad_res = quad(func, **params)
+        if full_output_requested:
+            value, error, message = quad_res
+        else:
+            value, error = quad_res
+            message = None
+        return IntegrationResult(value, error, message)

src/algorithms/support_algorithms/rqmc.py

@@ -5,6 +5,8 @@
 import scipy
 from numba import njit
 
+from src.algorithms.support_algorithms.integrator import IntegrationResult
+
 BITS = 30
 """Number of bits in XOR. Should be less than 64"""
 NUMBA_FAST_MATH = True
@@ -126,7 +128,8 @@ def _update(
 
         Returns:Updated mean of all rows
 
-        """
+
+                """
         values = []
         sum_of_new: float = 0.0
         for i in range(self.count):
@@ -212,9 +215,9 @@ def _xor_float(a: float, b: float) -> float:
         Returns: XOR float value
 
         """
-        a = int(a * (2**BITS))
-        b = int(b * (2**BITS))
-        return np.bitwise_xor(a, b) / 2**BITS
+        a = int(a * (2 ** BITS))
+        b = int(b * (2 ** BITS))
+        return np.bitwise_xor(a, b) / 2 ** BITS
 
     def __call__(self) -> tuple[float, float]:
         """Interface for users
@@ -223,3 +226,20 @@ def __call__(self) -> tuple[float, float]:
 
         """
         return self.rqmc()
+
+
+class RQMCIntegrator:
+
+    def compute_integral(self, func: Callable, params: dict) -> IntegrationResult:
+        """
+        Compute integral via RQMC integrator
+
+        Args:
+            func: integrated function
+            params: Parameters of integration algorithm
+
+        Returns: moment approximation and error tolerance
+        """
+
+        rqmc = RQMC(func, **params)()
+        return IntegrationResult(rqmc[0], rqmc[1])

src/mixtures/nm_mixture.py

@@ -2,13 +2,14 @@
 from typing import Any
 
 import numpy as np
-from scipy.integrate import quad
 from scipy.special import binom
 from scipy.stats import norm, rv_continuous
 from scipy.stats.distributions import rv_frozen
 
 from src.algorithms.support_algorithms.log_rqmc import LogRQMC
-from src.algorithms.support_algorithms.rqmc import RQMC
+from src.algorithms.support_algorithms.integrator import Integrator
+from src.algorithms.support_algorithms.rqmc import RQMCIntegrator
+from src.algorithms.support_algorithms.quad_integrator import QuadIntegrator
 from src.mixtures.abstract_mixture import AbstractMixtures
 
 
@@ -59,156 +60,175 @@ def _params_validation(self, data_collector: Any, params: dict[str, float | rv_c
             raise ValueError("Gamma cant be zero")
         return data_class
 
-    def _classical_moment(self, n: int, params: dict) -> tuple[float, float]:
+    def _classical_moment(self, n: int, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Compute n-th moment of classical NMM
 
         Args:
             n (): Moment ordinal
             params (): Parameters of integration algorithm
+            integrator (): type of integrator to computing
 
         Returns: moment approximation and error tolerance
 
         """
         mixture_moment = 0
         error_tolerance = 0
+        integrator = integrator or QuadIntegrator()
         for k in range(0, n + 1):
             for l in range(0, k + 1):
-                coefficient = binom(n, n - k) * binom(k, k - l) * (self.params.beta ** (k - l)) * (self.params.gamma**l)
-                mixing_moment = quad(lambda u: self.params.distribution.ppf(u) ** (k - l), 0, 1, **params)
-                error_tolerance += (self.params.beta ** (k - l)) * mixing_moment[1]
-                mixture_moment += coefficient * (self.params.alpha ** (n - k)) * mixing_moment[0] * norm.moment(l)
-        return mixture_moment, error_tolerance
+                coefficient = binom(n, n - k) * binom(k, k - l) * (self.params.beta ** (k - l)) * (
+                            self.params.gamma ** l)
+                mixing_moment = integrator.compute_integral(lambda u: self.params.distribution.ppf(u) ** (k - l),
+                                                            {"a": 0, "b": 1, **params})
+                error_tolerance += (self.params.beta ** (k - l)) * mixing_moment.error
+                mixture_moment += coefficient * (self.params.alpha ** (n - k)) * mixing_moment.value * norm.moment(l)
+            return mixture_moment, error_tolerance
+
 
-    def _canonical_moment(self, n: int, params: dict) -> tuple[float, float]:
+    def _canonical_moment(self, n: int, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Compute n-th moment of canonical NMM
 
         Args:
             n (): Moment ordinal
             params (): Parameters of integration algorithm
+            integrator (): type of integrator to computing
 
         Returns: moment approximation and error tolerance
 
         """
         mixture_moment = 0
         error_tolerance = 0
+        integrator = integrator or QuadIntegrator()
         for k in range(0, n + 1):
-            coefficient = binom(n, n - k) * (self.params.sigma**k)
-            mixing_moment = quad(lambda u: self.params.distribution.ppf(u) ** (n - k), 0, 1, **params)
-            error_tolerance += mixing_moment[1]
-            mixture_moment += coefficient * mixing_moment[0] * norm.moment(k)
+            coefficient = binom(n, n - k) * (self.params.sigma ** k)
+            mixing_moment = integrator.compute_integral(lambda u: self.params.distribution.ppf(u) ** (n - k),
+                                                        {"a": 0, "b": 1, **params})
+            error_tolerance += mixing_moment.error
+            mixture_moment += coefficient * mixing_moment.value * norm.moment(k)
         return mixture_moment, error_tolerance
 
-    def compute_moment(self, n: int, params: dict) -> tuple[float, float]:
+    def compute_moment(self, n: int, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Compute n-th moment of  NMM
 
         Args:
             n (): Moment ordinal
             params (): Parameters of integration algorithm
+            integrator (): type of integrator to computing
 
         Returns: moment approximation and error tolerance
 
         """
         if isinstance(self.params, _NMMClassicDataCollector):
-            return self._classical_moment(n, params)
-        return self._canonical_moment(n, params)
+            return self._classical_moment(n, params, integrator)
+        return self._canonical_moment(n, params, integrator)
 
-    def _canonical_compute_cdf(self, x: float, params: dict) -> tuple[float, float]:
+    def _canonical_compute_cdf(self, x: float, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Equation for canonical cdf
         Args:
             x (): point
             params (): parameters of RQMC algorithm
+            integrator (): type of integrator to computing
 
         Returns: computed cdf and error tolerance
 
         """
-        rqmc = RQMC(lambda u: norm.cdf((x - self.params.distribution.ppf(u)) / np.abs(self.params.sigma)), **params)
-        return rqmc()
+        integrator = integrator or RQMCIntegrator()
+        rqmc = integrator.compute_integral(func=lambda u: norm.cdf((x - self.params.distribution.ppf(u)) / np.abs(self.params.sigma)), **params)
+        return rqmc.value, rqmc.error
 
-    def _classical_compute_cdf(self, x: float, params: dict) -> tuple[float, float]:
+    def _classical_compute_cdf(self, x: float, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Equation for classic cdf
         Args:
             x (): point
             params (): parameters of RQMC algorithm
+            integrator (): type of integrator to computing
 
         Returns: computed cdf and error tolerance
 
         """
-        rqmc = RQMC(
+        integrator = integrator or RQMCIntegrator()
+        rqmc = integrator.compute_integral(func=
             lambda u: norm.cdf(
                 (x - self.params.alpha - self.params.beta * self.params.distribution.ppf(u)) / np.abs(self.params.gamma)
             ),
             **params
         )
-        return rqmc()
+        return rqmc.value, rqmc.error
 
-    def compute_cdf(self, x: float, params: dict) -> tuple[float, float]:
+    def compute_cdf(self, x: float, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Choose equation for cdf estimation depends on Mixture form
         Args:
             x (): point
             params (): parameters of RQMC algorithm
+            integrator (): type of integrator to computing
 
         Returns: Computed pdf and error tolerance
 
         """
         if isinstance(self.params, _NMMCanonicalDataCollector):
-            return self._canonical_compute_cdf(x, params)
-        return self._classical_compute_cdf(x, params)
+            return self._canonical_compute_cdf(x, params, integrator)
+        return self._classical_compute_cdf(x, params, integrator)
 
-    def _canonical_compute_pdf(self, x: float, params: dict) -> tuple[float, float]:
+    def _canonical_compute_pdf(self, x: float, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Equation for canonical pdf
         Args:
             x (): point
             params (): parameters of RQMC algorithm
+            integrator (): type of integrator to computing
 
         Returns: computed pdf and error tolerance
 
         """
-        rqmc = RQMC(
+        integrator = integrator or RQMCIntegrator()
+        rqmc = integrator.compute_integral(func=
             lambda u: (1 / np.abs(self.params.sigma))
             * norm.pdf((x - self.params.distribution.ppf(u)) / np.abs(self.params.sigma)),
             **params
         )
-        return rqmc()
+        return rqmc.value, rqmc.error
 
-    def _classical_compute_pdf(self, x: float, params: dict) -> tuple[float, float]:
+    def _classical_compute_pdf(self, x: float, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Equation for classic pdf
         Args:
             x (): point
             params (): parameters of RQMC algorithm
+            integrator (): type of integrator to computing
 
         Returns: computed pdf and error tolerance
 
         """
-        rqmc = RQMC(
+        integrator = integrator or RQMCIntegrator()
+        rqmc = integrator.compute_integral(func=
             lambda u: (1 / np.abs(self.params.gamma))
             * norm.pdf(
                 (x - self.params.alpha - self.params.beta * self.params.distribution.ppf(u)) / np.abs(self.params.gamma)
             ),
             **params
         )
-        return rqmc()
+        return rqmc.value, rqmc.error
 
-    def compute_pdf(self, x: float, params: dict) -> tuple[float, float]:
+    def compute_pdf(self, x: float, params: dict, integrator: Integrator = None) -> tuple[float, float]:
         """
         Choose equation for pdf estimation depends on Mixture form
         Args:
             x (): point
             params (): parameters of RQMC algorithm
+            integrator (): type of integrator to computing
 
         Returns: Computed pdf and error tolerance
 
         """
         if isinstance(self.params, _NMMCanonicalDataCollector):
-            return self._canonical_compute_pdf(x, params)
-        return self._classical_compute_pdf(x, params)
+            return self._canonical_compute_pdf(x, params, integrator)
+        return self._classical_compute_pdf(x, params, integrator)
 
     def _classical_compute_log_pdf(self, x: float, params: dict) -> tuple[float, float]:
         """
@@ -258,4 +278,4 @@ def compute_logpdf(self, x: float, params: dict) -> tuple[float, float]:
         """
         if isinstance(self.params, _NMMCanonicalDataCollector):
             return self._canonical_compute_log_pdf(x, params)
-        return self._classical_compute_log_pdf(x, params)
+        return self._classical_compute_log_pdf(x, params)