feat(estimators, optimizers): switch to generics

Author: xImoZA

rework_pysatl_mpest/core/mixture.py

@@ -347,6 +347,30 @@ def generate(self, size: int) -> NDArray[DType]:
         np.random.shuffle(samples)
         return samples
 
+    def astype(self, new_dtype: type[DType]) -> "MixtureModel[DType]":
+        """Creates a copy of the MixtureModel with a new data type.
+
+        If the specified `new_dtype` is the same as the instance's current `dtype`,
+        this method returns the original instance instead.
+
+        Parameters
+        ----------
+        new_dtype : type[DType]
+            The target NumPy data type for the new distribution instance.
+
+        Returns
+        -------
+        MixtureModel[DType]
+            A new MixtureModel instance with all components and weights converted to the
+            specified `new_dtype`, or the original instance if the `dtype` is
+            unchanged.
+        """
+        if self.dtype is new_dtype:
+            return self
+
+        new_mixture = MixtureModel(components=self.components, weights=self.weights.copy(), dtype=new_dtype)
+        return new_mixture
+
     def __getitem__(self, key: int) -> "ContinuousDistribution[DType]":
         """Retrieves components by index.
 

rework_pysatl_mpest/estimators/base_estimator.py

@@ -1,17 +1,19 @@
 """A module that provides an abstract class for implementing custom estimators."""
 
-__author__ = "Danil Totmyanin"
+__author__ = "Danil Totmyanin, Aleksandra Ri"
 __copyright__ = "Copyright (c) 2025 PySATL project"
 __license__ = "SPDX-License-Identifier: MIT"
 
 from abc import ABC, abstractmethod
+from typing import Generic
 
 from numpy.typing import ArrayLike
 
 from ..core import MixtureModel
+from ..typings import DType
 
 
-class BaseEstimator(ABC):
+class BaseEstimator(ABC, Generic[DType]):
     """Abstract class for a mixture model parameter estimator.
 
     This class defines the interface for all estimator algorithms. Estimators are responsible for
@@ -35,7 +37,7 @@ class BaseEstimator(ABC):
     """
 
     @abstractmethod
-    def fit(self, X: ArrayLike, mixture: MixtureModel) -> MixtureModel:
+    def fit(self, X: ArrayLike, mixture: MixtureModel[DType]) -> MixtureModel[DType]:
         """Fits the mixture model to the provided data.
 
         This method estimates the parameters of the model's components and their
@@ -45,11 +47,11 @@ def fit(self, X: ArrayLike, mixture: MixtureModel) -> MixtureModel:
         ----------
         X : ArrayLike
             The input data sample for fitting the model.
-        mixture : MixtureModel
+        mixture : MixtureModel[DType]
             The initial mixture model to be fitted.
 
         Returns
         -------
-        MixtureModel
+        MixtureModel[DType]
             The mixture model with estimated parameters.
         """

rework_pysatl_mpest/estimators/ecm.py

@@ -6,7 +6,7 @@
 to fit the parameters of a mixture model to data.
 """
 
-__author__ = "Danil Totmyanin"
+__author__ = "Danil Totmyanin, Aleksandra Ri"
 __copyright__ = "Copyright (c) 2025 PySATL project"
 __license__ = "SPDX-License-Identifier: MIT"
 
@@ -17,6 +17,7 @@
 
 from ..core import MixtureModel
 from ..optimizers import Optimizer
