feat: big base for distribution package

Author: LeonidElkin

.pre-commit-config.yaml

@@ -59,3 +59,7 @@ repos:
     hooks:
       - id: mypy
         args: [ "--config-file=pyproject.toml" ]
+        additional_dependencies:
+          # Иначе mypy не находит пакеты в венве
+          # Альтернатива сделать repo: local, но это менее чисто с точки зрения CI
+          - numpy>=2

pyproject.toml

@@ -43,6 +43,7 @@ urls.Issues = "https://github.com/PySATL/pysatl-core/issues"
 
 [tool.poetry.dependencies]
 python = ">=3.12"
+numpy = "^2.0.0"
 
 [tool.poetry.group.dev.dependencies]
 ruff = ">=0.6"

src/pysatl_core/distributions/__init__.py

@@ -0,0 +1,13 @@
+from typing import Any, Protocol
+
+from pysatl_core.types import (
+    GenericCharacteristicName,
+)
+
+from .distribution import Distribution
+
+
+class GenericCharacteristic[In, Out](Protocol):
+    name: GenericCharacteristicName
+
+    def __call__(self, distribution: Distribution, data: In, **options: Any) -> Out: ...

src/pysatl_core/distributions/characteristics.py

@@ -0,0 +1,62 @@
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from pysatl_core.distributions.distribution import Distribution
+from pysatl_core.types import (
+    GenericCharacteristicName,
+)
+
+
+@runtime_checkable
+class Computation[In, Out](Protocol):
+    @property
+    def target(self) -> GenericCharacteristicName: ...
+    def __call__(self, data: In) -> Out: ...
+
+
+@runtime_checkable
+class FittedComputationMethodProtocol[In, Out](Protocol):
+    @property
+    def target(self) -> GenericCharacteristicName: ...
+    @property
+    def sources(self) -> Sequence[GenericCharacteristicName]: ...
+    def __call__(self, data: In) -> Out: ...
+
+
+@runtime_checkable
+class ComputationMethodProtocol[In, Out](Protocol):
+    @property
+    def target(self) -> GenericCharacteristicName: ...
+    @property
+    def sources(self) -> Sequence[GenericCharacteristicName]: ...
+    def fit(self, distribution: Distribution) -> FittedComputationMethodProtocol[In, Out]: ...
+
+
+@dataclass(frozen=True, slots=True)
+class AnalyticalComputation[In, Out]:
+    target: GenericCharacteristicName
+    func: Callable[[In], Out]
+
+    def __call__(self, data: In) -> Out:
+        return self.func(data)
+
+
+@dataclass(frozen=True, slots=True)
+class FittedComputationMethod[In, Out]:
+    target: GenericCharacteristicName
+    sources: Sequence[GenericCharacteristicName]
+    func: Callable[[In], Out]
+
+    def __call__(self, data: In) -> Out:
+        return self.func(data)
+
+
+@dataclass(frozen=True, slots=True)
+class ComputationMethod[In, Out]:
+    target: GenericCharacteristicName
+    sources: Sequence[GenericCharacteristicName]
+    fitter: Callable[[Distribution], FittedComputationMethod[In, Out]]
+
+    def fit(self, distribution: Distribution) -> FittedComputationMethod[In, Out]:
+        return self.fitter(distribution)

src/pysatl_core/distributions/computation.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
+
+from pysatl_core.distributions.computation import AnalyticalComputation
+from pysatl_core.distributions.sampling import Sample
+from pysatl_core.distributions.strategies import (
+    ComputationStrategy,
+    DefaultComputationStrategy,
+    DefaultSamplingStrategy,
+    SamplingStrategy,
+)
+from pysatl_core.types import (
+    Dimension,
+    GenericCharacteristicName,
+    Kind,
+)
+
+
+@runtime_checkable
+class Distribution(Protocol):
+    @property
+    def dimension(self) -> Dimension: ...
+    @property
+    def kind(self) -> Kind: ...
+
+    @property
+    def analytical_computations(
+        self,
+    ) -> Mapping[GenericCharacteristicName, AnalyticalComputation[Any, Any]]: ...
+
+    @property
+    def sampling_strategy(self) -> SamplingStrategy: ...
+    @property
+    def computation_strategy(self) -> ComputationStrategy[Any, Any]: ...
+
+    def sample(self, n: int, **options: Any) -> Sample: ...
+    def log_likelihood(self, batch: Sample) -> float: ...
+
+
+@dataclass(slots=True)
+class StandaloneDistribution:
+    _dimension: Dimension
+    _kind: Kind
+    _analytical: dict[GenericCharacteristicName, AnalyticalComputation[Any, Any]]
+
+    def __init__(
+        self,
+        dimension: Dimension,
+        kind: Kind,
+        analytical_computations: Iterable[AnalyticalComputation[Any, Any]]
+        | Mapping[GenericCharacteristicName, AnalyticalComputation[Any, Any]] = (),
+    ):
+        self._dimension = dimension
+        self._kind = kind
+        if isinstance(analytical_computations, Mapping):
+            self._analytical = dict(analytical_computations)
+        else:
+            self._analytical = {ac.target: ac for ac in analytical_computations}
+
+    @property
+    def dimension(self) -> Dimension:
+        return self._dimension
+
+    @property
+    def kind(self) -> Kind:
+        return self._kind
+
+    @property
+    def analytical_computations(
+        self,
+    ) -> Mapping[GenericCharacteristicName, AnalyticalComputation[Any, Any]]:
+        return self._analytical
+
+    @property
+    def sampling_strategy(self) -> SamplingStrategy:
+        return DefaultSamplingStrategy()
+
+    @property
+    def computation_strategy(self) -> ComputationStrategy[Any, Any]:
+        return DefaultComputationStrategy()
+
+    def sample(self, n: int, **options: Any) -> Sample:
+        return self.sampling_strategy.sample(n, d=self.dimension, **options)
+
+    def log_likelihood(self, batch: Sample) -> float:
+        # TODO: Не ноль)
+        return 0.0

