diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..28ace4c --- /dev/null +++ b/.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" + } +} diff --git a/docs/_quarto.yml b/docs/_quarto.yml index a985521..2d4d840 100644 --- a/docs/_quarto.yml +++ b/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 \ No newline at end of file + - PoissonSample + - title: Stratification + desc: Make an allocation for stratified sampling. + contents: + - prop_allocation + - divisor_method \ No newline at end of file diff --git a/pysps/__init__.py b/pysps/__init__.py index af67c9d..ca067e7 100644 --- a/pysps/__init__.py +++ b/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" \ No newline at end of file +__version__ = "0.1.1.9002" diff --git a/pysps/allocate.py b/pysps/allocate.py new file mode 100644 index 0000000..7a592e5 --- /dev/null +++ b/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} diff --git a/pysps/inclusion_prob.py b/pysps/inclusion_prob.py index 4dd7593..c02bf44 100644 --- a/pysps/inclusion_prob.py +++ b/pysps/inclusion_prob.py @@ -7,9 +7,7 @@ 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: @@ -17,9 +15,7 @@ def _pi(x: np.ndarray, n: int) -> np.ndarray: 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 ---------- diff --git a/pysps/order_sample.py b/pysps/order_sample.py index 50b2561..4448f71 100644 --- a/pysps/order_sample.py +++ b/pysps/order_sample.py @@ -11,9 +11,7 @@ def _igpd(shape: float) -> Callable[[npt.ArrayLike], np.ndarray]: - """ - Inverse of the generalized Pareto distribution. - """ + """Inverse of the generalized Pareto distribution.""" if shape == 0.0: return lambda x: -np.log(1 - x) elif shape == 1.0: @@ -25,9 +23,7 @@ def _igpd(shape: float) -> Callable[[npt.ArrayLike], np.ndarray]: def _generate_random_deviates( prn: npt.ArrayLike | None, pi: InclusionProb ) -> np.ndarray: - """ - Generate a vector of random numbers for drawing a sample. - """ + """Generate a vector of random numbers for drawing a sample.""" if prn is None: u = np.random.default_rng().uniform(size=len(pi)) else: @@ -42,43 +38,31 @@ def _generate_random_deviates( class BaseSample: - """ - Interface for sample classes. Should not be used directly. - """ + """Interface for sample classes. Should not be used directly.""" @property def units(self) -> np.ndarray: - """ - Indices for units in the sample. - """ + """Indices for units in the sample.""" return self._units.copy() @property def weights(self) -> np.ndarray: - """ - Design weights for units in the sample. - """ + """Design weights for units in the sample.""" return 1 / self._pi._values[self._units] @property def take_all(self) -> np.ndarray: - """ - Take all units in the sample. - """ + """Take all units in the sample.""" return self._ta.copy() @property def take_some(self) -> np.ndarray: - """ - Take some units in the sample. - """ + """Take some units in the sample.""" return self._ts.copy() @property def prn(self) -> np.ndarray: - """ - Random numbers used for drawing the sample. - """ + """Random numbers used for drawing the sample.""" return self._prn.copy() def __len__(self) -> int: @@ -89,8 +73,7 @@ def __str__(self) -> str: class OrderSample(BaseSample): - """ - Order sampling scheme with fixed distribution shape. + """Order sampling with fixed distribution shape. Parameters ---------- @@ -102,11 +85,10 @@ class OrderSample(BaseSample): default draws a sample without permanent random numbers. shape : float, optional Shape parameter for the generalized Pareto distribution that is - used as the fixed order distribution shape. + used as the fixed order distribution shape. The default is sequential + Poisson sampling. Setting shape=0 gives Successive sampling and + shape=-1 gives Pareto order sampling. - shape=1 => Sequential Poisson sampling (the default) - shape=0 => Successive sampling - shape=-1 => Pareto order sampling sort_method : {'partial', 'stable'}, optional Sorting method to use for drawing the sample. The default uses a partial sort. Use 'stable' if ties should resolve in @@ -158,25 +140,25 @@ class OrderSample(BaseSample): ```{python} # Get the design weights. - + sample.weights ``` ```{python} # Units 0 to 2 are take-some units... - + sample.take_some ``` ```{python} # ... and units 3 to 5 are take-all units. - + sample.take_all ``` ```{python} # Draw a Pareto order sample using the same permanent random numbers. - + pysps.OrderSample(pi, prn, shape=-1).units ``` """ @@ -221,8 +203,7 @@ def __repr__(self) -> str: class PoissonSample(BaseSample): - """ - Ordinary Poisson sampling. + """Ordinary Poisson sampling. Parameters ---------- diff --git a/tests/test_allocate.py b/tests/test_allocate.py new file mode 100644 index 0000000..50c518b --- /dev/null +++ b/tests/test_allocate.py @@ -0,0 +1,88 @@ +import numpy as np +from pysps import prop_allocation +import pytest + + +def test_no_alabama_paradox(): + assert prop_allocation({1: 6, 2: 6, 3: 2}, 10) == {1: 5, 2: 4, 3: 1} + assert prop_allocation({1: 6, 2: 6, 3: 2}, 11) == {1: 5, 2: 5, 3: 1} + + +def test_corner_cases(): + assert prop_allocation({}, 0) == {} + assert prop_allocation({1: 0}, 0) == {1: 0} + assert prop_allocation({1: 1}, 0) == {1: 0} + assert prop_allocation({1: 1}, 5) == {1: 5} + assert prop_allocation({1: 0, 2: 1}, 5) == {1: 0, 2: 5} + + +def test_simple_allocations(): + assert prop_allocation({1: 1, 2: 100}, 5) == {1: 0, 2: 5} + assert prop_allocation({1: 1, 2: 100}, 5, initial={1: 1}) == {1: 1, 2: 4} + assert prop_allocation({1: 1, 2: 100}, 5, initial={1: 1}, available={2: 3}) == { + 1: 2, + 2: 3, + } + assert prop_allocation({1: 1, 2: 100}, 5, available={2: 0}) == {1: 5, 2: 0} + + +def test_voting_examples(): + # example from https://en.wikipedia.org/wiki/Highest_averages_method + x = {1: 47000, 2: 16000, 3: 15900, 4: 12000, 5: 6000, 6: 0, 7: 3100} + assert prop_allocation(x, 10) == {1: 5, 2: 2, 3: 2, 4: 1, 5: 0, 6: 0, 7: 0} + assert prop_allocation(x, 10, divisor=lambda a: a + 0.5) == { + 1: 4, + 2: 2, + 3: 2, + 4: 1, + 5: 1, + 6: 0, + 7: 0, + } + assert prop_allocation( + x, + 10, + divisor=lambda a: np.sqrt(a * (a + 1)), + initial={1: 1, 2: 1, 3: 1, 4: 1, 5: 1, 6: 0, 7: 1}, + available={6: 0}, + ) == {1: 4, 2: 2, 3: 1, 4: 1, 5: 1, 6: 0, 7: 1} + assert prop_allocation( + x, + 10, + divisor=lambda a: a, + initial={1: 1, 2: 1, 3: 1, 4: 1, 5: 1, 6: 0, 7: 1}, + available={6: 0}, + ) == {1: 3, 2: 2, 3: 2, 4: 1, 5: 1, 6: 0, 7: 1} + + +def test_ties(): + assert prop_allocation({1: 4, 2: 5}, 8) == {1: 3, 2: 5} + assert prop_allocation({1: 4, 2: 5}, 8, ties="first") == {1: 4, 2: 4} + assert prop_allocation({1: 1, 2: 1, 3: 1}, 2) == {1: 1, 2: 1, 3: 0} + assert prop_allocation({1: 1, 2: 1, 3: 1}, 2, ties="first") == {1: 1, 2: 1, 3: 0} + + +def test_exceptions(): + with pytest.raises(ValueError) as execinfo: + prop_allocation({1: 0}, 1) + assert execinfo.type is ValueError + + with pytest.raises(ValueError) as execinfo: + prop_allocation({1: 1}, 3, available={1: 2}) + assert execinfo.type is ValueError + + with pytest.raises(ValueError) as execinfo: + prop_allocation({1: 1}, 1, initial={1: 2}) + assert execinfo.type is ValueError + + with pytest.raises(ValueError) as execinfo: + prop_allocation({1: 1}, 3, initial={1: 2}, available={1: 1}) + assert execinfo.type is ValueError + + with pytest.raises(ValueError) as execinfo: + prop_allocation({1: 1}, 3, initial={2: 2}) + assert execinfo.type is ValueError + + with pytest.raises(ValueError) as execinfo: + prop_allocation({1: 1}, 3, available={2: 2}) + assert execinfo.type is ValueError diff --git a/tests/test_order_sample.py b/tests/test_order_sample.py index f8066a9..325422c 100644 --- a/tests/test_order_sample.py +++ b/tests/test_order_sample.py @@ -83,7 +83,7 @@ def test_sort_method(): x = [1, 2, 3, 4, 5, 6] prn = [0.1, 0.3, 0.2, 0.6, 0.7, 0.5] - samp1 = OrderSample(InclusionProb(x, 3)) - samp2 = OrderSample(InclusionProb(x, 3), sort_method="stable") + samp1 = OrderSample(InclusionProb(x, 3), prn=prn) + samp2 = OrderSample(InclusionProb(x, 3), prn=prn, sort_method="stable") - np.all(samp1.units == samp2.units) \ No newline at end of file + assert np.all(samp1.units == samp2.units)