Coverage for src/amisc/distribution.py: 81%

1"""Module for probability distribution functions (PDFs). 

2 

3Includes: 

4 

5- `Distribution` — an abstract interface for specifying a PDF. 

6- `Uniform` — a uniform distribution. 

7- `Normal` — a normal distribution. 

8- `Relative` — a relative distribution (i.e. uniform within a percentage of a nominal value). 

9- `Tolerance` — a tolerance distribution (i.e. uniform within a tolerance of a nominal value). 

10- `LogUniform` — a log-uniform distribution. 

11- `LogNormal` — a log-normal distribution. 

12 

13Distribution objects can be converted easily to/from strings for serialization. 

14""" 

15from __future__ import annotations 

16 

17from abc import ABC, abstractmethod 

18 

19import numpy as np 

20 

21from amisc.utils import parse_function_string 

22 

23__all__ = ['Distribution', 'Uniform', 'Normal', 'Relative', 'Tolerance', 'LogNormal', 'LogUniform'] 

24 

25 

26class Distribution(ABC): 

27 """Base class for PDF distributions that provide sample and pdf methods. Distributions should: 

28 

29 - `sample` - return samples from the distribution 

30 - `pdf` - return the probability density function of the distribution 

31 

32 :ivar dist_args: the arguments that define the distribution (e.g. mean and std for a normal distribution) 

33 :vartype dist_args: tuple 

34 """ 

35 

36 def __init__(self, dist_args: tuple): 

37 self.dist_args = dist_args 

38 

39 def __str__(self): 

40 """Serialize a `Distribution` object to/from string.""" 

41 return f'{type(self).__name__}{self.dist_args}' 

42 

43 def __repr__(self): 

44 return self.__str__() 

45 

46 def domain(self, dist_args: tuple = None) -> tuple: 

47 """Return the domain of this distribution. Defaults to `dist_args` 

48 

49 :param dist_args: overrides `self.dist_args` 

50 """ 

51 return dist_args or self.dist_args 

52 

53 def nominal(self, dist_args: tuple = None) -> float: 

54 """Return the nominal value of this distribution. Defaults to middle of domain. 

55 

56 :param dist_args: overrides `self.dist_args` 

57 """ 

58 lb, ub = self.domain(dist_args=dist_args) 

59 return (lb + ub) / 2 

60 

61 @classmethod 

62 def from_string(cls, dist_string: str) -> Distribution | None: 

63 """Convert a string to a `Distribution` object. 

64 

65 :param dist_string: specifies a PDF or distribution. Can be `Normal(mu, std)`, `Uniform(lb, ub)`, 

66 `LogUniform(lb, ub)`, `LogNormal(mu, std)`, 

67 `Relative(pct)`, or `Tolerance(tol)`. The shorthands `N(0, 1)`, `U(0, 1)`, `LU(0, 1)`, 

68 `LN(0, 1)`, `rel(5)`, or `tol(1)` are also accepted. 

69 :return: the corresponding `Distribution` object 

70 """ 

71 if not dist_string: 

72 return None 

73 

74 dist_name, args, kwargs = parse_function_string(dist_string) 

75 if dist_name in ['N', 'Normal', 'normal']: 

76 # Normal distribution like N(0, 1) 

77 try: 

78 mu = float(kwargs.get('mu', args[0])) 

79 std = float(kwargs.get('std', args[1])) 

80 return Normal((mu, std)) 

81 except Exception as e: 

82 raise ValueError(f'Normal distribution string "{dist_string}" is not valid: Try N(0, 1).') from e 

83 elif dist_name in ['U', 'Uniform', 'uniform']: 

84 # Uniform distribution like U(0, 1) 

85 try: 

86 lb = float(kwargs.get('lb', args[0])) 

87 ub = float(kwargs.get('ub', args[1])) 

88 return Uniform((lb, ub)) 

89 except Exception as e: 

90 raise ValueError(f'Uniform distribution string "{dist_string}" is not valid: Try U(0, 1).') from e 

91 elif dist_name in ['R', 'Relative', 'relative', 'rel']: 

92 # Relative uniform distribution like rel(+-5%) 

93 try: 

94 pct = float(kwargs.get('pct', args[0])) 

95 return Relative((pct,)) 

96 except Exception as e: 

97 raise ValueError(f'Relative distribution string "{dist_string}" is not valid: Try rel(5).') from e 

98 elif dist_name in ['T', 'Tolerance', 'tolerance', 'tol']: 

99 # Uniform distribution within a tolerance like tol(+-1) 

