feat(knowledge base): add optimization with knowledge base for ParametricFamilyDistribution

domosedy · domosedy · commit 280efce81c1c · 2026-05-17T20:36:22.000+03:00
diff --git a/src/pysatl_core/families/__init__.py b/src/pysatl_core/families/__init__.py
@@ -22,6 +22,7 @@
     parametrization,
 )
 from .registry import ParametricFamilyRegister
+from .registry_graph import BinaryOperationType
 
 __all__ = [
     "ParametricFamilyRegister",
@@ -32,6 +33,7 @@
     "constraint",
     "parametrization",
     "configure_families_register",
+    "BinaryOperationType",
     # builtins
     *_builtins_all,
 ]
diff --git a/src/pysatl_core/families/distribution.py b/src/pysatl_core/families/distribution.py
@@ -7,10 +7,15 @@
 
 from __future__ import annotations
 
+import dataclasses
+
+from pysatl_core.families.registry_graph import BinaryOperationType
+
 __author__ = "Leonid Elkin, Mikhail Mikhailov"
 __copyright__ = "Copyright (c) 2025 PySATL project"
 __license__ = "SPDX-License-Identifier: MIT"
 
+from types import NotImplementedType
 from typing import TYPE_CHECKING
 
 from pysatl_core.distributions.distribution import _KEEP, Distribution
