diff --git a/src/pysatl_core/families/__init__.py b/src/pysatl_core/families/__init__.py index ed30528..d0b3387 100644 --- 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/builtins/continuous/__init__.py b/src/pysatl_core/families/builtins/continuous/__init__.py index d7f9492..0bce521 100644 --- a/src/pysatl_core/families/builtins/continuous/__init__.py +++ b/src/pysatl_core/families/builtins/continuous/__init__.py @@ -10,6 +10,7 @@ from pysatl_core.families.builtins.continuous.exponential import configure_exponential_family +from pysatl_core.families.builtins.continuous.lognormal import configure_lognormal_family from pysatl_core.families.builtins.continuous.normal import configure_normal_family from pysatl_core.families.builtins.continuous.uniform import configure_uniform_family @@ -17,4 +18,5 @@ "configure_normal_family", "configure_uniform_family", "configure_exponential_family", + "configure_lognormal_family", ] diff --git a/src/pysatl_core/families/builtins/continuous/lognormal.py b/src/pysatl_core/families/builtins/continuous/lognormal.py new file mode 100644 index 0000000..e666168 --- /dev/null +++ b/src/pysatl_core/families/builtins/continuous/lognormal.py @@ -0,0 +1,149 @@ +""" +Log-Normal distribution family implementation. + +Contains the LogNormal family with multiple parameterizations. +""" + +from __future__ import annotations + +__author__ = "Fedor Myznikov" +__copyright__ = "Copyright (c) 2025" +__license__ = "SPDX-License-Identifier: MIT" + +from typing import cast + +import numpy as np +from scipy.special import erf, erfinv + +from pysatl_core.distributions.support import ContinuousSupport +from pysatl_core.families.parametric_family import ParametricFamily +from pysatl_core.families.parametrizations import ( + Parametrization, + constraint, + parametrization, +) +from pysatl_core.families.registry import ParametricFamilyRegister +from pysatl_core.types import ( + CharacteristicName, + FamilyName, + NumericArray, + UnivariateContinuous, +) + + +def configure_lognormal_family() -> None: + """Configure and register the LogNormal distribution family.""" + + if ParametricFamilyRegister.contains(FamilyName.LOGNORMAL): + return + + LOGNORMAL_DOC = """ + Log-Normal distribution. + + If a random variable Y is normally distributed with mean μ and standard deviation σ, + then X = exp(Y) follows a log‑normal distribution. Its probability density function is: + + f(x) = 1/(x σ √(2π)) * exp(-(ln x - μ)²/(2σ²)), for x > 0. + + The distribution is often used to model quantities that cannot be negative, + such as incomes, stock prices, or lifetimes. + """ + + def pdf(parameters: Parametrization, x: NumericArray) -> NumericArray: + """Probability density function of the log‑normal distribution.""" + params = cast(_MeanStd, parameters) + mu = params.mu + sigma = params.sigma + + if x <= 0: + raise ValueError("X must be in [0, +inf)") + + exponent = np.exp(-((np.log(x) - mu) ** 2) / (2 * sigma**2)) + coefficient = 1 / (x * sigma * np.sqrt(np.pi * 2)) + + return cast(NumericArray, coefficient * exponent) + + def cdf(parameters: Parametrization, x: NumericArray) -> NumericArray: + parameters = cast(_MeanStd, parameters) + if x <= 0: + raise ValueError("X must be in [0, +inf)") + + z = (np.log(x) - parameters.mu) / (parameters.sigma * np.sqrt(2)) + return cast(NumericArray, 0.5 * (1 + erf(z))) + + def ppf(parameters: Parametrization, p: NumericArray) -> NumericArray: + if np.any((p < 0) | (p > 1)): + raise ValueError("Probability must be in [0, 1]") + + parameters = cast(_MeanStd, parameters) + result = np.exp(parameters.mu + np.sqrt(2 * parameters.sigma**2) * erfinv(2 * p - 1)) + return cast(NumericArray, result) + + def lpdf(parameters: Parametrization, x: NumericArray) -> NumericArray: + return np.log(pdf(parameters, x)) + + def mean_func(parameters: Parametrization) -> float: + """Mean of normal distribution.""" + parameters = cast(_MeanStd, parameters) + return cast(float, np.exp(parameters.mu + parameters.sigma**2 / 2)) + + def var_func(parameters: Parametrization) -> float: + """Variance of normal distribution.""" + parameters = cast(_MeanStd, parameters) + return cast( + float, + (np.exp(parameters.sigma**2) - 1) * np.exp(2 * parameters.mu + parameters.sigma**2), + ) + + def skew_func(parameters: Parametrization) -> float: + """Skewness of normal distribution (always 0).""" + parameters = cast(_MeanStd, parameters) + return cast( + float, (np.exp(parameters.sigma**2) + 2) * np.sqrt(np.exp(parameters.sigma**2) - 1) + ) + + def kurt_func(parameters: Parametrization, excess: bool = False) -> float: + parameters = cast(_MeanStd, parameters) + if not excess: + return 0.0 + else: + return cast( + float, + np.exp(4 * parameters.sigma**2) + + 2 * np.exp(3 * parameters.sigma**2) + + 3 * np.exp(2 * parameters.sigma**2) + - 6, + ) + + def _support(_: Parametrization) -> ContinuousSupport: + """Support of the log‑normal distribution (0, ∞).""" + return ContinuousSupport(left=0.0, left_closed=False, right=np.inf, right_closed=False) + + LogNormal = ParametricFamily( + name=FamilyName.LOGNORMAL, + distr_type=UnivariateContinuous, + distr_parametrizations=["meanStd"], + distr_characteristics={ + CharacteristicName.PDF: pdf, + CharacteristicName.CDF: cdf, + CharacteristicName.PPF: ppf, + CharacteristicName.LPDF: lpdf, + CharacteristicName.MEAN: mean_func, + CharacteristicName.VAR: var_func, + CharacteristicName.SKEW: skew_func, + CharacteristicName.KURT: kurt_func, + }, + support_by_parametrization=_support, + ) + LogNormal.__doc__ = LOGNORMAL_DOC + + @parametrization(family=LogNormal, name="meanStd") # family will be set after Normal is created + class _MeanStd(Parametrization): + mu: float + sigma: float + + @constraint(description="sigma > 0") + def check_sigma_positive(self) -> bool: + return self.sigma > 0 + + ParametricFamilyRegister.register(LogNormal) diff --git a/src/pysatl_core/families/configuration.py b/src/pysatl_core/families/configuration.py index 7d0358a..e7e97bb 100644 --- a/src/pysatl_core/families/configuration.py +++ b/src/pysatl_core/families/configuration.py @@ -25,6 +25,7 @@ from pysatl_core.families.builtins import ( configure_exponential_family, + configure_lognormal_family, configure_normal_family, configure_uniform_family, ) @@ -48,6 +49,7 @@ def configure_families_register() -> ParametricFamilyRegister: configure_exponential_family() configure_uniform_family() configure_normal_family() + configure_lognormal_family() return ParametricFamilyRegister() diff --git a/src/pysatl_core/families/distribution.py b/src/pysatl_core/families/distribution.py index b2529c4..cd4d754 100644 --- 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/parametric_family.py b/src/pysatl_core/families/parametric_family.py index 1ed66e5..914646f 100644 --- a/src/pysatl_core/families/parametric_family.py +++ b/src/pysatl_core/families/parametric_family.py @@ -23,6 +23,7 @@ from pysatl_core.distributions.computations.computation import AnalyticalComputation from pysatl_core.families.distribution import ParametricFamilyDistribution from pysatl_core.families.parametrizations import Parametrization, ParametrizationConstraint +from pysatl_core.families.registry import ParametricFamilyRegister from pysatl_core.types import ( DEFAULT_ANALYTICAL_COMPUTATION_LABEL, ComputationFunc, @@ -384,6 +385,25 @@ def distribution( parameters = parametrization_class(**parameters_values) parameters.validate() + registry = ParametricFamilyRegister() + + optimal_result = registry.get_optimal_family(self.name, parameters) + if optimal_result is not None: + optimal_family, optimal_parametrization = optimal_result + distribution_type = optimal_family._distr_type(optimal_parametrization) + analytical_computations = optimal_family._build_analytical_computations( + optimal_parametrization + ) + return ParametricFamilyDistribution( + family_name=optimal_family.name, + distribution_type=distribution_type, + analytical_computations=analytical_computations, + parametrization=optimal_parametrization, + support=optimal_family.support_resolver(optimal_parametrization), + sampling_strategy=sampling_strategy, + computation_strategy=computation_strategy, + ) + base_parameters = self.to_base(parameters) distribution_type = self._distr_type(base_parameters) analytical_computations = self._build_analytical_computations(parameters) diff --git a/src/pysatl_core/families/registry.py b/src/pysatl_core/families/registry.py index 2b0ba6e..26acd9f 100644 --- a/src/pysatl_core/families/registry.py +++ b/src/pysatl_core/families/registry.py @@ -8,16 +8,24 @@ from __future__ import annotations +from pysatl_core.families.parametrizations import Parametrization +from pysatl_core.families.registry_graph import ( + BinaryOperationType, + RegistryGraphTransformations, +) + __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: from typing import ClassVar from pysatl_core.families.parametric_family import ParametricFamily + from pysatl_core.types import Number, NumericArray class ParametricFamilyRegister: @@ -30,14 +38,110 @@ class ParametricFamilyRegister: _instance: ClassVar[ParametricFamilyRegister | None] = None _registered_families: dict[str, ParametricFamily] + _registry_graph: RegistryGraphTransformations 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._registry_graph = RegistryGraphTransformations() return cls._instance + @classmethod + def get_optimal_family( + cls, name: str, parametrization: Parametrization + ) -> tuple[ParametricFamily, Parametrization] | None: + self = cls() + if name not in self._registered_families: + return None + + family_name, new_parametrization = self._registry_graph.get_optimal_parametrization( + name, parametrization + ) + return self.get(family_name), new_parametrization + + @classmethod + def get_optimal_density( + cls, name: str + ) -> None | tuple[ParametricFamily, Callable[[Number | NumericArray], Number | NumericArray]]: + self = cls() + if name not in self._registered_families: + return None + + 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, + 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 and tail_name in self._registered_families: + self._registry_graph.register_parametrization_transformation( + head_name, tail_name, transform_constraint, transform_function + ) + return True + + return False + + @classmethod + def register_density_transformation( + cls, + head_name: str, + tail_name: str, + transform_function: Callable[[Number | NumericArray], Number | NumericArray], + ) -> bool: + self = cls() + + if head_name in self._registered_families and tail_name in self._registered_families: + self._registry_graph.register_density_transformation( + head_name, tail_name, transform_function + ) + return True + + return False + @classmethod def get(cls, name: str) -> ParametricFamily: """ @@ -64,7 +168,7 @@ def get(cls, name: str) -> ParametricFamily: return self._registered_families[name] @classmethod - def register(cls, family: ParametricFamily) -> None: + def register(cls, family: ParametricFamily, temperature: int = 128) -> None: """ Register a new parametric family. @@ -79,10 +183,17 @@ def register(cls, family: ParametricFamily) -> None: If a family with the same name is already registered. """ self = cls() + 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: + self = cls() + self._registry_graph.register_family_temperature(family_name, new_temperature) + @classmethod def contains(cls, name: str) -> bool: """ diff --git a/src/pysatl_core/families/registry_graph.py b/src/pysatl_core/families/registry_graph.py new file mode 100644 index 0000000..8143ad9 --- /dev/null +++ b/src/pysatl_core/families/registry_graph.py @@ -0,0 +1,342 @@ +from __future__ import annotations + +from collections.abc import Callable +from enum import StrEnum +from queue import Queue +from typing import TYPE_CHECKING, cast + +from pysatl_core.families.parametrizations import Parametrization + +if TYPE_CHECKING: + from typing import ClassVar + + from pysatl_core.types import Number, NumericArray + + +class RegistryEdge: + def __init__(self, head_name: str, tail_name: str): + self._head_name = head_name + self._tail_name = tail_name + + @property + def tail_name(self) -> str: + return self._tail_name + + @property + def head_name(self) -> str: + return self._head_name + + +class TransformatedParametrizationEdge(RegistryEdge): + def __init__( + self, + head_name: str, + tail_name: str, + transform_constraint: Callable[[Parametrization], bool], + transform_function: Callable[[Parametrization], Parametrization], + ): + self._transform_function = transform_function + self._transform_constraint = transform_constraint + + RegistryEdge.__init__(self, head_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) + + +class TransformatedDensityEdge(RegistryEdge): + def __init__( + self, + head_name: str, + tail_name: str, + transform_function: Callable[[Number | NumericArray], Number | NumericArray], + ): + self._transform_function = transform_function + RegistryEdge.__init__(self, head_name, tail_name) + + def transform_density(self, argument: Number | NumericArray) -> Number | NumericArray: + 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, + start_vertex: str, + edge_representation: dict[str, list[RegistryEdge]], + edge_constraint: Callable[[RegistryEdge], bool], + visit_edge: Callable[[RegistryEdge], None], + path_callback: Callable[[RegistryEdge], None], + ) -> None: + self = cls() + + visited_vertexes = set() + previous_in_path: dict[str, tuple[str, RegistryEdge] | None] = {} + + visited_vertexes.add(start_vertex) + + previous_in_path[start_vertex] = None + queue: Queue[str] = Queue() + queue.put(start_vertex) + + while not queue.empty(): + current_vertex = queue.get() + + for next_vertex_edge in edge_representation.get(current_vertex, []): + next_name = next_vertex_edge.tail_name + if next_name not in visited_vertexes and edge_constraint(next_vertex_edge): + visited_vertexes.add(next_name) + previous_in_path[next_name] = current_vertex, next_vertex_edge + queue.put(next_name) + visit_edge(next_vertex_edge) + + best_choice = start_vertex + for visited_vertex in visited_vertexes: + if ( + self._registered_families_temperature[visited_vertex] + > self._registered_families_temperature[best_choice] + ): + best_choice = visited_vertex + + path = [] + while previous_in_path[best_choice] is not None: + previous = cast(tuple[str, RegistryEdge], previous_in_path[best_choice]) + path.append(previous[1]) + best_choice = previous[0] + + path = path[::-1] + for edge in path: + path_callback(edge) + + @classmethod + def get_optimal_parametrization( + cls, current_family: str, current_parametrization: Parametrization + ) -> tuple[str, Parametrization]: + self = cls() + + best_choice = current_family + result_parametrization = current_parametrization + + parametrizations = {} + parametrizations[current_family] = current_parametrization + + def path_callback(edge: RegistryEdge) -> None: + nonlocal best_choice + nonlocal result_parametrization + + edge = cast(TransformatedParametrizationEdge, edge) + new_parametrization = edge.transform_parametrization(result_parametrization) + result_parametrization = new_parametrization + best_choice = edge.tail_name + + def visit_edge(edge: RegistryEdge) -> None: + nonlocal parametrizations + edge = cast(TransformatedParametrizationEdge, edge) + + head = edge.head_name + tail = edge.tail_name + parametrizations[tail] = edge.transform_parametrization(parametrizations[head]) + + def edge_constraint(edge: RegistryEdge) -> bool: + nonlocal parametrizations + edge = cast(TransformatedParametrizationEdge, edge) + return edge.is_transoform_possible(parametrizations[edge.head_name]) + + RegistryGraphTransformations._run_bfs( + start_vertex=current_family, + edge_representation=self._registered_parametrzation_transformations, # type: ignore[arg-type] + edge_constraint=edge_constraint, + visit_edge=visit_edge, + path_callback=path_callback, + ) + + return best_choice, result_parametrization + + @classmethod + def get_optimal_transoformation( + cls, current_family: str + ) -> tuple[str, Callable[[Number | NumericArray], Number | NumericArray]]: + self = cls() + + def collected_function(x: Number | NumericArray) -> NumericArray | Number: + return x + + best_choice = current_family + + def edge_constraint(edge: RegistryEdge) -> bool: + return True + + def visit_edge(edge: RegistryEdge) -> None: + pass + + def path_callback(edge: RegistryEdge) -> None: + nonlocal collected_function + nonlocal best_choice + + edge = cast(TransformatedDensityEdge, edge) + _temp = collected_function + + def collected_function(x: Number | NumericArray) -> Number | NumericArray: + return edge.transform_density(_temp(x)) + + best_choice = edge.tail_name + + RegistryGraphTransformations._run_bfs( + current_family, + edge_representation=self._registered_transformations, # type: ignore[arg-type] + edge_constraint=edge_constraint, + visit_edge=visit_edge, + path_callback=path_callback, + ) + + return best_choice, collected_function + + @classmethod + def register_parametrization_transformation( + cls, + head_name: str, + tail_name: str, + transform_constraint: Callable[[Parametrization], bool], + transform_function: Callable[[Parametrization], Parametrization], + ) -> None: + self = cls() + self._registered_parametrzation_transformations.setdefault(head_name, []).append( + TransformatedParametrizationEdge( + head_name, tail_name, transform_constraint, transform_function + ) + ) + + @classmethod + def register_density_transformation( + cls, + head_name: str, + tail_name: str, + transform_function: Callable[[Number | NumericArray], Number | NumericArray], + ) -> None: + self = cls() + self._registered_transformations.setdefault(head_name, []).append( + TransformatedDensityEdge(head_name, tail_name, transform_function) + ) + + @classmethod + def register_family_temperature(cls, family_name: str, temperature: int) -> None: + self = cls() + self._registered_families_temperature[family_name] = temperature diff --git a/src/pysatl_core/types.py b/src/pysatl_core/types.py index 25aa63e..fa802bd 100644 --- a/src/pysatl_core/types.py +++ b/src/pysatl_core/types.py @@ -404,6 +404,7 @@ class FamilyName(StrEnum): NORMAL = "Normal" CONTINUOUS_UNIFORM = "ContinuousUniform" EXPONENTIAL = "Exponential" + LOGNORMAL = "LogNormal" # ============================================================================ diff --git a/tests/unit/families/test_optimization.py b/tests/unit/families/test_optimization.py new file mode 100644 index 0000000..0f7b407 --- /dev/null +++ b/tests/unit/families/test_optimization.py @@ -0,0 +1,131 @@ +from typing import cast + +import numpy as np +from numpy.testing import assert_allclose + +from pysatl_core.families.configuration import configure_families_register +from pysatl_core.families.distribution import ParametricFamilyDistribution +from pysatl_core.families.parametrizations import Parametrization +from pysatl_core.families.registry_graph import BinaryOperationType +from pysatl_core.types import CharacteristicName, FamilyName, Number, NumericArray + + +def test_parametrization_optimization(): + registry = configure_families_register() + exponential_fam = registry.get(FamilyName.EXPONENTIAL) + normal_fam = registry.get(FamilyName.NORMAL) + uniform_fam = registry.get(FamilyName.CONTINUOUS_UNIFORM) + + def transform_function_exponent(head_param: Parametrization) -> Parametrization: + head_param = head_param.transform_to_base_parametrization() + tail_param_type = normal_fam.get_parametrization(normal_fam.base_parametrization_name) + return tail_param_type(mu=0, sigma=1) # type: ignore[call-arg] + + def transform_constraint_exponent(head_param: Parametrization) -> bool: + head_param = head_param.transform_to_base_parametrization() + return head_param.lambda_ == 1.0 # type: ignore[attr-defined] + + registry.register_parametrization_transformation( + FamilyName.EXPONENTIAL, + FamilyName.NORMAL, + transform_constraint_exponent, + transform_function_exponent, + ) + + def transform_function_normal(head_param: Parametrization) -> Parametrization: + tail_param_type = uniform_fam.get_parametrization(uniform_fam.base_parametrization_name) + return tail_param_type(lower_bound=0.0, upper_bound=1.0) # type: ignore[call-arg] + + def transform_constraint_normal(head_param: Parametrization) -> bool: + head_param = head_param.transform_to_base_parametrization() + return head_param.mu == 0.0 and head_param.sigma == 1.0 # type: ignore[attr-defined] + + registry.register_parametrization_transformation( + FamilyName.NORMAL, + FamilyName.CONTINUOUS_UNIFORM, + transform_constraint_normal, + transform_function_normal, + ) + + registry._change_family_temperature(FamilyName.CONTINUOUS_UNIFORM, 129) + + 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 + registry._reset() + + +def test_density_optimization(): + registry = configure_families_register() + registry.get(FamilyName.NORMAL) + lognormal_fam = registry.get(FamilyName.LOGNORMAL) + + def transform_function(x: Number | NumericArray) -> Number | NumericArray: + return np.exp(x) + + registry.register_density_transformation( + FamilyName.NORMAL, FamilyName.LOGNORMAL, transform_function + ) + + def revert_function(x: Number | NumericArray) -> Number | NumericArray: + return np.log(x) + + registry.register_density_transformation( + FamilyName.LOGNORMAL, FamilyName.NORMAL, revert_function + ) + + registry._change_family_temperature(FamilyName.LOGNORMAL, 129) + + lognormal_result = registry.get_optimal_density(FamilyName.NORMAL) + if lognormal_result is None: + raise ValueError("The family is not in registry") + + lognormal_optimized, transformation = lognormal_result + + assert lognormal_optimized == lognormal_fam + + x = np.array(range(1, 10)) + assert_allclose( + np.array([transform_function(xx) for xx in x]), np.array([transformation(xx) for xx in x]) + ) + + +def test_transformations(): + registry = configure_families_register() + normal_fam = registry.get(FamilyName.NORMAL) + + def transform_function( + left_parametrization: Parametrization, right_parametrization: Parametrization + ) -> Parametrization: + param_type = normal_fam.get_parametrization(normal_fam.base_parametrization_name) + return param_type( # type: ignore[call-arg] + mu=left_parametrization.mu + right_parametrization.mu, # type: ignore[attr-defined] + sigma=left_parametrization.sigma + right_parametrization.sigma, # type: ignore[attr-defined] + ) + + registry.add_binary_transformation( + FamilyName.NORMAL, + FamilyName.NORMAL, + FamilyName.NORMAL, + BinaryOperationType.ADD, + transform_function, + ) + distribution_one = normal_fam(mu=1, sigma=1) + distribution_two = normal_fam(mu=2, sigma=3) + + distribution_add = cast(ParametricFamilyDistribution, distribution_one + distribution_two) + distribution_sub = distribution_one * distribution_two + + assert distribution_add.parametrization.mu == 3 # type: ignore[attr-defined] + assert distribution_add.parametrization.sigma == 4 # type: ignore[attr-defined] + + assert isinstance(distribution_add, ParametricFamilyDistribution) + assert not isinstance(distribution_sub, ParametricFamilyDistribution)