Coverage for HARK / distributions / continuous.py: 95%
139 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-07 05:16 +0000
1from typing import Any, Optional, Union
3import numpy as np
4from scipy.special import erfc
5from scipy.stats import (
6 rv_continuous,
7 norm,
8 lognorm,
9 uniform,
10 weibull_min,
11)
12from scipy.stats._distn_infrastructure import rv_continuous_frozen
14from HARK.distributions.base import Distribution
15from HARK.distributions.discrete import DiscreteDistribution
17# CONTINUOUS DISTRIBUTIONS
20class ContinuousFrozenDistribution(rv_continuous_frozen, Distribution):
21 """
22 Parameterized continuous distribution from scipy.stats with seed management.
23 """
25 def __init__(
26 self, dist: rv_continuous, *args: Any, seed: int = 0, **kwds: Any
27 ) -> None:
28 """
29 Parameterized continuous distribution from scipy.stats with seed management.
31 Parameters
32 ----------
33 dist : rv_continuous
34 Continuous distribution from scipy.stats.
35 seed : int, optional
36 Seed for random number generator, by default 0
37 """
38 rv_continuous_frozen.__init__(self, dist, *args, **kwds)
39 Distribution.__init__(self, seed=seed)
42class Normal(ContinuousFrozenDistribution):
43 """
44 A Normal distribution.
46 Parameters
47 ----------
48 mu : float or [float]
49 One or more means. Number of elements T in mu determines number
50 of rows of output.
51 sigma : float or [float]
52 One or more standard deviations. Number of elements T in sigma
53 determines number of rows of output.
54 seed : int
55 Seed for random number generator.
56 """
58 def __init__(self, mu=0.0, sigma=1.0, seed=0):
59 self.mu = np.asarray(mu)
60 self.sigma = np.asarray(sigma)
62 if self.mu.size != self.sigma.size:
63 raise AttributeError(
64 "mu and sigma must be of same size, are %s, %s"
65 % ((self.mu.size), (self.sigma.size))
66 )
68 super().__init__(norm, loc=mu, scale=sigma, seed=seed)
69 self.infimum = -np.inf * np.ones(self.mu.size)
70 self.supremum = np.inf * np.ones(self.mu.size)
72 def discretize(self, N, method="hermite", endpoints=False):
73 """
74 For normal distributions, the Gauss-Hermite quadrature rule is
75 used as the default method for discretization.
76 """
78 return super().discretize(N, method, endpoints)
80 def _approx_hermite(self, N, endpoints=False):
81 """
82 Returns a discrete approximation of this distribution
83 using the Gauss-Hermite quadrature rule.
85 Parameters
86 ----------
87 N : int
88 Number of discrete points to approximate the distribution.
90 Returns
91 -------
92 DiscreteDistribution
93 Discrete approximation of this distribution.
94 """
96 x, w = np.polynomial.hermite.hermgauss(N)
97 # normalize w
98 pmv = w * np.pi**-0.5
99 # correct x
100 atoms = np.sqrt(2.0) * self.sigma * x + self.mu
102 limit = {"dist": self, "method": "hermite", "N": N, "endpoints": endpoints}
104 return DiscreteDistribution(
105 pmv,
106 atoms,
107 seed=self._rng.integers(0, 2**31 - 1, dtype="int32"),
108 limit=limit,
109 )
111 def _approx_equiprobable(self, N, endpoints=False):
112 """
113 Returns a discrete equiprobable approximation of this distribution.
115 Parameters
116 ----------
117 N : int
118 Number of discrete points to approximate the distribution.
120 Returns
121 -------
122 DiscreteDistribution
123 Discrete approximation of this distribution.
124 """
126 CDF = np.linspace(0, 1, N + 1)
127 lims = norm.ppf(CDF)
128 pdf = norm.pdf(lims)
130 # Find conditional means using Mills's ratio
131 pmv = np.diff(CDF)
132 atoms = self.mu - np.diff(pdf) / pmv * self.sigma
134 limit = {"dist": self, "method": "equiprobable", "N": N, "endpoints": endpoints}
136 return DiscreteDistribution(
137 pmv,
138 atoms,
139 seed=self._rng.integers(0, 2**31 - 1, dtype="int32"),
140 limit=limit,
141 )
144class Lognormal(ContinuousFrozenDistribution):
145 """
146 A Lognormal distribution
148 Parameters
149 ----------
150 mu : float or [float]
151 One or more means of underlying normal distribution.
152 Number of elements T in mu determines number of rows of output.
153 sigma : float or [float]
154 One or more standard deviations of underlying normal distribution.
155 Number of elements T in sigma determines number of rows of output.
156 seed : int
157 Seed for random number generator.
158 """
160 def __new__(
161 cls,
162 mu: Union[float, np.ndarray] = 0.0,
163 sigma: Union[float, np.ndarray] = 1.0,
164 seed: Optional[int] = 0,
165 mean=None,
166 std=None,
167 ):
168 """
169 Create a new Lognormal distribution. If sigma is zero, return a
170 DiscreteDistribution with a single atom at exp(mu).
172 Parameters
173 ----------
174 mu : Union[float, np.ndarray], optional
175 Mean of underlying normal distribution, by default 0.0
176 sigma : Union[float, np.ndarray], optional
177 Standard deviation of underlying normal distribution, by default 1.0
178 seed : Optional[int], optional
179 Seed for random number generator, by default None
180 mean: optional
181 For alternative mean/std parameterization, the mean of the lognormal distribution
182 std: optional
183 For alternative mean/std parameterization, the standard deviation of the lognormal distribution
185 Returns
186 -------
187 Lognormal or DiscreteDistribution
188 Lognormal distribution or DiscreteDistribution with a single atom.
189 """
191 if sigma == 0:
192 # If sigma is zero, return a DiscreteDistribution with a single atom
193 if mean is not None:
194 return DiscreteDistribution([1.0], mean, seed=seed)
195 return DiscreteDistribution([1.0], [np.exp(mu)], seed=seed)
197 return super().__new__(cls)
199 def __init__(
200 self,
201 mu: Union[float, np.ndarray] = 0.0,
202 sigma: Union[float, np.ndarray] = 1.0,
203 seed: Optional[int] = 0,
204 mean=None,
205 std=None,
206 ):
207 if mean is not None and std is not None:
208 mean_squared = mean**2
209 variance = std**2
210 mu = np.log(mean / (np.sqrt(1.0 + variance / mean_squared)))
211 sigma = np.sqrt(np.log(1.0 + variance / mean_squared))
213 if mean is not None and std is None:
214 mu = np.log(mean) - sigma**2 / 2
216 self.mu = np.asarray(mu)
217 self.sigma = np.asarray(sigma)
219 if self.mu.size != self.sigma.size:
220 raise AttributeError(
221 "mu and sigma must be of same size, are %s, %s"
222 % ((self.mu.size), (self.sigma.size))
223 )
225 # Set up the RNG
226 super().__init__(lognorm, s=self.sigma, scale=np.exp(self.mu), loc=0, seed=seed)
227 self.infimum = np.array([0.0])
228 self.supremum = np.array([np.inf])
230 def _approx_equiprobable(
231 self, N, endpoints=False, tail_N=0, tail_bound=None, tail_order=np.e
232 ):
233 """
234 Construct a discrete approximation to a lognormal distribution with underlying
235 normal distribution N(mu,sigma). Makes an equiprobable distribution by
236 default, but user can optionally request augmented tails with exponentially
237 sized point masses. This can improve solution accuracy in some models.
239 TODO: add endpoints option
241 Parameters
242 ----------
243 N: int
244 Number of discrete points in the "main part" of the approximation.
245 tail_N: int
246 Number of points in each "tail part" of the approximation; 0 = no tail.
247 tail_bound: [float]
248 CDF boundaries of the tails vs main portion; tail_bound[0] is the lower
249 tail bound, tail_bound[1] is the upper tail bound. Inoperative when
250 tail_N = 0. Can make "one tailed" approximations with 0.0 or 1.0.
251 tail_order: float
252 Factor by which consecutive point masses in a "tail part" differ in
253 probability. Should be >= 1 for sensible spacing.
255 Returns
256 -------
257 d : DiscreteDistribution
258 Probability associated with each point in array of discrete
259 points for discrete probability mass function.
260 """
261 tail_bound = tail_bound if tail_bound is not None else [0.02, 0.98]
263 # Handle the trivial case first
264 if self.sigma == 0.0:
265 pmv = np.ones(N) / N
266 atoms = np.exp(self.mu) * np.ones(N)
268 else:
269 # Find the CDF boundaries of each segment
270 if tail_N > 0:
271 lo_cut = tail_bound[0]
272 hi_cut = tail_bound[1]
273 else:
274 lo_cut = 0.0
275 hi_cut = 1.0
276 inner_size = hi_cut - lo_cut
277 inner_CDF_vals = [
278 lo_cut + x * N ** (-1.0) * inner_size for x in range(1, N)
279 ]
280 if inner_size < 1.0:
281 scale = 1.0 / tail_order
282 mag = (1.0 - scale**tail_N) / (1.0 - scale)
283 lower_CDF_vals = [0.0]
284 if lo_cut > 0.0:
285 for x in range(tail_N - 1, -1, -1):
286 lower_CDF_vals.append(lower_CDF_vals[-1] + lo_cut * scale**x / mag)
287 upper_CDF_vals = [hi_cut]
288 if hi_cut < 1.0:
289 for x in range(tail_N):
290 upper_CDF_vals.append(
291 upper_CDF_vals[-1] + (1.0 - hi_cut) * scale**x / mag
292 )
293 CDF_vals = np.array(lower_CDF_vals + inner_CDF_vals + upper_CDF_vals)
294 CDF_vals[-1] = 1.0
295 CDF_vals[0] = 0.0 # somehow these need fixing sometimes
297 # Calculate probability masses for each node
298 pmv = CDF_vals[1:] - CDF_vals[:-1]
299 pmv /= np.sum(pmv)
301 # Translate the CDF values to z-scores (stdevs from mean), then get q-scores
302 z_cuts = norm.ppf(CDF_vals)
303 q_cuts = (z_cuts - self.sigma) / np.sqrt(2)
305 # Evaluate the (complementary) error function at the q values
306 erf_q = erfc(q_cuts)
307 erf_q_neg = erfc(-q_cuts)
309 # Evaluate the base for the conditional expectations
310 vals_base = erf_q[:-1] - erf_q[1:]
311 these = q_cuts[:-1] < -2.0 # flag low q values and use the *other* version
312 vals_base[these] = erf_q_neg[1:][these] - erf_q_neg[:-1][these]
314 # Make and apply the normalization factor and probability weights
315 norm_fac = 0.5 * np.exp(self.mu + 0.5 * self.sigma**2) / pmv
316 atoms = vals_base * norm_fac
318 limit = {
319 "dist": self,
320 "method": "equiprobable",
321 "N": N,
322 "endpoints": endpoints,
323 "tail_N": tail_N,
324 "tail_bound": tail_bound,
325 "tail_order": tail_order,
326 }
328 return DiscreteDistribution(
329 pmv,
330 atoms,
331 seed=self._rng.integers(0, 2**31 - 1, dtype="int32"),
332 limit=limit,
333 )
335 def _approx_hermite(self, N, endpoints=False):
336 """
337 Returns a discrete approximation of this distribution
338 using the Gauss-Hermite quadrature rule.
340 Parameters
341 ----------
342 N : int
343 Number of discrete points to approximate the distribution.
345 Returns
346 -------
347 DiscreteDistribution
348 Discrete approximation of this distribution.
349 """
351 x, w = np.polynomial.hermite.hermgauss(N)
352 # normalize w
353 pmv = w * np.pi**-0.5
354 # correct x
355 atoms = np.exp(np.sqrt(2.0) * self.sigma * x + self.mu)
357 limit = {"dist": self, "method": "hermite", "N": N, "endpoints": endpoints}
359 return DiscreteDistribution(
360 pmv,
361 atoms,
362 seed=self._rng.integers(0, 2**31 - 1, dtype="int32"),
363 limit=limit,
364 )
366 @classmethod
367 def from_mean_std(cls, mean, std, seed=0):
368 """
369 Construct a LogNormal distribution from its
370 mean and standard deviation.
372 This is unlike the normal constructor for the distribution,
373 which takes the mu and sigma for the normal distribution
374 that is the logarithm of the Log Normal distribution.
376 Parameters
377 ----------
378 mean : float or [float]
379 One or more means. Number of elements T in mu determines number
380 of rows of output.
381 std : float or [float]
382 One or more standard deviations. Number of elements T in sigma
383 determines number of rows of output.
384 seed : int
385 Seed for random number generator.
387 Returns
388 ---------
389 LogNormal
391 """
392 mean_squared = mean**2
393 variance = std**2
394 mu = np.log(mean / (np.sqrt(1.0 + variance / mean_squared)))
395 sigma = np.sqrt(np.log(1.0 + variance / mean_squared))
397 return cls(mu=mu, sigma=sigma, seed=seed)
400LogNormal = Lognormal
403class MeanOneLogNormal(Lognormal):
404 """
405 A Lognormal distribution with mean 1.
406 """
408 def __init__(self, sigma=1.0, seed=0):
409 mu = -0.5 * sigma**2
410 super().__init__(mu=mu, sigma=sigma, seed=seed)
413class Uniform(ContinuousFrozenDistribution):
414 """
415 A Uniform distribution.
417 Parameters
418 ----------
419 bot : float or [float]
420 One or more bottom values.
421 Number of elements T in mu determines number
422 of rows of output.
423 top : float or [float]
424 One or more top values.
425 Number of elements T in top determines number of
426 rows of output.
427 seed : int
428 Seed for random number generator.
429 """
431 def __init__(self, bot=0.0, top=1.0, seed=0):
432 self.bot = np.asarray(bot)
433 self.top = np.asarray(top)
435 super().__init__(uniform, loc=self.bot, scale=self.top - self.bot, seed=seed)
436 self.infimum = np.array([0.0])
437 self.supremum = np.array([np.inf])
439 def _approx_equiprobable(self, N, endpoints=False):
440 """
441 Makes a discrete approximation to this uniform distribution.
443 Parameters
444 ----------
445 N : int
446 The number of points in the discrete approximation.
447 endpoints : bool
448 Whether to include the endpoints in the approximation.
450 Returns
451 -------
452 d : DiscreteDistribution
453 Probability associated with each point in array of discrete
454 points for discrete probability mass function.
455 """
456 pmv = np.ones(N) / float(N)
458 center = (self.top + self.bot) / 2.0
459 width = (self.top - self.bot) / 2.0
460 atoms = center + width * np.linspace(-(N - 1.0) / 2.0, (N - 1.0) / 2.0, N) / (
461 N / 2.0
462 )
464 if endpoints: # insert endpoints with infinitesimally small mass
465 atoms = np.concatenate(([self.bot], atoms, [self.top]))
466 pmv = np.concatenate(([0.0], pmv, [0.0]))
468 limit = {
469 "dist": self,
470 "method": "equiprobable",
471 "N": N,
472 "endpoints": endpoints,
473 }
475 return DiscreteDistribution(
476 pmv,
477 atoms,
478 seed=self._rng.integers(0, 2**31 - 1, dtype="int32"),
479 limit=limit,
480 )
483class Weibull(ContinuousFrozenDistribution):
484 """
485 A Weibull distribution.
487 Parameters
488 ----------
489 scale : float or [float]
490 One or more scales. Number of elements T in scale
491 determines number of
492 rows of output.
493 shape : float or [float]
494 One or more shape parameters. Number of elements T in scale
495 determines number of rows of output.
496 seed : int
497 Seed for random number generator.
498 """
500 def __init__(self, scale=1.0, shape=1.0, seed=0):
501 self.scale = np.asarray(scale)
502 self.shape = np.asarray(shape)
504 # Set up the RNG
505 super().__init__(weibull_min, c=shape, scale=scale, seed=seed)
506 self.infimum = np.array([0.0])
507 self.supremum = np.array([np.inf])