refactor(distribution): now there can be several analytical calculations for one characteristic

Author: LeonidElkin

src/pysatl_core/distributions/distribution.py

@@ -12,27 +12,28 @@
 __license__ = "SPDX-License-Identifier: MIT"
 
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Mapping
+from collections.abc import Mapping
 from copy import deepcopy
 from typing import TYPE_CHECKING, Self, cast
 
 from pysatl_core.distributions.strategies import (
     ComputationStrategy,
     SamplingStrategy,
 )
-from pysatl_core.types import NumericArray
+from pysatl_core.types import DEFAULT_ANALYTICAL_COMPUTATION_LABEL, NumericArray
 
 _KEEP: object = object()
 
 
 if TYPE_CHECKING:
     from typing import Any
 
-    from pysatl_core.distributions.computation import AnalyticalComputation, Method
+    from pysatl_core.distributions.computation import AnalyticalComputation
     from pysatl_core.distributions.support import Support
     from pysatl_core.types import (
         DistributionType,
         GenericCharacteristicName,
+        LabelName,
         Method,
     )
 
@@ -49,7 +50,13 @@ class Distribution(ABC):
     ----------
     distribution_type : DistributionType
         Type information about the distribution (kind, dimension, etc.).
-    analytical_computations : Mapping[str, AnalyticalComputation]
+    analytical_computations : Mapping[
+        GenericCharacteristicName,
+        (
+            AnalyticalComputation[Any, Any]
+            | Mapping[LabelName, AnalyticalComputation[Any, Any]]
+        ),
+    ]
         Direct analytical computations provided by the distribution.
     sampling_strategy : SamplingStrategy
         Strategy for generating random samples.
