Multistage Selection in One Call with Python

Every stage drawn from one frame, with per-stage probabilities and weights

Tutorials
Sample Selection
Cluster Sampling
Python
Draw a multistage sample (districts, enumeration areas, households) in one call with svy’s multistage(): PPS or SRS at each stage, per-stage and overall selection probabilities, design weights and the variance design.
Author

Mamadou S. Diallo, Ph.D.

Published

October 17, 2026

Modified

October 17, 2026

Keywords

multistage sampling Python, three-stage sampling Python, two-stage cluster sampling Python, PPS sampling Python, sampling from a register Python, design weights Python, survey sample selection Python

The sample selection tutorial draws a two-stage sample the way most household surveys do: select enumeration areas (EAs), list their households in the field, then select households from that listing with add_stage(). When every stage’s frame already exists, for example a population or address register, a school list with classes and students, or a simulated population, sample.sampling.multistage() draws all the stages in one call.

When to Use multistage() or add_stage()

multistage() add_stage()
Frames One frame with every stage’s ids (one row per final unit) One frame per stage; later frames built after the earlier stages are drawn
Typical use Registers, school and business lists, simulation, teaching Household surveys with field listing
Stages Any number, in one call One call per stage, chained

Both give the same kind of result: each stage’s conditional probability as a column, the overall probability and weight, and a design ready for estimation.

Setting Up

The bundled World Bank synthetic data has an EA frame and the full household population of those EAs. Joining them gives a register: one row per household, with its region (geo1), district (geo2), residence (urbrur), EA and the EA’s census household count.

import polars as pl
import svy
from svy.selection import Stage

ea_frame = svy.datasets.load(name="ea_frame_wb_2023", source="bundled")
households = svy.datasets.load(name="hld_pop_wb_2023", source="bundled")

register = households.join(ea_frame.select("ea", "n_hlds_census"), on="ea")
print(register.select("geo1", "geo2", "urbrur", "ea", "n_hlds_census", "hid").head(3))
shape: (3, 6)
┌────────┬───────────┬────────┬───────┬───────────────┬─────────────┐
│ geo1   ┆ geo2      ┆ urbrur ┆ ea    ┆ n_hlds_census ┆ hid         │
│ ---    ┆ ---       ┆ ---    ┆ ---   ┆ ---           ┆ ---         │
│ str    ┆ str       ┆ str    ┆ i64   ┆ u32           ┆ str         │
╞════════╪═══════════╪════════╪═══════╪═══════════════╪═════════════╡
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12012 ┆ 521           ┆ 003607f75c3 │
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12012 ┆ 521           ┆ 009631410df │
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12012 ┆ 521           ┆ 011fd2d09ce │
└────────┴───────────┴────────┴───────┴───────────────┴─────────────┘

Two Stages in One Call

The design of the selection tutorial, drawn from the register: in each region × residence stratum, 4 EAs with probability proportional to their census household count, then 21 households in each selected EA.

frame = svy.Sample(register, svy.Design(stratum=("geo1", "urbrur")))

hh_sample = frame.sampling.multistage(
    [
        Stage("ea", n=4, mos="n_hlds_census"),  # PPS systematic, 4 EAs per stratum
        Stage(None, n=21),                       # 21 households per selected EA
    ],
    rstate=12345,
)

print(f"psu={hh_sample.design.psu!r}, households selected: {hh_sample.n_records}")
psu='ea', households selected: 588

Each stage’s probability, given that its parent was selected, is a column, next to the overall weight:

