Coverage for src/amisc/rv.py: 83%
161 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-08-29 21:38 +0000
1"""Provides small classes for random variables.
3Includes:
5- `BaseRV`: Abstract wrapper class of a random variable
6- `UniformRV`: a uniformly-distributed random variable
7- `NormalRV`: a normally-distributed random variable
8- `ScalarRV`: a stand-in class for a variable with no uncertainty or pdf
9- `LogUniformRV`: base 10 log-uniform
10- `LogNormalRV`: base 10 log-normal
11"""
12import random
13import string
14from abc import ABC, abstractmethod
16import numpy as np
19class BaseRV(ABC):
20 """Small wrapper class similar to `scipy.stats` random variables (RVs).
22 :ivar id: an identifier for the variable
23 :ivar bds: the explicit domain bounds of the variable (limits of where you expect to use it)
24 :ivar nominal: a typical value for this variable (within `bds`)
25 :ivar tex: latex format for the random variable, i.e. r"$x_i$"
26 :ivar description: a lengthier description of the variable
27 :ivar units: assumed units for the variable (if applicable)
28 :ivar param_type: an additional descriptor for how this rv is used, e.g. calibration, operating, design, etc.
30 :vartype id: str
31 :vartype bds: tuple[float, float]
32 :vartype nominal: float
33 :vartype tex: str
34 :vartype description: str
35 :vartype units: str
36 :vartype param_type: str
37 """
39 def __init__(self, id='', *, tex='', description='Random variable', units='-',
40 param_type='calibration', nominal=1, domain=(0, 1)):
41 """Child classes must implement `sample` and `pdf` methods."""
42 self.bds = tuple(domain)
43 self.nominal = nominal
44 self.id = id if id != '' else 'X_' + ''.join(random.choices(string.ascii_uppercase + string.digits, k=4))
45 self.is_custom_id = id != '' # Whether a custom id was assigned to this variable
46 self.tex = tex
47 self.description = description
48 self.units = units
49 self.param_type = param_type
51 def __repr__(self):
52 return r'{}'.format(f"{self.id} - {self.description} ({self.units})")
54 def __str__(self):
55 return self.id
57 def __eq__(self, other):
58 """Consider two RVs equal if they share the same string id.
60 Also returns true when checking if this RV is equal to a string id by itself.
61 """
62 if isinstance(other, BaseRV):
63 return self.id == other.id
64 elif isinstance(other, str):
65 return self.id == other
66 else:
67 return NotImplemented
69 def __hash__(self):
70 return hash(self.id)
72 def to_tex(self, units=False, symbol=True):
73 """Return a raw string that is well-formatted for plotting (with tex).
75 :param units: whether to include the units in the string
76 :param symbol: just latex symbol if true, otherwise the full description
77 """
78 s = self.tex if symbol else self.description
79 if s == '':
80 s = str(self)
81 return r'{} [{}]'.format(s, self.units) if units else r'{}'.format(s)
83 def bounds(self):
84 """Return a tuple of the defined domain of this RV."""
85 return self.bds
87 def update_bounds(self, lb, ub):
88 """Update the defined domain of this RV to `(lb, ub)`."""
89 self.bds = (lb, ub)
91 def sample_domain(self, shape: tuple | int) -> np.ndarray:
92 """Return an array of the given `shape` for random samples over the domain of this RV.
94 :param shape: the shape of samples to return
95 :returns samples: random samples over the domain of the random variable
96 """
97 if isinstance(shape, int):
98 shape = (shape, )
99 return np.random.rand(*shape) * (self.bds[1] - self.bds[0]) + self.bds[0]
101 @abstractmethod
102 def pdf(self, x: np.ndarray) -> np.ndarray:
103 """Compute the PDF of the RV at the given `x` locations.
105 :param x: locations to compute the PDF at
106 :returns f: the PDF evaluations at `x`
107 """
108 pass
110 @abstractmethod
111 def sample(self, shape: tuple | int, nominal: float = None) -> np.ndarray:
112 """Draw samples from the PDF.
114 :param shape: the shape of the returned samples
115 :param nominal: a nominal value to use if applicable (i.e. a center for relative Uniform or Normal)
116 :returns: samples from the PDF of this random variable
117 """
118 pass
121class ScalarRV(BaseRV):
122 """A stand-in variable with no uncertainty/pdf, just scalars."""
124 def pdf(self, x):
125 return np.ones(x.shape)
127 def sample(self, shape, nominal=None):
128 if isinstance(shape, int):
129 shape = (shape, )
130 if nominal is not None:
131 return np.ones(shape)*nominal
132 else:
133 return self.sample_domain(shape)
136class UniformRV(BaseRV):
137 """A uniformly-distributed random variable.
139 Can be uniformly distributed in one of three ways: between global bounds, relative within a percent, or relative
140 within a set absolute tolerance.
142 :ivar type: specifies the type of uniform distribution, either 'bds', 'pct', or 'tol' as described above
143 :ivar value: the absolute tolerance or percent uncertainty if type is 'tol' or 'pct'
145 :vartype type: str
146 :vartype value: float
147 """
149 def __init__(self, arg1: float, arg2: float | str, id='', **kwargs):
150 """Construct a uniformly-distributed random variable.
152 :param arg1: lower bound if specifying U(lb, ub), otherwise a tol or pct if specifying U(+/- tol/pct)
153 :param arg2: upper bound if specifying U(lb, ub), otherwise a str of either 'tol' or 'pct'
154 """
155 domain = kwargs.get('domain', None)
156 if isinstance(arg2, str):
157 self.value = arg1
158 self.type = arg2
159 else:
160 self.value = None
161 self.type = 'bds'
162 if self.type == 'bds':
163 domain = (arg1, arg2) if domain is None else tuple(domain) # This means domain overrides (arg1, arg2)
164 else:
165 domain = (0, 1) if domain is None else tuple(domain)
166 kwargs['domain'] = domain
167 super().__init__(id, **kwargs)
169 # Set default nominal value as middle of the domain if not specified
170 if kwargs.get('nominal', None) is None:
171 self.nominal = (self.bds[1] + self.bds[0]) / 2
173 def __str__(self):
174 return self.id if self.is_custom_id else f'U({self.bds[0]}, {self.bds[1]})'
176 def get_uniform_bounds(self, nominal: float = None) -> tuple[float, float]:
177 """Return the correct set of bounds based on type of uniform distribution.
179 :param nominal: the center value for relative uniform distributions
180 :returns: the uniform bounds
181 """
182 match self.type:
183 case 'bds':
184 return self.bds
185 case 'pct':
186 if nominal is None:
187 return self.bds
188 return nominal * (1 - self.value), nominal * (1 + self.value)
189 case 'tol':
190 if nominal is None:
191 return self.bds
192 return nominal - self.value, nominal + self.value
193 case other:
194 raise NotImplementedError(f'self.type = {other} not known. Choose from ["pct, "tol", "bds"]')
196 def pdf(self, x: np.ndarray, nominal: float = None) -> np.ndarray:
197 """Compute the pdf for a uniform distribution.
199 :param x: locations to compute the pdf at
200 :param nominal: center location for relative uniform rvs
201 :returns: the evaluated PDF at `x`
202 """
203 x = np.atleast_1d(x)
204 bds = self.get_uniform_bounds(nominal)
205 den = bds[1] - bds[0]
206 den = 1 if np.isclose(den, 0) else den
207 y = np.broadcast_to(1 / den, x.shape).copy()
208 y[np.where(x > bds[1])] = 0
209 y[np.where(x < bds[0])] = 0
210 return y
212 def sample(self, shape, nominal=None):
213 if isinstance(shape, int):
214 shape = (shape, )
215 bds = self.get_uniform_bounds(nominal)
216 return np.random.rand(*shape) * (bds[1] - bds[0]) + bds[0]
219class LogUniformRV(BaseRV):
220 """A base 10 log-uniform distributed random variable, only supports absolute bounds."""
222 def __init__(self, log10_a: float, log10_b: float, id='', **kwargs):
223 """Construct the log-uniform random variable.
225 :param log10_a: the lower bound in log10 space
226 :param log10_b: the upper bound in log10 space
227 """
228 super().__init__(id, **kwargs)
229 self.bds = (10**log10_a, 10**log10_b)
231 def __str__(self):
232 return self.id if self.is_custom_id else f'LU({np.log10(self.bds[0]): .2f}, {np.log10(self.bds[1]): .2f})'
234 def pdf(self, x):
235 return np.log10(np.e) / (x * (np.log10(self.bds[1]) - np.log10(self.bds[0])))
237 def sample(self, shape, nominal=None):
238 if isinstance(shape, int):
239 shape = (shape, )
240 lb = np.log10(self.bds[0])
241 ub = np.log10(self.bds[1])
242 return 10 ** (np.random.rand(*shape) * (ub - lb) + lb)
245class LogNormalRV(BaseRV):
246 """A base 10 log-normal distributed random variable.
248 :ivar mu: the center of the log-normal distribution
249 :ivar std: the standard deviation of the log-normal distribution
251 :vartype mu: float
252 :vartype std: float
253 """
255 def __init__(self, mu: float, std: float, id='', **kwargs):
256 """Construct the RV with the mean and std of the underlying distribution,
257 i.e. $\\log_{10}(x) \\sim N(\\mu, \\sigma)$.
259 :param mu: the center of the log-normal distribution
260 :param std: the standard deviation of the log-normal distribution
261 """
262 domain = kwargs.get('domain', None)
263 if domain is None:
264 domain = (10 ** (mu - 3*std), 10 ** (mu + 3*std)) # Use a default domain of +- 3std
265 kwargs['domain'] = domain
266 super().__init__(id, **kwargs)
267 self.std = std
268 self.mu = mu
270 def recenter(self, mu: float, std: float = None):
271 """Move the center of the distribution to `mu` with standard deviation `std` (optional)
273 :param mu: the new center of the distribution
274 :param std: (optional) new standard deviation
275 """
276 self.mu = mu
277 if std is not None:
278 self.std = std
280 def __str__(self):
281 return self.id if self.is_custom_id else f'LN_10({self.mu}, {self.std})'
283 def pdf(self, x):
284 return (np.log10(np.e) / (x * self.std * np.sqrt(2 * np.pi))) * \
285 np.exp(-0.5 * ((np.log10(x) - self.mu) / self.std) ** 2)
287 def sample(self, shape, nominal=None):
288 if isinstance(shape, int):
289 shape = (shape, )
290 scale = np.log10(np.e)
291 center = self.mu if nominal is None else nominal
292 return np.random.lognormal(mean=(1 / scale) * center, sigma=(1 / scale) * self.std, size=shape)
293 # return 10 ** (np.random.randn(*size)*self.std + center) # Alternatively
296class NormalRV(BaseRV):
297 """A normally-distributed random variable.
299 :ivar mu: float, the mean of the normal distribution
300 :ivar std: float, the standard deviation of the normal distribution
302 :vartype mu: float
303 :vartype std: float
304 """
306 def __init__(self, mu, std, id='', **kwargs):
307 domain = kwargs.get('domain', None)
308 if domain is None:
309 domain = (mu - 2.5*std, mu + 2.5*std) # Use a default domain of +- 2.5std
310 kwargs['domain'] = domain
311 super().__init__(id, **kwargs)
312 self.mu = mu
313 self.std = std
315 # Set default nominal value as the provided mean
316 if kwargs.get('nominal', None) is None:
317 self.nominal = mu
319 def recenter(self, mu: float, std: float =None):
320 """Move the center of the distribution to `mu` with standard deviation `std` (optional)
322 :param mu: the new center of the distribution
323 :param std: (optional) new standard deviation
324 """
325 self.mu = mu
326 if std is not None:
327 self.std = std
329 def __str__(self):
330 return self.id if self.is_custom_id else f'N({self.mu}, {self.std})'
332 def pdf(self, x):
333 x = np.atleast_1d(x)
334 return (1 / (np.sqrt(2 * np.pi) * self.std)) * np.exp(-0.5 * ((x - self.mu) / self.std) ** 2)
336 def sample(self, shape, nominal=None):
337 if isinstance(shape, int):
338 shape = (shape, )
339 center = self.mu if nominal is None else nominal
340 return np.random.randn(*shape) * self.std + center