@@ -62,10 +69,10 @@ class Distribution(ABC):
     def __init__(
         self,
         distribution_type: DistributionType,
-        analytical_computations: (
-            Iterable[AnalyticalComputation[Any, Any]]
-            | Mapping[GenericCharacteristicName, AnalyticalComputation[Any, Any]]
-        ),
+        analytical_computations: Mapping[
+            GenericCharacteristicName,
+            (AnalyticalComputation[Any, Any] | Mapping[LabelName, AnalyticalComputation[Any, Any]]),
+        ],
         support: Support | None = None,
         sampling_strategy: SamplingStrategy | None = None,
         computation_strategy: ComputationStrategy | None = None,
@@ -78,7 +85,13 @@ def __init__(
         distribution_type : DistributionType
             Type information about the distribution (kind, dimension, etc.).
         analytical_computations :
-            Iterable[AnalyticalComputation] | Mapping[str, AnalyticalComputation]
+            Mapping[
+                GenericCharacteristicName,
+                (
+                    AnalyticalComputation[Any, Any]
+                    | Mapping[LabelName, AnalyticalComputation[Any, Any]]
+                ),
+            ]
             Analytical computations provided by the distribution.
         support : Support or None, default=None
             Support of the distribution.
@@ -93,14 +106,27 @@ def __init__(
         )
 
         self._distribution_type = distribution_type
-        if isinstance(analytical_computations, Mapping):
-            normalized_analytical = dict(analytical_computations)
-        else:
-            normalized_analytical = {ac.target: ac for ac in analytical_computations}
+        normalized_analytical: dict[
+            GenericCharacteristicName, dict[LabelName, AnalyticalComputation[Any, Any]]
+        ] = {}
+        for characteristic_name, methods in analytical_computations.items():
+            if isinstance(methods, Mapping):
+                normalized_analytical[characteristic_name] = dict(methods)
+            else:
+                normalized_analytical[characteristic_name] = {
+                    DEFAULT_ANALYTICAL_COMPUTATION_LABEL: methods
+                }
 
         if not normalized_analytical:
             raise ValueError("Distribution requires at least one analytical computation.")
 
+        for characteristic_name, labeled_methods in normalized_analytical.items():
+            if not labeled_methods:
+                raise ValueError(
+                    f"Characteristic '{characteristic_name}' must provide at least one "
+                    "analytical computation."
+                )
+
         self._analytical_computations = normalized_analytical
         self._support = support
         self._sampling_strategy = sampling_strategy or DefaultSamplingUnivariateStrategy()
@@ -114,7 +140,7 @@ def distribution_type(self) -> DistributionType:
     @property
     def analytical_computations(
         self,
-    ) -> Mapping[GenericCharacteristicName, AnalyticalComputation[Any, Any]]:
+    ) -> Mapping[GenericCharacteristicName, Mapping[LabelName, AnalyticalComputation[Any, Any]]]:
         """Return analytical computations provided directly by this distribution."""
         return self._analytical_computations
 

src/pysatl_core/distributions/strategies.py

@@ -19,11 +19,16 @@
 from pysatl_core.types import CharacteristicName, Method, NumericArray
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
     from typing import Any
 
-    from pysatl_core.distributions.computation import FittedComputationMethod
+    from pysatl_core.distributions.computation import (
+        AnalyticalComputation,
+        FittedComputationMethod,
+        Method,
+    )
     from pysatl_core.distributions.distribution import Distribution
-    from pysatl_core.types import GenericCharacteristicName
+    from pysatl_core.types import GenericCharacteristicName, LabelName
 
 
 class ComputationStrategy(Protocol):
@@ -101,6 +106,26 @@ def _pop_guard(self, distr: Distribution, state: GenericCharacteristicName) -> N
             if not seen:
                 self._resolving.pop(key, None)
 
+    @staticmethod
+    def _pick_analytical_method(
+        state: GenericCharacteristicName,
+        methods: Mapping[LabelName, AnalyticalComputation[Any, Any]],
+    ) -> AnalyticalComputation[Any, Any]:
+        """
+        Pick the first available analytical method for a characteristic.
+
+        Raises
+        ------
+        RuntimeError
+            If no labeled analytical methods are available for the characteristic.
+        """
+        try:
+            return next(iter(methods.values()))
+        except StopIteration as exc:
+            raise RuntimeError(
+                f"Characteristic '{state}' provides no labeled analytical computations."
+            ) from exc
+
     def query_method(
         self, state: GenericCharacteristicName, distr: Distribution, **options: Any
     ) -> Method[Any, Any]:
@@ -134,7 +159,7 @@ def query_method(
         """
         # 1. Check for analytical implementation
         if state in distr.analytical_computations:
-            return distr.analytical_computations[state]
+            return self._pick_analytical_method(state, distr.analytical_computations[state])
 
         # 2. Check cache if enabled
         if self._enable_caching:
@@ -156,7 +181,7 @@ def query_method(
             # 5. Try each analytical characteristic as a source
             for src in distr.analytical_computations:
                 if src == state:
-                    return distr.analytical_computations[src]
+                    return self._pick_analytical_method(src, distr.analytical_computations[src])
 
                 # Find conversion path in the graph
                 path = reg.find_path(src, state)

src/pysatl_core/families/distribution.py

@@ -35,6 +35,7 @@
     from pysatl_core.types import (
         DistributionType,
         GenericCharacteristicName,
+        LabelName,
         ParametrizationName,
     )
 
@@ -69,7 +70,8 @@ def __init__(
         family_name: str,
         distribution_type: DistributionType,
         analytical_computations: Mapping[
-            GenericCharacteristicName, AnalyticalComputation[Any, Any]
+            GenericCharacteristicName,
+            Mapping[LabelName, AnalyticalComputation[Any, Any]],
         ],
         parametrization: Parametrization,
         support: Support | None,

src/pysatl_core/families/parametric_family.py

@@ -19,7 +19,11 @@
 
 from pysatl_core.distributions.computation import AnalyticalComputation
 from pysatl_core.families.distribution import ParametricFamilyDistribution
-from pysatl_core.types import ComputationFunc, DistributionType
+from pysatl_core.types import (
+    DEFAULT_ANALYTICAL_COMPUTATION_LABEL,
+    ComputationFunc,
+    DistributionType,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -31,18 +35,23 @@
     )
     from pysatl_core.types import (
         GenericCharacteristicName,
+        LabelName,
         ParametrizationName,
     )
 
     type SupportArg = Callable[[Parametrization], Support | None] | None
     type SupportResolver = Callable[[Parametrization], Support | None]
+    type LabeledCharacteristicProvider = (
+        Mapping[LabelName, ParametricFamilyCharacteristic[Any, Any]]
+        | ParametricFamilyCharacteristic[Any, Any]
+    )
     type CharacteristicProvider = (
-        Mapping[ParametrizationName, CharacteristicFunction[Any, Any]]
-        | CharacteristicFunction[Any, Any]
+        Mapping[ParametrizationName, LabeledCharacteristicProvider]
+        | ParametricFamilyCharacteristic[Any, Any]
     )
     type CharacteristicsMap = Mapping[GenericCharacteristicName, CharacteristicProvider]
     type NonParametrizedCharacteristic[In, Out] = Callable[[], Out]
-    type CharacteristicFunction[In, Out] = (
+    type ParametricFamilyCharacteristic[In, Out] = (
         NonParametrizedCharacteristic[In, Out] | ParametrizedCharacteristic[In, Out]
     )
 
@@ -78,7 +87,9 @@ class ParametricFamily:
         - nullary characteristics (e.g., mean, var): provider(params, **kwargs) -> Any
         - pointwise characteristics (e.g., pdf, cdf, ppf): provider(params, x, **kwargs) -> Any
 
-        If a single callable is provided, it is treated as defined for the base parametrization.
+        Providers are grouped by parametrization and may define multiple labeled methods.
+        If a single callable is provided, it is treated as the base-parametrization method
+        under ``DEFAULT_ANALYTICAL_COMPUTATION_LABEL``.
     support_by_parametrization : Callable or None, optional
         Function that returns support for given parameters.
     """
@@ -109,31 +120,63 @@ def __init__(
         # Runtime registry of parametrization classes
         self._parametrizations: dict[ParametrizationName, type[Parametrization]] = {}
 
-        def _normalize_characteristic(
-            value: Mapping[ParametrizationName, CharacteristicFunction[Any, Any]]
-            | CharacteristicFunction[Any, Any],
-        ) -> dict[ParametrizationName, CharacteristicFunction[Any, Any]]:
-            return (
-                dict(value)
-                if isinstance(value, Mapping)
-                else {self.base_parametrization_name: value}
+        def _normalize_labeled_provider(
+            characteristic_name: GenericCharacteristicName,
+            parametrization_name: ParametrizationName,
+            provider: LabeledCharacteristicProvider,
+        ) -> dict[LabelName, ParametricFamilyCharacteristic[Any, Any]]:
+            normalized = (
+                dict(provider)
+                if isinstance(provider, Mapping)
+                else {DEFAULT_ANALYTICAL_COMPUTATION_LABEL: provider}
             )
+            if not normalized:
+                raise ValueError(
+                    f"Characteristic '{characteristic_name}' has no labeled providers for "
+                    f"parametrization '{parametrization_name}'."
+                )
+            return normalized
+
+        def _normalize_characteristic(
+            characteristic_name: GenericCharacteristicName,
+            value: CharacteristicProvider,
+        ) -> dict[ParametrizationName, dict[LabelName, ParametricFamilyCharacteristic[Any, Any]]]:
+            if not isinstance(value, Mapping):
+                base_name = self.base_parametrization_name
+                return {
+                    base_name: _normalize_labeled_provider(characteristic_name, base_name, value)
+                }
+
+            normalized_by_parametrization: dict[
+                ParametrizationName, dict[LabelName, ParametricFamilyCharacteristic[Any, Any]]
+            ] = {}
+            for parametrization_name, provider in value.items():
+                normalized_by_parametrization[parametrization_name] = _normalize_labeled_provider(
+                    characteristic_name,
+                    parametrization_name,
+                    provider,
+                )
+            return normalized_by_parametrization
 
         self.distr_characteristics: dict[
-            GenericCharacteristicName, dict[ParametrizationName, CharacteristicFunction[Any, Any]]
-        ] = {k: _normalize_characteristic(v) for k, v in distr_characteristics.items()}
+            GenericCharacteristicName,
+            dict[ParametrizationName, dict[LabelName, ParametricFamilyCharacteristic[Any, Any]]],
+        ] = {
+            characteristic_name: _normalize_characteristic(characteristic_name, provider)
+            for characteristic_name, provider in distr_characteristics.items()
+        }
 
         # Validate characteristic providers
         valid_names = set(self.parametrization_names)
         for char_name, forms in self.distr_characteristics.items():
+            if not forms:
+                raise ValueError(f"Characteristic '{char_name}' has no providers.")
             unknown = set(forms) - valid_names
             if unknown:
                 raise ValueError(
                     f"Characteristic '{char_name}' has providers for unknown parametrizations: "
                     f"{sorted(unknown)}."
                 )
-            if self.base_parametrization_name not in forms and len(forms) == 0:
-                raise ValueError(f"Characteristic '{char_name}' has no providers.")
 
         # Precompute analytical plan: for each parametrization pick provider (self or base)
         self._analytical_plan: dict[
@@ -236,7 +279,7 @@ def to_base(self, parameters: Parametrization) -> Parametrization:
 
     @staticmethod
     def _bind_parametrization[In, Out](
-        func: CharacteristicFunction[In, Out], params_obj: Parametrization
+        func: ParametricFamilyCharacteristic[In, Out], params_obj: Parametrization
     ) -> ComputationFunc[In, Out]:
         """Bind ``params_obj`` to ``func`` only when ``func`` can accept positional arguments.
 
@@ -265,14 +308,16 @@ def _bind_parametrization[In, Out](
 
     def _build_analytical_computations(
         self, parameters: Parametrization
-    ) -> dict[GenericCharacteristicName, AnalyticalComputation[Any, Any]]:
+    ) -> dict[GenericCharacteristicName, dict[LabelName, AnalyticalComputation[Any, Any]]]:
         """
         Build analytical computations for given parameters.
 
         Uses precomputed provider plan for efficient computation.
         """
         plan = self._analytical_plan.get(parameters.name, {})
-        result: dict[GenericCharacteristicName, AnalyticalComputation[Any, Any]] = {}
+        result: dict[
+            GenericCharacteristicName, dict[LabelName, AnalyticalComputation[Any, Any]]
+        ] = {}
         base_params: Parametrization | None = None
 
         for characteristic, provider_name in plan.items():
@@ -282,11 +327,14 @@ def _build_analytical_computations(
                 base_params = base_params or self.to_base(parameters)
                 params_obj = base_params
 
-            func_factory = self.distr_characteristics[characteristic][provider_name]
-            result[characteristic] = AnalyticalComputation(
-                target=characteristic,
-                func=self._bind_parametrization(func_factory, params_obj),
-            )
+            labeled_providers = self.distr_characteristics[characteristic][provider_name]
+            result[characteristic] = {
+                label_name: AnalyticalComputation(
+                    target=characteristic,
+                    func=self._bind_parametrization(func_factory, params_obj),
+                )
+                for label_name, func_factory in labeled_providers.items()
+            }
 
         return result
 

src/pysatl_core/types.py

@@ -243,6 +243,12 @@ def shape(self) -> ContinuousSupportShape1D:
 type GenericCharacteristicName = str
 """Type alias for characteristic names (e.g., 'pdf', 'cdf')."""
 
+type LabelName = str
+"""Type alias for labels that distinguish computation variants."""
+
+DEFAULT_ANALYTICAL_COMPUTATION_LABEL: LabelName = "PySATL_default_analytical_computation"
+"""Default label for analytical methods when a label is not explicitly provided."""
+
 type ParametrizationName = str
 """Type alias for parametrization names."""
 
@@ -302,6 +308,8 @@ class FamilyName(StrEnum):
     "UnivariateContinuous",
     "UnivariateDiscrete",
     "GenericCharacteristicName",
+    "LabelName",
+    "DEFAULT_ANALYTICAL_COMPUTATION_LABEL",
     "ParametrizationName",
     "ComputationFunc",
     "DistributionType",