Sample Selection: Two-Stage Cluster Sampling in Python

Two-stage cluster sampling with PPS and simple random sampling

Tutorials

Sample Selection

Cluster Sampling

Python

Learn how to implement two-stage cluster sampling using Python and the svy library. Step-by-step tutorial covering PSU selection with PPS sampling and household selection with SRS.

Author

Mamadou S. Diallo, Ph.D.

Published

January 18, 2026

Modified

July 21, 2026

Keywords

two-stage cluster sampling Python, probability proportional to size sampling Python, PPS sampling Python, household survey sampling Python, cluster sampling Python, simple random sampling Python, stratified cluster sampling Python, PSU selection Python, enumeration area sampling, sampling frame Python, survey sample selection Python

Two-stage cluster sampling is one of the most widely used survey sampling designs in practice. This tutorial walks you through implementing a complete two-stage cluster sampling workflow in Python using the svy library, covering everything from defining your target population to selecting primary sampling units (PSUs) and households, and linking the two stages with correct probability chaining.

Understanding Target Populations and Sampling Frames

In survey sampling, the target population is the entire group of individuals or units about which you want to draw conclusions. Defining the target population is one of the first steps in designing any survey. You’ll need to answer several key questions:

What units are of interest? (e.g., individuals, households, schools, businesses)
Where are they located? (e.g., a country, region, administrative district, or economic sector)
When should they be observed? (e.g., during a specific year, quarter, or survey cycle)

What Is a Sampling Frame?

A sampling frame is a list of all units in the target population from which a sample can be drawn. For practical purposes, researchers use sampling frames to operationalize the target population.

Ideally, the sampling frame should fully represent the target population—including all units that belong to it and excluding any that do not. However, in practice, sampling frames are often:

Incomplete (missing some population members)
Outdated (based on older census data)
Containing duplicates or ineligible units

For a comprehensive treatment of sampling frame construction and assessment, see Section 1.2 of Lohr (2021), Sampling: Design and Analysis.

Note

This tutorial focuses on how to use the svy library to select samples, rather than on constructing or validating sampling frames.

Two-Stage Cluster Sampling Design Overview

This tutorial demonstrates a two-stage cluster sampling design, which is commonly used in household surveys, health surveys, and demographic research.

Stage 1: Selecting Primary Sampling Units (PSUs)

In the first stage, we select Primary Sampling Units (PSUs)—geographic clusters of households that cover the entire country without overlap. These PSUs are selected using probability proportional to size (PPS) sampling, meaning larger clusters have a higher probability of being included in the sample.

The first-stage inclusion probability for EA \(i\) in stratum \(h\) is:

\[\pi_{hi} = \frac{m_h \, M_{hi}}{M_h}\]

where \(m_h\) is the number of EAs selected in stratum \(h\), \(M_{hi}\) is the census household count for EA \(i\), and \(M_h = \sum_{i=1}^{N_h} M_{hi}\) is the stratum total.

Stage 2: Selecting Secondary Sampling Units (SSUs)

In the second stage, within each selected PSU, we select Secondary Sampling Units (SSUs)—typically households—using simple random sampling (SRS) without replacement. The conditional inclusion probability for household \(j\) in EA \(i\) is:

\[\pi_{j|hi} = \frac{k}{M_{hi}^*}\]

where \(k\) is the number of households selected per EA and \(M_{hi}^*\) is the listing count (which may differ from the census count \(M_{hi}\)).

Linking the Two Stages

The overall inclusion probability for household \(j\) in EA \(i\) of stratum \(h\) is the product:

\[\pi_{hij} = \pi_{hi} \times \pi_{j|hi}\]

The design weight is the reciprocal: \(w_{hij} = 1 / \pi_{hij}\). svy computes this chaining automatically through the add_stage() method, which links the two stages before the second selection call.

Setting Up

import numpy as np
import svy

seed = 12345
rng = np.random.default_rng(seed)

Stage 1: Selecting Primary Sampling Units (PSUs)

Loading the Enumeration Area Frame

We begin by loading the enumeration area (EA) frame — a listing of all primary sampling units in the target population.

Note

Sample selection needs a realistic frame, so this tutorial runs on the small bundled version of the World Bank synthetic frame — it works offline. For a national-scale frame, load the same slugs with source="remote" and scale the selection numbers up.

ea_frame_df = svy.datasets.load(name="ea_frame_wb_2023", source="bundled")
print(ea_frame_df.head(5))

