blob: 4b172662770fb029605665f198f82c835257484d [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 Hettingere4810b22019-09-05 00:18:47 -0700627 The *data* can be any iterable containing sample data or it can be
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700628 an instance of a class that defines an inv_cdf() method. For sample
629 data, the cut points are linearly interpolated between data points.
630
Raymond Hettingere4810b22019-09-05 00:18:47 -0700631 If *method* is set to *inclusive*, *data* is treated as population
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700632 data. The minimum value is treated as the 0th percentile and the
633 maximum value is treated as the 100th percentile.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700634 """
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700635 if n < 1:
636 raise StatisticsError('n must be at least 1')
Raymond Hettingere4810b22019-09-05 00:18:47 -0700637 if hasattr(data, 'inv_cdf'):
638 return [data.inv_cdf(i / n) for i in range(1, n)]
639 data = sorted(data)
Raymond Hettinger9013ccf2019-04-23 00:06:35 -0700640 ld = len(data)
641 if ld < 2:
642 raise StatisticsError('must have at least two data points')
643 if method == 'inclusive':
644 m = ld - 1
645 result = []
646 for i in range(1, n):
647 j = i * m // n
648 delta = i*m - j*n
649 interpolated = (data[j] * (n - delta) + data[j+1] * delta) / n
650 result.append(interpolated)
651 return result
652 if method == 'exclusive':
653 m = ld + 1
654 result = []
655 for i in range(1, n):
656 j = i * m // n # rescale i to m/n
657 j = 1 if j < 1 else ld-1 if j > ld-1 else j # clamp to 1 .. ld-1
658 delta = i*m - j*n # exact integer math
659 interpolated = (data[j-1] * (n - delta) + data[j] * delta) / n
660 result.append(interpolated)
661 return result
662 raise ValueError(f'Unknown method: {method!r}')
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700663
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700664
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700665# === Measures of spread ===
666
667# See http://mathworld.wolfram.com/Variance.html
668# http://mathworld.wolfram.com/SampleVariance.html
669# http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
670#
671# Under no circumstances use the so-called "computational formula for
672# variance", as that is only suitable for hand calculations with a small
673# amount of low-precision data. It has terrible numeric properties.
674#
675# See a comparison of three computational methods here:
676# http://www.johndcook.com/blog/2008/09/26/comparing-three-methods-of-computing-standard-deviation/
677
678def _ss(data, c=None):
679 """Return sum of square deviations of sequence data.
680
681 If ``c`` is None, the mean is calculated in one pass, and the deviations
682 from the mean are calculated in a second pass. Otherwise, deviations are
683 calculated from ``c`` as given. Use the second case with care, as it can
684 lead to garbage results.
685 """
686 if c is None:
687 c = mean(data)
Steven D'Apranob28c3272015-12-01 19:59:53 +1100688 T, total, count = _sum((x-c)**2 for x in data)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700689 # The following sum should mathematically equal zero, but due to rounding
690 # error may not.
Steven D'Apranob28c3272015-12-01 19:59:53 +1100691 U, total2, count2 = _sum((x-c) for x in data)
692 assert T == U and count == count2
693 total -= total2**2/len(data)
694 assert not total < 0, 'negative sum of square deviations: %f' % total
695 return (T, total)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700696
697
698def variance(data, xbar=None):
699 """Return the sample variance of data.
700
701 data should be an iterable of Real-valued numbers, with at least two
702 values. The optional argument xbar, if given, should be the mean of
703 the data. If it is missing or None, the mean is automatically calculated.
704
705 Use this function when your data is a sample from a population. To
706 calculate the variance from the entire population, see ``pvariance``.
707
708 Examples:
709
710 >>> data = [2.75, 1.75, 1.25, 0.25, 0.5, 1.25, 3.5]
711 >>> variance(data)
712 1.3720238095238095
713
714 If you have already calculated the mean of your data, you can pass it as
715 the optional second argument ``xbar`` to avoid recalculating it:
716
717 >>> m = mean(data)
718 >>> variance(data, m)
719 1.3720238095238095
720
721 This function does not check that ``xbar`` is actually the mean of
722 ``data``. Giving arbitrary values for ``xbar`` may lead to invalid or
723 impossible results.
724
725 Decimals and Fractions are supported:
726
727 >>> from decimal import Decimal as D
728 >>> variance([D("27.5"), D("30.25"), D("30.25"), D("34.5"), D("41.75")])
729 Decimal('31.01875')
730
731 >>> from fractions import Fraction as F
732 >>> variance([F(1, 6), F(1, 2), F(5, 3)])
733 Fraction(67, 108)
734
735 """
736 if iter(data) is data:
737 data = list(data)
738 n = len(data)
739 if n < 2:
740 raise StatisticsError('variance requires at least two data points')
Steven D'Apranob28c3272015-12-01 19:59:53 +1100741 T, ss = _ss(data, xbar)
742 return _convert(ss/(n-1), T)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700743
744
745def pvariance(data, mu=None):
746 """Return the population variance of ``data``.
747
Raymond Hettingere4810b22019-09-05 00:18:47 -0700748 data should be a sequence or iterator of Real-valued numbers, with at least one
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700749 value. The optional argument mu, if given, should be the mean of
750 the data. If it is missing or None, the mean is automatically calculated.
751
752 Use this function to calculate the variance from the entire population.
753 To estimate the variance from a sample, the ``variance`` function is
754 usually a better choice.
755
756 Examples:
757
758 >>> data = [0.0, 0.25, 0.25, 1.25, 1.5, 1.75, 2.75, 3.25]
759 >>> pvariance(data)
760 1.25
761
762 If you have already calculated the mean of the data, you can pass it as
763 the optional second argument to avoid recalculating it:
764
765 >>> mu = mean(data)
766 >>> pvariance(data, mu)
767 1.25
768
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700769 Decimals and Fractions are supported:
770
771 >>> from decimal import Decimal as D
772 >>> pvariance([D("27.5"), D("30.25"), D("30.25"), D("34.5"), D("41.75")])
773 Decimal('24.815')
774
775 >>> from fractions import Fraction as F
776 >>> pvariance([F(1, 4), F(5, 4), F(1, 2)])
777 Fraction(13, 72)
778
779 """
780 if iter(data) is data:
781 data = list(data)
782 n = len(data)
783 if n < 1:
784 raise StatisticsError('pvariance requires at least one data point')
Steven D'Apranob28c3272015-12-01 19:59:53 +1100785 T, ss = _ss(data, mu)
786 return _convert(ss/n, T)
Larry Hastingsf5e987b2013-10-19 11:50:09 -0700787
788
789def stdev(data, xbar=None):
790 """Return the square root of the sample variance.
791
792 See ``variance`` for arguments and other details.
793
794 >>> stdev([1.5, 2.5, 2.5, 2.75, 3.25, 4.75])
795 1.0810874155219827
796
797 """
798 var = variance(data, xbar)
799 try:
800 return var.sqrt()
801 except AttributeError:
802 return math.sqrt(var)
803
804
805def pstdev(data, mu=None):
806 """Return the square root of the population variance.
807
808 See ``pvariance`` for arguments and other details.
809
810 >>> pstdev([1.5, 2.5, 2.5, 2.75, 3.25, 4.75])
811 0.986893273527251
812
813 """
814 var = pvariance(data, mu)
815 try:
816 return var.sqrt()
817 except AttributeError:
818 return math.sqrt(var)
Raymond Hettinger11c79532019-02-23 14:44:07 -0800819
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700820
Raymond Hettinger11c79532019-02-23 14:44:07 -0800821## Normal Distribution #####################################################
822
Dong-hee Na0a18ee42019-08-24 07:20:30 +0900823
824def _normal_dist_inv_cdf(p, mu, sigma):
825 # There is no closed-form solution to the inverse CDF for the normal
826 # distribution, so we use a rational approximation instead:
827 # Wichura, M.J. (1988). "Algorithm AS241: The Percentage Points of the
828 # Normal Distribution". Applied Statistics. Blackwell Publishing. 37
829 # (3): 477–484. doi:10.2307/2347330. JSTOR 2347330.
830 q = p - 0.5
831 if fabs(q) <= 0.425:
832 r = 0.180625 - q * q
833 # Hash sum: 55.88319_28806_14901_4439
834 num = (((((((2.50908_09287_30122_6727e+3 * r +
835 3.34305_75583_58812_8105e+4) * r +
836 6.72657_70927_00870_0853e+4) * r +
837 4.59219_53931_54987_1457e+4) * r +
838 1.37316_93765_50946_1125e+4) * r +
839 1.97159_09503_06551_4427e+3) * r +
840 1.33141_66789_17843_7745e+2) * r +
841 3.38713_28727_96366_6080e+0) * q
842 den = (((((((5.22649_52788_52854_5610e+3 * r +
843 2.87290_85735_72194_2674e+4) * r +
844 3.93078_95800_09271_0610e+4) * r +
845 2.12137_94301_58659_5867e+4) * r +
846 5.39419_60214_24751_1077e+3) * r +
847 6.87187_00749_20579_0830e+2) * r +
848 4.23133_30701_60091_1252e+1) * r +
849 1.0)
850 x = num / den
851 return mu + (x * sigma)
852 r = p if q <= 0.0 else 1.0 - p
853 r = sqrt(-log(r))
854 if r <= 5.0:
855 r = r - 1.6
856 # Hash sum: 49.33206_50330_16102_89036
857 num = (((((((7.74545_01427_83414_07640e-4 * r +
858 2.27238_44989_26918_45833e-2) * r +
859 2.41780_72517_74506_11770e-1) * r +
860 1.27045_82524_52368_38258e+0) * r +
861 3.64784_83247_63204_60504e+0) * r +
862 5.76949_72214_60691_40550e+0) * r +
863 4.63033_78461_56545_29590e+0) * r +
864 1.42343_71107_49683_57734e+0)
865 den = (((((((1.05075_00716_44416_84324e-9 * r +
866 5.47593_80849_95344_94600e-4) * r +
867 1.51986_66563_61645_71966e-2) * r +
868 1.48103_97642_74800_74590e-1) * r +
869 6.89767_33498_51000_04550e-1) * r +
870 1.67638_48301_83803_84940e+0) * r +
871 2.05319_16266_37758_82187e+0) * r +
872 1.0)
873 else:
874 r = r - 5.0
875 # Hash sum: 47.52583_31754_92896_71629
876 num = (((((((2.01033_43992_92288_13265e-7 * r +
877 2.71155_55687_43487_57815e-5) * r +
878 1.24266_09473_88078_43860e-3) * r +
879 2.65321_89526_57612_30930e-2) * r +
880 2.96560_57182_85048_91230e-1) * r +
881 1.78482_65399_17291_33580e+0) * r +
882 5.46378_49111_64114_36990e+0) * r +
883 6.65790_46435_01103_77720e+0)
884 den = (((((((2.04426_31033_89939_78564e-15 * r +
885 1.42151_17583_16445_88870e-7) * r +
886 1.84631_83175_10054_68180e-5) * r +
887 7.86869_13114_56132_59100e-4) * r +
888 1.48753_61290_85061_48525e-2) * r +
889 1.36929_88092_27358_05310e-1) * r +
890 5.99832_20655_58879_37690e-1) * r +
891 1.0)
892 x = num / den
893 if q < 0.0:
894 x = -x
895 return mu + (x * sigma)
896
897
Raymond Hettinger11c79532019-02-23 14:44:07 -0800898class NormalDist:
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700899 "Normal distribution of a random variable"
Raymond Hettinger11c79532019-02-23 14:44:07 -0800900 # https://en.wikipedia.org/wiki/Normal_distribution
901 # https://en.wikipedia.org/wiki/Variance#Properties
902
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700903 __slots__ = {
904 '_mu': 'Arithmetic mean of a normal distribution',
905 '_sigma': 'Standard deviation of a normal distribution',
906 }
Raymond Hettinger11c79532019-02-23 14:44:07 -0800907
908 def __init__(self, mu=0.0, sigma=1.0):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700909 "NormalDist where mu is the mean and sigma is the standard deviation."
Raymond Hettinger11c79532019-02-23 14:44:07 -0800910 if sigma < 0.0:
911 raise StatisticsError('sigma must be non-negative')
Raymond Hettingere4810b22019-09-05 00:18:47 -0700912 self._mu = float(mu)
913 self._sigma = float(sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -0800914
915 @classmethod
916 def from_samples(cls, data):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700917 "Make a normal distribution instance from sample data."
Raymond Hettinger11c79532019-02-23 14:44:07 -0800918 if not isinstance(data, (list, tuple)):
919 data = list(data)
920 xbar = fmean(data)
921 return cls(xbar, stdev(data, xbar))
922
Raymond Hettingerfb8c7d52019-04-23 01:46:18 -0700923 def samples(self, n, *, seed=None):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700924 "Generate *n* samples for a given mean and standard deviation."
Raymond Hettinger11c79532019-02-23 14:44:07 -0800925 gauss = random.gauss if seed is None else random.Random(seed).gauss
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700926 mu, sigma = self._mu, self._sigma
Raymond Hettinger11c79532019-02-23 14:44:07 -0800927 return [gauss(mu, sigma) for i in range(n)]
928
929 def pdf(self, x):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700930 "Probability density function. P(x <= X < x+dx) / dx"
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700931 variance = self._sigma ** 2.0
Raymond Hettinger11c79532019-02-23 14:44:07 -0800932 if not variance:
933 raise StatisticsError('pdf() not defined when sigma is zero')
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700934 return exp((x - self._mu)**2.0 / (-2.0*variance)) / sqrt(tau*variance)
Raymond Hettinger11c79532019-02-23 14:44:07 -0800935
936 def cdf(self, x):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700937 "Cumulative distribution function. P(X <= x)"
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700938 if not self._sigma:
Raymond Hettinger11c79532019-02-23 14:44:07 -0800939 raise StatisticsError('cdf() not defined when sigma is zero')
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700940 return 0.5 * (1.0 + erf((x - self._mu) / (self._sigma * sqrt(2.0))))
Raymond Hettinger11c79532019-02-23 14:44:07 -0800941
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700942 def inv_cdf(self, p):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700943 """Inverse cumulative distribution function. x : P(X <= x) = p
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700944
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700945 Finds the value of the random variable such that the probability of
946 the variable being less than or equal to that value equals the given
947 probability.
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700948
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700949 This function is also called the percent point function or quantile
950 function.
951 """
952 if p <= 0.0 or p >= 1.0:
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700953 raise StatisticsError('p must be in the range 0.0 < p < 1.0')
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700954 if self._sigma <= 0.0:
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700955 raise StatisticsError('cdf() not defined when sigma at or below zero')
Dong-hee Na0a18ee42019-08-24 07:20:30 +0900956 return _normal_dist_inv_cdf(p, self._mu, self._sigma)
Raymond Hettinger714c60d2019-03-18 20:17:14 -0700957
Raymond Hettinger318d5372019-03-06 22:59:40 -0800958 def overlap(self, other):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700959 """Compute the overlapping coefficient (OVL) between two normal distributions.
Raymond Hettinger318d5372019-03-06 22:59:40 -0800960
961 Measures the agreement between two normal probability distributions.
962 Returns a value between 0.0 and 1.0 giving the overlapping area in
963 the two underlying probability density functions.
964
965 >>> N1 = NormalDist(2.4, 1.6)
966 >>> N2 = NormalDist(3.2, 2.0)
967 >>> N1.overlap(N2)
968 0.8035050657330205
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700969 """
Raymond Hettinger318d5372019-03-06 22:59:40 -0800970 # See: "The overlapping coefficient as a measure of agreement between
971 # probability distributions and point estimation of the overlap of two
972 # normal densities" -- Henry F. Inman and Edwin L. Bradley Jr
973 # http://dx.doi.org/10.1080/03610928908830127
974 if not isinstance(other, NormalDist):
975 raise TypeError('Expected another NormalDist instance')
976 X, Y = self, other
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700977 if (Y._sigma, Y._mu) < (X._sigma, X._mu): # sort to assure commutativity
Raymond Hettinger318d5372019-03-06 22:59:40 -0800978 X, Y = Y, X
979 X_var, Y_var = X.variance, Y.variance
980 if not X_var or not Y_var:
981 raise StatisticsError('overlap() not defined when sigma is zero')
982 dv = Y_var - X_var
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700983 dm = fabs(Y._mu - X._mu)
Raymond Hettinger318d5372019-03-06 22:59:40 -0800984 if not dv:
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700985 return 1.0 - erf(dm / (2.0 * X._sigma * sqrt(2.0)))
986 a = X._mu * Y_var - Y._mu * X_var
987 b = X._sigma * Y._sigma * sqrt(dm**2.0 + dv * log(Y_var / X_var))
Raymond Hettinger318d5372019-03-06 22:59:40 -0800988 x1 = (a + b) / dv
989 x2 = (a - b) / dv
990 return 1.0 - (fabs(Y.cdf(x1) - X.cdf(x1)) + fabs(Y.cdf(x2) - X.cdf(x2)))
991
Raymond Hettinger11c79532019-02-23 14:44:07 -0800992 @property
Raymond Hettinger9e456bc2019-02-24 11:44:55 -0800993 def mean(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700994 "Arithmetic mean of the normal distribution."
Raymond Hettinger02c91f52019-07-21 00:34:47 -0700995 return self._mu
Raymond Hettinger9e456bc2019-02-24 11:44:55 -0800996
997 @property
998 def stdev(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -0700999 "Standard deviation of the normal distribution."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001000 return self._sigma
Raymond Hettinger9e456bc2019-02-24 11:44:55 -08001001
1002 @property
Raymond Hettinger11c79532019-02-23 14:44:07 -08001003 def variance(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001004 "Square of the standard deviation."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001005 return self._sigma ** 2.0
Raymond Hettinger11c79532019-02-23 14:44:07 -08001006
1007 def __add__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001008 """Add a constant or another NormalDist instance.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001009
1010 If *other* is a constant, translate mu by the constant,
1011 leaving sigma unchanged.
1012
1013 If *other* is a NormalDist, add both the means and the variances.
1014 Mathematically, this works only if the two distributions are
1015 independent or if they are jointly normally distributed.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001016 """
Raymond Hettinger11c79532019-02-23 14:44:07 -08001017 if isinstance(x2, NormalDist):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001018 return NormalDist(x1._mu + x2._mu, hypot(x1._sigma, x2._sigma))
1019 return NormalDist(x1._mu + x2, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001020
1021 def __sub__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001022 """Subtract a constant or another NormalDist instance.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001023
1024 If *other* is a constant, translate by the constant mu,
1025 leaving sigma unchanged.
1026
1027 If *other* is a NormalDist, subtract the means and add the variances.
1028 Mathematically, this works only if the two distributions are
1029 independent or if they are jointly normally distributed.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001030 """
Raymond Hettinger11c79532019-02-23 14:44:07 -08001031 if isinstance(x2, NormalDist):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001032 return NormalDist(x1._mu - x2._mu, hypot(x1._sigma, x2._sigma))
1033 return NormalDist(x1._mu - x2, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001034
1035 def __mul__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001036 """Multiply both mu and sigma by a constant.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001037
1038 Used for rescaling, perhaps to change measurement units.
1039 Sigma is scaled with the absolute value of the constant.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001040 """
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001041 return NormalDist(x1._mu * x2, x1._sigma * fabs(x2))
Raymond Hettinger11c79532019-02-23 14:44:07 -08001042
1043 def __truediv__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001044 """Divide both mu and sigma by a constant.
Raymond Hettinger5f1e8b42019-03-18 22:24:15 -07001045
1046 Used for rescaling, perhaps to change measurement units.
1047 Sigma is scaled with the absolute value of the constant.
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001048 """
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001049 return NormalDist(x1._mu / x2, x1._sigma / fabs(x2))
Raymond Hettinger11c79532019-02-23 14:44:07 -08001050
1051 def __pos__(x1):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001052 "Return a copy of the instance."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001053 return NormalDist(x1._mu, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001054
1055 def __neg__(x1):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001056 "Negates mu while keeping sigma the same."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001057 return NormalDist(-x1._mu, x1._sigma)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001058
1059 __radd__ = __add__
1060
1061 def __rsub__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001062 "Subtract a NormalDist from a constant or another NormalDist."
Raymond Hettinger11c79532019-02-23 14:44:07 -08001063 return -(x1 - x2)
1064
1065 __rmul__ = __mul__
1066
1067 def __eq__(x1, x2):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001068 "Two NormalDist objects are equal if their mu and sigma are both equal."
Raymond Hettinger11c79532019-02-23 14:44:07 -08001069 if not isinstance(x2, NormalDist):
1070 return NotImplemented
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001071 return (x1._mu, x2._sigma) == (x2._mu, x2._sigma)
1072
1073 def __hash__(self):
Raymond Hettinger1c0e9bb2019-07-21 12:13:07 -07001074 "NormalDist objects hash equal if their mu and sigma are both equal."
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001075 return hash((self._mu, self._sigma))
Raymond Hettinger11c79532019-02-23 14:44:07 -08001076
1077 def __repr__(self):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001078 return f'{type(self).__name__}(mu={self._mu!r}, sigma={self._sigma!r})'
Raymond Hettinger11c79532019-02-23 14:44:07 -08001079
Dong-hee Na0a18ee42019-08-24 07:20:30 +09001080# If available, use C implementation
1081try:
1082 from _statistics import _normal_dist_inv_cdf
1083except ImportError:
1084 pass
1085
Raymond Hettinger11c79532019-02-23 14:44:07 -08001086
1087if __name__ == '__main__':
1088
1089 # Show math operations computed analytically in comparsion
1090 # to a monte carlo simulation of the same operations
1091
1092 from math import isclose
1093 from operator import add, sub, mul, truediv
1094 from itertools import repeat
Raymond Hettingerfc06a192019-03-12 00:43:27 -07001095 import doctest
Raymond Hettinger11c79532019-02-23 14:44:07 -08001096
1097 g1 = NormalDist(10, 20)
1098 g2 = NormalDist(-5, 25)
1099
1100 # Test scaling by a constant
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001101 assert (g1 * 5 / 5).mean == g1.mean
1102 assert (g1 * 5 / 5).stdev == g1.stdev
Raymond Hettinger11c79532019-02-23 14:44:07 -08001103
1104 n = 100_000
1105 G1 = g1.samples(n)
1106 G2 = g2.samples(n)
1107
1108 for func in (add, sub):
1109 print(f'\nTest {func.__name__} with another NormalDist:')
1110 print(func(g1, g2))
1111 print(NormalDist.from_samples(map(func, G1, G2)))
1112
1113 const = 11
1114 for func in (add, sub, mul, truediv):
1115 print(f'\nTest {func.__name__} with a constant:')
1116 print(func(g1, const))
1117 print(NormalDist.from_samples(map(func, G1, repeat(const))))
1118
1119 const = 19
1120 for func in (add, sub, mul):
1121 print(f'\nTest constant with {func.__name__}:')
1122 print(func(const, g1))
1123 print(NormalDist.from_samples(map(func, repeat(const), G1)))
1124
1125 def assert_close(G1, G2):
Raymond Hettinger02c91f52019-07-21 00:34:47 -07001126 assert isclose(G1.mean, G1.mean, rel_tol=0.01), (G1, G2)
1127 assert isclose(G1.stdev, G2.stdev, rel_tol=0.01), (G1, G2)
Raymond Hettinger11c79532019-02-23 14:44:07 -08001128
1129 X = NormalDist(-105, 73)
1130 Y = NormalDist(31, 47)
1131 s = 32.75
1132 n = 100_000
1133
1134 S = NormalDist.from_samples([x + s for x in X.samples(n)])
1135 assert_close(X + s, S)
1136
1137 S = NormalDist.from_samples([x - s for x in X.samples(n)])
1138 assert_close(X - s, S)
1139
1140 S = NormalDist.from_samples([x * s for x in X.samples(n)])
1141 assert_close(X * s, S)
1142
1143 S = NormalDist.from_samples([x / s for x in X.samples(n)])
1144 assert_close(X / s, S)
1145
1146 S = NormalDist.from_samples([x + y for x, y in zip(X.samples(n),
1147 Y.samples(n))])
1148 assert_close(X + Y, S)
1149
1150 S = NormalDist.from_samples([x - y for x, y in zip(X.samples(n),
1151 Y.samples(n))])
1152 assert_close(X - Y, S)
Raymond Hettingerfc06a192019-03-12 00:43:27 -07001153
1154 print(doctest.testmod())