100 try: 

101 tol = float(kwargs.get('tol', args[0])) 

102 return Tolerance((tol,)) 

103 except Exception as e: 

104 raise ValueError(f'Tolerance distribution string "{dist_string}" is not valid: Try tol(1).') from e 

105 elif dist_name in ['LogUniform', 'LU']: 

106 # LogUniform distribution like LU(1e-3, 1e-1) 

107 try: 

108 lb = float(kwargs.get('lb', args[0])) 

109 ub = float(kwargs.get('ub', args[1])) 

110 base = float(kwargs.get('base', args[2] if len(args) > 2 else 10)) 

111 return LogUniform((lb, ub), base=base) 

112 except Exception as e: 

113 raise ValueError(f'LogUniform distr string "{dist_string}" is not valid: Try LU(1e-3, 1e-1).') from e 

114 elif dist_name in ['LogNormal', 'LN']: 

115 # LogNniform distribution like LN(-2, 1) 

116 try: 

117 mu = float(kwargs.get('mu', args[0])) 

118 std = float(kwargs.get('std', args[1])) 

119 base = float(kwargs.get('base', args[2] if len(args) > 2 else 10)) 

120 return LogNormal((mu, std), base=base) 

121 except Exception as e: 

122 raise ValueError(f'LogNormal distr string "{dist_string}" is not valid: Try LN(-2, 1).') from e 

123 else: 

124 raise NotImplementedError(f'The distribution "{dist_string}" is not recognized.') 

125 

126 @abstractmethod 

127 def sample(self, shape: int | tuple, nominal: float | np.ndarray = None, dist_args: tuple = None) -> np.ndarray: 

128 """Sample from the distribution. 

129 

130 :param shape: shape of the samples to return 

131 :param nominal: a nominal value(s) for sampling (e.g. for relative distributions) 

132 :param dist_args: overrides `Distribution.dist_args` 

133 :return: the samples of the given shape 

134 """ 

135 raise NotImplementedError 

136 

137 @abstractmethod 

138 def pdf(self, x: np.ndarray, dist_args: tuple = None) -> np.ndarray: 

139 """Evaluate the pdf of this distribution at the `x` locations. 

140 

141 :param x: the locations at which to evaluate the pdf 

142 :param dist_args: overrides `Distribution.dist_args` 

143 :return: the pdf evaluations 

144 """ 

145 raise NotImplementedError 

146 

147 

148class Uniform(Distribution): 

149 """A Uniform distribution. Specify by string as "Uniform(lb, ub)" or "U(lb, ub)" in shorthand.""" 

150 

151 def __str__(self): 

152 return f'U({self.dist_args[0]}, {self.dist_args[1]})' 

153 

154 def sample(self, shape, nominal=None, dist_args=None): 

155 lb, ub = dist_args or self.dist_args 

156 return np.random.rand(*shape) * (ub - lb) + lb 

157 

158 def pdf(self, x, dist_args=None): 

159 x = np.atleast_1d(x) 

160 lb, ub = dist_args or self.dist_args 

161 pdf = np.broadcast_to(1 / (ub - lb), x.shape).copy() 

162 pdf[np.where(x > ub)] = 0 

163 pdf[np.where(x < lb)] = 0 

164 return pdf 

165 

166 

167class LogUniform(Distribution): 

168 """A LogUniform distribution. Specify by string as `LogUniform(lb, ub)` or `LU(lb, ub)` in shorthand. Uses 

169 base-10 by default. 

170 

171 !!! Example 

172 ```python 

173 x = LogUniform((1e-3, 1e-1)) # log10(x) ~ U(-3, -1) 

174 ``` 

175 """ 

176 def __init__(self, dist_args: tuple, base=10): 

177 self.base = base 

178 super().__init__(dist_args) 

179 

180 def __str__(self): 

181 return f'LU({self.dist_args[0]}, {self.dist_args[1]}, base={self.base})' 

182 

183 def sample(self, shape, nominal=None, dist_args=None): 

184 lb, ub = dist_args or self.dist_args 

185 c = 1 / np.log(self.base) 

186 return self.base ** (np.random.rand(*shape) * c * (np.log(ub) - np.log(lb)) + c * np.log(lb)) 

187 

188 def pdf(self, x, dist_args=None): 

189 x = np.atleast_1d(x) 

190 lb, ub = dist_args or self.dist_args 

191 c = 1 / np.log(self.base) 

192 const = 1 / (c * (np.log(ub) - np.log(lb)) * np.log(self.base)) 

