Started prop_allocation()

.vscode/settings.json

@@ -0,0 +1,11 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true,
+    "[python]": {
+        "editor.formatOnSave": true,
+        "editor.defaultFormatter": "charliermarsh.ruff"
+    }
+}

docs/_quarto.yml

@@ -39,7 +39,12 @@ quartodoc:
         - InclusionProb
         - becomes_ta
     - title: Draw a sample
-      desc: Classes to draw sequential and ordinary Poisson samples
+      desc: Classes to draw sequential and ordinary Poisson samples.
       contents:
         - OrderSample
-        - PoissonSample
+        - PoissonSample
+    - title: Stratification
+      desc: Make an allocation for stratified sampling.
+      contents:
+        - prop_allocation
+        - divisor_method

pysps/__init__.py

@@ -4,7 +4,15 @@
 
 from pysps.inclusion_prob import InclusionProb, becomes_ta
 from pysps.order_sample import OrderSample, PoissonSample
+from pysps.allocate import prop_allocation, divisor_method
 
-__all__ = ["InclusionProb", "OrderSample", "PoissonSample", "becomes_ta"]
+__all__ = [
+    "InclusionProb",
+    "OrderSample",
+    "PoissonSample",
+    "becomes_ta",
+    "prop_allocation",
+    "divisor_method",
+]
 
-__version__ = "0.1.1"
+__version__ = "0.1.1.9002"

pysps/allocate.py

@@ -0,0 +1,148 @@
+import math
+from typing import Callable
+
+
+def divisor_method(name: str) -> Callable[[int], float]:
+    """Divisor function for highest averages apportionment method.
+
+    Parameters
+    ----------
+    name : str
+        Name of the divisor method, one of:
+        "Jefferson/D'Hondt",
+        "Webster/Sainte-Lague",
+        "Imperiali",
+        "Huntington-Hill",
+        "Danish",
+        "Adams",
+        "Dean"
+
+    Returns
+    -------
+    Callable[[int], float]
+        The divisor function.
+    """
+    name = str(name).lower()
+    if name == "jefferson/d'hondt":
+        return lambda x: x + 1.0
+    elif name == "webster/sainte-laguë":
+        return lambda x: x + 0.5
+    elif name == "imperiali":
+        return lambda x: x + 2.0
+    elif name == "huntington-hill":
+        return lambda x: math.sqrt(x * (x + 1.0))
+    elif name == "danish":
+        return lambda x: x + 1 / 3
+    elif name == "adams":
+        return lambda x: x
+    elif name == "Dean":
+        return lambda x: x * (x + 1.0) / (x + 0.5)
+    else:
+        return ValueError("'name' is not implemented")
+
+
+def prop_allocation(
+    x: dict,
+    n: int,
+    *,
+    initial: dict = {},
+    available: dict = {},
+    divisor: Callable[[int], float] = divisor_method("Jefferson/D'Hondt"),
+    ties: str = "largest",
+) -> dict:
+    """Proportional to size allocation for stratified sampling.
+
+    Parameters
+    ----------
+    x : dict
+        The size of each stratum.
+    n : int
+        Sample size.
+    initial : dict, optional
+        Initial allocation for strata in 'x'.
+    available : dict, optional
+        Number of available units for each strata in `x`.
+    divisor : Callable[[int], float], optional
+        A function for the divisor (highest-averages)
+        apportionment method. The default uses the Jefferson/D'Hondt method.
+    ties : {'largest', 'first'}, optional
+        Either 'largest' to break ties in favor of the stratum with the
+        largest size (the default), or 'first' to break ties in favor of the
+        ordering of 'x'.
+
+    Returns
+    -------
+    dict
+        Allocation for each stratum in 'x'.
+
+    References
+    ----------
+    Balinksi, M. L. and Young, H. P. (1982).
+    *Fair Representation: Meeting the Ideal of One Man, One Vote*.
+    Yale University Press.
+
+    Examples
+    --------
+    ```{python}
+    import pysps
+
+    pysps.prop_allocation({"a": 1, "b": 2, "c": 3}, 5)
+    ```
+
+    ```{python}
+    pysps.prop_allocation(
+        {"a": 1, "b": 2, "c": 3},
+        5,
+        divisor=pysps.divisor_method("Danish")
+    )
+    ```
+    """
+    if not set(x.keys()).issuperset(initial.keys()):
+        raise ValueError("all keys in 'initial' must also be in 'x'")
+    if not set(x.keys()).issuperset(available.keys()):
+        raise ValueError("all keys in 'available' must also be in 'x'")
+    if any(x < 0 for x in x.values()):
+        raise ValueError("all elements of 'x' must be positive")
+    if any(x < 0 for x in initial.values()):
+        raise ValueError("all elements of 'initial' must be positive")
+
+    res = dict.fromkeys(x.keys(), 0)
+    res.update(initial)
+
+    n = int(n)
+    if n < 0:
+        raise ValueError("'n' must be greater than or equal to 0")
+
+    upper = dict.fromkeys(x.keys(), n)
+    upper.update(available | {k: 0 for k, v in x.items() if v == 0})
+
+    if n < sum(res.values()):
+        raise ValueError("'n' is smaller than initial allocation")
+    if n > sum(upper.values()):
+        raise ValueError("'n' is larger than number of available units")
+    for k in res:
+        if res[k] > upper[k]:
+            raise ValueError(
+                "initial allocation must be smaller that number of available units"
+            )
+
+    ord = list(x.keys())
+    if ties == "largest":
+        x = {k: x[k] for k in sorted(x, key=x.get, reverse=True)}
+    elif ties == "first":
+        pass
+    else:
+        raise ValueError("'ties' must be either 'largest' or 'first'")
+
+    s = {k: v / divisor(res[k]) for k, v in x.items() if res[k] < upper[k]}
+    n -= sum(res.values())
+    while n > 0:
+        k = max(s, key=s.get)
+        res[k] += 1
+        n -= 1
+        if res[k] < upper[k]:
+            s[k] = x[k] / divisor(res[k])
+        else:
+            del s[k]
+
+    return {k: res[k] for k in ord}

