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

April 18, 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 polars as pl
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. This example uses the World Bank synthetic population dataset for an imaginary country.

Important

We intentionally perturbed the number of EAs per cluster to mirror a common real-world challenge: sampling frames built from an earlier census are often outdated by the time a survey is fielded.

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

shape: (5, 5)
┌────────┬───────────┬────────┬───────┬───────────────┐
│ geo1   ┆ geo2      ┆ urbrur ┆ ea    ┆ n_hlds_census │
│ ---    ┆ ---       ┆ ---    ┆ ---   ┆ ---           │
│ str    ┆ str       ┆ str    ┆ i64   ┆ i64           │
╞════════╪═══════════╪════════╪═══════╪═══════════════╡
│ geo_01 ┆ geo_01_01 ┆ Urban  ┆ 11001 ┆ 265           │
│ geo_01 ┆ geo_01_01 ┆ Urban  ┆ 11002 ┆ 466           │
│ geo_01 ┆ geo_01_01 ┆ Urban  ┆ 11003 ┆ 443           │
│ geo_01 ┆ geo_01_01 ┆ Urban  ┆ 11004 ┆ 500           │
│ geo_01 ┆ geo_01_01 ┆ Urban  ┆ 11005 ┆ 455           │
└────────┴───────────┴────────┴───────┴───────────────┘

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 5940
The average cluster size is 374.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 imaginary country, the first region (geo_01) contains only urban areas. Other regions include both urban and rural areas, yielding 19 strata.

Implementing PPS Systematic Sampling

From the sample size calculations in the planning tutorial, we need 30 urban and 20 rural EAs per province (assuming 21 households per EA). The n parameter accepts a dictionary that svy broadcasts automatically across strata. Here, the keys match levels of urbrur, and by="geo1" tells svy to apply these within each province:

ea_sample = ea_frame.sampling.pps_sys(
    n={"Urban": 30, "Rural": 20},
    by="geo1",
    rstate=rng,
)

print(ea_sample)

╭──────────────── Sample ─────────────────╮
│ Survey Data                             │
│   Rows     : 477                        │
│   Columns  : 12                         │
│   Strata   : 19                         │
│   PSUs     : 477                        │
│                                         │
│ 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.

Note

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

Handling Certainty Selections

When an EA’s inclusion probability exceeds 1 (i.e., the stratum has fewer EAs than requested), svy selects it with certainty and flags it with svy_certainty = True. You can identify such units:

certainty_eas = ea_sample.show_data(
    where={"svy_certainty": True},
    columns=["geo1", "urbrur", "ea", "svy_prob_selection"],
)
print(certainty_eas)

shape: (5, 4)
┌────────┬────────┬───────┬────────────────────┐
│ geo1   ┆ urbrur ┆ ea    ┆ svy_prob_selection │
│ ---    ┆ ---    ┆ ---   ┆ ---                │
│ str    ┆ str    ┆ i64   ┆ f64                │
╞════════╪════════╪═══════╪════════════════════╡
│ geo_08 ┆ Urban  ┆ 81002 ┆ 1.0                │
│ geo_08 ┆ Urban  ┆ 81043 ┆ 1.0                │
│ geo_08 ┆ Urban  ┆ 82002 ┆ 1.0                │
│ geo_08 ┆ Urban  ┆ 82041 ┆ 1.0                │
│ geo_08 ┆ Urban  ┆ 82048 ┆ 1.0                │
└────────┴────────┴───────┴────────────────────┘

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    ┆ i64           ┆ f64                ┆ f64               │
╞════════╪════════╪════════╪═══════════════╪════════════════════╪═══════════════════╡
│ 53013  ┆ geo_05 ┆ Rural  ┆ 244           ┆ 0.037201           ┆ 26.880738         │
│ 98060  ┆ geo_09 ┆ Urban  ┆ 329           ┆ 0.090319           ┆ 11.071834         │
│ 72074  ┆ geo_07 ┆ Rural  ┆ 360           ┆ 0.098063           ┆ 10.1975           │
│ 103004 ┆ geo_10 ┆ Rural  ┆ 118           ┆ 0.020125           ┆ 49.688559         │
│ 97054  ┆ geo_09 ┆ Rural  ┆ 529           ┆ 0.072329           ┆ 13.825709         │
└────────┴────────┴────────┴───────────────┴────────────────────┴───────────────────┘

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 create a listing frame. We simulate this by loading the census households for the selected EAs only. The where argument of datasets.load() avoids loading the full 2.5 million household census into memory:

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

hld_frame = svy.datasets.load(
    name="hld_pop_wb_2023",
    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: 477
Households in the listing frame: 222708

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     : 10017                      │
│   Columns  : 61                         │
│   Strata   : 19                         │
│   PSUs     : 477                        │
│                                         │
│ 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        │
╞═══════╪════════╪════════╪════════════════╪════════════╡
│ 44097 ┆ geo_04 ┆ Urban  ┆ 0.002694       ┆ 371.23784  │
│ 31011 ┆ geo_03 ┆ Rural  ┆ 0.005103       ┆ 195.98018  │
│ 43169 ┆ geo_04 ┆ Urban  ┆ 0.00171        ┆ 584.71681  │
│ 55067 ┆ geo_05 ┆ Urban  ┆ 0.010005       ┆ 99.951334  │
│ 98035 ┆ geo_09 ┆ Urban  ┆ 0.006293       ┆ 158.902753 │
└───────┴────────┴────────┴────────────────┴────────────┘

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

Next Steps

References

Linking the Two Stages with `add_stage()`