1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
|
import copy
import math
import random
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
class Distribution:
def __init__(self, prob_table=None):
self.probabilities = prob_table
@property
def probabilities(self):
return self._prob_table
@probabilities.setter
def probabilities(self, value):
if value is None:
self._prob_table = defaultdict(lambda: 0.0)
elif isinstance(value, list) or isinstance(value, tuple):
self._prob_table = defaultdict(lambda: 0.0, {value[0], value[1]})
elif isinstance(value, dict) or isinstance(value, defaultdict):
self._prob_table = defaultdict(lambda: 0.0, value)
else:
raise ValueError("Unknown value type", type(value))
def expected(self):
ex = 0
for val, prob in self.probabilities.items():
ex += val * prob
return ex
def variance(self):
# Squared deviation from the mean
return (self + (-self.expected())).squared()
def standard_deviation(self):
return math.sqrt(self.variance().expected())
def squared(self):
return self.apply(lambda x: x*x)
@classmethod
def convert(cls, value):
if isinstance(value, cls):
return value
res = cls()
if isinstance(value, int) or isinstance(value, float):
res.probabilities = {value: 1.0}
else:
res.probabilities = value
return res
def prob_table(self):
return tuple(zip(*self._prob_table.items()))
def normalize(self):
summed = sum(self.probabilities.values())
self.probabilities = {k: v/summed for k, v in self.probabilities.items()}
def roll(self, count=1):
values, probs = self.prob_table()
res = random.choices(values, weights=probs, k=count)
return sum(res)
def show(self):
values, probs = self.prob_table()
self.__show_plot(values, probs, "Value", "Probability")
def test(self, sample_count):
values = list(self.probabilities.keys())
counts = defaultdict(lambda: 0)
for _ in range(sample_count):
counts[self.roll()] += 1
self.__show_plot(values, [counts[v] for v in values], "Value", "Times rolled")
@staticmethod
def __show_plot(xax, yax, xlb, ylb):
fig, ax = plt.subplots()
ax.bar(xax, yax)
ax.set_ylabel(ylb)
ax.set_xlabel(xlb)
ax.xaxis.set_major_locator(ticker.FixedLocator(xax))
# ax.yaxis.set_major_locator(ticker.FixedLocator(list(set(yax))))
plt.grid(True)
plt.show()
def insert(self, value, probability=1.0):
for val, prob in self.convert(value).probabilities.items():
self._prob_table[val] += probability * prob
def accumulate_prob(self, pred):
res = 0
for val, prob in self.probabilities.items():
if pred(val):
res += prob
return res
def apply(self, operator, *args, **kwargs):
res = Distribution()
for lv, lprob in self.probabilities.items():
res.insert(operator(lv, *args, **kwargs), lprob)
return res
def apply2(self, operator, other, *args, **kwargs):
res = Distribution()
for lv, lprob in self.probabilities.items():
for rv, rprob in self.convert(other).probabilities.items():
res.insert(operator(lv, rv, *args, **kwargs), lprob*rprob)
return res
def __add__(self, other):
return self.apply2(lambda x, y: x + y, other)
def __neg__(self):
return self.apply(lambda x: -x)
def __rmul__(self, count):
if not isinstance(count, int):
raise ValueError("Multiplication allowed only for int")
if count == 0:
return Distribution({0, 1.0})
res = copy.deepcopy(self)
if count < 0:
res = -res
count = -count
for _ in range(count-1):
res += self
return res
|
