Cómo realizar un muestreo estratificado en Pandas

Usa Pandas para tomar muestras dentro de grupos definidos por una columna de estratificación. Elige la misma cantidad de cada grupo cuando te interesen las comparaciones entre grupos, o reparte un total fijo en proporción al tamaño de los grupos cuando quieras que la muestra refleje la composición de la población. Una semilla aleatoria fija permite reproducir una muestra ilustrativa.

Muestreo estratificado en estadística

En el muestreo estratificado, la población se divide en grupos mutuamente excluyentes llamados estratos y luego se seleccionan observaciones de cada grupo. Un estrato puede corresponder a un curso, una región, una franja de edad u otra característica pertinente para la pregunta. El muestreo dentro de estos grupos puede asegurar que los grupos importantes aparezcan en una muestra, pero por sí solo no elimina el sesgo de selección ni garantiza que la muestra represente a la población. El marco muestral y la forma de elegir las observaciones también importan.

Los investigadores recurren a la estratificación cuando los grupos difieren en aspectos relevantes para una estimación. Quienes trabajan con datos también pueden usarla al crear conjuntos de entrenamiento y prueba para que una categoría esté representada en ambos. Decide qué columna define los estratos y si cada grupo debe aportar la misma cantidad o una proporción basada en su tamaño dentro de la población. Estas opciones responden a preguntas distintas: las cantidades iguales facilitan las comparaciones entre grupos, mientras que las proporcionales mantienen sus tamaños relativos.

Cómo realizar un muestreo estratificado en Pandas

El ejemplo crea un DataFrame de estudiantes con diez filas repartidas en tres grupos de cursos. Toma dos estudiantes de cada curso para obtener una muestra de igual tamaño y luego distribuye proporcionalmente un total de seis estudiantes. Esta distribución proporcional usa el método de restos mayores: se toma la parte entera inferior de la cuota exacta de cada grupo y después se asignan las plazas restantes a los grupos con los mayores restos fraccionarios. Aquí los grupos tienen 5, 3 y 2 estudiantes, por lo que la distribución de las seis filas es 3, 2 y 1.

import pandas as pd

students = {
    "Name": [
        "sanay", "shivesh", "rutwik", "preet", "yash",
        "mann", "pritesh", "hritesh", "raj", "tarun",
    ],
    "ID": ["001", "002", "003", "004", "005", "006", "007", "008", "009", "010"],
    "Grade": ["A", "A", "C", "B", "B", "B", "C", "A", "A", "A"],
    "Category": [2, 3, 1, 3, 2, 3, 3, 1, 2, 1],
}
df = pd.DataFrame(students)

# Take the same number from each grade.
equal_sample = df.groupby("Grade", group_keys=False).sample(n=2, random_state=7)
print("Equal-sized counts:")
print(equal_sample["Grade"].value_counts().sort_index().to_dict())

# Allocate a fixed total in proportion to each grade's population size.
target = 6
group_sizes = df["Grade"].value_counts().sort_index()
quotas = group_sizes / len(df) * target
allocation = quotas.astype(int)
places_left = target - int(allocation.sum())
remainders = (quotas - allocation).sort_values(ascending=False)
for grade in remainders.index[:places_left]:
    allocation.loc[grade] += 1

proportional_sample = pd.concat(
    [
        group.sample(n=int(allocation.loc[grade]), random_state=7)
        for grade, group in df.groupby("Grade", sort=True)
    ],
    ignore_index=True,
)
print("Proportional counts:")
print(proportional_sample["Grade"].value_counts().sort_index().to_dict())

DataFrame.groupby() forma los grupos que se usan para el muestreo. Consulta la referencia de groupby() de Pandas para obtener más detalles.

Para una muestra de igual tamaño, sample(n=2) solicita dos filas de cada curso. Por tanto, el grupo más pequeño debe tener al menos dos filas; si no las tiene, elige una cantidad menor por grupo o modifica los estratos. En una muestra proporcional, target es el total solicitado y el paso de distribución gestiona el redondeo para que la suma de las cantidades por grupo coincida con ese total. Cuando hay restos fraccionarios empatados, distintas reglas de redondeo pueden asignar una plaza restante a grupos diferentes; para que el análisis sea repetible, documenta la regla que uses.

Ambas llamadas usan random_state=7 para que los mismos datos de entrada y el mismo entorno de Pandas produzcan una muestra repetible. Cambia la semilla para obtener otra muestra aleatoria. La estratificación controla cuántas observaciones se toman de cada grupo, pero no hace que las observaciones sean aleatorias dentro del grupo si los datos de origen o el proceso de muestreo están sesgados. Si necesitas reproducir la muestra, indica la columna de estratificación, la regla de distribución, el tamaño objetivo y la semilla.