''' Tool creating and manipulating normal distributions of random variable. Features a composite class that treats the mean and standard deviation of measurements as single entity. ''' from math import hypot, sqrt, fabs, exp, erf, tau from statistics import mean as fmean, stdev import random class NormalDist: 'Normal distribution of a random variable' # https://en.wikipedia.org/wiki/Normal_distribution # https://en.wikipedia.org/wiki/Variance#Properties __slots__ = ('mu', 'sigma') def __init__(x1, mu, sigma=0.0): assert sigma >= 0.0 x1.mu = mu x1.sigma = sigma @classmethod def from_samples(cls, data): 'Make a normal distribution from a sample' data = list(data) xbar = fmean(data) return cls(xbar, stdev(data, xbar)) @property def variance(self): 'Square of the standard deviation' return self.sigma ** 2.0 def examples(self, n, seed=None): 'Generate *n* samples for given mean and standard deviation' gauss = random.gauss if seed is None else random.Random(seed).gauss mu, sigma = self.mu, self.sigma return [gauss(mu, sigma) for i in range(n)] def pdf(self, x): 'Probability density function: P(x <= X < x+dx) / dx' variance = self.sigma ** 2.0 return exp((x - self.mu)**2.0 / (-2.0*variance)) / sqrt(tau * variance) def cdf(self, x): 'Cumulative density function: P(X <= x)' return 0.5 * (1.0 + erf((x - self.mu) / (self.sigma * sqrt(2.0)))) def __repr__(self): return f'{type(self).__name__}(mu={self.mu!r}, sigma={self.sigma!r})' def __add__(x1, x2): if isinstance(x2, NormalDist): return NormalDist(x1.mu + x2.mu, hypot(x1.sigma, x2.sigma)) return NormalDist(x1.mu + x2, x1.sigma) def __sub__(x1, x2): if isinstance(x2, NormalDist): return NormalDist(x1.mu - x2.mu, hypot(x1.sigma, x2.sigma)) return NormalDist(x1.mu - x2, x1.sigma) def __mul__(x1, x2): return NormalDist(x1.mu * x2, x1.sigma * fabs(x2)) def __truediv__(num, den): return NormalDist(num.mu / den, num.sigma / fabs(den)) def __pos__(x1): return x1 def __neg__(x1): return NormalDist(-x1.mu, x1.sigma) __radd__ = __add__ def __rsub__(x1, x2): return -(x1 - x2) __rmul__ = __mul__ if __name__ == '__main__': from math import isclose from operator import add, sub, mul, truediv from itertools import repeat g1 = NormalDist(10, 20) g2 = NormalDist(-5, 25) # Test scaling by a constant assert (g1 * 5 / 5).mu == g1.mu assert (g1 * 5 / 5).sigma == g1.sigma n = 100_000 G1 = g1.examples(n) G2 = g2.examples(n) for func in (add, sub): print(f'\nTest {func.__name__} with another NormalDist:') print(func(g1, g2)) print(NormalDist.from_samples(map(func, G1, G2))) const = 11 for func in (add, sub, mul, truediv): print(f'\nTest {func.__name__} with a constant:') print(func(g1, const)) print(NormalDist.from_samples(map(func, G1, repeat(const)))) const = 19 for func in (add, sub, mul): print(f'\nTest constant with {func.__name__}:') print(func(const, g1)) print(NormalDist.from_samples(map(func, repeat(const), G1))) def assert_close(G1, G2): assert isclose(G1.mu, G1.mu, rel_tol=0.01), (G1, G2) assert isclose(G1.sigma, G2.sigma, rel_tol=0.01), (G1, G2) X = NormalDist(-105, 73) Y = NormalDist(31, 47) s = 32.75 n = 100_000 S = NormalDist.from_samples([x + s for x in X.examples(n)]) assert_close(X + s, S) S = NormalDist.from_samples([x - s for x in X.examples(n)]) assert_close(X - s, S) S = NormalDist.from_samples([x * s for x in X.examples(n)]) assert_close(X * s, S) S = NormalDist.from_samples([x / s for x in X.examples(n)]) assert_close(X / s, S) S = NormalDist.from_samples([x + y for x, y in zip(X.examples(n), Y.examples(n))]) assert_close(X + Y, S) S = NormalDist.from_samples([x - y for x, y in zip(X.examples(n), Y.examples(n))]) assert_close(X - Y, S)