antecedent.population

  1"""Named predicates and custom target-distribution weights for analyze()."""
  2
  3from __future__ import annotations
  4
  5from dataclasses import dataclass, field
  6from typing import Mapping, Sequence
  7
  8
  9@dataclass
 10class PopulationRegistry:
 11    """Caller bindings for named row predicates and custom target weights.
 12
 13    Stratification estimators reject ``CustomDistribution``; use IPW/matching
 14    (``estimator=\"propensity.weighting\"``) for weighted target populations.
 15    """
 16
 17    predicates: dict[str, list[int]] = field(default_factory=dict)
 18    distributions: dict[int, list[float]] = field(default_factory=dict)
 19
 20    def insert_predicate(self, name: str, rows: Sequence[int]) -> None:
 21        self.predicates[str(name)] = [int(r) for r in rows]
 22
 23    def insert_distribution(self, distribution_id: int, weights: Sequence[float]) -> None:
 24        self.distributions[int(distribution_id)] = [float(w) for w in weights]
 25
 26
 27def target_all() -> dict[str, str]:
 28    return {"kind": "all"}
 29
 30
 31def target_treated() -> dict[str, str]:
 32    return {"kind": "treated"}
 33
 34
 35def target_untreated() -> dict[str, str]:
 36    return {"kind": "untreated"}
 37
 38
 39def target_named(name: str) -> dict[str, str]:
 40    return {"kind": "named", "name": str(name)}
 41
 42
 43def target_rows(rows: Sequence[int]) -> dict[str, object]:
 44    return {"kind": "rows", "rows": [int(r) for r in rows]}
 45
 46
 47def target_custom_distribution(distribution_id: int) -> dict[str, object]:
 48    return {"kind": "custom_distribution", "id": int(distribution_id)}
 49
 50
 51def coerce_target_population(spec) -> dict[str, object] | None:
 52    """Normalize AverageEffect.target_population / analyze kwargs to a wire dict."""
 53    if spec is None:
 54        return None
 55    if isinstance(spec, str):
 56        key = spec.strip().lower().replace("-", "_")
 57        if key in {"all", "all_observed", "observed"}:
 58            return target_all()
 59        if key == "treated":
 60            return target_treated()
 61        if key in {"untreated", "control"}:
 62            return target_untreated()
 63        raise ValueError(
 64            f"unknown target_population string {spec!r}; "
 65            "use all|treated|untreated or a target_* helper"
 66        )
 67    if isinstance(spec, Mapping):
 68        kind = str(spec.get("kind", "")).lower()
 69        if kind in {"all", "all_observed"}:
 70            return target_all()
 71        if kind == "treated":
 72            return target_treated()
 73        if kind == "untreated":
 74            return target_untreated()
 75        if kind == "named":
 76            return target_named(str(spec["name"]))
 77        if kind == "rows":
 78            return target_rows(spec["rows"])  # type: ignore[arg-type]
 79        if kind in {"custom_distribution", "custom"}:
 80            return target_custom_distribution(int(spec["id"]))  # type: ignore[arg-type]
 81        raise ValueError(f"unknown target_population mapping {spec!r}")
 82    raise TypeError(f"unsupported target_population type: {type(spec)!r}")
 83
 84
 85def registry_wire(registry: PopulationRegistry | None) -> tuple[dict[str, list[int]], dict[int, list[float]]]:
 86    if registry is None:
 87        return {}, {}
 88    return dict(registry.predicates), dict(registry.distributions)
 89
 90
 91__all__ = [
 92    "PopulationRegistry",
 93    "coerce_target_population",
 94    "registry_wire",
 95    "target_all",
 96    "target_custom_distribution",
 97    "target_named",
 98    "target_rows",
 99    "target_treated",
100    "target_untreated",
101]