blob: 70c48d605d1969b5ab158ae54035f26a703c5da1 [file] [log] [blame]
Larry Hastingsf5e987b2013-10-19 11:50:09 -07001"""
2Basic statistics module.
3
4This module provides functions for calculating statistics of data, including
5averages, variance, and standard deviation.
6
7Calculating averages
8--------------------
9
Raymond Hettinger9013ccf2019-04-23 00:06:35 -070010================== ==================================================
Larry Hastingsf5e987b2013-10-19 11:50:09 -070011Function Description
Raymond Hettinger9013ccf2019-04-23 00:06:35 -070012================== ==================================================
Larry Hastingsf5e987b2013-10-19 11:50:09 -070013mean Arithmetic mean (average) of data.
Raymond Hettinger72800482019-04-23 01:35:16 -070014fmean Fast, floating point arithmetic mean.
Raymond Hettinger6463ba32019-04-07 09:20:03 -070015geometric_mean Geometric mean of data.
Steven D'Apranoa474afd2016-08-09 12:49:01 +100016harmonic_mean Harmonic mean of data.
Larry Hastingsf5e987b2013-10-19 11:50:09 -070017median Median (middle value) of data.
18median_low Low median of data.
19median_high High median of data.
20median_grouped Median, or 50th percentile, of grouped data.
21mode Mode (most common value) of data.
Raymond Hettinger6463ba32019-04-07 09:20:03 -070022multimode List of modes (most common values of data).
Raymond Hettinger9013ccf2019-04-23 00:06:35 -070023quantiles Divide data into intervals with equal probability.
24================== ==================================================
Larry Hastingsf5e987b2013-10-19 11:50:09 -070025
26Calculate the arithmetic mean ("the average") of data:
27
28>>> mean([-1.0, 2.5, 3.25, 5.75])
292.625
30
31
32Calculate the standard median of discrete data:
33
34>>> median([2, 3, 4, 5])
353.5
36
37
38Calculate the median, or 50th percentile, of data grouped into class intervals
39centred on the data values provided. E.g. if your data points are rounded to
40the nearest whole number:
41
42>>> median_grouped([2, 2, 3, 3, 3, 4]) #doctest: +ELLIPSIS
432.8333333333...
44
45This should be interpreted in this way: you have two data points in the class
46interval 1.5-2.5, three data points in the class interval 2.5-3.5, and one in
47the class interval 3.5-4.5. The median of these data points is 2.8333...
48
49
50Calculating variability or spread
51---------------------------------
52
53================== =============================================
54Function Description
55================== =============================================
56pvariance Population variance of data.
57variance Sample variance of data.
58pstdev Population standard deviation of data.
59stdev Sample standard deviation of data.
60================== =============================================
61
62Calculate the standard deviation of sample data:
63
64>>> stdev([2.5, 3.25, 5.5, 11.25, 11.75]) #doctest: +ELLIPSIS
654.38961843444...
66
67If you have previously calculated the mean, you can pass it as the optional
68second argument to the four "spread" functions to avoid recalculating it:
69
70>>> data = [1, 2, 2, 4, 4, 4, 5, 6]
71>>> mu = mean(data)
72>>> pvariance(data, mu)
732.5
74
75
76Exceptions
77----------
78
79A single exception is defined: StatisticsError is a subclass of ValueError.
80
81"""
82
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -070083__all__ = [
84 'NormalDist',
85 'StatisticsError',
86 'fmean',
87 'geometric_mean',
88 'harmonic_mean',
89 'mean',
90 'median',
91 'median_grouped',
92 'median_high',
93 'median_low',
94 'mode',
95 'multimode',
96 'pstdev',
97 'pvariance',
98 'quantiles',
99 'stdev',
100 'variance',
101]
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700102
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700103import math
Steven D'Apranoa474afd2016-08-09 12:49:01 +1000104import numbers
Raymond Hettinger11c79532019-02-23 14:44:07 -0800105import random
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700106
107from fractions import Fraction
108from decimal import Decimal
Victor Stinnerd6debb22017-03-27 16:05:26 +0200109from itertools import groupby
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000110from bisect import bisect_left, bisect_right
Raymond Hettinger318d5372019-03-06 22:59:40 -0800111from math import hypot, sqrt, fabs, exp, erf, tau, log, fsum
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700112from operator import itemgetter
113from collections import Counter
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700114
115# === Exceptions ===
116
117class StatisticsError(ValueError):
118 pass
119
120
121# === Private utilities ===
122
123def _sum(data, start=0):
Steven D'Apranob28c3272015-12-01 19:59:53 +1100124 """_sum(data [, start]) -> (type, sum, count)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700125
Steven D'Apranob28c3272015-12-01 19:59:53 +1100126 Return a high-precision sum of the given numeric data as a fraction,
127 together with the type to be converted to and the count of items.
128
129 If optional argument ``start`` is given, it is added to the total.
130 If ``data`` is empty, ``start`` (defaulting to 0) is returned.
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700131
132
133 Examples
134 --------
135
136 >>> _sum([3, 2.25, 4.5, -0.5, 1.0], 0.75)
Benjamin Petersonab078e92016-07-13 21:13:29 -0700137 (<class 'float'>, Fraction(11, 1), 5)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700138
139 Some sources of round-off error will be avoided:
140
Steven D'Apranoa474afd2016-08-09 12:49:01 +1000141 # Built-in sum returns zero.
142 >>> _sum([1e50, 1, -1e50] * 1000)
Benjamin Petersonab078e92016-07-13 21:13:29 -0700143 (<class 'float'>, Fraction(1000, 1), 3000)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700144
145 Fractions and Decimals are also supported:
146
147 >>> from fractions import Fraction as F
148 >>> _sum([F(2, 3), F(7, 5), F(1, 4), F(5, 6)])
Benjamin Petersonab078e92016-07-13 21:13:29 -0700149 (<class 'fractions.Fraction'>, Fraction(63, 20), 4)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700150
151 >>> from decimal import Decimal as D
152 >>> data = [D("0.1375"), D("0.2108"), D("0.3061"), D("0.0419")]
153 >>> _sum(data)
Benjamin Petersonab078e92016-07-13 21:13:29 -0700154 (<class 'decimal.Decimal'>, Fraction(6963, 10000), 4)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700155
Nick Coghlan73afe2a2014-02-08 19:58:04 +1000156 Mixed types are currently treated as an error, except that int is
157 allowed.
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700158 """
Steven D'Apranob28c3272015-12-01 19:59:53 +1100159 count = 0
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700160 n, d = _exact_ratio(start)
Steven D'Apranob28c3272015-12-01 19:59:53 +1100161 partials = {d: n}
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700162 partials_get = partials.get
Steven D'Apranob28c3272015-12-01 19:59:53 +1100163 T = _coerce(int, type(start))
164 for typ, values in groupby(data, type):
165 T = _coerce(T, typ) # or raise TypeError
166 for n,d in map(_exact_ratio, values):
167 count += 1
168 partials[d] = partials_get(d, 0) + n
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700169 if None in partials:
Steven D'Apranob28c3272015-12-01 19:59:53 +1100170 # The sum will be a NAN or INF. We can ignore all the finite
171 # partials, and just look at this special one.
172 total = partials[None]
173 assert not _isfinite(total)
174 else:
175 # Sum all the partial sums using builtin sum.
176 # FIXME is this faster if we sum them in order of the denominator?
177 total = sum(Fraction(n, d) for d, n in sorted(partials.items()))
178 return (T, total, count)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700179
180
Steven D'Apranob28c3272015-12-01 19:59:53 +1100181def _isfinite(x):
182 try:
183 return x.is_finite() # Likely a Decimal.
184 except AttributeError:
185 return math.isfinite(x) # Coerces to float first.
186
187
188def _coerce(T, S):
189 """Coerce types T and S to a common type, or raise TypeError.
190
191 Coercion rules are currently an implementation detail. See the CoerceTest
192 test class in test_statistics for details.
193 """
194 # See http://bugs.python.org/issue24068.
195 assert T is not bool, "initial type T is bool"
196 # If the types are the same, no need to coerce anything. Put this
197 # first, so that the usual case (no coercion needed) happens as soon
198 # as possible.
199 if T is S: return T
200 # Mixed int & other coerce to the other type.
201 if S is int or S is bool: return T
202 if T is int: return S
203 # If one is a (strict) subclass of the other, coerce to the subclass.
204 if issubclass(S, T): return S
205 if issubclass(T, S): return T
206 # Ints coerce to the other type.
207 if issubclass(T, int): return S
208 if issubclass(S, int): return T
209 # Mixed fraction & float coerces to float (or float subclass).
210 if issubclass(T, Fraction) and issubclass(S, float):
211 return S
212 if issubclass(T, float) and issubclass(S, Fraction):
213 return T
214 # Any other combination is disallowed.
215 msg = "don't know how to coerce %s and %s"
216 raise TypeError(msg % (T.__name__, S.__name__))
Nick Coghlan73afe2a2014-02-08 19:58:04 +1000217
218
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700219def _exact_ratio(x):
Steven D'Apranob28c3272015-12-01 19:59:53 +1100220 """Return Real number x to exact (numerator, denominator) pair.
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700221
222 >>> _exact_ratio(0.25)
223 (1, 4)
224
225 x is expected to be an int, Fraction, Decimal or float.
226 """
227 try:
Steven D'Apranob28c3272015-12-01 19:59:53 +1100228 # Optimise the common case of floats. We expect that the most often
229 # used numeric type will be builtin floats, so try to make this as
230 # fast as possible.
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000231 if type(x) is float or type(x) is Decimal:
Steven D'Apranob28c3272015-12-01 19:59:53 +1100232 return x.as_integer_ratio()
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700233 try:
Steven D'Apranob28c3272015-12-01 19:59:53 +1100234 # x may be an int, Fraction, or Integral ABC.
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700235 return (x.numerator, x.denominator)
236 except AttributeError:
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700237 try:
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000238 # x may be a float or Decimal subclass.
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700239 return x.as_integer_ratio()
240 except AttributeError:
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000241 # Just give up?
242 pass
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700243 except (OverflowError, ValueError):
Steven D'Apranob28c3272015-12-01 19:59:53 +1100244 # float NAN or INF.
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000245 assert not _isfinite(x)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700246 return (x, None)
Steven D'Apranob28c3272015-12-01 19:59:53 +1100247 msg = "can't convert type '{}' to numerator/denominator"
248 raise TypeError(msg.format(type(x).__name__))
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700249
250
Steven D'Apranob28c3272015-12-01 19:59:53 +1100251def _convert(value, T):
252 """Convert value to given numeric type T."""
253 if type(value) is T:
254 # This covers the cases where T is Fraction, or where value is
255 # a NAN or INF (Decimal or float).
256 return value
257 if issubclass(T, int) and value.denominator != 1:
258 T = float
259 try:
260 # FIXME: what do we do if this overflows?
261 return T(value)
262 except TypeError:
263 if issubclass(T, Decimal):
264 return T(value.numerator)/T(value.denominator)
265 else:
266 raise
267
268
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000269def _find_lteq(a, x):
270 'Locate the leftmost value exactly equal to x'
271 i = bisect_left(a, x)
272 if i != len(a) and a[i] == x:
273 return i
274 raise ValueError
275
276
277def _find_rteq(a, l, x):
278 'Locate the rightmost value exactly equal to x'
279 i = bisect_right(a, x, lo=l)
280 if i != (len(a)+1) and a[i-1] == x:
281 return i-1
282 raise ValueError
283
Steven D'Apranoa474afd2016-08-09 12:49:01 +1000284
285def _fail_neg(values, errmsg='negative value'):
286 """Iterate over values, failing if any are less than zero."""
287 for x in values:
288 if x < 0:
289 raise StatisticsError(errmsg)
290 yield x
291
292
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700293# === Measures of central tendency (averages) ===
294
295def mean(data):
296 """Return the sample arithmetic mean of data.
297
298 >>> mean([1, 2, 3, 4, 4])
299 2.8
300
301 >>> from fractions import Fraction as F
302 >>> mean([F(3, 7), F(1, 21), F(5, 3), F(1, 3)])
303 Fraction(13, 21)
304
305 >>> from decimal import Decimal as D
306 >>> mean([D("0.5"), D("0.75"), D("0.625"), D("0.375")])
307 Decimal('0.5625')
308
309 If ``data`` is empty, StatisticsError will be raised.
310 """
311 if iter(data) is data:
312 data = list(data)
313 n = len(data)
314 if n < 1:
315 raise StatisticsError('mean requires at least one data point')
Steven D'Apranob28c3272015-12-01 19:59:53 +1100316 T, total, count = _sum(data)
317 assert count == n
318 return _convert(total/n, T)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700319
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700320
Raymond Hettinger47d99872019-02-21 15:06:29 -0800321def fmean(data):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700322 """Convert data to floats and compute the arithmetic mean.
Raymond Hettinger47d99872019-02-21 15:06:29 -0800323
324 This runs faster than the mean() function and it always returns a float.
Raymond Hettinger47d99872019-02-21 15:06:29 -0800325 If the input dataset is empty, it raises a StatisticsError.
326
327 >>> fmean([3.5, 4.0, 5.25])
328 4.25
Raymond Hettinger47d99872019-02-21 15:06:29 -0800329 """
330 try:
331 n = len(data)
332 except TypeError:
333 # Handle iterators that do not define __len__().
334 n = 0
Raymond Hettinger6c01ebc2019-06-05 07:39:38 -0700335 def count(iterable):
Raymond Hettinger47d99872019-02-21 15:06:29 -0800336 nonlocal n
Raymond Hettinger6c01ebc2019-06-05 07:39:38 -0700337 for n, x in enumerate(iterable, start=1):
338 yield x
339 total = fsum(count(data))
Raymond Hettinger47d99872019-02-21 15:06:29 -0800340 else:
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700341 total = fsum(data)
Raymond Hettinger47d99872019-02-21 15:06:29 -0800342 try:
343 return total / n
344 except ZeroDivisionError:
345 raise StatisticsError('fmean requires at least one data point') from None
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700346
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700347
Raymond Hettinger6463ba32019-04-07 09:20:03 -0700348def geometric_mean(data):
349 """Convert data to floats and compute the geometric mean.
350
351 Raises a StatisticsError if the input dataset is empty,
352 if it contains a zero, or if it contains a negative value.
353
354 No special efforts are made to achieve exact results.
355 (However, this may change in the future.)
356
357 >>> round(geometric_mean([54, 24, 36]), 9)
358 36.0
359 """
360 try:
361 return exp(fmean(map(log, data)))
362 except ValueError:
363 raise StatisticsError('geometric mean requires a non-empty dataset '
364 ' containing positive numbers') from None
365
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700366
Steven D'Apranoa474afd2016-08-09 12:49:01 +1000367def harmonic_mean(data):
368 """Return the harmonic mean of data.
369
370 The harmonic mean, sometimes called the subcontrary mean, is the
371 reciprocal of the arithmetic mean of the reciprocals of the data,
372 and is often appropriate when averaging quantities which are rates
373 or ratios, for example speeds. Example:
374
375 Suppose an investor purchases an equal value of shares in each of
376 three companies, with P/E (price/earning) ratios of 2.5, 3 and 10.
377 What is the average P/E ratio for the investor's portfolio?
378
379 >>> harmonic_mean([2.5, 3, 10]) # For an equal investment portfolio.
380 3.6
381
382 Using the arithmetic mean would give an average of about 5.167, which
383 is too high.
384
385 If ``data`` is empty, or any element is less than zero,
386 ``harmonic_mean`` will raise ``StatisticsError``.
387 """
388 # For a justification for using harmonic mean for P/E ratios, see
389 # http://fixthepitch.pellucid.com/comps-analysis-the-missing-harmony-of-summary-statistics/
390 # http://papers.ssrn.com/sol3/papers.cfm?abstract_id=2621087
391 if iter(data) is data:
392 data = list(data)
393 errmsg = 'harmonic mean does not support negative values'
394 n = len(data)
395 if n < 1:
396 raise StatisticsError('harmonic_mean requires at least one data point')
397 elif n == 1:
398 x = data[0]
399 if isinstance(x, (numbers.Real, Decimal)):
400 if x < 0:
401 raise StatisticsError(errmsg)
402 return x
403 else:
404 raise TypeError('unsupported type')
405 try:
406 T, total, count = _sum(1/x for x in _fail_neg(data, errmsg))
407 except ZeroDivisionError:
408 return 0
409 assert count == n
410 return _convert(n/total, T)
411
412
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700413# FIXME: investigate ways to calculate medians without sorting? Quickselect?
414def median(data):
415 """Return the median (middle value) of numeric data.
416
417 When the number of data points is odd, return the middle data point.
418 When the number of data points is even, the median is interpolated by
419 taking the average of the two middle values:
420
421 >>> median([1, 3, 5])
422 3
423 >>> median([1, 3, 5, 7])
424 4.0
425
426 """
427 data = sorted(data)
428 n = len(data)
429 if n == 0:
430 raise StatisticsError("no median for empty data")
431 if n%2 == 1:
432 return data[n//2]
433 else:
434 i = n//2
435 return (data[i - 1] + data[i])/2
436
437
438def median_low(data):
439 """Return the low median of numeric data.
440
441 When the number of data points is odd, the middle value is returned.
442 When it is even, the smaller of the two middle values is returned.
443
444 >>> median_low([1, 3, 5])
445 3
446 >>> median_low([1, 3, 5, 7])
447 3
448
449 """
450 data = sorted(data)
451 n = len(data)
452 if n == 0:
453 raise StatisticsError("no median for empty data")
454 if n%2 == 1:
455 return data[n//2]
456 else:
457 return data[n//2 - 1]
458
459
460def median_high(data):
461 """Return the high median of data.
462
463 When the number of data points is odd, the middle value is returned.
464 When it is even, the larger of the two middle values is returned.
465
466 >>> median_high([1, 3, 5])
467 3
468 >>> median_high([1, 3, 5, 7])
469 5
470
471 """
472 data = sorted(data)
473 n = len(data)
474 if n == 0:
475 raise StatisticsError("no median for empty data")
476 return data[n//2]
477
478
479def median_grouped(data, interval=1):
Zachary Waredf2660e2015-10-27 22:00:41 -0500480 """Return the 50th percentile (median) of grouped continuous data.
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700481
482 >>> median_grouped([1, 2, 2, 3, 4, 4, 4, 4, 4, 5])
483 3.7
484 >>> median_grouped([52, 52, 53, 54])
485 52.5
486
487 This calculates the median as the 50th percentile, and should be
488 used when your data is continuous and grouped. In the above example,
489 the values 1, 2, 3, etc. actually represent the midpoint of classes
490 0.5-1.5, 1.5-2.5, 2.5-3.5, etc. The middle value falls somewhere in
491 class 3.5-4.5, and interpolation is used to estimate it.
492
493 Optional argument ``interval`` represents the class interval, and
494 defaults to 1. Changing the class interval naturally will change the
495 interpolated 50th percentile value:
496
497 >>> median_grouped([1, 3, 3, 5, 7], interval=1)
498 3.25
499 >>> median_grouped([1, 3, 3, 5, 7], interval=2)
500 3.5
501
502 This function does not check whether the data points are at least
503 ``interval`` apart.
504 """
505 data = sorted(data)
506 n = len(data)
507 if n == 0:
508 raise StatisticsError("no median for empty data")
509 elif n == 1:
510 return data[0]
511 # Find the value at the midpoint. Remember this corresponds to the
512 # centre of the class interval.
513 x = data[n//2]
514 for obj in (x, interval):
515 if isinstance(obj, (str, bytes)):
516 raise TypeError('expected number but got %r' % obj)
517 try:
518 L = x - interval/2 # The lower limit of the median interval.
519 except TypeError:
520 # Mixed type. For now we just coerce to float.
521 L = float(x) - float(interval)/2
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000522
523 # Uses bisection search to search for x in data with log(n) time complexity
Martin Panterf1579822016-05-26 06:03:33 +0000524 # Find the position of leftmost occurrence of x in data
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000525 l1 = _find_lteq(data, x)
Martin Panterf1579822016-05-26 06:03:33 +0000526 # Find the position of rightmost occurrence of x in data[l1...len(data)]
Steven D'Aprano3b06e242016-05-05 03:54:29 +1000527 # Assuming always l1 <= l2
528 l2 = _find_rteq(data, l1, x)
529 cf = l1
530 f = l2 - l1 + 1
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700531 return L + interval*(n/2 - cf)/f
532
533
534def mode(data):
535 """Return the most common data point from discrete or nominal data.
536
537 ``mode`` assumes discrete data, and returns a single value. This is the
538 standard treatment of the mode as commonly taught in schools:
539
Raymond Hettingere4810b22019-09-05 00:18:47 -0700540 >>> mode([1, 1, 2, 3, 3, 3, 3, 4])
541 3
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700542
543 This also works with nominal (non-numeric) data:
544
Raymond Hettingere4810b22019-09-05 00:18:47 -0700545 >>> mode(["red", "blue", "blue", "red", "green", "red", "red"])
546 'red'
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700547
Raymond Hettingere4810b22019-09-05 00:18:47 -0700548 If there are multiple modes with same frequency, return the first one
549 encountered:
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700550
551 >>> mode(['red', 'red', 'green', 'blue', 'blue'])
552 'red'
553
554 If *data* is empty, ``mode``, raises StatisticsError.
555
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700556 """
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700557 data = iter(data)
558 try:
559 return Counter(data).most_common(1)[0][0]
560 except IndexError:
561 raise StatisticsError('no mode for empty data') from None
562
563
564def multimode(data):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700565 """Return a list of the most frequently occurring values.
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700566
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700567 Will return more than one result if there are multiple modes
568 or an empty list if *data* is empty.
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700569
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700570 >>> multimode('aabbbbbbbbcc')
571 ['b']
572 >>> multimode('aabbbbccddddeeffffgg')
573 ['b', 'd', 'f']
574 >>> multimode('')
575 []
Raymond Hettingerfc06a192019-03-12 00:43:27 -0700576 """
577 counts = Counter(iter(data)).most_common()
578 maxcount, mode_items = next(groupby(counts, key=itemgetter(1)), (0, []))
579 return list(map(itemgetter(0), mode_items))
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700580
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700581
Raymond Hettingercba9f842019-06-02 21:07:43 -0700582# Notes on methods for computing quantiles
583# ----------------------------------------
584#
585# There is no one perfect way to compute quantiles. Here we offer
586# two methods that serve common needs. Most other packages
587# surveyed offered at least one or both of these two, making them
588# "standard" in the sense of "widely-adopted and reproducible".
589# They are also easy to explain, easy to compute manually, and have
590# straight-forward interpretations that aren't surprising.
591
592# The default method is known as "R6", "PERCENTILE.EXC", or "expected
593# value of rank order statistics". The alternative method is known as
594# "R7", "PERCENTILE.INC", or "mode of rank order statistics".
595
596# For sample data where there is a positive probability for values
597# beyond the range of the data, the R6 exclusive method is a
598# reasonable choice. Consider a random sample of nine values from a
599# population with a uniform distribution from 0.0 to 100.0. The
600# distribution of the third ranked sample point is described by
601# betavariate(alpha=3, beta=7) which has mode=0.250, median=0.286, and
602# mean=0.300. Only the latter (which corresponds with R6) gives the
603# desired cut point with 30% of the population falling below that
604# value, making it comparable to a result from an inv_cdf() function.
605
606# For describing population data where the end points are known to
607# be included in the data, the R7 inclusive method is a reasonable
608# choice. Instead of the mean, it uses the mode of the beta
609# distribution for the interior points. Per Hyndman & Fan, "One nice
610# property is that the vertices of Q7(p) divide the range into n - 1
611# intervals, and exactly 100p% of the intervals lie to the left of
612# Q7(p) and 100(1 - p)% of the intervals lie to the right of Q7(p)."
613
Raymond Hettingereed5e9a2019-07-19 01:57:22 -0700614# If needed, other methods could be added. However, for now, the
615# position is that fewer options make for easier choices and that
616# external packages can be used for anything more advanced.
Raymond Hettingercba9f842019-06-02 21:07:43 -0700617
Raymond Hettingere4810b22019-09-05 00:18:47 -0700618def quantiles(data, /, *, n=4, method='exclusive'):
619 """Divide *data* into *n* continuous intervals with equal probability.
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700620
621 Returns a list of (n - 1) cut points separating the intervals.
622
623 Set *n* to 4 for quartiles (the default). Set *n* to 10 for deciles.
624 Set *n* to 100 for percentiles which gives the 99 cuts points that
Raymond Hettingere4810b22019-09-05 00:18:47 -0700625 separate *data* in to 100 equal sized groups.
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700626
Raymond Hettinger4db25d52019-09-08 16:57:58 -0700627 The *data* can be any iterable containing sample.
628 The cut points are linearly interpolated between data points.
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700629
Raymond Hettingere4810b22019-09-05 00:18:47 -0700630 If *method* is set to *inclusive*, *data* is treated as population
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700631 data. The minimum value is treated as the 0th percentile and the
632 maximum value is treated as the 100th percentile.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700633 """
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700634 if n < 1:
635 raise StatisticsError('n must be at least 1')
Raymond Hettingere4810b22019-09-05 00:18:47 -0700636 data = sorted(data)
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700637 ld = len(data)
638 if ld < 2:
639 raise StatisticsError('must have at least two data points')
640 if method == 'inclusive':
641 m = ld - 1
642 result = []
643 for i in range(1, n):
644 j = i * m // n
645 delta = i*m - j*n
646 interpolated = (data[j] * (n - delta) + data[j+1] * delta) / n
647 result.append(interpolated)
648 return result
649 if method == 'exclusive':
650 m = ld + 1
651 result = []
652 for i in range(1, n):
653 j = i * m // n # rescale i to m/n
654 j = 1 if j < 1 else ld-1 if j > ld-1 else j # clamp to 1 .. ld-1
655 delta = i*m - j*n # exact integer math
656 interpolated = (data[j-1] * (n - delta) + data[j] * delta) / n
657 result.append(interpolated)
658 return result
659 raise ValueError(f'Unknown method: {method!r}')
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700660
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700661
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700662# === Measures of spread ===
663
664# See http://mathworld.wolfram.com/Variance.html
665# http://mathworld.wolfram.com/SampleVariance.html
666# http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
667#
668# Under no circumstances use the so-called "computational formula for
669# variance", as that is only suitable for hand calculations with a small
670# amount of low-precision data. It has terrible numeric properties.
671#
672# See a comparison of three computational methods here:
673# http://www.johndcook.com/blog/2008/09/26/comparing-three-methods-of-computing-standard-deviation/
674
675def _ss(data, c=None):
676 """Return sum of square deviations of sequence data.
677
678 If ``c`` is None, the mean is calculated in one pass, and the deviations
679 from the mean are calculated in a second pass. Otherwise, deviations are
680 calculated from ``c`` as given. Use the second case with care, as it can
681 lead to garbage results.
682 """
683 if c is None:
684 c = mean(data)
Steven D'Apranob28c3272015-12-01 19:59:53 +1100685 T, total, count = _sum((x-c)**2 for x in data)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700686 # The following sum should mathematically equal zero, but due to rounding
687 # error may not.
Steven D'Apranob28c3272015-12-01 19:59:53 +1100688 U, total2, count2 = _sum((x-c) for x in data)
689 assert T == U and count == count2
690 total -= total2**2/len(data)
691 assert not total < 0, 'negative sum of square deviations: %f' % total
692 return (T, total)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700693
694
695def variance(data, xbar=None):
696 """Return the sample variance of data.
697
698 data should be an iterable of Real-valued numbers, with at least two
699 values. The optional argument xbar, if given, should be the mean of
700 the data. If it is missing or None, the mean is automatically calculated.
701
702 Use this function when your data is a sample from a population. To
703 calculate the variance from the entire population, see ``pvariance``.
704
705 Examples:
706
707 >>> data = [2.75, 1.75, 1.25, 0.25, 0.5, 1.25, 3.5]
708 >>> variance(data)
709 1.3720238095238095
710
711 If you have already calculated the mean of your data, you can pass it as
712 the optional second argument ``xbar`` to avoid recalculating it:
713
714 >>> m = mean(data)
715 >>> variance(data, m)
716 1.3720238095238095
717
718 This function does not check that ``xbar`` is actually the mean of
719 ``data``. Giving arbitrary values for ``xbar`` may lead to invalid or
720 impossible results.
721
722 Decimals and Fractions are supported:
723
724 >>> from decimal import Decimal as D
725 >>> variance([D("27.5"), D("30.25"), D("30.25"), D("34.5"), D("41.75")])
726 Decimal('31.01875')
727
728 >>> from fractions import Fraction as F
729 >>> variance([F(1, 6), F(1, 2), F(5, 3)])
730 Fraction(67, 108)
731
732 """
733 if iter(data) is data:
734 data = list(data)
735 n = len(data)
736 if n < 2:
737 raise StatisticsError('variance requires at least two data points')
Steven D'Apranob28c3272015-12-01 19:59:53 +1100738 T, ss = _ss(data, xbar)
739 return _convert(ss/(n-1), T)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700740
741
742def pvariance(data, mu=None):
743 """Return the population variance of ``data``.
744
Raymond Hettingere4810b22019-09-05 00:18:47 -0700745 data should be a sequence or iterator of Real-valued numbers, with at least one
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700746 value. The optional argument mu, if given, should be the mean of
747 the data. If it is missing or None, the mean is automatically calculated.
748
749 Use this function to calculate the variance from the entire population.
750 To estimate the variance from a sample, the ``variance`` function is
751 usually a better choice.
752
753 Examples:
754
755 >>> data = [0.0, 0.25, 0.25, 1.25, 1.5, 1.75, 2.75, 3.25]
756 >>> pvariance(data)
757 1.25
758
759 If you have already calculated the mean of the data, you can pass it as
760 the optional second argument to avoid recalculating it:
761
762 >>> mu = mean(data)
763 >>> pvariance(data, mu)
764 1.25
765
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700766 Decimals and Fractions are supported:
767
768 >>> from decimal import Decimal as D
769 >>> pvariance([D("27.5"), D("30.25"), D("30.25"), D("34.5"), D("41.75")])
770 Decimal('24.815')
771
772 >>> from fractions import Fraction as F
773 >>> pvariance([F(1, 4), F(5, 4), F(1, 2)])
774 Fraction(13, 72)
775
776 """
777 if iter(data) is data:
778 data = list(data)
779 n = len(data)
780 if n < 1:
781 raise StatisticsError('pvariance requires at least one data point')
Steven D'Apranob28c3272015-12-01 19:59:53 +1100782 T, ss = _ss(data, mu)
783 return _convert(ss/n, T)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700784
785
786def stdev(data, xbar=None):
787 """Return the square root of the sample variance.
788
789 See ``variance`` for arguments and other details.
790
791 >>> stdev([1.5, 2.5, 2.5, 2.75, 3.25, 4.75])
792 1.0810874155219827
793
794 """
795 var = variance(data, xbar)
796 try:
797 return var.sqrt()
798 except AttributeError:
799 return math.sqrt(var)
800
801
802def pstdev(data, mu=None):
803 """Return the square root of the population variance.
804
805 See ``pvariance`` for arguments and other details.
806
807 >>> pstdev([1.5, 2.5, 2.5, 2.75, 3.25, 4.75])
808 0.986893273527251
809
810 """
811 var = pvariance(data, mu)
812 try:
813 return var.sqrt()
814 except AttributeError:
815 return math.sqrt(var)
Raymond Hettinger11c79532019-02-23 14:44:07 -0800816
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700817
Raymond Hettinger11c79532019-02-23 14:44:07 -0800818## Normal Distribution #####################################################
819
Dong-hee Na0a18ee42019-08-24 07:20:30 +0900820
821def _normal_dist_inv_cdf(p, mu, sigma):
822 # There is no closed-form solution to the inverse CDF for the normal
823 # distribution, so we use a rational approximation instead:
824 # Wichura, M.J. (1988). "Algorithm AS241: The Percentage Points of the
825 # Normal Distribution". Applied Statistics. Blackwell Publishing. 37
826 # (3): 477–484. doi:10.2307/2347330. JSTOR 2347330.
827 q = p - 0.5
828 if fabs(q) <= 0.425:
829 r = 0.180625 - q * q
830 # Hash sum: 55.88319_28806_14901_4439
831 num = (((((((2.50908_09287_30122_6727e+3 * r +
832 3.34305_75583_58812_8105e+4) * r +
833 6.72657_70927_00870_0853e+4) * r +
834 4.59219_53931_54987_1457e+4) * r +
835 1.37316_93765_50946_1125e+4) * r +
836 1.97159_09503_06551_4427e+3) * r +
837 1.33141_66789_17843_7745e+2) * r +
838 3.38713_28727_96366_6080e+0) * q
839 den = (((((((5.22649_52788_52854_5610e+3 * r +
840 2.87290_85735_72194_2674e+4) * r +
841 3.93078_95800_09271_0610e+4) * r +
842 2.12137_94301_58659_5867e+4) * r +
843 5.39419_60214_24751_1077e+3) * r +
844 6.87187_00749_20579_0830e+2) * r +
845 4.23133_30701_60091_1252e+1) * r +
846 1.0)
847 x = num / den
848 return mu + (x * sigma)
849 r = p if q <= 0.0 else 1.0 - p
850 r = sqrt(-log(r))
851 if r <= 5.0:
852 r = r - 1.6
853 # Hash sum: 49.33206_50330_16102_89036
854 num = (((((((7.74545_01427_83414_07640e-4 * r +
855 2.27238_44989_26918_45833e-2) * r +
856 2.41780_72517_74506_11770e-1) * r +
857 1.27045_82524_52368_38258e+0) * r +
858 3.64784_83247_63204_60504e+0) * r +
859 5.76949_72214_60691_40550e+0) * r +
860 4.63033_78461_56545_29590e+0) * r +
861 1.42343_71107_49683_57734e+0)
862 den = (((((((1.05075_00716_44416_84324e-9 * r +
863 5.47593_80849_95344_94600e-4) * r +
864 1.51986_66563_61645_71966e-2) * r +
865 1.48103_97642_74800_74590e-1) * r +
866 6.89767_33498_51000_04550e-1) * r +
867 1.67638_48301_83803_84940e+0) * r +
868 2.05319_16266_37758_82187e+0) * r +
869 1.0)
870 else:
871 r = r - 5.0
872 # Hash sum: 47.52583_31754_92896_71629
873 num = (((((((2.01033_43992_92288_13265e-7 * r +
874 2.71155_55687_43487_57815e-5) * r +
875 1.24266_09473_88078_43860e-3) * r +
876 2.65321_89526_57612_30930e-2) * r +
877 2.96560_57182_85048_91230e-1) * r +
878 1.78482_65399_17291_33580e+0) * r +
879 5.46378_49111_64114_36990e+0) * r +
880 6.65790_46435_01103_77720e+0)
881 den = (((((((2.04426_31033_89939_78564e-15 * r +
882 1.42151_17583_16445_88870e-7) * r +
883 1.84631_83175_10054_68180e-5) * r +
884 7.86869_13114_56132_59100e-4) * r +
885 1.48753_61290_85061_48525e-2) * r +
886 1.36929_88092_27358_05310e-1) * r +
887 5.99832_20655_58879_37690e-1) * r +
888 1.0)
889 x = num / den
890 if q < 0.0:
891 x = -x
892 return mu + (x * sigma)
893
894
Raymond Hettinger11c79532019-02-23 14:44:07 -0800895class NormalDist:
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700896 "Normal distribution of a random variable"
Raymond Hettinger11c79532019-02-23 14:44:07 -0800897 # https://en.wikipedia.org/wiki/Normal_distribution
898 # https://en.wikipedia.org/wiki/Variance#Properties
899
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700900 __slots__ = {
901 '_mu': 'Arithmetic mean of a normal distribution',
902 '_sigma': 'Standard deviation of a normal distribution',
903 }
Raymond Hettinger11c79532019-02-23 14:44:07 -0800904
905 def __init__(self, mu=0.0, sigma=1.0):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700906 "NormalDist where mu is the mean and sigma is the standard deviation."
Raymond Hettinger11c79532019-02-23 14:44:07 -0800907 if sigma < 0.0:
908 raise StatisticsError('sigma must be non-negative')
Raymond Hettingere4810b22019-09-05 00:18:47 -0700909 self._mu = float(mu)
910 self._sigma = float(sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -0800911
912 @classmethod
913 def from_samples(cls, data):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700914 "Make a normal distribution instance from sample data."
Raymond Hettinger11c79532019-02-23 14:44:07 -0800915 if not isinstance(data, (list, tuple)):
916 data = list(data)
917 xbar = fmean(data)
918 return cls(xbar, stdev(data, xbar))
919
Raymond Hettingerfb8c7d52019-04-23 01:46:18 -0700920 def samples(self, n, *, seed=None):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700921 "Generate *n* samples for a given mean and standard deviation."
Raymond Hettinger11c79532019-02-23 14:44:07 -0800922 gauss = random.gauss if seed is None else random.Random(seed).gauss
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700923 mu, sigma = self._mu, self._sigma
Raymond Hettinger11c79532019-02-23 14:44:07 -0800924 return [gauss(mu, sigma) for i in range(n)]
925
926 def pdf(self, x):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700927 "Probability density function. P(x <= X < x+dx) / dx"
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700928 variance = self._sigma ** 2.0
Raymond Hettinger11c79532019-02-23 14:44:07 -0800929 if not variance:
930 raise StatisticsError('pdf() not defined when sigma is zero')
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700931 return exp((x - self._mu)**2.0 / (-2.0*variance)) / sqrt(tau*variance)
Raymond Hettinger11c79532019-02-23 14:44:07 -0800932
933 def cdf(self, x):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700934 "Cumulative distribution function. P(X <= x)"
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700935 if not self._sigma:
Raymond Hettinger11c79532019-02-23 14:44:07 -0800936 raise StatisticsError('cdf() not defined when sigma is zero')
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700937 return 0.5 * (1.0 + erf((x - self._mu) / (self._sigma * sqrt(2.0))))
Raymond Hettinger11c79532019-02-23 14:44:07 -0800938
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700939 def inv_cdf(self, p):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700940 """Inverse cumulative distribution function. x : P(X <= x) = p
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700941
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700942 Finds the value of the random variable such that the probability of
943 the variable being less than or equal to that value equals the given
944 probability.
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700945
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700946 This function is also called the percent point function or quantile
947 function.
948 """
949 if p <= 0.0 or p >= 1.0:
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700950 raise StatisticsError('p must be in the range 0.0 < p < 1.0')
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700951 if self._sigma <= 0.0:
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700952 raise StatisticsError('cdf() not defined when sigma at or below zero')
Dong-hee Na0a18ee42019-08-24 07:20:30 +0900953 return _normal_dist_inv_cdf(p, self._mu, self._sigma)
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700954
Raymond Hettinger4db25d52019-09-08 16:57:58 -0700955 def quantiles(self, n=4):
956 """Divide into *n* continuous intervals with equal probability.
957
958 Returns a list of (n - 1) cut points separating the intervals.
959
960 Set *n* to 4 for quartiles (the default). Set *n* to 10 for deciles.
961 Set *n* to 100 for percentiles which gives the 99 cuts points that
962 separate the normal distribution in to 100 equal sized groups.
963 """
964 return [self.inv_cdf(i / n) for i in range(1, n)]
965
Raymond Hettinger318d5372019-03-06 22:59:40 -0800966 def overlap(self, other):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700967 """Compute the overlapping coefficient (OVL) between two normal distributions.
Raymond Hettinger318d5372019-03-06 22:59:40 -0800968
969 Measures the agreement between two normal probability distributions.
970 Returns a value between 0.0 and 1.0 giving the overlapping area in
971 the two underlying probability density functions.
972
973 >>> N1 = NormalDist(2.4, 1.6)
974 >>> N2 = NormalDist(3.2, 2.0)
975 >>> N1.overlap(N2)
976 0.8035050657330205
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700977 """
Raymond Hettinger318d5372019-03-06 22:59:40 -0800978 # See: "The overlapping coefficient as a measure of agreement between
979 # probability distributions and point estimation of the overlap of two
980 # normal densities" -- Henry F. Inman and Edwin L. Bradley Jr
981 # http://dx.doi.org/10.1080/03610928908830127
982 if not isinstance(other, NormalDist):
983 raise TypeError('Expected another NormalDist instance')
984 X, Y = self, other
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700985 if (Y._sigma, Y._mu) < (X._sigma, X._mu): # sort to assure commutativity
Raymond Hettinger318d5372019-03-06 22:59:40 -0800986 X, Y = Y, X
987 X_var, Y_var = X.variance, Y.variance
988 if not X_var or not Y_var:
989 raise StatisticsError('overlap() not defined when sigma is zero')
990 dv = Y_var - X_var
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700991 dm = fabs(Y._mu - X._mu)
Raymond Hettinger318d5372019-03-06 22:59:40 -0800992 if not dv:
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700993 return 1.0 - erf(dm / (2.0 * X._sigma * sqrt(2.0)))
994 a = X._mu * Y_var - Y._mu * X_var
995 b = X._sigma * Y._sigma * sqrt(dm**2.0 + dv * log(Y_var / X_var))
Raymond Hettinger318d5372019-03-06 22:59:40 -0800996 x1 = (a + b) / dv
997 x2 = (a - b) / dv
998 return 1.0 - (fabs(Y.cdf(x1) - X.cdf(x1)) + fabs(Y.cdf(x2) - X.cdf(x2)))
999
Raymond Hettinger11c79532019-02-23 14:44:07 -08001000 @property
Raymond Hettinger9e456bc2019-02-24 11:44:55 -08001001 def mean(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001002 "Arithmetic mean of the normal distribution."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001003 return self._mu
Raymond Hettinger9e456bc2019-02-24 11:44:55 -08001004
1005 @property
Raymond Hettinger4db25d52019-09-08 16:57:58 -07001006 def median(self):
1007 "Return the median of the normal distribution"
1008 return self._mu
1009
1010 @property
1011 def mode(self):
1012 """Return the mode of the normal distribution
1013
1014 The mode is the value x where which the probability density
1015 function (pdf) takes its maximum value.
1016 """
1017 return self._mu
1018
1019 @property
Raymond Hettinger9e456bc2019-02-24 11:44:55 -08001020 def stdev(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001021 "Standard deviation of the normal distribution."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001022 return self._sigma
Raymond Hettinger9e456bc2019-02-24 11:44:55 -08001023
1024 @property
Raymond Hettinger11c79532019-02-23 14:44:07 -08001025 def variance(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001026 "Square of the standard deviation."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001027 return self._sigma ** 2.0
Raymond Hettinger11c79532019-02-23 14:44:07 -08001028
1029 def __add__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001030 """Add a constant or another NormalDist instance.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001031
1032 If *other* is a constant, translate mu by the constant,
1033 leaving sigma unchanged.
1034
1035 If *other* is a NormalDist, add both the means and the variances.
1036 Mathematically, this works only if the two distributions are
1037 independent or if they are jointly normally distributed.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001038 """
Raymond Hettinger11c79532019-02-23 14:44:07 -08001039 if isinstance(x2, NormalDist):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001040 return NormalDist(x1._mu + x2._mu, hypot(x1._sigma, x2._sigma))
1041 return NormalDist(x1._mu + x2, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001042
1043 def __sub__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001044 """Subtract a constant or another NormalDist instance.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001045
1046 If *other* is a constant, translate by the constant mu,
1047 leaving sigma unchanged.
1048
1049 If *other* is a NormalDist, subtract the means and add the variances.
1050 Mathematically, this works only if the two distributions are
1051 independent or if they are jointly normally distributed.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001052 """
Raymond Hettinger11c79532019-02-23 14:44:07 -08001053 if isinstance(x2, NormalDist):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001054 return NormalDist(x1._mu - x2._mu, hypot(x1._sigma, x2._sigma))
1055 return NormalDist(x1._mu - x2, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001056
1057 def __mul__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001058 """Multiply both mu and sigma by a constant.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001059
1060 Used for rescaling, perhaps to change measurement units.
1061 Sigma is scaled with the absolute value of the constant.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001062 """
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001063 return NormalDist(x1._mu * x2, x1._sigma * fabs(x2))
Raymond Hettinger11c79532019-02-23 14:44:07 -08001064
1065 def __truediv__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001066 """Divide both mu and sigma by a constant.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001067
1068 Used for rescaling, perhaps to change measurement units.
1069 Sigma is scaled with the absolute value of the constant.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001070 """
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001071 return NormalDist(x1._mu / x2, x1._sigma / fabs(x2))
Raymond Hettinger11c79532019-02-23 14:44:07 -08001072
1073 def __pos__(x1):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001074 "Return a copy of the instance."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001075 return NormalDist(x1._mu, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001076
1077 def __neg__(x1):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001078 "Negates mu while keeping sigma the same."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001079 return NormalDist(-x1._mu, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001080
1081 __radd__ = __add__
1082
1083 def __rsub__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001084 "Subtract a NormalDist from a constant or another NormalDist."
Raymond Hettinger11c79532019-02-23 14:44:07 -08001085 return -(x1 - x2)
1086
1087 __rmul__ = __mul__
1088
1089 def __eq__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001090 "Two NormalDist objects are equal if their mu and sigma are both equal."
Raymond Hettinger11c79532019-02-23 14:44:07 -08001091 if not isinstance(x2, NormalDist):
1092 return NotImplemented
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001093 return (x1._mu, x2._sigma) == (x2._mu, x2._sigma)
1094
1095 def __hash__(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001096 "NormalDist objects hash equal if their mu and sigma are both equal."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001097 return hash((self._mu, self._sigma))
Raymond Hettinger11c79532019-02-23 14:44:07 -08001098
1099 def __repr__(self):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001100 return f'{type(self).__name__}(mu={self._mu!r}, sigma={self._sigma!r})'
Raymond Hettinger11c79532019-02-23 14:44:07 -08001101
Dong-hee Na0a18ee42019-08-24 07:20:30 +09001102# If available, use C implementation
1103try:
1104 from _statistics import _normal_dist_inv_cdf
1105except ImportError:
1106 pass
1107
Raymond Hettinger11c79532019-02-23 14:44:07 -08001108
1109if __name__ == '__main__':
1110
1111 # Show math operations computed analytically in comparsion
1112 # to a monte carlo simulation of the same operations
1113
1114 from math import isclose
1115 from operator import add, sub, mul, truediv
1116 from itertools import repeat
Raymond Hettingerfc06a192019-03-12 00:43:27 -07001117 import doctest
Raymond Hettinger11c79532019-02-23 14:44:07 -08001118
1119 g1 = NormalDist(10, 20)
1120 g2 = NormalDist(-5, 25)
1121
1122 # Test scaling by a constant
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001123 assert (g1 * 5 / 5).mean == g1.mean
1124 assert (g1 * 5 / 5).stdev == g1.stdev
Raymond Hettinger11c79532019-02-23 14:44:07 -08001125
1126 n = 100_000
1127 G1 = g1.samples(n)
1128 G2 = g2.samples(n)
1129
1130 for func in (add, sub):
1131 print(f'\nTest {func.__name__} with another NormalDist:')
1132 print(func(g1, g2))
1133 print(NormalDist.from_samples(map(func, G1, G2)))
1134
1135 const = 11
1136 for func in (add, sub, mul, truediv):
1137 print(f'\nTest {func.__name__} with a constant:')
1138 print(func(g1, const))
1139 print(NormalDist.from_samples(map(func, G1, repeat(const))))
1140
1141 const = 19
1142 for func in (add, sub, mul):
1143 print(f'\nTest constant with {func.__name__}:')
1144 print(func(const, g1))
1145 print(NormalDist.from_samples(map(func, repeat(const), G1)))
1146
1147 def assert_close(G1, G2):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001148 assert isclose(G1.mean, G1.mean, rel_tol=0.01), (G1, G2)
1149 assert isclose(G1.stdev, G2.stdev, rel_tol=0.01), (G1, G2)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001150
1151 X = NormalDist(-105, 73)
1152 Y = NormalDist(31, 47)
1153 s = 32.75
1154 n = 100_000
1155
1156 S = NormalDist.from_samples([x + s for x in X.samples(n)])
1157 assert_close(X + s, S)
1158
1159 S = NormalDist.from_samples([x - s for x in X.samples(n)])
1160 assert_close(X - s, S)
1161
1162 S = NormalDist.from_samples([x * s for x in X.samples(n)])
1163 assert_close(X * s, S)
1164
1165 S = NormalDist.from_samples([x / s for x in X.samples(n)])
1166 assert_close(X / s, S)
1167
1168 S = NormalDist.from_samples([x + y for x, y in zip(X.samples(n),
1169 Y.samples(n))])
1170 assert_close(X + Y, S)
1171
1172 S = NormalDist.from_samples([x - y for x, y in zip(X.samples(n),
1173 Y.samples(n))])
1174 assert_close(X - Y, S)
Raymond Hettingerfc06a192019-03-12 00:43:27 -07001175
1176 print(doctest.testmod())