Cómo generar muestras estratificadas en Python para datos: paso a paso

Generar muestras estratificadas en Python para datos ayuda a garantizar que las proporciones de clases o categorías relevantes se mantienen en la muestra, haciendo los análisis y los modelos más representativos. Esta guía muestra cómo crear muestras estratificadas simples y prácticas, por ejemplo para entrenamiento/prueba o validación, y explica errores comunes a evitar.

Requisitos previos

• Python 3.8+ instalado
• Bibliotecas: pandas y scikit-learn (sklearn)
• Archivo CSV o DataFrame con al menos una columna categórica para estratificar

Paso 1: ¿Por qué usar muestreo estratificado?

El muestreo estratificado garantiza que la distribución de las categorías (por ejemplo, etiquetas de clase) en la muestra refleja la distribución en la población. Es útil cuando hay clases desbalanceadas y cuando queremos evitar sesgos en la evaluación de modelos o en la exploración de datos.

Paso 2: Instalar e importar bibliotecas

Instale pandas y scikit-learn si aún no están disponibles. Después importe lo necesario. Evite instalar en cada ejecución; utilice un entorno virtual para reproducibilidad.

pip install pandas scikit-learn

import pandas as pd
from sklearn.model_selection import train_test_split

Paso 3: Cargar datos e identificar el estrato

Lectura del CSV en un DataFrame y elección de la columna que define los estratos (por ejemplo, 'target' o 'categoria'). Verifique la distribución original antes de muestrear.

# Ejemplo mínimo
df = pd.read_csv('dados.csv')  # o use un DataFrame ya cargado
print(df['target'].value_counts(normalize=True))

Paso 4: Muestra estratificada para entrenamiento/prueba con sklearn

Use train_test_split con el argumento stratify para mantener proporciones. Defina test_size o train_size según necesite. Este método es el más simple y robusto para la división inicial.

# Separar X y y si es necesario
X = df.drop(columns=['target'])
y = df['target']

# Muestra estratificada: 80% entrenamiento, 20% prueba
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# Reconstruimos DataFrames completos si es útil
train_df = X_train.copy()
train_df['target'] = y_train

test_df = X_test.copy()
test_df['target'] = y_test

Paso 5: Muestreo estratificado por grupo o múltiples columnas

Si el estrato es la combinación de varias columnas (por ejemplo, 'sexo' + 'idade_cat') cree primero una columna compuesta y luego estratifique por ella. Para estratificar por grupos grandes, atención al número mínimo de muestras por estrato.

# Crear un estrato compuesto
df['estrato'] = df['sexo'].astype(str) + '_' + df['idade_cat'].astype(str)

# Verificar tamaños mínimos por estrato
print(df['estrato'].value_counts().head())

# Usar stratify con la columna compuesta
X = df.drop(columns=['target'])
y = df['target']
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=1, stratify=df['estrato']
)

Paso 6: Muestrear una fracción estratificada (ej.: 10% de cada estrato)

Si desea una muestra que contenga el 10% de cada estrato, use groupby + sample en pandas. Esto permite controlar la fracción por estrato, pero atención a estratos con pocas observaciones.

# Muestra del 10% de cada estrato (con replace=False)
sample_frac = 0.10
sampled = df.groupby('estrato', group_keys=False).apply(
    lambda x: x.sample(frac=sample_frac, random_state=42)
).reset_index(drop=True)

# Si existen estratos muy pequeños puede usar replace=True o filtrar

Paso 7: Errores comunes y cómo evitarlos

Errores frecuentes incluyen stratify con estratos que tienen solo 1 observación (train_test_split falla) y olvidar mantener el mismo random_state para reproducibilidad. Verifique siempre value_counts antes y después del muestreo.

# Verificar distribuciones
print('Original:', df['target'].value_counts(normalize=True))
print('Entrenamiento:', train_df['target'].value_counts(normalize=True))
print('Prueba:', test_df['target'].value_counts(normalize=True))

Verificar el resultado

Confirme que las proporciones por estrato en la muestra corresponden a las de la población. Compare value_counts(normalize=True) entre original, entrenamiento y prueba o entre original y la muestra estratificada.

def comparar_proporciones(col, *dfs):
    for i, d in enumerate(dfs, 1):
        print(f'Dataset {i}:')
        print(d[col].value_counts(normalize=True))
        print()

comparar_proporcoes('target', df, train_df, test_df)

Conclusión

El muestreo estratificado en Python con pandas y sklearn garantiza representatividad por categorías/estratos y reduce el sesgo en análisis y validación de modelos. Pasos siguientes: experimentar k-fold estratificado con sklearn.model_selection.StratifiedKFold y tratar estratos muy pequeños (agregar o usar oversampling). Consejo: confirme siempre las distribuciones antes y después del muestreo.