pysps/inclusion_prob.py

@@ -7,19 +7,15 @@
 
 
 def _pi(x: np.ndarray, n: int) -> np.ndarray:
-    """
-    Unbounded first-order inclusion probabilities.
-    """
+    """Unbounded first-order inclusion probabilities."""
     if n == 0:
         return np.repeat(0.0, len(x))
     else:
         return x * (n / np.sum(x))
 
 
 def _which_ta(x: np.ndarray, n: int, alpha: float, sort_method: str) -> np.ndarray:
-    """
-    Indices for take-all units.
-    """
+    """Indices for take-all units."""
     # Sorting should be stable if there are ties in x.
     # Sorting in reverse order then flipping means ties resolve
     # according to the order of x.
@@ -38,9 +34,7 @@ def _which_ta(x: np.ndarray, n: int, alpha: float, sort_method: str) -> np.ndarr
 
 
 def _validate_input(x: npt.ArrayLike, n: int, alpha: float, cutoff: float) -> None:
-    """
-    Validate inputs for inclusion probabilities.
-    """
+    """Validate inputs for inclusion probabilities."""
     if np.any(x < 0.0):
         raise ValueError("elements of x must be greater than or equal to 0")
     if not np.all(np.isfinite(x)):
@@ -62,8 +56,7 @@ def _validate_input(x: npt.ArrayLike, n: int, alpha: float, cutoff: float) -> No
 
 
 class InclusionProb:
-    """
-    First-order inclusion probabilities for units in the population.
+    """First-order inclusion probabilities for units in the population.
 
     Parameters
     ----------
@@ -95,7 +88,7 @@ class InclusionProb:
     take_all : Array
         Take-all units.
     take_some : Array
-        Take-some units.    
+        Take-some units.
 
     References
     ----------
@@ -118,14 +111,14 @@ class InclusionProb:
     ```{python}
     # Units 1-4 belong to the take-some stratum, and units 5 belongs to
     # the take-all stratum
-    
+
     pi.take_some
     pi.take_all
     ```
 
     ```{python}
     # Calculate design weights for a PPS sampling scheme
-    
+
     1 / pi.values
     ```
     """
@@ -172,30 +165,22 @@ def __init__(
 
     @property
     def values(self) -> np.ndarray:
-        """
-        Vector of inclusion probabilties.
-        """
+        """Vector of inclusion probabilties."""
         return self._values.copy()
 
     @property
     def n(self) -> int:
-        """
-        Sample size.
-        """
+        """Sample size."""
         return self._n
 
     @property
     def take_all(self) -> np.ndarray:
-        """
-        Take all units.
-        """
+        """Take all units."""
         return self._ta.copy()
 
     @property
     def take_some(self) -> np.ndarray:
-        """
-        Take some units.
-        """
+        """Take some units."""
         return self._ts.copy()
 
     # Inclusion probabilities are a fixed point for the inclusion
@@ -213,8 +198,7 @@ def __len__(self) -> int:
 def becomes_ta(
     x: npt.ArrayLike, *, alpha: float = 0.001, cutoff: float = np.inf
 ) -> np.ndarray:
-    """
-    Calculate the sample size at which a unit enters the take-all stratum.
+    """Sample size at which a unit enters the take-all stratum.
 
     Parameters
     ----------