shape: (5, 5)
┌────────┬───────────┬────────┬───────┬───────────────┐
│ geo1   ┆ geo2      ┆ urbrur ┆ ea    ┆ n_hlds_census │
│ ---    ┆ ---       ┆ ---    ┆ ---   ┆ ---           │
│ str    ┆ str       ┆ str    ┆ i64   ┆ u32           │
╞════════╪═══════════╪════════╪═══════╪═══════════════╡
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12012 ┆ 521           │
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12022 ┆ 563           │
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12030 ┆ 488           │
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12037 ┆ 478           │
│ geo_01 ┆ geo_01_02 ┆ Urban  ┆ 12042 ┆ 510           │
└────────┴───────────┴────────┴───────┴───────────────┘

Let’s examine the frame’s basic characteristics:

print(f"The number of EAs is {ea_frame_df.shape[0]}")
print(f"The average cluster size is {round(ea_frame_df['n_hlds_census'].mean(), 1)}")

The number of EAs is 133
The average cluster size is 435.1

Defining the Design

We declare the sampling design by specifying the measure of size, stratification variables, and PSU identifier. We stratify by administrative region (geo1) and residence type (urbrur), and use the census household count as the measure of size for PPS selection.

ea_design = svy.Design(
    stratum=("geo1", "urbrur"),
    psu="ea",
    mos="n_hlds_census",
)

ea_frame = svy.Sample(data=ea_frame_df, design=ea_design)

In this bundled example, the frame covers four regions. The first region (geo_01) contains only urban areas; the others include both urban and rural areas, yielding 7 strata.

Implementing PPS Systematic Sampling

We’ll select 6 urban and 4 rural EAs per province (the bundled frame holds 15–20 EAs per stratum). The design already stratifies by geo1 × urbrur, so a dict keyed by urbrur is all we need: svy broadcasts {"Urban": 6, "Rural": 4} across every region automatically. (With the full remote frame you’d scale these up — e.g. 30 urban / 20 rural.)

ea_sample = ea_frame.sampling.pps_sys(
    n={"Urban": 6, "Rural": 4},
    rstate=rng,
)

print(ea_sample)

╭──────────────── Sample ─────────────────╮
│ Survey Data                             │
│   Rows     : 36                         │
│   Columns  : 12                         │
│   Strata   : 7                          │
│   PSUs     : 36                         │
│                                         │
│ Survey Design                           │
│   Row index          svy_row_index      │
│   Stratum            (geo1, urbrur)     │
│   PSU                ea                 │
│   SSU                None               │
│   Weight             svy_sample_weight  │
│   With replacement   False              │
│   Prob               svy_prob_selection │
│   Hit                svy_number_of_hits │
│   MOS                n_hlds_census      │
│   Population size    None               │
│   Replicate weights  None               │
╰─────────────────────────────────────────╯

Sample size broadcasting

The n parameter is flexible. You can pass:

A scalar (e.g., n=20): same size for every stratum
A dict keyed by one stratification variable (e.g., n={"Urban": 30, "Rural": 20}): broadcast across all levels of the other stratification variables
A dict keyed by the full stratum combination (e.g., n={("geo_01", "Urban"): 23, ("geo_02", "Urban"): 20, ...}): exact per-stratum sizes

The keys must match stratum labels exactly in spelling, case, and type.

When to use by

by adds an extra grouping for the draw on top of the design’s strata — svy selects n independently within each stratum × by combination. Reach for it when you need to group by something that is not a stratum variable, as we do in stage 2 below (by="ea" selects households within each selected EA). Passing a variable that is already in the stratum (e.g. by="geo1" here) is redundant and changes nothing, so we leave it off.

Note

Variables starting with svy_ are automatically created by the library to track selection probabilities and weights.

Handling Certainty Selections

A very large EA can have a first-stage inclusion probability of 1 — its measure of size is a big enough share of the stratum that it is selected with certainty. svy flags such units with svy_certainty = True and handles them automatically (a certain EA contributes no first-stage variance). This well-sized selection produces none, but you can always check:

n_certain = ea_sample.wrangling.filter_records(svy.col("svy_certainty")).n_records
print(f"Certainty EAs: {n_certain}")

Certainty EAs: 0

By default a unit becomes certain only when its inclusion probability reaches 1 (certainty_threshold=1.0). Lower the threshold to pull dominant units in earlier: any EA whose would-be inclusion probability meets or exceeds the threshold is taken with certainty, and the remaining sample is redistributed over the rest of the stratum. Every PPS method accepts it:

