diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 133 |
1 files changed, 1 insertions, 132 deletions
@@ -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): |