python/rwsims/distributions.py

from __future__ import annotations

from abc import ABC, abstractmethod

import numpy as np
# from numpy.typing import ArrayLike
from numpy.random import Generator


class BaseDistribution(ABC):
    def __init__(self, tau: float, rng: Generator | None = None):
        self._tau = tau
        self._rng = rng

        self.tau_jump = tau

    @property
    def name(self) -> str:
        return self.__class__.__name__

    @abstractmethod
    def __repr__(self):
        pass

    @abstractmethod
    def start(self):
        pass

    @property
    @abstractmethod
    def mean_tau(self):
        pass

    def wait(self, size: int = 1) -> ArrayLike:
        return self._rng.exponential(self.tau_jump, size=size)


class DeltaDistribution(BaseDistribution):

    def __repr__(self):
        return f'Delta Distribution (tau={self._tau})'

    def start(self):
        self.tau_jump = self._tau

    @property
    def mean_tau(self):
        return self._tau


class LogGaussianDistribution(BaseDistribution):
    def __init__(self, tau: float, sigma: float, rng: Generator):
        super().__init__(tau=tau, rng=rng)

        self._sigma = sigma

    def __repr__(self):
        return f'Log-Gaussian(tau={self._tau}, sigma={self._sigma})'

    def start(self):
        self.tau_jump = self._rng.lognormal(np.log(self._tau), self._sigma)

    @property
    def mean_tau(self):
        return self._tau * np.exp(self._sigma**2 / 2)