Como gerar amostras estratificadas em Python para dados: passo a passo

Gerar amostras estratificadas em Python para dados ajuda a garantir que as proporções de classes ou categorias relevantes se mantêm na amostra, tornando as análises e os modelos mais representativos. Este guia mostra como criar amostras estratificadas simples e práticas, por exemplo para treino/teste ou validação, e explica erros comuns a evitar.

Pré-requisitos

• Python 3.8+ instalado
• Bibliotecas: pandas e scikit-learn (sklearn)
• Ficheiro CSV ou DataFrame com pelo menos uma coluna categórica para estratificar

Passo 1: Por que usar amostragem estratificada?

Amostragem estratificada garante que a distribuição das categorias (por exemplo, rótulos de classe) na amostra reflete a distribuição na população. É útil quando há classes desbalanceadas e quando queremos evitar viés na avaliação de modelos ou na exploração de dados.

Passo 2: Instalar e importar bibliotecas

Instale pandas e scikit-learn se ainda não estiverem disponíveis. Depois importe o que é necessário. Evite instalar em cada execução; utilize um ambiente virtual para reprodutibilidade.

pip install pandas scikit-learn

import pandas as pd
from sklearn.model_selection import train_test_split

Passo 3: Carregar dados e identificar o estrato

Leitura do CSV para um DataFrame e escolha da coluna que define os estratos (por exemplo, 'target' ou 'categoria'). Verifique a distribuição original antes de amostrar.

# Exemplo mínimo
df = pd.read_csv('dados.csv')  # ou use um DataFrame já carregado
print(df['target'].value_counts(normalize=True))

Passo 4: Amostra estratificada para treino/teste com sklearn

Use train_test_split com o argumento stratify para manter proporções. Defina test_size ou train_size conforme necessitar. Este método é o mais simples e robusto para divisão inicial.

# Separar X e y se precisar
X = df.drop(columns=['target'])
y = df['target']

# Amostra estratificada: 80% treino, 20% teste
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# Recriamos DataFrames completos se for útil
train_df = X_train.copy()
train_df['target'] = y_train

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

Passo 5: Amostragem estratificada por grupo ou múltiplas colunas

Se o estrato for a combinação de várias colunas (por exemplo, 'sexo' + 'idade_cat') crie primeiro uma coluna composta e depois estratifique por ela. Para estratificar por grupos grandes, atenção ao número mínimo de amostras por estrato.

# Criar um estrato composto
df['estrato'] = df['sexo'].astype(str) + '_' + df['idade_cat'].astype(str)

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

# Usar stratify com a coluna composta
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']
)

Passo 6: Amostrar uma fracção estratificada (ex.: 10% de cada estrato)

Se pretender uma amostra que contenha 10% de cada estrato, use groupby + sample em pandas. Isto permite controlar a fracção por estrato, mas atenção a estratos com poucas observações.

# Amostra de 10% de cada estrato (com 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)

# Se existirem estratos muito pequenos pode usar replace=True ou filtrar

Passo 7: Erros comuns e como os evitar

Erros frequentes incluem stratify com estratos que têm apenas 1 observação (train_test_split falha) e esquecer de manter o mesmo random_state para reprodutibilidade. Verifique sempre value_counts antes e depois da amostragem.

# Verificar distribuições
print('Original:', df['target'].value_counts(normalize=True))
print('Treino:', train_df['target'].value_counts(normalize=True))
print('Teste:', test_df['target'].value_counts(normalize=True))

Verificar o resultado

Confirme que as proporções por estrato na amostra correspondem às da população. Compare value_counts(normalize=True) entre original, treino e teste ou entre original e a amostra estratificada.

def comparar_proporcoes(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)

Conclusão

Amostragem estratificada em Python com pandas e sklearn garante representatividade por categorias/estratos e reduz o viés em análises e validação de modelos. Próximos passos: experimentar k-fold estratificado com sklearn.model_selection.StratifiedKFold e tratar estratos muito pequenos (agregar ou usar oversampling). Dica: confirme sempre as distribuições antes e depois da amostragem.