add Population.pad_denominator() method #197

move `zero()`, `zero_outcome()` from `Die` to `Population`

src/icepool/population/base.py

@@ -146,6 +146,32 @@ def nearest(self, comparison: Literal['<=', '<', '>=', '>'], outcome,
             case _:
                 raise ValueError(f'Invalid comparison {comparison}')
 
+    @staticmethod
+    def _zero(x):
+        return x * 0
+
+    def zero(self: C) -> C:
+        """Zeros all outcomes of this population.
+
+        This is done by multiplying all outcomes by `0`.
+
+        The result will have the same denominator.
+
+        Raises:
+            ValueError: If the zeros did not resolve to a single outcome.
+        """
+        result = self._unary_operator(Population._zero)
+        if len(result) != 1:
+            raise ValueError('zero() did not resolve to a single outcome.')
+        return result
+
+    def zero_outcome(self) -> T_co:
+        """A zero-outcome for this population.
+
+        E.g. `0` for a `Population` whose outcomes are `int`s.
+        """
+        return self.zero().outcomes()[0]
+
     # Quantities.
 
     @overload
@@ -279,6 +305,35 @@ def has_zero_quantities(self) -> bool:
         """`True` iff `self` contains at least one outcome with zero quantity. """
         return 0 in self.values()
 
+    def pad_denominator(self: C, target: int, /, outcome: Hashable) -> C:
+        """Changes the denominator to a target number by changing the quantity of a specified outcome.
+        
+        Args:
+            `target`: The denominator of the result.
+            `outcome`: The outcome whose quantity will be adjusted.
+
+        Returns:
+            A `Population` like `self` but with the quantity of `outcome`
+            adjusted so that the overall denominator is equal to  `target`.
+            If the denominator is reduced to zero, it will be removed.
+
+        Raises:
+            `ValueError` if this would require the quantity of the specified
+            outcome to be negative.
+        """
+        adjustment = target - self.denominator()
+        data = {outcome: quantity for outcome, quantity in self.items()}
+        new_quantity = data.get(outcome, 0) + adjustment
+        if new_quantity > 0:
+            data[outcome] = new_quantity
+        elif new_quantity == 0:
+            del data[outcome]
+        else:
+            raise ValueError(
+                f'Padding to denominator of {target} would require a negative quantity of {new_quantity} for {outcome}'
+            )
+        return self._new_type(data)
+
     # Probabilities.
 
     @overload

src/icepool/population/die.py

@@ -1272,32 +1272,6 @@ def ceil(self) -> 'Die':
 
     __ceil__ = ceil
 
-    @staticmethod
-    def _zero(x):
-        return x * 0
-
-    def zero(self) -> 'Die[T_co]':
-        """Zeros all outcomes of this die.
-
-        This is done by multiplying all outcomes by `0`.
-
-        The result will have the same denominator as this die.
-
-        Raises:
-            ValueError: If the zeros did not resolve to a single outcome.
-        """
-        result = self.unary_operator(Die._zero)
-        if len(result) != 1:
-            raise ValueError('zero() did not resolve to a single outcome.')
-        return result
-
-    def zero_outcome(self) -> T_co:
-        """A zero-outcome for this die.
-
-        E.g. `0` for a `Die` whose outcomes are `int`s.
-        """
-        return self.zero().outcomes()[0]
-
     # Binary operators.
 
     def __add__(self, other) -> 'Die':

tests/statistics_test.py

@@ -1,7 +1,7 @@
 import icepool
 import pytest
 
-from icepool import Die, d6
+from icepool import Die, d6, Deck
 from fractions import Fraction
 
 
@@ -53,3 +53,23 @@ def test_percent(comparison):
     die = 3 @ d6
     assert die.probability(
         comparison, 4, percent=True) == die.probability(comparison, 4) * 100.0
+
+
+def test_pad_to_denominator_add():
+    deck = Deck([0, 1, 2, 3]).pad_denominator(6, 0)
+    assert deck == Deck([0, 0, 0, 1, 2, 3])
+
+
+def test_pad_to_denominator_remove():
+    deck = Deck([0, 0, 0, 1, 2, 3]).pad_denominator(4, 0)
+    assert deck == Deck([0, 1, 2, 3])
+
+
+def test_pad_to_denominator_zero():
+    deck = Deck([0, 0, 0, 1, 2, 3]).pad_denominator(3, 0)
+    assert deck == Deck([1, 2, 3])
+
+
+def test_pad_to_denominator_negative_error():
+    with pytest.raises(ValueError):
+        deck = Deck([0, 0, 0, 1, 2, 3]).pad_denominator(2, 0)