+from ..typings import DType
 from .base_estimator import BaseEstimator
 from .iterative import (
     Breakpointer,
@@ -30,7 +31,7 @@
 from .iterative._logger import IterationsHistory
 
 
-class ECM(BaseEstimator):
+class ECM(BaseEstimator[DType]):
     """An estimator that implements the Expectation-Conditional Maximization (ECM) algorithm.
 
     This class encapsulates the logic for the ECM algorithm, a variant of the
@@ -87,10 +88,10 @@ def __init__(self, breakpointers: Sequence[Breakpointer], pruners: Sequence[Prun
         self.breakpointers = list(breakpointers)
         self.pruners = list(pruners)
         self.optimizer = optimizer
-        self._logger: IterationsHistory | None = None
+        self._logger: IterationsHistory[DType] | None = None
 
     @property
-    def logger(self) -> IterationsHistory:
+    def logger(self) -> IterationsHistory[DType]:
         """An object that collects information about each iteration.
 
         Raises
@@ -103,7 +104,7 @@ def logger(self) -> IterationsHistory:
             raise AttributeError("Logger is not available. Call the 'fit' method first.")
         return self._logger
 
-    def fit(self, X: ArrayLike, mixture: MixtureModel, once_in_iterations: int = 1) -> MixtureModel:
+    def fit(self, X: ArrayLike, mixture: MixtureModel[DType], once_in_iterations: int = 1) -> MixtureModel[DType]:
         """Fits the mixture model to the data using the ECM algorithm.
 
         This method sets up and runs an iterative pipeline to estimate the
@@ -115,15 +116,15 @@ def fit(self, X: ArrayLike, mixture: MixtureModel, once_in_iterations: int = 1)
         ----------
         X : ArrayLike
             The input dataset for fitting the model.
-        mixture : MixtureModel
+        mixture : MixtureModel[DType]
             The initial mixture model to be fitted.
         once_in_iterations : int, optional
             The logging frequency. A value of `n` means logging occurs every
             `n` iterations. Defaults to 1.
 
         Returns
         -------
-        MixtureModel
+        MixtureModel[DType]
             The mixture model with the estimated parameters.
         """
 
@@ -132,7 +133,7 @@ def fit(self, X: ArrayLike, mixture: MixtureModel, once_in_iterations: int = 1)
             block = OptimizationBlock(i, comp.params_to_optimize, MaximizationStrategy.QFUNCTION)
             blocks.append(block)
 
-        pipeline = Pipeline(
+        pipeline: Pipeline[DType] = Pipeline(
             [ExpectationStep(), MaximizationStep(blocks, self.optimizer)],
             self.breakpointers,
             self.pruners,

rework_pysatl_mpest/estimators/iterative/_logger.py

@@ -6,22 +6,22 @@
 and errors across iterations.
 """
 
-__author__ = "Maksim Pastukhov"
+__author__ = "Maksim Pastukhov, Aleksandra Ri"
 __copyright__ = "Copyright (c) 2025 PySATL project"
 __license__ = "SPDX-License-Identifier: MIT"
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Generic, Optional
 
-from numpy import float64
 from numpy.typing import NDArray
 
 from ...core import MixtureModel
+from ...typings import DType
 from .pruner import Pruner
 
 
 @dataclass
-class IterationRecord:
+class IterationRecord(Generic[DType]):
     """Data class representing a single pipeline iteration snapshot.
 
     This class captures the complete state of a pipeline iteration after
@@ -32,15 +32,15 @@ class IterationRecord:
     ----------
     iteration : int
         The iteration number (0-based index).
-    mixture : MixtureModel
+    mixture : MixtureModel[DType]
         The state of the mixture model after pruning in this iteration.
-    X : NDArray[float64]
+    X : NDArray[DType]
         The input data sample being processed (conventionally named `X`).
-    H : Optional[NDArray[float64]]
+    H : Optional[NDArray[DType]]
         The responsibility matrix (posterior probabilities) if available,
         where `H[i, j]` represents the probability that data point `i`
         belongs to component `j`. May be `None` if not computed.
-    pruners_used : Optional[list[Pruner]]
+    pruners_used : Optional[list[Pruner[DType]]]
         List of pruner instances that were applied during this iteration.
         `None` or empty if no pruning occurred.
     error : Optional[Exception]
@@ -49,14 +49,14 @@ class IterationRecord:
     """
 
     iteration: int
-    mixture: MixtureModel
-    X: NDArray[float64]
-    H: Optional[NDArray[float64]]
-    pruners_used: Optional[list[Pruner]]
+    mixture: MixtureModel[DType]
+    X: NDArray[DType]
+    H: Optional[NDArray[DType]]
+    pruners_used: Optional[list[Pruner[DType]]]
     error: Optional[Exception]
 
 
-class IterationsHistory:
+class IterationsHistory(Generic[DType]):
     """A container for storing and accessing pipeline iteration history.
 
     `IterationsHistory` collects and stores snapshots of each pipeline iteration
@@ -98,7 +98,7 @@ class IterationsHistory:
     _counter : int
         Internal counter tracking the total number of `log()` calls
         (i.e., total iterations processed, not just recorded ones).
-    _logs : list[IterationRecord]
+    _logs : list[IterationRecord[DType]]
         List of stored iteration records. Only iterations matching the
         recording frequency are appended.
 
@@ -123,11 +123,11 @@ def __init__(self, once_in_iterations: int = 1) -> None:
         if once_in_iterations < 1:
             raise ValueError("once_in_iterations must be a positive integer")
 
-        self._logs: list[IterationRecord] = []
+        self._logs: list[IterationRecord[DType]] = []
         self._counter: int = 0
         self.once_in_iterations = once_in_iterations
 
-    def log(self, record: IterationRecord) -> None:
+    def log(self, record: IterationRecord[DType]) -> None:
         """Store an iteration record based on the configured frequency.
 
         The record is stored only if the current internal counter is divisible
@@ -136,7 +136,7 @@ def log(self, record: IterationRecord) -> None:
 
         Parameters
         ----------
-        record : IterationRecord
+        record : IterationRecord[DType]
             The iteration snapshot to potentially store. The `record.iteration`
             should ideally match the logger's internal state, though this is
             not enforced.
@@ -168,7 +168,7 @@ def __len__(self) -> int:
         """
         return len(self._logs)
 
-    def __getitem__(self, index: int) -> IterationRecord:
+    def __getitem__(self, index: int) -> IterationRecord[DType]:
         """Access a stored iteration record by index.
 
         Supports both positive (0-based) and negative indexing (e.g., `-1` for last).
@@ -181,7 +181,7 @@ def __getitem__(self, index: int) -> IterationRecord:
 
         Returns
         -------
-        IterationRecord
+        IterationRecord[DType]
             The recorded state of the specified iteration.
 
         Raises

rework_pysatl_mpest/estimators/iterative/_strategies/__init__.py

@@ -9,12 +9,13 @@
 
 from ....distributions import ContinuousDistribution
 from ....optimizers import Optimizer
+from ....typings import DType
 from ..pipeline_state import PipelineState
 from ..steps import OptimizationBlock
 from .q_function import q_function_strategy as _q_function_strategy
 
 q_function_strategy: Callable[
-    [ContinuousDistribution, PipelineState, OptimizationBlock, Optimizer], tuple[int, dict[str, float]]
+    [ContinuousDistribution, PipelineState, OptimizationBlock, Optimizer], tuple[int, dict[str, DType]]
 ] = _q_function_strategy