73 lines
1.5 KiB
Python
73 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
import numpy as np
|
|
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
|
|
|
|
def header(self):
|
|
return f'tau = {self._tau}'
|
|
|
|
@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 header(self) -> str:
|
|
return (
|
|
f'tau = {self._tau}\n'
|
|
f'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) |