diff options
| -rw-r--r-- | distribution.py | 132 | ||||
| -rw-r--r-- | main.py | 133 |
2 files changed, 133 insertions, 132 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 + @@ -1,136 +1,5 @@ -import copy import math -import random -import numpy as np -from collections import defaultdict -import matplotlib.pyplot as plt -import matplotlib.ticker as ticker -from enum import Enum - - -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 +from distribution import Distribution def d(edge_count): |