ea_frame.sampling.pps_sys(
    n={"Urban": 6, "Rural": 4},
    certainty_threshold=0.9,   # take near-certain EAs (π ≥ 0.9) as certainty units
    rstate=rng,
)

Viewing Sample Data with Selection Probabilities

Use the show_data() method to examine a slice of the selected sample, including selection probabilities and sample weights:

print(
    ea_sample.show_data(
        columns=[
            "ea",
            "geo1",
            "urbrur",
            "n_hlds_census",
            "svy_prob_selection",
            "svy_sample_weight",
        ],
        order_type="random",
        rstate=seed,
    )
)

shape: (5, 6)
┌───────┬────────┬────────┬───────────────┬────────────────────┬───────────────────┐
│ ea    ┆ geo1   ┆ urbrur ┆ n_hlds_census ┆ svy_prob_selection ┆ svy_sample_weight │
│ ---   ┆ ---    ┆ ---    ┆ ---           ┆ ---                ┆ ---               │
│ i64   ┆ str    ┆ str    ┆ u32           ┆ f64                ┆ f64               │
╞═══════╪════════╪════════╪═══════════════╪════════════════════╪═══════════════════╡
│ 26049 ┆ geo_02 ┆ Urban  ┆ 408           ┆ 0.271247           ┆ 3.686683          │
│ 31078 ┆ geo_03 ┆ Rural  ┆ 437           ┆ 0.30699            ┆ 3.257437          │
│ 35018 ┆ geo_03 ┆ Urban  ┆ 517           ┆ 0.350627           ┆ 2.852031          │
│ 41277 ┆ geo_04 ┆ Rural  ┆ 337           ┆ 0.182804           ┆ 5.470326          │
│ 33087 ┆ geo_03 ┆ Rural  ┆ 257           ┆ 0.180541           ┆ 5.538911          │
└───────┴────────┴────────┴───────────────┴────────────────────┴───────────────────┘

Stage 2: Selecting Households

Loading the Household Listing Frame

In practice, after selecting EAs, field teams visit each selected EA and enumerate all households to build a listing frame. We simulate this by loading the census households for the selected EAs. The where argument of datasets.load() filters as it reads:

eas_in_sample = ea_sample.show_data(columns="ea", n=None).unique()

hld_frame = svy.datasets.load(
    name="hld_pop_wb_2023",
    source="bundled",
    where=svy.col("ea").is_in(eas_in_sample.to_series()),
)

print(f"Number of selected EAs: {len(eas_in_sample)}")
print(f"Households in the listing frame: {hld_frame.shape[0]}")

Number of selected EAs: 36
Households in the listing frame: 17721

Linking the Two Stages with `add_stage()`

Before selecting households, we link the two stages using add_stage(). This method joins the stage-1 design columns (selection probabilities, weights) onto the second-stage frame, renames them with a _stage1 suffix to preserve the audit trail, and configures the combined Sample so that the next selection call chains probabilities automatically.

The prob_name argument names the column that will store the final unconditional inclusion probability \(\pi_{hij} = \pi_{hi} \times \pi_{j|hi}\); the corresponding design weight is derived automatically:

hh_sample = ea_sample.sampling.add_stage(
    next_stage=hld_frame,
    prob_name="prob_inclusion",
).sampling.srs(
    n=21,
    by="ea",
    wgt_name="ds_wgt",
    rstate=rng,
)

print(hh_sample)

╭──────────────── Sample ─────────────────╮
│ Survey Data                             │
│   Rows     : 756                        │
│   Columns  : 29                         │
│   Strata   : 7                          │
│   PSUs     : 36                         │
│                                         │
│ Survey Design                           │
│   Row index          svy_row_index      │
│   Stratum            (geo1, urbrur)     │
│   PSU                ea                 │
│   SSU                None               │
│   Weight             ds_wgt             │
│   With replacement   False              │
│   Prob               prob_inclusion     │
│   Hit                svy_number_of_hits │
│   MOS                None               │
│   Population size    None               │
│   Replicate weights  None               │
╰─────────────────────────────────────────╯

This single chained call does three things:

add_stage() joins stage-1 probabilities onto the household frame and sets up the chaining
srs(n=21, by="ea") selects 21 households within each EA using simple random sampling
The final probabilities are computed as \(\pi_{hij} = \pi_{hi} \times \pi_{j|hi}\) and stored in prob_inclusion, with the design weight ds_wgt = \(1 / \pi_{hij}\)

Note

