refactor: fit-evaluate strategy now work more clearly

Author: LeonidElkin

src/pysatl_core/distributions/computation.py

@@ -13,7 +13,7 @@
 
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Protocol, overload, runtime_checkable
+from typing import TYPE_CHECKING, Protocol, cast, overload, runtime_checkable
 
 from pysatl_core.types import ComputationFunc
 
@@ -25,6 +25,11 @@
     from pysatl_core.distributions.distribution import Distribution
     from pysatl_core.types import GenericCharacteristicName
 
+type Fitter[In, Out] = Callable[[Distribution, KwArg(Any)], FittedComputationMethod[In, Out]]
+type Evaluator[In, Out] = (
+    Callable[[Distribution, KwArg(Any)], Out] | Callable[[Distribution, In, KwArg(Any)], Out]
+)
+
 
 @runtime_checkable
 class Computation[In, Out](Protocol):
@@ -58,7 +63,7 @@ class AnalyticalComputation[In, Out]:
     ----------
     target : str
         Characteristic name (e.g., "pdf", "cdf").
-    func : Callable[..., Out]
+    func : ComputationFunc[In, Out]
         Analytical function that computes the characteristic.
     """
 
@@ -87,7 +92,7 @@ class FittedComputationMethod[In, Out]:
         Destination characteristic name.
     sources : Sequence[str]
         Source characteristic names (typically length 1 for unary conversions).
-    func : Callable[..., Out]
+    func : ComputationFunc[In, Out]
         Callable implementing the fitted conversion.
     """
 
@@ -120,13 +125,72 @@ class ComputationMethod[In, Out]:
         Destination characteristic name.
     sources : Sequence[str]
         Source characteristic names (typically length 1 for unary conversions).
-    fitter : Callable[[Distribution, **options], FittedComputationMethod]
+    fitter : Fitter[In, Out] | None
         Function that fits the computation method to a distribution.
+        If provided, the method is considered *cacheable* (fitting may perform
+        expensive precomputation).
+    evaluator : Evaluator[In, Out] | None
+        Direct evaluator that performs the computation in one step, without
+        a separate fitting stage. If provided, the method is considered
+        *non-cacheable* at the strategy level.
     """
 
     target: GenericCharacteristicName
     sources: Sequence[GenericCharacteristicName]
-    fitter: Callable[[Distribution, KwArg(Any)], FittedComputationMethod[In, Out]]
+    fitter: Fitter[In, Out] | None = None
+    evaluator: Evaluator[In, Out] | None = None
+
+    def __post_init__(self) -> None:
+        has_fitter = self.fitter is not None
+        has_eval = self.evaluator is not None
+        if has_fitter == has_eval:
+            raise ValueError(
+                "ComputationMethod must define exactly one of 'fitter' or 'evaluator'."
+            )
+
+    @property
+    def cacheable(self) -> bool:
+        """Whether it makes sense to cache the prepared method at strategy level."""
+        return self.fitter is not None
+
+    def prepare(
+        self, distribution: Distribution, **options: Any
+    ) -> FittedComputationMethod[In, Out]:
+        """Prepare a callable method for a specific distribution.
+
+        - If ``fitter`` is provided, run the fitting stage and return the fitted method.
+        - If ``evaluator`` is provided, bind the distribution and return a lightweight
+          fitted wrapper.
+        """
+        if self.fitter is not None:
+            return self.fitter(distribution, **options)
+
+        def _bound(*args: Any, **kwargs: Any) -> Out:
+            return cast(Evaluator[In, Out], self.evaluator)(distribution, *args, **kwargs)
+
+        return FittedComputationMethod[In, Out](
+            target=self.target,
+            sources=self.sources,
+            func=_bound,
+        )
+
+    @overload
+    def __call__(self, distribution: Distribution, **options: Any) -> Out: ...
+
+    @overload
+    def __call__(self, distribution: Distribution, data: In, **options: Any) -> Out: ...
+
+    def __call__(self, distribution: Distribution, *args: Any, **options: Any) -> Out:
+        """Evaluate *direct* computation methods.
+
+        This is only available for methods defined via ``evaluator``.
+        """
+        if self.evaluator is None:
+            raise RuntimeError(
+                "This ComputationMethod requires fitting. "
+                "Call .fit(...) / .prepare(...) to obtain a callable."
+            )
+        return self.evaluator(distribution, *args, **options)
 
     def fit(self, distribution: Distribution, **options: Any) -> FittedComputationMethod[In, Out]:
         """
@@ -144,6 +208,11 @@ def fit(self, distribution: Distribution, **options: Any) -> FittedComputationMe
         FittedComputationMethod
             Fitted method ready for evaluation.
         """
+        if self.fitter is None:
+            raise RuntimeError(
+                "This ComputationMethod is evaluator-based and does not support .fit(). "
+                "Use .prepare(...) or call the method directly."
+            )
         return self.fitter(distribution, **options)
 
 

src/pysatl_core/distributions/strategies.py

@@ -164,8 +164,8 @@ def query_method(
                 # Fit each edge along the path
                 last_fitted: FittedComputationMethod[Any, Any] | None = None
                 for edge in path:
-                    fitted = edge.fit(distr, **options)
-                    if self.enable_caching:
+                    fitted = edge.prepare(distr, **options)
+                    if self.enable_caching and edge.cacheable:
                         self._cache[edge.target] = fitted
                     last_fitted = fitted