@@ -193,3 +198,112 @@ def sample(self, n: int, **options: Any) -> NumericArray:
             the sampling strategy.
         """
         return self.sampling_strategy.sample(n, distr=self, **options)
+
+    def _try_to_transform_with_optimization(
+        self, other: ParametricFamilyDistribution, kind: BinaryOperationType
+    ) -> None | ParametricFamilyDistribution:
+        registry = ParametricFamilyRegister()
+        transform_result = registry.find_binary_transformation(
+            self.family_name, other.family_name, kind
+        )
+        if transform_result is None:
+            return None
+
+        family, transform_parametrization = transform_result
+        new_parametrization = transform_parametrization(
+            self.parametrization.transform_to_base_parametrization(),
+            other.parametrization.transform_to_base_parametrization(),
+        )
+        return family(new_parametrization.name, **dataclasses.asdict(new_parametrization))  # type:ignore[call-overload]
+
+    def __add__(
+        self, other: object
+    ) -> ParametricFamilyDistribution | Distribution | NotImplementedType:
+        """Return ``self + other`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = self._try_to_transform_with_optimization(
+                other, BinaryOperationType.ADD
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__add__(self, other)
+
+    def __radd__(
+        self, other: object
+    ) -> ParametricFamilyDistribution | Distribution | NotImplementedType:
+        """Return ``other + self`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = other._try_to_transform_with_optimization(
+                self, BinaryOperationType.ADD
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__radd__(self, other)
+
+    def __sub__(self, other: object) -> Distribution | NotImplementedType:
+        """Return ``self - other`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = self._try_to_transform_with_optimization(
+                other, BinaryOperationType.SUB
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__sub__(self, other)
+
+    def __rsub__(self, other: object) -> Distribution | NotImplementedType:
+        """Return ``other - self`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = other._try_to_transform_with_optimization(
+                self, BinaryOperationType.SUB
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__rsub__(self, other)
+
+    def __mul__(self, other: object) -> Distribution | NotImplementedType:
+        """Return ``self * other`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = self._try_to_transform_with_optimization(
+                other, BinaryOperationType.MUL
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__mul__(self, other)
+
+    def __rmul__(self, other: object) -> Distribution | NotImplementedType:
+        """Return ``other * self`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = other._try_to_transform_with_optimization(
+                self, BinaryOperationType.MUL
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__rmul__(self, other)
+
+    def __truediv__(self, other: object) -> Distribution | NotImplementedType:
+        """Return ``self / other`` for scalar or distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = self._try_to_transform_with_optimization(
+                other, BinaryOperationType.DIV
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__truediv__(self, other)
+
+    def __rtruediv__(self, other: object) -> Distribution | NotImplementedType:
+        """Return ``other / self`` for distribution operands."""
+        if isinstance(other, ParametricFamilyDistribution):
+            transformation_result = other._try_to_transform_with_optimization(
+                self, BinaryOperationType.DIV
+            )
+            if transformation_result is not None:
+                return transformation_result
+
+        return TransformationOperatorsMixin.__rtruediv__(self, other)
diff --git a/src/pysatl_core/families/registry.py b/src/pysatl_core/families/registry.py
@@ -9,7 +9,10 @@
 from __future__ import annotations
 
 from pysatl_core.families.parametrizations import Parametrization
-from pysatl_core.families.registry_graph import RegistryGraphTransformations
+from pysatl_core.families.registry_graph import (
+    BinaryOperationType,
+    RegistryGraphTransformations,
+)
 
 __author__ = "Leonid Elkin, Mikhail Mikhailov, Fedor Myznikov"
 __copyright__ = "Copyright (c) 2025 PySATL project"
@@ -69,6 +72,41 @@ def get_optimal_density(
         family_name, transform_function = self._registry_graph.get_optimal_transoformation(name)
         return self.get(family_name), transform_function
 
+    @classmethod
+    def add_binary_transformation(
+        cls,
+        left: str,
+        right: str,
+        result: str,
+        operation: BinaryOperationType,
+        parametrization_transformation: Callable[
+            [Parametrization, Parametrization], Parametrization
+        ],
+    ) -> bool:
+        self = cls()
+        return self._registry_graph.add_binary_transformation(
+            left, right, result, operation, parametrization_transformation
+        )
+
+    @classmethod
+    def find_binary_transformation(
+        cls, left: str, right: str, operation: BinaryOperationType
+    ) -> (
+        None
+        | tuple[ParametricFamily, Callable[[Parametrization, Parametrization], Parametrization]]
+    ):
+        self = cls()
+        transformation = self._registry_graph.find_binary_transform(left, right, operation)
+
+        if transformation is None:
+            return transformation
+
+        family = self._registered_families.get(transformation.result, None)
+        if family is None:
+            return family
+
+        return family, transformation.transformation
+
     @classmethod
     def register_parametrization_transformation(
         cls,
@@ -145,14 +183,14 @@ def register(cls, family: ParametricFamily, temperature: int = 128) -> None:
             If a family with the same name is already registered.
         """
         self = cls()
-        self.change_family_temperature(family.name, temperature)
+        self._change_family_temperature(family.name, temperature)
 
         if family.name in self._registered_families:
             raise ValueError(f"Family {family.name} already found in register")
         self._registered_families[family.name] = family
 
     @classmethod
-    def change_family_temperature(cls, family_name: str, new_temperature: int) -> None:
+    def _change_family_temperature(cls, family_name: str, new_temperature: int) -> None:
         self = cls()
         self._registry_graph.register_family_temperature(family_name, new_temperature)
 
diff --git a/src/pysatl_core/families/registry_graph.py b/src/pysatl_core/families/registry_graph.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from enum import StrEnum
 from queue import Queue
 from typing import TYPE_CHECKING, cast
 
@@ -60,20 +61,122 @@ def transform_density(self, argument: Number | NumericArray) -> Number | Numeric
         return self._transform_function(argument)
 
 
+class BinaryOperationType(StrEnum):
+    ADD = "add"
+    SUB = "sub"
+    MUL = "multiply"
+    DIV = "divide"
+
+
+class FamilyBinaryOperationTransformationRecord:
+    """
+    Record for a bainray transformation between 2 families
+    It doesn't accept constraints on transformation, just for now
+    """
+
+    def __init__(
+        self,
+        left: str,
+        right: str,
+        result: str,
+        kind: BinaryOperationType,
+        parametrization_transformation: Callable[
+            [Parametrization, Parametrization], Parametrization
+        ],
+    ):
+        self._left = left
+        self._right = right
+        self._kind = kind
+        self._result = result
+        self._transformation = parametrization_transformation
+
+    def accepts(self, left: str, right: str, kind: BinaryOperationType) -> bool:
+        return (self._left, self._right, self._kind) == (left, right, kind)
+
+    def transform(
+        self, left_parametrization: Parametrization, right_parametrization: Parametrization
+    ) -> tuple[str, Parametrization]:
+        return self._result, self._transformation(left_parametrization, right_parametrization)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, FamilyBinaryOperationTransformationRecord):
+            return False
+
+        return self.accepts(other._left, other._right, other._kind)
+
+    @property
+    def result(self) -> str:
+        return self._result
+
+    @property
+    def transformation(self) -> Callable[[Parametrization, Parametrization], Parametrization]:
+        return self._transformation
+
+    def __hash__(self) -> int:
+        return (
+            self._left.__hash__()
+            ^ self._right.__hash__()
+            ^ self._kind.__hash__()
+            ^ self._result.__hash__()
+        )
+
+
 class RegistryGraphTransformations:
+    """
+    Registry with transformations, such as transformation, when densitry equals or smth like this
+    TODO: find a way to merge ways between density and parametrization transformations
+    """
+
     _instance: ClassVar[RegistryGraphTransformations | None] = None
     _registered_families_temperature: dict[str, int]
     _registered_parametrzation_transformations: dict[str, list[TransformatedParametrizationEdge]]
     _registered_transformations: dict[str, list[TransformatedDensityEdge]]
+    _binary_transformations: list[FamilyBinaryOperationTransformationRecord]
 
     def __new__(cls) -> RegistryGraphTransformations:
         if cls._instance is None:
             cls._instance = super().__new__(cls)
             cls._registered_parametrzation_transformations = {}
             cls._registered_families_temperature = {}
             cls._registered_transformations = {}
+            cls._binary_transformations = []
         return cls._instance
 
+    @classmethod
+    def find_binary_transform(
+        cls, left: str, right: str, operation: BinaryOperationType
+    ) -> None | FamilyBinaryOperationTransformationRecord:
+        self = cls()
+
+        for transformation in self._binary_transformations:
+            if transformation.accepts(left, right, operation):
+                return transformation
+
+        return None
+
+    @classmethod
+    def add_binary_transformation(
+        cls,
+        left: str,
+        right: str,
+        result: str,
+        operation: BinaryOperationType,
+        parametrization_transformation: Callable[
+            [Parametrization, Parametrization], Parametrization
+        ],
+    ) -> bool:
+        transformation = FamilyBinaryOperationTransformationRecord(
+            left, right, result, operation, parametrization_transformation
+        )
+        self = cls()
+
+        for registered_transformation in self._binary_transformations:
+            if transformation == registered_transformation:
+                return False
+
+        self._binary_transformations.append(transformation)
+        return True
+
     @classmethod
     def _run_bfs(
         cls,
diff --git a/tests/unit/families/test_optimization.py b/tests/unit/families/test_optimization.py

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,7 @@`
`22`	`22`	`parametrization,`
`23`	`23`	`)`
`24`	`24`	`from .registry import ParametricFamilyRegister`
	`25`	`+from .registry_graph import BinaryOperationType`
`25`	`26`
`26`	`27`	`__all__ = [`
`27`	`28`	`"ParametricFamilyRegister",`
`@@ -32,6 +33,7 @@`
`32`	`33`	`"constraint",`
`33`	`34`	`"parametrization",`
`34`	`35`	`"configure_families_register",`
	`36`	`+ "BinaryOperationType",`
`35`	`37`	`# builtins`
`36`	`38`	`*_builtins_all,`
`37`	`39`	`]`