The census count \(M_{hi}\) (used for stage-1 PPS probabilities) and the listing count \(M_{hi}^*\) (used for stage-2 SRS probabilities) may differ due to population change between the census and the field listing. svy handles this correctly because each stage uses its own frame.

Examining the Combined Sample

print(
    hh_sample.show_data(
        columns=[
            "ea",
            "geo1",
            "urbrur",
            "prob_inclusion",
            "ds_wgt",
        ],
        order_type="random",
        rstate=seed,
    )
)

shape: (5, 5)
┌───────┬────────┬────────┬────────────────┬───────────┐
│ ea    ┆ geo1   ┆ urbrur ┆ prob_inclusion ┆ ds_wgt    │
│ ---   ┆ ---    ┆ ---    ┆ ---            ┆ ---       │
│ i64   ┆ str    ┆ str    ┆ f64            ┆ f64       │
╞═══════╪════════╪════════╪════════════════╪═══════════╡
│ 26037 ┆ geo_02 ┆ Rural  ┆ 0.012969       ┆ 77.107143 │
│ 44108 ┆ geo_04 ┆ Urban  ┆ 0.011939       ┆ 83.761905 │
│ 28071 ┆ geo_02 ┆ Urban  ┆ 0.013961       ┆ 71.626984 │
│ 43159 ┆ geo_04 ┆ Rural  ┆ 0.011391       ┆ 87.785714 │
│ 32063 ┆ geo_03 ┆ Urban  ┆ 0.014242       ┆ 70.214286 │
└───────┴────────┴────────┴────────────────┴───────────┘

More Selection Options

The two-stage workflow above uses the common path, but every selection method (srs and all the pps_* variants) exposes a few more controls worth knowing.

Restrict who’s eligible with where. Draw only from a subset of the frame; ineligible rows stay in the output with null selection columns. It composes with by:

# 6 urban / 4 rural per region, but only among EAs with at least 40 households
ea_frame.sampling.pps_sys(
    n={"Urban": 6, "Rural": 4},
    where=svy.col("n_hlds_census") >= 40,
    rstate=rng,
)

Name the output columns with prob_name, wgt_name, hit_name. Selection writes svy_prob_selection, svy_sample_weight, and svy_number_of_hits by default; override any of them to match your own conventions (we used wgt_name="ds_wgt" at stage 2):

ea_frame.sampling.pps_sys(
    n={"Urban": 6, "Rural": 4},
    prob_name="pi_1",
    wgt_name="w_1",
    rstate=rng,
)

Control the sort with order_by / order_type. PPS systematic selection walks the frame in order, so the ordering affects which units land in the sample. Sort by a variable, or set order_type="random" to shuffle first (the default is "ascending"):

ea_frame.sampling.pps_sys(n={"Urban": 6, "Rural": 4}, order_type="random", rstate=rng)

Compute the allocation instead of hand-writing n. group_sizes() counts frame units per group and allocate() turns a total sample size into a per-stratum allocation (proportional, neyman, equal, or rate); the result drops straight back into n=:

sizes = ea_frame.sampling.group_sizes(by="urbrur")
n_alloc = ea_frame.sampling.allocate(sizes, method="proportional", n_total=100)
ea_frame.sampling.pps_sys(n=n_alloc, rstate=rng)

Reach for other PPS estimators when you need them. pps_sys is systematic PPS. svy also ships pps_wr (with replacement), pps_brewer, pps_murphy (for n = 2), and pps_rs (Rao–Sampford) — same signature, different joint-inclusion behavior for variance estimation.

Next Steps

After selecting your sample, the next step is to derive and adjust sample weights—including nonresponse adjustments, poststratification, and calibration.

Ready to continue?
Learn how to calculate and adjust weights in Survey Weighting →

References

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

Understanding Target Populations and Sampling Frames

What Is a Sampling Frame?

Two-Stage Cluster Sampling Design Overview

Stage 1: Selecting Primary Sampling Units (PSUs)

Stage 2: Selecting Secondary Sampling Units (SSUs)

Linking the Two Stages

Setting Up

Stage 1: Selecting Primary Sampling Units (PSUs)

Loading the Enumeration Area Frame

Defining the Design

Implementing PPS Systematic Sampling

Handling Certainty Selections

Viewing Sample Data with Selection Probabilities

Stage 2: Selecting Households

Loading the Household Listing Frame

Linking the Two Stages with add_stage()

Examining the Combined Sample

More Selection Options

Next Steps

References

Linking the Two Stages with `add_stage()`