Fabricating finite populations

fabricate() evaluates columns in declaration order. A column callable receives the current Polars frame and a NumPy generator, so later variables can depend on earlier variables without string evaluation.

import numpy as np
import polars as pl
import declaredesign as dd

population = dd.fabricate(
    2_000,
    rng=44,
    X=lambda data, rng: rng.normal(size=data.height),
    sector=lambda data, rng: dd.draw_categorical(
        data.height,
        [0.45, 0.35, 0.20],
        categories=["services", "manufacturing", "public"],
        rng=rng,
    ),
    earnings=lambda data, rng: np.exp(
        9.5 + 0.25 * data["X"].to_numpy() + rng.normal(scale=0.5, size=data.height)
    ),
)
population.head()

shape: (5, 4)

ID	X	sector	earnings
i64	f64	str	f64
1	1.445963	"services"	22365.442614
2	0.102305	"manufacturing"	14918.505095
3	0.326753	"services"	13916.566996
4	1.136462	"services"	13629.444446
5	0.824244	"services"	32887.180235

Nested levels and intraclass correlation

schools = dd.fabricate(50, rng=5, school_quality=lambda data, rng: rng.normal(size=data.height))
students = dd.add_level(
    schools,
    N=30,
    level="student",
    rng=6,
    baseline=lambda data, rng: (
        data["school_quality"].to_numpy() + rng.normal(size=data.height)
    ),
)

students.select("ID", "student_ID", "student_parent_ID", "school_quality", "baseline").head()

shape: (5, 5)

ID	student_ID	student_parent_ID	school_quality	baseline
i64	i64	i64	f64	f64
1	1	1	-0.801931	0.251184
1	2	1	-0.801931	0.97456
1	3	1	-0.801931	-3.355223
1	4	1	-0.801931	-0.939896
1	5	1	-0.801931	0.211788

draw_normal_icc() and draw_binary_icc() expose the within-cluster dependence directly. draw_multivariate() returns a Polars frame with a requested covariance matrix. Other helpers cover binary, binomial, categorical, count, ordered, Likert, and quantile-valued variables.

Potential outcomes and reveal

population = dd.potential_outcomes(
    population,
    outcome="Y",
    assignment="Z",
    conditions=[0, 1],
    model=lambda data, z, rng: (
        0.2 * data["X"].to_numpy() + 1.25 * z + rng.normal(scale=0.2, size=data.height)
    ),
    rng=10,
)
population = population.with_columns(pl.Series("Z", dd.complete_ra(population.height, rng=11)))
observed = dd.reveal_outcomes(population, outcome="Y", assignment="Z")
observed.select("Z", "Y_Z_0", "Y_Z_1", "Y").head()

shape: (5, 4)

Z	Y_Z_0	Y_Z_1	Y
i64	f64	f64	f64
1	0.068525	1.82455	1.82455
1	-0.124544	1.147167	1.147167
1	-0.09101	1.64837	1.64837
1	0.280688	1.710156	1.710156
1	0.115133	1.493938	1.493938

Potential outcomes are ordinary columns named by outcome, assignment variable, and condition. Reveal is therefore transparent and testable.