Source code for momba.utils.distribution

```# -*- coding:utf-8 -*-
#
# Copyright (C) 2020, Maximilian Köhl <koehl@cs.uni-saarland.de>

from __future__ import annotations

import typing as t

import fractions
import random

from mxu.maps import FrozenMap

ElementT = t.TypeVar("ElementT", bound=t.Hashable)

[docs]class Distribution(t.Generic[ElementT]):
"""A probability distribution."""

_mapping: FrozenMap[ElementT, fractions.Fraction]

@classmethod
def create_dirac(cls, element: ElementT) -> Distribution[ElementT]:
return cls({element: 1})

@classmethod
def create_uniform(cls, *elements: ElementT) -> Distribution[ElementT]:
probability = fractions.Fraction(1, len(elements))
return cls({element: probability for element in elements})

def __init__(
self, mapping: t.Mapping[ElementT, t.Union[int, float, fractions.Fraction]]
) -> None:
self._mapping = FrozenMap.transfer_ownership(
{
element: fractions.Fraction(probability)
for element, probability in mapping.items()
}
)
assert all(probability >= 0 for probability in self._mapping.values())
# assert sum(self._mapping.values()) == 1

def __str__(self) -> str:
return f"Distribution({self._mapping})"

@property
def support(self) -> t.List[ElementT]:
return list(
element for element, probability in self._mapping.items() if probability > 0
)

@property
def is_dirac(self) -> bool:
return len(self.support) == 1

def get_probability(self, element: ElementT) -> fractions.Fraction:
return self._mapping.get(element, fractions.Fraction(0))

[docs]    def pick(self) -> ElementT:
"""Picks an element at random according to the distribution."""
max_denominator = max(
probability.denominator for probability in self._mapping.values()
)
outcome = fractions.Fraction(
random.randint(0, max_denominator), max_denominator
)
total = fractions.Fraction(0)
for element in self.support:
total += self.get_probability(element)
if outcome <= total:
return element
raise RuntimeError("empty distribution is not possible")
```