src/pysatl_core/distributions/distribution.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Any, ClassVar, Self
+
+from pysatl_core.distributions.computation import ComputationMethod
+from pysatl_core.types import (
+    Dimension,
+    GenericCharacteristicName,
+    Kind,
+)
+
+
+@dataclass
+class GenericCharacteristicRegister:
+    dimension: Dimension
+    kind: Kind
+
+    __registered_indefinitive_chars: dict[
+        GenericCharacteristicName,
+        list[ComputationMethod[Any, Any]],
+    ] = field(default_factory=dict)
+
+    __registered_definitive_chars: dict[
+        GenericCharacteristicName, list[ComputationMethod[Any, Any]]
+    ] = field(default_factory=dict)
+
+    def register_indefinitive_characteristic(
+        self,
+        default_computation_method: ComputationMethod[Any, Any],
+    ) -> None:
+        self.__registered_indefinitive_chars[default_computation_method.target].append(
+            default_computation_method
+        )
+
+    def register_definitive_characteristic(
+        self,
+        default_computation_method: ComputationMethod[Any, Any],
+        default_inversion_method: ComputationMethod[Any, Any],
+    ) -> None:
+        if (
+            default_computation_method.target not in default_inversion_method.sources
+            or default_inversion_method.target not in default_computation_method.sources
+        ):
+            raise AttributeError(
+                "New characteristic always must be in targets for default "
+                "computation method and in sources for default inversion method"
+            )
+        self.__registered_definitive_chars[default_computation_method.target].append(
+            default_computation_method
+        )
+        self.__registered_definitive_chars[default_inversion_method.target].append(
+            default_inversion_method
+        )
+
+    def get_available_indefinitive_characteristics(
+        self,
+    ) -> dict[GenericCharacteristicName, list[ComputationMethod[Any, Any]]]:
+        return self.__registered_indefinitive_chars
+
+    def get_available_definitive_characteristics(
+        self,
+    ) -> dict[GenericCharacteristicName, list[ComputationMethod[Any, Any]]]:
+        return self.__registered_definitive_chars
+
+    def get_all_available_characteristics_keys(self) -> list[GenericCharacteristicName]:
+        return list(self.__registered_definitive_chars.keys()) + list(
+            self.__registered_indefinitive_chars.keys()
+        )
+
+
+class DistributionTypeRegister:
+    _instance: ClassVar[Self | None] = None
+    _register_kinds: dict[tuple[Dimension, Kind], GenericCharacteristicRegister]
+
+    def __new__(cls) -> Self:
+        if cls._instance is None:
+            self = super().__new__(cls)
+            self._register_kinds = {}
+            cls._instance = self
+        return cls._instance
+
+    def get(self, dimension: Dimension, kind: Kind) -> GenericCharacteristicRegister:
+        key = (dimension, kind)
+        reg = self._register_kinds.get(key)
+        if reg is None:
+            reg = GenericCharacteristicRegister(dimension=dimension, kind=kind)
+            self._register_kinds[key] = reg
+        if reg.dimension != dimension or reg.kind != kind:
+            raise TypeError(
+                f"Inconsistent registry under key ({dimension}, {kind}): "
+                f"got ({reg.dimension}, {reg.kind}) inside"
+            )
+        return reg
+
+    __call__ = get
+
+
+def _configure(reg: DistributionTypeRegister) -> None:
+    reg.get(1, Kind.DISCRETE)
+    reg.get(1, Kind.CONTINUOUS)
+
+    # Тут много заполнений дефолтами
+
+
+@lru_cache(maxsize=1)
+def distribution_type_register() -> DistributionTypeRegister:
+    reg = DistributionTypeRegister()
+    _configure(reg)
+    return reg
+
+
+def _reset_distribution_type_register_for_tests() -> None:
+    distribution_type_register.cache_clear()

src/pysatl_core/distributions/registry.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any, Protocol
+
+import numpy as np
+import numpy.typing as npt
+
+from pysatl_core.types import Dimension
+
+
+class Sample(Protocol):
+    def __len__(self) -> int: ...
+    @property
+    def dim(self) -> Dimension: ...
+    @property
+    def array(self) -> npt.NDArray[np.floating[Any]]: ...
+    @property
+    def shape(self) -> tuple[int, ...]: ...
+
+
+class ArraySample(Sample):
+    dimension: Dimension
+    data: npt.NDArray[np.floating[Any]]
+
+    def __init__(self, data: npt.NDArray[np.floating[Any]]) -> None:
+        self.data = data
+        self.dimension = data.shape[1]
+
+    def __len__(self) -> int:
+        return int(self.data.shape[0])
+
+    @property
+    def dim(self) -> Dimension:
+        return self.dimension
+
+    def __iter__(self) -> Iterator[npt.NDArray[np.floating[Any]]]:
+        yield from self.data
+
+    @property
+    def array(self) -> npt.NDArray[np.floating[Any]]:
+        return self.data
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        n, d = self.data.shape
+        return int(n), int(d)