Skip to content

Commit

Permalink
add sequence() method to Die and Deck
Browse files Browse the repository at this point in the history
  • Loading branch information
HighDiceRoller committed Aug 30, 2024
1 parent c4ed32e commit 463c271
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 2 deletions.
27 changes: 26 additions & 1 deletion src/icepool/population/deck.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

from collections import Counter, defaultdict
from functools import cached_property
from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence, Type, overload
from typing import Any, Callable, Iterable, Iterator, Mapping, MutableSequence, Sequence, Type, overload


class Deck(Population[T_co]):
Expand Down Expand Up @@ -292,6 +292,31 @@ def transition_function(outcome):
[transition_function(outcome) for outcome in self.outcomes()],
times=self.quantities())

@cached_property
def _sequence_cache(
self) -> 'MutableSequence[icepool.Die[tuple[T_co, ...]]]':
return [icepool.Die([()])]

def sequence(self, deals: int, /) -> 'icepool.Die[tuple[T_co, ...]]':
"""Possible sequences produced by dealing from this deck a number of times.
This is extremely expensive computationally. If you don't care about
order, use `deal()` instead.
"""
if deals < 0:
raise ValueError('The number of deals cannot be negative.')
for i in range(len(self._sequence_cache), deals + 1):

def transition(curr):
remaining = icepool.Die(self - curr)
return icepool.map(lambda curr, next: curr + (next, ), curr,
remaining)

result: 'icepool.Die[tuple[T_co, ...]]' = self._sequence_cache[
i - 1].map(transition)
self._sequence_cache.append(result)
return result

@cached_property
def _hash_key(self) -> tuple:
return Deck, tuple(self.items())
Expand Down
9 changes: 9 additions & 0 deletions src/icepool/population/die.py
Original file line number Diff line number Diff line change
Expand Up @@ -794,6 +794,15 @@ def __rmatmul__(self, other: 'int | Die[int]') -> 'Die':
other = implicit_convert_to_die(other)
return other.__matmul__(self)

def sequence(self, rolls: int) -> 'icepool.Die[tuple[T_co, ...]]':
"""Possible sequences produced by rolling this die a number of times.
This is extremely expensive computationally. If possible, use `reduce()`
instead; if you don't care about order, `Die.pool()` is better.
"""
return icepool.cartesian_product(*(self for _ in range(rolls)),
outcome_type=tuple) # type: ignore

def pool(self, rolls: int | Sequence[int] = 1, /) -> 'icepool.Pool[T_co]':
"""Creates a `Pool` from this `Die`.
Expand Down
23 changes: 22 additions & 1 deletion tests/vector_test.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import icepool
import pytest

from icepool import d, d4, d6, d8, vectorize, Die, Vector
from icepool import d, d4, d6, d8, vectorize, Die, Vector, map, Deck
from collections import namedtuple


Expand Down Expand Up @@ -150,3 +150,24 @@ def test_named_tuple():

def test_auto_tupleize():
assert Die([(d6, d6, 6)]).map(sum) == 2 @ d6 + 6


def test_die_sequence():
assert d6.sequence(3) == map(lambda a, b, c: (a, b, c), d6, d6, d6)


def test_deck_sequence():
result = Deck(range(5)).sequence(3)
assert len(result) == 5 * 4 * 3
assert result.map(
lambda a, b, c: a != b and a != c and b != c).probability(True) == 1


def test_deck_sequence_with_dups():
result = Deck([0, 0, 1, 1]).sequence(2)
assert result == Die({
(0, 0): 2,
(0, 1): 4,
(1, 0): 4,
(1, 1): 2,
})

0 comments on commit 463c271

Please sign in to comment.