print(
    hh_sample.show_data(
        columns=[
            "geo1",
            "urbrur",
            "ea",
            "svy_prob_selection_stage1",
            "svy_prob_selection_stage2",
            "svy_sample_weight",
        ],
        n=5,
        order_type="random",
        rstate=1,
    )
)
shape: (5, 6)
┌────────┬────────┬───────┬──────────────────────────┬─────────────────────────┬───────────────────┐
│ geo1   ┆ urbrur ┆ ea    ┆ svy_prob_selection_stage ┆ svy_prob_selection_stag ┆ svy_sample_weight │
│ ---    ┆ ---    ┆ ---   ┆ 1                        ┆ e2                      ┆ ---               │
│ str    ┆ str    ┆ i64   ┆ ---                      ┆ ---                     ┆ f64               │
│        ┆        ┆       ┆ f64                      ┆ f64                     ┆                   │
╞════════╪════════╪═══════╪══════════════════════════╪═════════════════════════╪═══════════════════╡
│ geo_04 ┆ Urban  ┆ 42121 ┆ 0.189881                 ┆ 0.041916                ┆ 125.642857        │
│ geo_03 ┆ Urban  ┆ 32063 ┆ 0.232847                 ┆ 0.040777                ┆ 105.321429        │
│ geo_02 ┆ Urban  ┆ 28065 ┆ 0.199003                 ┆ 0.046771                ┆ 107.440476        │
│ geo_03 ┆ Rural  ┆ 32036 ┆ 0.232525                 ┆ 0.063444                ┆ 67.785714         │
│ geo_01 ┆ Urban  ┆ 14078 ┆ 0.198464                 ┆ 0.04277                 ┆ 117.809524        │
└────────┴────────┴───────┴──────────────────────────┴─────────────────────────┴───────────────────┘

How multistage() Works

Each Stage describes one stage:

Field Meaning
unit Id column of the stage’s units (or a tuple of columns); None, at the last stage only, draws the frame’s rows
n Units per stratum at stage 1 (an int, or a mapping keyed by stratum), per selected parent after; None takes all
mos Size measure, constant within a unit; given, the stage is PPS
method "srs" without mos; "sys" (default), "brewer", "murphy" or "rs" with it
certainty_threshold PPS only: units whose probability reaches it are taken with certainty

Stages are drawn in the order listed. Stage 1 draws within the design’s strata; each later stage draws within every unit the stage before selected. Units are keyed by the chain of ids above them, so an id only has to be unique within its parent: EA 1 of one district and EA 1 of another are different units. One rstate drives all stages, so the same seed gives the same sample.

WarningOrder matters

List the stages from the coarsest to the finest, as the sampling plan nests them. svy cannot tell from the frame that districts contain EAs rather than the reverse: listing Stage("ea") before Stage("geo2") runs, but treats every EA number as one unit across districts, and the probabilities and design are wrong.

Design Weights

A stage’s probability is conditional on its parent:

  • SRS: \(n / N\), with \(N\) the parent’s number of units;
  • PPS: \(n \, M_i / \sum_j M_j\) within the parent, as in the pps_* methods (1 for a certainty unit);
  • 1 when the parent has no more than \(n\) units, or with n=None.

The overall inclusion probability is their product, and the design weight its reciprocal, so the overall weight is the product of the stage weights:

\[\pi_{i} = \prod_{k} \pi_{i,k \mid k-1}, \qquad w_i = \frac{1}{\pi_i} = \prod_k w_{i,k}\]

check = hh_sample.data.select(
    (
        pl.col("svy_sample_weight_stage1") * pl.col("svy_sample_weight_stage2")
        - pl.col("svy_sample_weight")
    ).abs().max().alias("largest difference")
)
print(check)
shape: (1, 1)
┌────────────────────┐
│ largest difference │
│ ---                │
│ f64                │
╞════════════════════╡
│ 2.8422e-14         │
└────────────────────┘

These are base weights. Nonresponse adjustment, trimming and calibration come next, in Survey Weighting.

Three Stages

Districts first: in each region, 3 districts by simple random sampling, then 2 EAs per district by PPS, then 10 households per EA.

frame3 = svy.Sample(register, svy.Design(stratum="geo1"))

