Dice pools

CHANGES.md

@@ -3,6 +3,7 @@
 * **[Breaking change]** This package now only supports Python 3.9 and higher.
 * **[Breaking change]** `get_percentiles` and `get_log_percentiles` now always return a dictionary, even if there's only one element.
 * **[Breaking change]** `.type` is now removed from distribution objects.
+* **[Breaking change]** `sq.flip_coin` and `sq.roll_die` as functions are now functionally replaced by sampling from `sq.coin` or `sq.die` distribution objects.
 * Package load time is now ~2x faster.
 * Pandas and matplotlib as removed as required dependencies, but their related features are lazily enabled when the modules are available. These packages are still available for install  as extras, installable with `pip install squigglepy[plots]` (for plotting-related functionality, matplotlib for now), `pip install squigglepy[ecosystem]` (for pandas, and in the future other related packages), or `pip install squigglepy[all]` (for all extras).
 * Multicore distribution now does extra checks to avoid crashing from race conditions.

README.md

@@ -181,7 +181,8 @@ roll_die(sides=6, n=10)
 # [2, 6, 5, 2, 6, 2, 3, 1, 5, 2]
 ```
 
-This is already included standard in the utils of this package. Use `sq.roll_die`.
+(This is a toy example. For better dice rolling, use the built-in `sq.die`)
+
 
 ### Bayesian inference
 
@@ -396,10 +397,10 @@ from squigglepy.numbers import K, M, B, T
 from squigglepy import bayes
 
 def define_event():
-    if sq.flip_coin() == 'heads': # Blue bag
-        return sq.roll_die(6)
+    if ~sq.coin() == 'heads': # Blue bag
+        return sq.die(6)
     else: # Red bag
-        return sq.discrete([4, 6, 10, 20]) >> sq.roll_die
+        return sq.discrete([4, 6, 10, 20]) >> sq.die
 
 
 bayes.bayesnet(define_event,

squigglepy/__init__.py

@@ -1,4 +1,5 @@
 from .distributions import *  # noqa ignore=F405
+from .dice import *  # noqa ignore=F405
 from .numbers import *  # noqa ignore=F405
 from .samplers import *  # noqa ignore=F405
 from .utils import *  # noqa ignore=F405

squigglepy/dice.py

@@ -0,0 +1,82 @@
+from .utils import is_dist
+from .distributions import OperableDistribution
+
+
+class Die(OperableDistribution):
+    def __init__(self, sides=None, explode_on=None):
+        super().__init__()
+
+        if is_dist(sides) or callable(sides):
+            from .samplers import sample
+
+            sides = sample(sides)
+        elif sides < 2:
+            raise ValueError("cannot roll less than a 2-sided die.")
+        elif not isinstance(sides, int):
+            raise ValueError("can only roll an integer number of sides")
+        elif explode_on is not None:
+            if not isinstance(explode_on, list):
+                explode_on = [explode_on]
+            if len(explode_on) >= sides:
+                raise ValueError("cannot explode on every value")
+        self.sides = sides
+        self.explode_on = explode_on
+
+    def __str__(self):
+        explode_out = (
+            "" if self.explode_on is None else ", explodes on {}".format(str(self.explode_on))
+        )
+        out = "<Distribution> Die({}{})".format(self.sides, explode_out)
+        return out
+
+
+def die(sides, explode_on=None):
+    """
+    Create a distribution for a die.
+
+    Parameters
+    ----------
+    sides : int
+        The number of sides of the die that is rolled.
+
+    explode_on : list or int
+        An additional die will be rolled if the initial die rolls any of these values.
+        Implements "exploding dice".
+
+    Returns
+    -------
+    Distribution
+        A distribution that models a coin flip, returning either "heads" or "tails"
+
+    Examples
+    --------
+    >>> die(6)
+    <Distribution> Die(6)
+    """
+    return Die(sides=sides, explode_on=explode_on)
+
+
+class Coin(OperableDistribution):
+    def __init__(self):
+        super().__init__()
+
+    def __str__(self):
+        out = "<Distribution> Coin"
+        return out
+
+
+def coin():
+    """
+    Create a distribution for a coin
+
+    Returns
+    -------
+    Distribution
+        A distribution that models a coin flip, returning either "heads" or "tails"
+
+    Examples
+    --------
+    >>> coin()
+    <Distribution> Coin
+    """
+    return Coin()

squigglepy/samplers.py

@@ -42,6 +42,8 @@
     const,
 )
 
+from .dice import Coin, Die
+
 _squigglepy_internal_sample_caches = {}
 
 
@@ -914,6 +916,34 @@ def run_dist(dist, pbar=None, tick=1):
                 _multicore_tqdm_cores=_multicore_tqdm_cores,
             )
 
+        elif isinstance(dist, Die):
+            samples = discrete_sample(
+                list(range(1, dist.sides + 1)),
+                samples=n,
+                _multicore_tqdm_n=_multicore_tqdm_n,
+                _multicore_tqdm_cores=_multicore_tqdm_cores,
+            )
+            if dist.explode_on is not None:
+                samples = _enlist(samples)
+                explosion_samples = samples
+                n_explosions = 1
+                while n_explosions > 0:
+                    explosion_samples = _enlist(explosion_samples)
+                    n_explosions = sum([s in dist.explode_on for s in explosion_samples])
+                    if n_explosions > 0:
+                        explosion_samples = discrete_sample(
+                            list(range(1, dist.sides + 1)), samples=n_explosions
+                        )
+                        samples.append(explosion_samples)
+
+        elif isinstance(dist, Coin):
+            samples = discrete_sample(
+                ["heads", "tails"],
+                samples=n,
+                _multicore_tqdm_n=_multicore_tqdm_n,
+                _multicore_tqdm_cores=_multicore_tqdm_cores,
+            )
+
         elif isinstance(dist, NormalDistribution):
             samples = normal_sample(mean=dist.mean, sd=dist.sd, samples=n)
 

squigglepy/utils.py

@@ -799,72 +799,6 @@ def doubling_time_to_growth_rate(doubling_time):
         return math.exp(math.log(2) / doubling_time) - 1
 
 
-def roll_die(sides, n=1):
-    """
-    Roll a die.
-
-    Parameters
-    ----------
-    sides : int
-        The number of sides of the die that is rolled.
-    n : int
-        The number of dice to be rolled.
-
-    Returns
-    -------
-    int or list
-        Returns the value of each die roll.
-
-    Examples
-    --------
-    >>> set_seed(42)
-    >>> roll_die(6)
-    5
-    """
-    if is_dist(sides) or callable(sides):
-        from .samplers import sample
-
-        sides = sample(sides)
-    if not isinstance(n, int):
-        raise ValueError("can only roll an integer number of times")
-    elif sides < 2:
-        raise ValueError("cannot roll less than a 2-sided die.")
-    elif not isinstance(sides, int):
-        raise ValueError("can only roll an integer number of sides")
-    else:
-        from .samplers import sample
-        from .distributions import discrete
-
-        return sample(discrete(list(range(1, sides + 1))), n=n) if sides > 0 else None
-
-
-def flip_coin(n=1):
-    """
-    Flip a coin.
-
-    Parameters
-    ----------
-    n : int
-        The number of coins to be flipped.
-
-    Returns
-    -------
-    str or list
-        Returns the value of each coin flip, as either "heads" or "tails"
-
-    Examples
-    --------
-    >>> set_seed(42)
-    >>> flip_coin()
-    'heads'
-    """
-    rolls = roll_die(2, n=n)
-    if isinstance(rolls, int):
-        rolls = [rolls]
-    flips = ["heads" if d == 2 else "tails" for d in rolls]
-    return flips[0] if len(flips) == 1 else flips
-
-
 def kelly(my_price, market_price, deference=0, bankroll=1, resolve_date=None, current=0):
     """
     Calculate the Kelly criterion.