Coverage for HARK / distributions / continuous.py: 90%
155 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-25 05:22 +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 __new__(
59 cls,
60 mu: Union[float, np.ndarray] = 0.0,
61 sigma: Union[float, np.ndarray] = 1.0,
62 seed: Optional[int] = None,
63 ):
64 """
65 Create a new Normal distribution. If sigma is zero, return a
66 DiscreteDistribution with a single atom at mu.
67 Parameters
68 ----------
69 mu : Union[float, np.ndarray], optional
70 Mean of normal distribution, by default 0.0
71 sigma : Union[float, np.ndarray], optional
72 Standard deviation of normal distribution, by default 1.0
73 seed : Optional[int], optional
74 Seed for random number generator, by default None
75 Returns
76 -------
77 Normal or DiscreteDistribution
78 Normal distribution or DiscreteDistribution with a single atom.
79 """
81 if sigma == 0:
82 # If sigma is zero, return a DiscreteDistribution with a single atom
83 return DiscreteDistribution([1.0], mu, seed=seed)
85 return super().__new__(cls)
87 def __init__(self, mu=0.0, sigma=1.0, seed=None):
88 self.mu = np.asarray(mu)
89 self.sigma = np.asarray(sigma)
91 if self.mu.size != self.sigma.size:
92 raise AttributeError(
93 f"'mu' and 'sigma' must be of the same size. Instead 'mu' is of "
94 f"size {self.mu.size} and 'sigma' is of size {self.sigma.size}."
95 )
97 super().__init__(norm, loc=mu, scale=sigma, seed=seed)
98 self.infimum = -np.inf * np.ones(self.mu.size)
99 self.supremum = np.inf * np.ones(self.mu.size)
101 def discretize(self, N, method="hermite", endpoints=False):
102 """
103 For normal distributions, the Gauss-Hermite quadrature rule is
104 used as the default method for discretization.
105 """
107 return super().discretize(N, method, endpoints)
109 def _approx_hermite(self, N, endpoints=False):
110 """
111 Returns a discrete approximation of this distribution
112 using the Gauss-Hermite quadrature rule.
114 Parameters
115 ----------
116 N : int
117 Number of discrete points to approximate the distribution.
119 Returns
120 -------
121 DiscreteDistribution
122 Discrete approximation of this distribution.
123 """
125 x, w = np.polynomial.hermite.hermgauss(N)
126 # normalize w
127 pmv = w * np.pi**-0.5
128 # correct x
129 atoms = np.sqrt(2.0) * self.sigma * x + self.mu
131 if endpoints:
132 atoms = np.r_[-np.inf, atoms, np.inf]
133 pmv = np.r_[0.0, pmv, 0.0]
135 limit = {"dist": self, "method": "hermite", "N": N, "endpoints": endpoints}
137 return DiscreteDistribution(
138 pmv,
139 atoms,
140 seed=self.random_seed(),
141 limit=limit,
142 )
144 def _approx_equiprobable(self, N, endpoints=False):
145 """
146 Returns a discrete equiprobable approximation of this distribution.
148 Parameters
149 ----------
150 N : int
151 Number of discrete points to approximate the distribution.
153 Returns
154 -------
155 DiscreteDistribution
156 Discrete approximation of this distribution.
157 """
159 CDF = np.linspace(0, 1, N + 1)
160 lims = norm.ppf(CDF)
161 pdf = norm.pdf(lims)
163 # Find conditional means using Mills's ratio
164 pmv = np.diff(CDF)
165 atoms = self.mu - np.diff(pdf) / pmv * self.sigma
167 if endpoints:
168 atoms = np.r_[-np.inf, atoms, np.inf]
169 pmv = np.r_[0.0, pmv, 0.0]
171 limit = {"dist": self, "method": "equiprobable", "N": N, "endpoints": endpoints}
173 return DiscreteDistribution(
174 pmv,
175 atoms,
176 seed=self._rng.integers(0, 2**31 - 1, dtype="int32"),
177 limit=limit,
178 )
181class Lognormal(ContinuousFrozenDistribution):
182 """
183 A Lognormal distribution
185 Parameters
186 ----------
187 mu : float or [float]
188 One or more means of underlying normal distribution.
189 Number of elements T in mu determines number of rows of output.
190 sigma : float or [float]
191 One or more standard deviations of underlying normal distribution.
192 Number of elements T in sigma determines number of rows of output.
193 seed : int
194 Seed for random number generator.
195 """
197 def __new__(
198 cls,
199 mu: Union[float, np.ndarray] = 0.0,
200 sigma: Union[float, np.ndarray] = 1.0,
201 seed: Optional[int] = None,
202 mean=None,
203 std=None,
204 ):
205 """
206 Create a new Lognormal distribution. If sigma is zero, return a
207 DiscreteDistribution with a single atom at exp(mu).
209 Parameters
210 ----------
211 mu : Union[float, np.ndarray], optional
212 Mean of underlying normal distribution, by default 0.0
213 sigma : Union[float, np.ndarray], optional
214 Standard deviation of underlying normal distribution, by default 1.0
215 seed : Optional[int], optional
216 Seed for random number generator, by default None
217 mean: optional
218 For alternative mean/std parameterization, the mean of the lognormal distribution
219 std: optional
220 For alternative mean/std parameterization, the standard deviation of the lognormal distribution
222 Returns
223 -------
224 Lognormal or DiscreteDistribution
225 Lognormal distribution or DiscreteDistribution with a single atom.
226 """
228 if sigma == 0:
229 # If sigma is zero, return a DiscreteDistribution with a single atom
230 if mean is not None:
231 return DiscreteDistribution([1.0], mean, seed=seed)
232 return DiscreteDistribution([1.0], [np.exp(mu)], seed=seed)
234 return super().__new__(cls)
236 def __init__(
237 self,
238 mu: Union[float, np.ndarray] = 0.0,
239 sigma: Union[float, np.ndarray] = 1.0,
240 seed: Optional[int] = None,
241 mean=None,
242 std=None,
243 ):
244 if mean is not None and std is not None:
245 mean_squared = mean**2
246 variance = std**2
247 mu = np.log(mean / (np.sqrt(1.0 + variance / mean_squared)))
248 sigma = np.sqrt(np.log(1.0 + variance / mean_squared))
250 if mean is not None and std is None:
251 mu = np.log(mean) - sigma**2 / 2
253 self.mu = np.asarray(mu)
254 self.sigma = np.asarray(sigma)
256 if self.mu.size != self.sigma.size:
257 raise AttributeError(
258 "mu and sigma must be of same size, are %s, %s"
259 % ((self.mu.size), (self.sigma.size))
260 )
262 # Set up the RNG
263 super().__init__(lognorm, s=self.sigma, scale=np.exp(self.mu), loc=0, seed=seed)
264 self.infimum = np.array([0.0])
265 self.supremum = np.array([np.inf])
267 def _approx_equiprobable(
268 self, N, endpoints=False, tail_N=0, tail_bound=None, tail_order=np.e
269 ):
270 """
271 Construct a discrete approximation to a lognormal distribution with underlying
272 normal distribution N(mu,sigma). Makes an equiprobable distribution by
273 default, but user can optionally request augmented tails with exponentially
274 sized point masses. This can improve solution accuracy in some models.
276 Parameters
277 ----------
278 N: int
279 Number of discrete points in the "main part" of the approximation.
280 tail_N: int
281 Number of points in each "tail part" of the approximation; 0 = no tail.
282 tail_bound: [float]
283 CDF boundaries of the tails vs main portion; tail_bound[0] is the lower
284 tail bound, tail_bound[1] is the upper tail bound. Inoperative when
285 tail_N = 0. Can make "one tailed" approximations with 0.0 or 1.0.
286 tail_order: float
287 Factor by which consecutive point masses in a "tail part" differ in
288 probability. Should be >= 1 for sensible spacing.
290 Returns
291 -------
292 d : DiscreteDistribution
293 Probability associated with each point in array of discrete
294 points for discrete probability mass function.
295 """
296 tail_bound = tail_bound or [0.02, 0.98]
298 # Handle the trivial case first
299 if self.sigma == 0.0:
300 pmv = np.ones(N) / N
301 atoms = np.exp(self.mu) * np.ones(N)
303 else:
304 # Find the CDF boundaries of each segment
305 if tail_N > 0:
306 lo_cut = tail_bound[0]
307 hi_cut = tail_bound[1]
308 else:
309 lo_cut = 0.0
310 hi_cut = 1.0
311 inner_size = hi_cut - lo_cut
312 inner_CDF_vals = [
313 lo_cut + x * N ** (-1.0) * inner_size for x in range(1, N)
314 ]
315 if inner_size < 1.0:
316 scale = 1.0 / tail_order
317 mag = (1.0 - scale**tail_N) / (1.0 - scale)
318 lower_CDF_vals = [0.0]
319 if lo_cut > 0.0:
320 for x in range(tail_N - 1, -1, -1):
321 lower_CDF_vals.append(lower_CDF_vals[-1] + lo_cut * scale**x / mag)
322 upper_CDF_vals = [hi_cut]
323 if hi_cut < 1.0:
324 for x in range(tail_N):
325 upper_CDF_vals.append(
326 upper_CDF_vals[-1] + (1.0 - hi_cut) * scale**x / mag
327 )
328 CDF_vals = np.array(lower_CDF_vals + inner_CDF_vals + upper_CDF_vals)
329 CDF_vals[-1] = 1.0
330 CDF_vals[0] = 0.0 # somehow these need fixing sometimes
332 # Calculate probability masses for each node
333 pmv = CDF_vals[1:] - CDF_vals[:-1]
334 pmv /= np.sum(pmv)
336 # Translate the CDF values to z-scores (stdevs from mean), then get q-scores
337 z_cuts = norm.ppf(CDF_vals)
338 q_cuts = (z_cuts - self.sigma) / np.sqrt(2)
340 # Evaluate the (complementary) error function at the q values
341 erf_q = erfc(q_cuts)
342 erf_q_neg = erfc(-q_cuts)
344 # Evaluate the base for the conditional expectations
345 vals_base = erf_q[:-1] - erf_q[1:]
346 these = q_cuts[:-1] < -2.0 # flag low q values and use the *other* version
347 vals_base[these] = erf_q_neg[1:][these] - erf_q_neg[:-1][these]
349 # Make and apply the normalization factor and probability weights
350 norm_fac = 0.5 * np.exp(self.mu + 0.5 * self.sigma**2) / pmv
351 atoms = vals_base * norm_fac
353 if endpoints:
354 atoms = np.r_[0.0, atoms, np.inf]
355 pmv = np.r_[0.0, pmv, 0.0]
357 limit = {
358 "dist": self,
359 "method": "equiprobable",
360 "N": N,
361 "endpoints": endpoints,
362 "tail_N": tail_N,
363 "tail_bound": tail_bound,
364 "tail_order": tail_order,
365 }
367 return DiscreteDistribution(
368 pmv,
369 atoms,
370 seed=self.random_seed(),
371 limit=limit,
372 )
374 def _approx_hermite(self, N, endpoints=False):
375 """
376 Returns a discrete approximation of this distribution
377 using the Gauss-Hermite quadrature rule.
379 Parameters
380 ----------
381 N : int
382 Number of discrete points to approximate the distribution.
384 Returns
385 -------
386 DiscreteDistribution
387 Discrete approximation of this distribution.
388 """
390 x, w = np.polynomial.hermite.hermgauss(N)
391 # normalize w
392 pmv = w * np.pi**-0.5
393 # correct x
394 atoms = np.exp(np.sqrt(2.0) * self.sigma * x + self.mu)
396 if endpoints:
397 atoms = np.r_[0.0, atoms, np.inf]
398 pmv = np.r_[0.0, pmv, 0.0]
400 limit = {"dist": self, "method": "hermite", "N": N, "endpoints": endpoints}
402 return DiscreteDistribution(
403 pmv,
404 atoms,
405 seed=self.random_seed(),
406 limit=limit,
407 )
409 @classmethod
410 def from_mean_std(cls, mean, std, seed=None):
411 """
412 Construct a LogNormal distribution from its
413 mean and standard deviation.
415 This is unlike the normal constructor for the distribution,
416 which takes the mu and sigma for the normal distribution
417 that is the logarithm of the Log Normal distribution.
419 Parameters
420 ----------
421 mean : float or [float]
422 One or more means. Number of elements T in mu determines number
423 of rows of output.
424 std : float or [float]
425 One or more standard deviations. Number of elements T in sigma
426 determines number of rows of output.
427 seed : int
428 Seed for random number generator.
430 Returns
431 ---------
432 LogNormal
434 """
435 mean_squared = mean**2
436 variance = std**2
437 mu = np.log(mean / (np.sqrt(1.0 + variance / mean_squared)))
438 sigma = np.sqrt(np.log(1.0 + variance / mean_squared))
440 return cls(mu=mu, sigma=sigma, seed=seed)
443LogNormal = Lognormal
446class MeanOneLogNormal(Lognormal):
447 """
448 A Lognormal distribution with mean 1.
449 """
451 def __init__(self, sigma=1.0, seed=0):
452 mu = -0.5 * sigma**2
453 super().__init__(mu=mu, sigma=sigma, seed=seed)
456class Uniform(ContinuousFrozenDistribution):
457 """
458 A Uniform distribution.
460 Parameters
461 ----------
462 bot : float or [float]
463 One or more bottom values.
464 Number of elements T in mu determines number
465 of rows of output.
466 top : float or [float]
467 One or more top values.
468 Number of elements T in top determines number of
469 rows of output.
470 seed : int
471 Seed for random number generator.
472 """
474 def __init__(self, bot=0.0, top=1.0, seed=None):
475 self.bot = np.asarray(bot)
476 self.top = np.asarray(top)
478 super().__init__(uniform, loc=self.bot, scale=self.top - self.bot, seed=seed)
479 self.infimum = np.array([0.0])
480 self.supremum = np.array([np.inf])
482 def _approx_equiprobable(self, N, endpoints=False):
483 """
484 Makes a discrete approximation to this uniform distribution.
486 Parameters
487 ----------
488 N : int
489 The number of points in the discrete approximation.
490 endpoints : bool
491 Whether to include the endpoints in the approximation.
493 Returns
494 -------
495 d : DiscreteDistribution
496 Probability associated with each point in array of discrete
497 points for discrete probability mass function.
498 """
499 pmv = np.ones(N) / float(N)
501 center = (self.top + self.bot) / 2.0
502 width = (self.top - self.bot) / 2.0
503 atoms = center + width * np.linspace(-(N - 1.0) / 2.0, (N - 1.0) / 2.0, N) / (
504 N / 2.0
505 )
507 if endpoints: # insert endpoints with infinitesimally small mass
508 atoms = np.concatenate(([self.bot], atoms, [self.top]))
509 pmv = np.concatenate(([0.0], pmv, [0.0]))
511 limit = {
512 "dist": self,
513 "method": "equiprobable",
514 "N": N,
515 "endpoints": endpoints,
516 }
518 return DiscreteDistribution(
519 pmv,
520 atoms,
521 seed=self.random_seed(),
522 limit=limit,
523 )
526class Weibull(ContinuousFrozenDistribution):
527 """
528 A Weibull distribution.
530 Parameters
531 ----------
532 scale : float or [float]
533 One or more scales. Number of elements T in scale
534 determines number of
535 rows of output.
536 shape : float or [float]
537 One or more shape parameters. Number of elements T in scale
538 determines number of rows of output.
539 seed : int
540 Seed for random number generator.
541 """
543 def __init__(self, scale=1.0, shape=1.0, seed=None):
544 self.scale = np.asarray(scale)
545 self.shape = np.asarray(shape)
547 # Set up the RNG
548 super().__init__(weibull_min, c=shape, scale=scale, seed=seed)
549 self.infimum = np.array([0.0])
550 self.supremum = np.array([np.inf])