sample3 = frame3.sampling.multistage(
    [
        Stage("geo2", n=3),                     # 3 districts per region
        Stage("ea", n=2, mos="n_hlds_census"),  # 2 EAs per district, PPS
        Stage(None, n=10),                      # 10 households per EA
    ],
    rstate=12345,
)

d = sample3.design
print(f"stratum={d.stratum!r}, psu={d.psu!r}, ssu={d.ssu!r}")
stratum='geo1', psu='geo2', ssu='ea'

Some districts hold a single EA. A stage that asks for more units than a parent holds takes all of them, with probability 1, and says so once:

print(sample3.warnings)
╭───────────────────────────────── Findings (1 info) ──────────────────────────────────╮
│ Level   INFO                                                                         │
│ Code    STAGE_TAKES_ALL                                                              │
│ Detail  Stage 2 (ea) asks for more units than 2 parent(s) hold, so all their units   │
│         are taken, with probability 1: geo1=geo_02, geo2=geo_02_02; geo1=geo_03,     │
│         geo2=geo_03_05                                                               │
╰──────────────────────────────────────────────────────────────────────────────────────╯

Stratification

  • Stage 1 is stratified by the design’s strata, given on the input Design. n is per stratum: an int, or a mapping keyed by stratum (by the full tuple when the strata are several columns).
  • Later stages draw within each selected parent: the parents act as strata.
  • The result’s design keeps the stage-1 strata, with psu the stage-1 unit and ssu the stage-2 unit. Variance is computed from PSU totals within these strata (the ultimate-cluster approach), so stratification below the first stage enters the weights but not the variance strata.

Strata inside a later stage, such as EAs split into urban and rural within each selected district, are drawn with add_stage() and by= instead.

Refining the Design

The design can be updated afterwards, for example to finer variance strata. This changes how the variance is computed, not the weights, so the new strata should describe the actual design. The PSUs must still nest within them, which check().nesting reports:

sample3.update_design(stratum=("geo1", "urbrur"))
print(sample3.check().nesting)
╭───────────────────────── Nesting check: geo2 in geo1, urbrur ──────────────────────────╮
│ Rows                           220                                                     │
│ Strata                         6                                                       │
│ PSU codes                      12                                                      │
│ PSU codes in several strata    3                                                       │
│ Examples                       'geo_03_01' in ('geo_03', 'Rural'), ('geo_03', 'Urban') │
│                                'geo_03_02' in ('geo_03', 'Rural'), ('geo_03', 'Urban') │
│                                'geo_04_01' in ('geo_04', 'Rural'), ('geo_04', 'Urban') │
╰────────────────────────────────────────────────────────────────────────────────────────╯

Here some districts span urban and rural areas, so they cannot be PSUs within residence strata: keep the regional strata for this sample.

sample3 = sample3.update_design(stratum="geo1")

Estimating from the Sample

The result is an ordinary Sample:

print(sample3.estimation.mean("pc_exp", by="geo1"))
╭────────────────────── Estimate: MEAN (TAYLOR) ──────────────────────╮
│ y: pc_exp                                                           │
│                                                                     │
│  geo1            est         se          lci          uci   cv (%)  │
│  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  │
│  geo_01   4,631.3163   471.0177   2,604.6908   6,657.9418    10.17  │
│  geo_02   2,646.5368   232.1593   1,647.6359   3,645.4378     8.77  │
│  geo_03   2,909.8188   430.8220   1,056.1412   4,763.4965    14.81  │
│  geo_04   3,743.2540   347.0735   2,249.9174   5,236.5907     9.27  │
│                                                                     │
╰─────────────────────────────────────────────────────────────────────╯

Limits

  • No stratification within a stage beyond the parent; use add_stage() with by=.
  • Stages are without replacement (pps_wr is not available as a stage).
  • The input design must not already have selection weights: multistage() starts from the frame.

Next Steps

Ready to continue?
Adjust the design weights in Survey Weighting →

References

  • Lohr, S. L. (2021). Sampling: Design and Analysis (3rd ed.). CRC Press.