1 files changed, 132 insertions, 0 deletions
diff --git a/distribution.py b/distribution.py
new file mode 100644
index 0000000..d2503fa
--- /dev/null
+++ b/distribution.py
@@ -0,0 +1,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
+