feat(knowledge base): add prototype of knowledge graph

Author: domosedy

src/pysatl_core/families/fixed_parametrization_edge.py

@@ -0,0 +1,25 @@
+from collections.abc import Callable
+
+from pysatl_core.families.parametrizations import Parametrization
+
+
+class EdgeWithFixedParametrization:
+    def __init__(
+        self,
+        tail_name: str,
+        transform_constraint: Callable[[Parametrization], bool],
+        transform_function: Callable[[Parametrization], Parametrization],
+    ):
+        self._transform_function = transform_function
+        self._transform_constraint = transform_constraint
+        self._tail_name = tail_name
+
+    def is_transoform_possible(self, parametrization: Parametrization) -> bool:
+        return self._transform_constraint(parametrization)
+
+    def transform_parametrization(self, parametrization: Parametrization) -> Parametrization:
+        return self._transform_function(parametrization)
+
+    @property
+    def tail_name(self) -> str:
+        return self._tail_name

src/pysatl_core/families/parametric_family.py

@@ -21,6 +21,7 @@
 
 from pysatl_core.distributions.computation import AnalyticalComputation
 from pysatl_core.families.distribution import ParametricFamilyDistribution
+from pysatl_core.families.registry import ParametricFamilyRegister
 from pysatl_core.types import (
     DEFAULT_ANALYTICAL_COMPUTATION_LABEL,
     ComputationFunc,
@@ -386,6 +387,19 @@ def distribution(
         parameters = parametrization_class(**parameters_values)
         parameters.validate()
         base_parameters = self.to_base(parameters)
+        registry = ParametricFamilyRegister()
+        optimized_family = registry.get_optimal_family(self.name, parameters)
+        if optimized_family is not None:
+            edge, new_family = optimized_family
+            new_parametrization = edge.transform_parametrization(parameters)
+
+            return new_family.distribution(
+                parametrization_name=None,
+                sampling_strategy=sampling_strategy,
+                computation_strategy=computation_strategy,
+                **new_parametrization.parameters,
+            )
+
         distribution_type = self._distr_type(base_parameters)
         analytical_computations = self._build_analytical_computations(parameters)
         return ParametricFamilyDistribution(

src/pysatl_core/families/registry.py

@@ -8,10 +8,14 @@
 
 from __future__ import annotations
 
+from pysatl_core.families.fixed_parametrization_edge import EdgeWithFixedParametrization
+from pysatl_core.families.parametrizations import Parametrization
+
 __author__ = "Leonid Elkin, Mikhail Mikhailov, Fedor Myznikov"
 __copyright__ = "Copyright (c) 2025 PySATL project"
 __license__ = "SPDX-License-Identifier: MIT"
 
+from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -30,14 +34,49 @@ class ParametricFamilyRegister:
 
     _instance: ClassVar[ParametricFamilyRegister | None] = None
     _registered_families: dict[str, ParametricFamily]
+    _registered_transformations: dict[str, list[EdgeWithFixedParametrization]]
 
     def __new__(cls) -> ParametricFamilyRegister:
         """Create or return the singleton instance."""
         if cls._instance is None:
             cls._instance = super().__new__(cls)
             cls._instance._registered_families = {}
+            cls._instance._registered_transformations = {}
         return cls._instance
 
+    @classmethod
+    def get_optimal_family(
+        cls, name: str, parametrization: Parametrization
+    ) -> None | tuple[EdgeWithFixedParametrization, ParametricFamily]:
+        self = cls()
+
+        for optimization_edge in self._registered_transformations.get(name, []):
+            if (
+                optimization_edge.tail_name in self._registered_families
+                and optimization_edge.is_transoform_possible(parametrization)
+            ):
+                return optimization_edge, self._registered_families[optimization_edge.tail_name]
+
+        return None
+
+    @classmethod
+    def add_optimization_edge(
+        cls,
+        head_name: str,
+        tail_name: str,
+        transform_constraint: Callable[[Parametrization], bool],
+        transform_function: Callable[[Parametrization], Parametrization],
+    ) -> bool:
+        self = cls()
+
+        if head_name in self._registered_families:
+            self._registered_transformations.setdefault(head_name, []).append(
+                EdgeWithFixedParametrization(tail_name, transform_constraint, transform_function)
+            )
+            return True
+
+        return False
+
     @classmethod
     def get(cls, name: str) -> ParametricFamily:
         """

tests/unit/families/test_optimization.py

@@ -0,0 +1,38 @@
+from pysatl_core.families.configuration import configure_families_register
+from pysatl_core.families.parametrizations import Parametrization
+from pysatl_core.types import CharacteristicName, FamilyName
+
+
+def test_optimization():
+    registry = configure_families_register()
+    exponential_fam = registry.get(FamilyName.EXPONENTIAL)
+    uniform_fam = registry.get(FamilyName.CONTINUOUS_UNIFORM)
+
+    def transform_function(head_param: Parametrization) -> Parametrization:
+        tail_param_type = uniform_fam.get_parametrization(uniform_fam.base_parametrization_name)
+
+        return tail_param_type(lower_bound=0, upper_bound=1)  # type: ignore[call-arg]
+
+    def transform_constraint(head_param: Parametrization) -> bool:
+        exponential_fam.get_parametrization(exponential_fam.base_parametrization_name)
+
+        return head_param.lambda_ == 1.0  # type: ignore[attr-defined]
+
+    registry.add_optimization_edge(
+        FamilyName.EXPONENTIAL,
+        FamilyName.CONTINUOUS_UNIFORM,
+        transform_constraint,
+        transform_function,
+    )
+
+    exponential = exponential_fam(lambda_=1.0)
+    exponential_parametrization = exponential.parametrization
+
+    assert exponential.family_name == FamilyName.CONTINUOUS_UNIFORM
+    assert exponential_parametrization.lower_bound == 0  # type: ignore[attr-defined]
+    assert exponential_parametrization.upper_bound == 1  # type: ignore[attr-defined]
+
+    pdf = exponential.query_method(CharacteristicName.PDF)
+
+    assert pdf(0.5) == 1
+    assert pdf(10) == 0