feat(families): added ability for partial fixing parameters in parametric families

Author: MyznikovFD

src/pysatl_core/families/parametric_family.py

@@ -15,7 +15,7 @@
 import inspect
 from collections.abc import Mapping
 from functools import partial
-from typing import TYPE_CHECKING, Any, cast, dataclass_transform
+from typing import TYPE_CHECKING, Any, cast, dataclass_transform, overload
 
 import numpy as np
 
@@ -451,4 +451,224 @@ def score(self, parameters: Parametrization, x: NumericArray) -> NumericArray:
         base_grad = self._base_score(base_params, x_arr)
         return parameters.gradient_transform(base_grad)
 
+    def view(
+        self,
+        *,
+        parametrization_name: str | None = None,
+        **fixed_params: Any,
+    ) -> PartialParametricFamily:
+        """
+        Create a view of this family with partially fixed parameters.
+
+        Parameters
+        ----------
+        parametrization_name : str, optional
+            Name of the parametrization in which the fixed parameters are given.
+            If not provided, the base parametrization of the family is used.
+        **fixed_params : Any
+            Parameter names and values to fix.
+
+        Returns
+        -------
+        PartialParametricFamily
+            A view that behaves like the original family but with the specified
+            parameters fixed.
+
+        Examples
+        --------
+        >>> uniform = ParametricFamilyRegister.get("uniform")
+        >>> uniform_lower0 = uniform.view(lower_bound=0)
+        >>> dist = uniform_lower0.distribution(upper_bound=1)  # Uniform(0,1)
+        """
+        return PartialParametricFamily(
+            base_family=self,
+            fixed_params=fixed_params,
+            parametrization_name=parametrization_name,
+        )
+
+    __call__ = distribution
+
+
+class PartialParametricFamily(ParametricFamily):
+    """
+    View on a parametric family with partially fixed parameters.
+
+    This class represents a parametric family where some parameters have been
+    fixed to specific values. It inherits all behaviour from `ParametricFamily`
+    but restricts the available parametrization to the one in which parameters
+    are fixed. All analytical characteristics are preserved via delegation to
+    the base parametrization of the original family.
+
+    Parameters
+    ----------
+    base_family : ParametricFamily
+        The original parametric family.
+    fixed_params : dict[str, Any]
+        Dictionary of fixed parameter names and their values.
+    parametrization_name : str, optional
+        Name of the parametrization in which the fixed parameters are specified.
+        If not provided, the base parametrization of the family is used.
+
+    Raises
+    ------
+    ValueError
+        If all parameters of the chosen parametrization are fixed (use `.distribution()` directly),
+        or if any fixed parameter name is unknown for that parametrization,
+        or if the parametrization name is not registered in the family.
+    """
+
+    def __init__(
+        self,
+        base_family: ParametricFamily,
+        fixed_params: dict[str, Any],
+        parametrization_name: str | None = None,
+    ) -> None:
+        self._fixed_in_param = parametrization_name or base_family.base_parametrization_name
+        param_class = base_family.get_parametrization(self._fixed_in_param)
+
+        # Check completeness: if all parameters are fixed, raise an error
+        required_fields = set(getattr(param_class, "__dataclass_fields__", {}).keys())
+        if required_fields.issubset(fixed_params):
+            raise ValueError(
+                f"All parameters of parametrization '{self._fixed_in_param}' are already fixed. "
+                "Use `.distribution()` directly to create a distribution instance."
+            )
+
+        # Keep only characteristics that have an implementation for fixed_in
+        filtered_chars = {}
+        for char_name, char_map in base_family.distr_characteristics.items():
+            if self._fixed_in_param in char_map:
+                filtered_chars[char_name] = {self._fixed_in_param: char_map[self._fixed_in_param]}
+
+        super().__init__(
+            name=base_family._name,
+            distr_type=base_family._distr_type,
+            distr_parametrizations=[self._fixed_in_param],
+            distr_characteristics=filtered_chars,
+            support_by_parametrization=base_family._support_resolver,
+            base_score=base_family._base_score,
+        )
+
+        # Register the parametrization (needed for parent methods)
+        self.register_parametrization(self._fixed_in_param, param_class)
+
+        self._base_family = base_family
+        self._param_class = param_class
+        self._fixed_params = fixed_params.copy()
+
+        # Validate that fixed parameters exist
+        unknown = set(self._fixed_params) - required_fields
+        if unknown:
+            raise ValueError(
+                f"Unknown parameters for parametrization '{self._fixed_in_param}': {unknown}"
+            )
+
+    @property
+    def base(self) -> type[Parametrization]:
+        """Return the only parametrization class available in this view."""
+        return self._param_class
+
+    @property
+    def parametrizations(self) -> dict[str, type[Parametrization]]:
+        """Return a dictionary containing only the fixed parametrization."""
+        return {self._fixed_in_param: self._param_class}
+
+    @overload
+    def get_parametrization(self) -> type[Parametrization]: ...
+
+    @overload
+    def get_parametrization(self, name: ParametrizationName) -> type[Parametrization]: ...
+
+    def get_parametrization(self, name: ParametrizationName | None = None) -> type[Parametrization]:
+        if name is None:
+            return self._param_class
+        if name != self._fixed_in_param:
+            raise KeyError(
+                f"Parametrization '{name}' is not available in this view. "
+                f"Only '{self._fixed_in_param}' is available."
+            )
+        return self._param_class
+
+    def _build_analytical_computations(
+        self, parameters: Parametrization
+    ) -> dict[GenericCharacteristicName, dict[LabelName, AnalyticalComputation[Any, Any]]]:
+        """
+        Build analytical computations using the parent family's plan.
+        Characteristics not defined in the fixed parametrization are computed
+        via the base parametrization of the original family.
+        """
+        # Use the parent family's plan (full, not filtered)
+        base_plan = self._base_family._analytical_plan.get(parameters.name, {})
+        result: dict[
+            GenericCharacteristicName,
+            dict[LabelName, AnalyticalComputation[Any, Any]],
+        ] = {}
+        base_params: Parametrization | None = None
+
+        for characteristic, provider_name in base_plan.items():
+            if provider_name == self._fixed_in_param:
+                params_obj = parameters
+            else:
+                if base_params is None:
+                    base_params = self._base_family.to_base(parameters)
+                params_obj = base_params
+
+            labeled_providers = self._base_family.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
+
+    def is_complete(self) -> bool:
+        """Check whether all parameters of the fixed parametrization are already fixed."""
+        required_fields = set(getattr(self._param_class, "__dataclass_fields__", {}).keys())
+        return required_fields.issubset(self._fixed_params)
+
+    def distribution(
+        self,
+        parametrization_name: str | None = None,
+        sampling_strategy: SamplingStrategy | None = None,
+        computation_strategy: ComputationStrategy | None = None,
+        **kwargs: Any,
+    ) -> ParametricFamilyDistribution:
+        target = parametrization_name or self._fixed_in_param
+        if target != self._fixed_in_param:
+            raise ValueError(
+                f"Only parametrization '{self._fixed_in_param}' is available in this view. "
+                "Please omit 'parametrization_name' or use the fixed one."
+            )
+        for key, fixed_val in self._fixed_params.items():
+            if key in kwargs and kwargs[key] != fixed_val:
+                raise ValueError(
+                    f"Parameter '{key}' is fixed to {fixed_val}, but got {kwargs[key]}"
+                )
+        all_params = {**self._fixed_params, **kwargs}
+        return super().distribution(
+            parametrization_name=target,
+            sampling_strategy=sampling_strategy,
+            computation_strategy=computation_strategy,
+            **all_params,
+        )
+
+    def view(
+        self,
+        *,
+        parametrization_name: str | None = None,
+        **additional_params: Any,
+    ) -> PartialParametricFamily:
+        if parametrization_name is not None and parametrization_name != self._fixed_in_param:
+            raise ValueError(
+                f"Cannot change parametrization. Current fixed parametrization is "
+                f"'{self._fixed_in_param}'. Use the same or omit the argument."
+            )
+        new_fixed = {**self._fixed_params, **additional_params}
+        return PartialParametricFamily(self._base_family, new_fixed, self._fixed_in_param)
+
     __call__ = distribution