193 pdf = const / x 

194 pdf[np.where(x > ub)] = 0 

195 pdf[np.where(x < lb)] = 0 

196 return pdf 

197 

198 

199class Normal(Distribution): 

200 """A Normal distribution. Specify by string as `Normal(mu, std)` or `N(mu, std)` in shorthand.""" 

201 

202 def __str__(self): 

203 return f'N({self.dist_args[0]}, {self.dist_args[1]})' 

204 

205 def domain(self, dist_args=None): 

206 """Defaults the domain of the distribution to 3 standard deviations above and below the mean.""" 

207 mu, std = dist_args or self.dist_args 

208 return mu - 3 * std, mu + 3 * std 

209 

210 def sample(self, shape, nominal=None, dist_args=None): 

211 mu, std = dist_args or self.dist_args 

212 return np.random.randn(*shape) * std + mu 

213 

214 def pdf(self, x, dist_args=None): 

215 mu, std = dist_args or self.dist_args 

216 return (1 / (np.sqrt(2 * np.pi) * std)) * np.exp(-0.5 * ((x - mu) / std) ** 2) 

217 

218 

219class LogNormal(Distribution): 

220 """A LogNormal distribution. Specify by string as `LogNormal(mu, sigma)` or `LN(mu, sigma)` in shorthand. 

221 Uses base-10 by default. 

222 

223 !!! Example 

224 ```python 

225 x = LogNormal((-2, 1)) # log10(x) ~ N(-2, 1) 

226 ``` 

227 """ 

228 def __init__(self, dist_args: tuple, base=10): 

229 self.base = base 

230 super().__init__(dist_args) 

231 

232 def __str__(self): 

233 return f'LN({self.dist_args[0]}, {self.dist_args[1]}, base={self.base})' 

234 

235 def domain(self, dist_args=None): 

236 """Defaults the domain of the distribution to 3 standard deviations above and below the mean.""" 

237 mu, std = dist_args or self.dist_args 

238 return self.base ** (mu - 3 * std), self.base ** (mu + 3 * std) 

239 

240 def sample(self, shape, nominal=None, dist_args=None): 

241 mu, std = dist_args or self.dist_args 

242 return self.base ** (np.random.randn(*shape) * std + mu) 

243 

244 def pdf(self, x, dist_args=None): 

245 mu, std = dist_args or self.dist_args 

246 const = (1 / (np.sqrt(2 * np.pi) * std * x * np.log(self.base))) 

247 return const * np.exp(-0.5 * ((np.log(x)/np.log(self.base) - mu) / std) ** 2) 

248 

249 

250class Relative(Distribution): 

251 """A Relative distribution. Specify by string as `Relative(pct)` or `rel(pct%)` in shorthand. 

252 Will attempt to sample uniformly within the given percent of a nominal value. 

253 """ 

254 

255 def __str__(self): 

256 return rf'rel({self.dist_args[0]})' 

257 

258 def domain(self, dist_args=None): 

259 return None 

260 

261 def nominal(self, dist_args=None): 

262 return None 

263 

264 def sample(self, shape, nominal=None, dist_args=None): 

265 if nominal is None: 

266 raise ValueError('Cannot sample relative distribution when no nominal value is provided.') 

267 dist_args = dist_args or self.dist_args 

268 tol = abs((dist_args[0] / 100) * nominal) 

269 return np.random.rand(*shape) * 2 * tol - tol + nominal 

270 

271 def pdf(self, x, dist_args=None): 

272 return np.ones(x.shape) 

273 

274 

275class Tolerance(Distribution): 

276 """A Tolerance distribution. Specify by string as `Tolerance(tol)` or `tol(tol)` in shorthand. 

277 Will attempt to sample uniformly within a given absolute tolerance of a nominal value. 

278 """ 

279 

280 def __str__(self): 

281 return rf'tol({self.dist_args[0]})' 

282 

283 def domain(self, dist_args=None): 

284 return None 

285 

286 def nominal(self, dist_args=None): 

287 return None 

288 

289 def sample(self, shape, nominal=None, dist_args=None): 

290 if nominal is None: 

291 raise ValueError('Cannot sample tolerance distribution when no nominal value is provided.') 

292 dist_args = dist_args or self.dist_args 

293 tol = abs(dist_args[0]) 

294 return np.random.rand(*shape) * 2 * tol - tol + nominal 

295 

296 def pdf(self, x, dist_args=None): 

297 return np.ones(np.atleast_1d(x).shape)