audit-labs/audit-tools

  1"""Sampling methods."""
  2
  3from __future__ import annotations
  4
  5import math
  6import random
  7
  8import pandas as pd
  9
 10from .io import AuditSamplingError
 11
 12
 13def ensure_seed(seed: int | None) -> tuple[int, bool]:
 14    if seed is not None:
 15        return int(seed), False
 16    return random.SystemRandom().randint(1, 2_147_483_647), True
 17
 18
 19def parse_key_ints(value: str | None, label: str) -> dict[str, int]:
 20    if not value:
 21        return {}
 22    parsed: dict[str, int] = {}
 23    for part in value.split(","):
 24        if "=" not in part:
 25            raise AuditSamplingError(f"Invalid {label} entry '{part}'. Use Name=Count.")
 26        key, raw_count = part.split("=", 1)
 27        key = key.strip()
 28        try:
 29            count = int(raw_count.strip())
 30        except ValueError as exc:
 31            raise AuditSamplingError(
 32                f"Invalid {label} count for '{key}': {raw_count}"
 33            ) from exc
 34        if count <= 0:
 35            raise AuditSamplingError(f"{label} count for '{key}' must be positive.")
 36        parsed[key] = count
 37    return parsed
 38
 39
 40def parse_key_floats(value: str | None, label: str) -> dict[str, float]:
 41    if not value:
 42        return {}
 43    parsed: dict[str, float] = {}
 44    for part in value.split(","):
 45        if "=" not in part:
 46            raise AuditSamplingError(
 47                f"Invalid {label} entry '{part}'. Use Name=Proportion."
 48            )
 49        key, raw_proportion = part.split("=", 1)
 50        key = key.strip()
 51        try:
 52            proportion = float(raw_proportion.strip())
 53        except ValueError as exc:
 54            raise AuditSamplingError(
 55                f"Invalid {label} proportion for '{key}': {raw_proportion}"
 56            ) from exc
 57        if proportion <= 0:
 58            raise AuditSamplingError(
 59                f"{label} proportion for '{key}' must be positive."
 60            )
 61        parsed[key] = proportion
 62    if not math.isclose(sum(parsed.values()), 1.0, rel_tol=1e-9, abs_tol=1e-9):
 63        raise AuditSamplingError(f"{label} proportions must sum to 1.0.")
 64    return parsed
 65
 66
 67def largest_remainder_allocation(
 68    sample_size: int, proportions: dict[str, float]
 69) -> dict[str, int]:
 70    raw = {
 71        stratum: {
 72            "floor": math.floor(sample_size * proportion),
 73            "remainder": sample_size * proportion
 74            - math.floor(sample_size * proportion),
 75        }
 76        for stratum, proportion in proportions.items()
 77    }
 78    allocation = {stratum: values["floor"] for stratum, values in raw.items()}
 79    remaining = sample_size - sum(allocation.values())
 80    ranked = sorted(
 81        raw,
 82        key=lambda stratum: (-raw[stratum]["remainder"], stratum),
 83    )
 84    for stratum in ranked[:remaining]:
 85        allocation[stratum] += 1
 86    return allocation
 87
 88
 89def random_sample(
 90    population: pd.DataFrame, sample_size: int, seed: int
 91) -> pd.DataFrame:
 92    return population.sample(n=sample_size, random_state=seed)
 93
 94
 95def stratified_sample(
 96    population: pd.DataFrame,
 97    stratify_column: str,
 98    counts: dict[str, int],
 99    seed: int,
100    allow_shortfall: bool,
101) -> tuple[pd.DataFrame, list[dict[str, object]]]:
102    samples: list[pd.DataFrame] = []
103    summary: list[dict[str, object]] = []
104    for index, (stratum, requested) in enumerate(counts.items()):
105        stratum_population = population[population[stratify_column] == stratum]
106        actual_count = min(requested, len(stratum_population))
107        if actual_count < requested and not allow_shortfall:
108            raise AuditSamplingError(
109                f"Stratum '{stratum}' has {len(stratum_population)} rows; "
110                f"requested {requested}. Use --allow-shortfall to continue."
111            )
112        if actual_count:
113            sampled = stratum_population.sample(
114                n=actual_count, random_state=seed + index
115            )
116            samples.append(sampled)
117        summary.append(
118            {
119                "Stratum": stratum,
120                "Population Count": len(stratum_population),
121                "Requested Sample Count": requested,
122                "Actual Sample Count": actual_count,
123                "Shortfall": requested - actual_count,
124            }
125        )
126
127    if samples:
128        return pd.concat(samples), summary
129    return population.iloc[0:0].copy(), summary