def test_appearances_in_rolls(self) -> None:
def _sum_method(p: P, outcome: _OutcomeT) -> H:
return H((sum(1 for v in roll if v == outcome), count)
for roll, count in p.rolls_with_counts())
p_empty = P()
assert p_empty.appearances_in_rolls(0) == _sum_method(p_empty, 0)
assert p_empty.appearances_in_rolls(0) == H({}).eq(0)
p_4d6 = 4 @ P(6)
assert p_4d6.appearances_in_rolls(2) == _sum_method(p_4d6, 2)
assert p_4d6.appearances_in_rolls(2) == 4 @ H(6).eq(2)
assert p_4d6.appearances_in_rolls(7) == _sum_method(p_4d6, 7)
assert p_4d6.appearances_in_rolls(7) == 4 @ H(6).eq(7)
p_mixed = P(4, 4, 6, 6, 6, H({}))
assert p_mixed.appearances_in_rolls(3) == _sum_method(p_mixed, 3)
assert p_mixed.appearances_in_rolls(3) == (2 @ H(4).eq(3) +
3 @ H(6).eq(3) +
H({}).eq(3))
assert p_mixed.appearances_in_rolls(5) == _sum_method(p_mixed, 5)
assert p_mixed.appearances_in_rolls(5) == (2 @ H(4).eq(5) +
3 @ H(6).eq(5) +
H({}).eq(5))
assert p_mixed.appearances_in_rolls(7) == _sum_method(p_mixed, 7)
assert p_mixed.appearances_in_rolls(7) == (2 @ H(4).eq(7) +
3 @ H(6).eq(7) +
H({}).eq(7))
def do_it(_: str) -> None:
import matplotlib.pyplot
p_4d6 = 4 @ P(6)
res1 = p_4d6.h(slice(1, None))
d6_reroll_first_one = H(6).substitute(lambda h, outcome: H(6)
if outcome == 1 else outcome)
p_4d6_reroll_first_one = 4 @ P(d6_reroll_first_one)
res2 = p_4d6_reroll_first_one.h(slice(1, None))
p_4d6_reroll_all_ones = 4 @ P(H(range(2, 7)))
res3 = p_4d6_reroll_all_ones.h(slice(1, None))
matplotlib.pyplot.plot(
*res1.distribution_xy(),
marker=".",
label="Discard lowest",
)
matplotlib.pyplot.plot(
*res2.distribution_xy(),
marker=".",
label="Re-roll first 1; discard lowest",
)
matplotlib.pyplot.plot(
*res3.distribution_xy(),
marker=".",
label="Re-roll all 1s; discard lowest",
)
matplotlib.pyplot.legend()
# Should match the corresponding img[alt] text
matplotlib.pyplot.title(r"Comparing various take-three-of-4d6 methods")
def test_op_sub_num(self) -> None:
p_d6 = P(6)
p_minus_d6 = P(H(range(0, -6, -1)))
p_d8 = P(8)
p_d8_minus = P(H(range(0, 8)))
assert 1 - p_d6 == p_minus_d6
assert p_d8_minus == p_d8 - 1
def test_op_add_num(self) -> None:
p_d6 = P(6)
p_d6_plus = P(H(range(2, 8)))
p_d8 = P(8)
p_d8_plus = P(H(range(2, 10)))
assert 1 + p_d6 == p_d6_plus
assert p_d8_plus == p_d8 + 1
def test_op_unary(self) -> None:
p = P(H(-v if v % 2 else v for v in range(10, 20)))
assert isinstance(+p, type(p))
assert (+p) == P(H((10, -11, 12, -13, 14, -15, 16, -17, 18, -19)))
assert isinstance(-p, type(p))
assert (-p) == P(H((-10, 11, -12, 13, -14, 15, -16, 17, -18, 19)))
assert isinstance(abs(p), type(p))
assert abs(p) == P(H((10, 11, 12, 13, 14, 15, 16, 17, 18, 19)))
def test_init_multiple_histograms(self) -> None:
d6 = H(6)
p_d6 = P(6)
p_2d6 = P(d6, d6)
assert p_2d6.h() == H(sum(v) for v in itertools.product(d6, d6))
assert P(p_d6, p_d6) == p_2d6
assert P(p_d6, d6) == p_2d6
assert P(d6, p_d6) == p_2d6
def test_getitem_slice(self) -> None:
d4n = H(-4)
d8 = H(8)
p_3d4n_3d8 = 3 @ P(d4n, d8)
assert p_3d4n_3d8[:] == p_3d4n_3d8
assert p_3d4n_3d8[:0] == P()
assert p_3d4n_3d8[6:] == P()
assert p_3d4n_3d8[2:4] == P(d4n, d8)
def test_h_flatten(self) -> None:
r_d6 = range(1, 7)
r_d8 = range(1, 9)
d6_d8 = H(sum(v) for v in itertools.product(r_d6, r_d8) if v)
p_d6 = P(6)
p_d8 = P(8)
p_d6_d8 = P(p_d6, p_d8)
assert p_d6_d8.h() == d6_d8
assert P().h() == H({})
def test_equivalence(self) -> None:
p_d6 = P(6)
p_d6n = P(-6)
assert -p_d6 == p_d6n
assert p_d6 - p_d6 == p_d6 + p_d6n
assert -p_d6 + p_d6 == p_d6n + p_d6
assert -p_d6 - p_d6 == p_d6n - p_d6
assert p_d6 + p_d6 == p_d6 - p_d6n
assert P(p_d6, -p_d6) == p_d6 + p_d6n
assert P(p_d6n, -p_d6n) == p_d6n + p_d6
assert 2 @ p_d6 - p_d6 == p_d6 + p_d6 + p_d6n
assert -(2 @ p_d6) == p_d6n + p_d6n
def test_op_truediv_h(self) -> None:
d2 = H(2)
d3n = H(-3)
p_d2 = P(d2)
p_d3n = P(d3n)
d2_truediv_d3n = d2 / d3n
d3n_truediv_d2 = d3n / d2
assert p_d2 / p_d3n == d2_truediv_d3n
assert p_d2 / d3n == d2_truediv_d3n
assert d2 / p_d3n == d2_truediv_d3n
assert p_d3n / p_d2 == d3n_truediv_d2
assert p_d3n / d2 == d3n_truediv_d2
assert d3n / p_d2 == d3n_truediv_d2
assert p_d2 / p_d3n != p_d3n / p_d2
def test_rolls_with_counts_take_heterogeneous_dice_vs_known_correct(
self) -> None:
p_d3 = P(3)
p_d4n = P(-4)
p_3d3_4d4n = P(3 @ p_d3, 4 @ p_d4n)
for which in (
# All outcomes
slice(None),
# 4 outcomes
slice(4),
slice(-4, None),
):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_3d3_4d4n, which)
karonen.assert_not_called()
multinomial_coefficient.assert_called()
for which in (
# 3 outcomes
slice(3),
slice(-3, None),
):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_3d3_4d4n, which)
karonen.assert_called() # called for 4d4n where k < n
multinomial_coefficient.assert_called(
) # called for 3d3 where k == n
for which in (
# 2 outcomes
slice(2),
slice(-2, None),
# 1 outcome
slice(1),
slice(-1, None),
):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_3d3_4d4n, which)
karonen.assert_called()
multinomial_coefficient.assert_not_called()
for which in (
# No outcomes
slice(0, 0), ):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_3d3_4d4n, which)
karonen.assert_not_called()
multinomial_coefficient.assert_not_called()
def test_rolls_with_counts_take_homogeneous_dice_vs_known_correct(
self) -> None:
p_df = P(H((-1, 0, 1)))
p_4df = 4 @ p_df
for which in (
# All outcomes
slice(None),
slice(0, 4),
):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_4df, which)
karonen.assert_not_called()
multinomial_coefficient.assert_called()
for which in (
# 1 Outcome
slice(0, 1),
slice(1, 2),
slice(2, 3),
slice(3, 4),
):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_4df, which)
karonen.assert_called()
multinomial_coefficient.assert_not_called()
for which in (
# No outcomes
slice(0, 0), ):
karonen, multinomial_coefficient = _rwc_validation_helper(
p_4df, which)
karonen.assert_not_called()
multinomial_coefficient.assert_not_called()
def test_op_truediv_num(self) -> None:
p_d10 = P(10)
p1 = P(H(range(100, 0, -10)))
assert p_d10 == p1 / 10
assert (2 * 2 * 2 * 3 * 3 * 5 * 7) / p_d10 == H({
252.0: 1,
280.0: 1,
315.0: 1,
360.0: 1,
420.0: 1,
504.0: 1,
630.0: 1,
840.0: 1,
1260.0: 1,
2520.0: 1,
})
def dupes(p: P):
for roll, count in p.rolls_with_counts():
dupes = 0
for i in range(1, len(roll)):
if roll[i] == roll[i - 1]:
dupes += 1
yield dupes, count
def do_it(_: str) -> None:
import matplotlib.pyplot
p_2d6 = 2 @ P(H(6))
outcomes, probabilities = p_2d6.h(0).distribution_xy()
matplotlib.pyplot.bar(
[v - 0.125 for v in outcomes],
probabilities,
alpha=0.75,
width=0.5,
label="Lowest",
)
outcomes, probabilities = p_2d6.h(-1).distribution_xy()
matplotlib.pyplot.bar(
[v + 0.125 for v in outcomes],
probabilities,
alpha=0.75,
width=0.5,
label="Highest",
)
matplotlib.pyplot.legend()
# Should match the corresponding img[alt] text
matplotlib.pyplot.title(r"Taking the lowest or highest die of 2d6")
def test_h_take_heterogeneous_dice_vs_known_correct(self) -> None:
p_d3 = P(3)
p_d3n = -p_d3
p_d4 = P(4)
p_d4n = -p_d4
p_4d3_4d4 = 2 @ P(p_d3, p_d3n, p_d4n, p_d4)
for which in (
slice(0, 0),
slice(-1, None),
slice(-2, None),
slice(2),
slice(1),
):
assert p_4d3_4d4.h(which) == H(
(sum(roll), count)
for roll, count in _brute_force_combinations_with_counts(
tuple(p_4d3_4d4), which))
def do_it(_: str) -> None:
import matplotlib.pyplot
res = (10 @ P(H(10).explode(max_depth=3))).h(slice(-3, None))
matplotlib.pyplot.plot(*res.distribution_xy(), marker=".")
# Should match the corresponding img[alt] text
matplotlib.pyplot.title(
r"Modeling taking the three highest of ten exploding d10s")
def test_repr(self) -> None:
assert repr(P()) == "P()"
assert repr(P(0)) == "P()"
assert repr(P(-6)) == "P(-6)"
assert repr(P(6)) == "P(6)"
assert (repr(P(P(6), P(8), P(H({
3: 1,
2: 2,
1: 3,
0: 1
})))) == "P(H({0: 1, 1: 3, 2: 2, 3: 1}), 6, 8)")
def test_getitem_int(self) -> None:
d4n = H(-4)
d8 = H(8)
p_3d4n_3d8 = 3 @ P(d4n, d8)
assert p_3d4n_3d8[0] == d4n
assert p_3d4n_3d8[2] == d4n
assert p_3d4n_3d8[-3] == d8
assert p_3d4n_3d8[-1] == d8
with pytest.raises(IndexError):
_ = p_3d4n_3d8[6]
def test_op_eq(self) -> None:
p_d6 = P(6)
p_d6_2 = P(H(range(6, 0, -1)))
p_d6_3 = P(p_d6_2)
assert p_d6_2 == p_d6
assert p_d6_3 == p_d6_2
p_d4n = P(-4)
p_d6 = P(6)
p_d4n_d6 = P(p_d4n, p_d6)
p_d6_d4n = P(p_d6, p_d4n)
assert p_d4n_d6 == p_d6_d4n
assert p_d4n_d6.h() == p_d6_d4n
assert p_d4n_d6 == p_d6_d4n.h()
assert p_d4n_d6.h() == p_d6_d4n.h()
assert P(p_d6, -4) == p_d4n_d6
assert P(p_d4n, 6) == p_d4n_d6
def test_rolls_with_counts_heterogeneous(self) -> None:
assert sorted(P(2, 3).rolls_with_counts()) == [
((1, 1), 1),
((1, 2), 1),
(
(1, 2),
1,
), # originated as ((2, 1), 1), but outcomes get sorted in each roll
((1, 3), 1),
((2, 2), 1),
((2, 3), 1),
]
def test_op_floordiv_h(self) -> None:
d2 = H(2)
d3n = H(-3)
p_d2 = P(d2)
p_d3n = P(d3n)
d2_floordiv_d3n = d2 // d3n
d3n_floordiv_d2 = d3n // d2
assert p_d2 // p_d3n == d2_floordiv_d3n
assert p_d2 // d3n == d2_floordiv_d3n
assert d2 // p_d3n == d2_floordiv_d3n
assert p_d3n // p_d2 == d3n_floordiv_d2
assert p_d3n // d2 == d3n_floordiv_d2
assert d3n // p_d2 == d3n_floordiv_d2
assert p_d2 // p_d3n != p_d3n // p_d2
assert p_d2 // P() == P()
assert P() // p_d2 == P()
assert P() // P() == P()
def test_op_mod_h(self) -> None:
d2 = H(2)
d3n = H(-3)
p_d2 = P(d2)
p_d3n = P(d3n)
d2_mod_d3n = d2 % d3n
d3n_mod_d2 = d3n % d2
assert p_d2 % p_d3n == d2_mod_d3n
assert p_d2 % d3n == d2_mod_d3n
assert d2 % p_d3n == d2_mod_d3n
assert p_d3n % p_d2 == d3n_mod_d2
assert p_d3n % d2 == d3n_mod_d2
assert d3n % p_d2 == d3n_mod_d2
assert p_d2 % p_d3n != p_d3n % p_d2
assert p_d2 % P() == P()
assert P() % p_d2 == P()
assert P() % P() == P()
def test_op_pow_h(self) -> None:
d2 = H(2)
d3 = H(3)
p_d2 = P(d2)
p_d3 = P(d3)
d2_pow_d3 = d2**d3
d3_pow_d2 = d3**d2
assert p_d2**p_d3 == d2_pow_d3
assert p_d2**d3 == d2_pow_d3
assert d2**p_d3 == d2_pow_d3
assert p_d3**p_d2 == d3_pow_d2
assert p_d3**d2 == d3_pow_d2
assert d3**p_d2 == d3_pow_d2
assert p_d2**p_d3 != p_d3**p_d2
assert p_d2**P() == P()
assert P()**p_d2 == P()
assert P()**P() == P()
def test_op_mul_h(self) -> None:
d2 = H(2)
d3n = H(-3)
p_d2 = P(d2)
p_d3n = P(d3n)
d2_mul_d3n = d2 * d3n
d3n_mul_d2 = d3n * d2
assert p_d2 * p_d3n == d2_mul_d3n
assert p_d2 * d3n == d2_mul_d3n
assert d2 * p_d3n == d2_mul_d3n
assert p_d3n * p_d2 == d3n_mul_d2
assert p_d3n * d2 == d3n_mul_d2
assert d3n * p_d2 == d3n_mul_d2
assert d2_mul_d3n == d3n_mul_d2
assert p_d2 * p_d3n == p_d3n * p_d2
assert p_d2 * P() == P()
assert P() * p_d2 == P()
assert P() * P() == P()
def test_h_take_heterogeneous_dice(self) -> None:
p_d3 = P(3)
p_d3n = -p_d3
p_d4 = P(4)
p_d4n = -p_d4
p_4d3_4d4 = 2 @ P(p_d3, p_d3n, p_d4n, p_d4)
with pytest.raises(IndexError):
_ = p_4d3_4d4.h(len(p_4d3_4d4))
assert p_4d3_4d4.h(slice(0, 0)) == {}
assert p_4d3_4d4.h(slice(None)) == p_4d3_4d4.h()
assert p_4d3_4d4.h(*range(len(p_4d3_4d4))) == p_4d3_4d4.h()
assert p_4d3_4d4.h(slice(1)) == p_4d3_4d4.h(0)
assert p_4d3_4d4.h(slice(-1, None)) == p_4d3_4d4.h(-1)
assert p_4d3_4d4.h(0, 2) == p_4d3_4d4.h(slice(None, 3, 2))
assert p_4d3_4d4.h(0, slice(2, None,
2)) == p_4d3_4d4.h(slice(None, None, 2))
assert p_4d3_4d4.h(0, 2, 4, 6) == p_4d3_4d4.h(slice(None, None, 2))
assert p_4d3_4d4.h(*range(0, len(p_4d3_4d4), 2)) == p_4d3_4d4.h(
slice(None, None, 2))
assert p_4d3_4d4.h(*(i for i in range(len(p_4d3_4d4))
if i & 0x1)) == p_4d3_4d4.h(slice(1, None, 2))
def test_h_take_homogeneous_dice_vs_known_correct(self) -> None:
# Use the brute-force mechanism to validate our harder-to-understand
# implementation
p_df = P(H((-1, 0, 1)))
p_4df = 4 @ p_df
for which in (
slice(0, 0),
slice(2, 3),
slice(1, 2),
slice(0, 1),
):
assert p_4df.h(which) == H(
(sum(roll), count)
for roll, count in _brute_force_combinations_with_counts(
tuple(p_4df), which))
def test_len(self) -> None:
p_d0_1 = P()
p_d0_2 = P(H({}))
p_d6 = P(6)
p_d8 = P(8)
assert len(p_d0_1) == 0
assert len(p_d0_2) == 0
assert len(p_d6) == 1
assert len(p_d8) == 1
assert len(P(p_d6, p_d8)) == 2
assert len(P(p_d6, p_d8, p_d6, p_d8)) == 4
def test_op_sub_h(self) -> None:
d2 = H(2)
d3n = H(-3)
p_d2 = P(d2)
p_d3n = P(d3n)
d2_sub_d3n = d2 - d3n
d3n_sub_d2 = d3n - d2
assert p_d2 - p_d3n == d2_sub_d3n
assert p_d2 - d3n == d2_sub_d3n
assert d2 - p_d3n == d2_sub_d3n
assert p_d3n - p_d2 == d3n_sub_d2
assert p_d3n - d2 == d3n_sub_d2
assert d3n - p_d2 == d3n_sub_d2
assert p_d2 - p_d3n != p_d3n - p_d2
assert p_d2 - P() == p_d2
assert P() - p_d2 == -p_d2
assert P() - P() == P()
def do_it(style: str) -> None:
import matplotlib.pyplot
def dupes(p: P):
for roll, count in p.rolls_with_counts():
dupes = 0
for i in range(1, len(roll)):
if roll[i] == roll[i - 1]:
dupes += 1
yield dupes, count
res = H(dupes(8 @ P(10)))
plot_burst(
res,
# Should match the corresponding img[alt] text
desc=r"Chances of rolling $n$ duplicates in 8d10",
text_color="white" if style == "dark" else "black",
)
matplotlib.pyplot.tight_layout()