Что такое стратификация?

Стратификация: разделение данных по слоям

Стратификация — это процесс разделения совокупности (датасета) на слои (страты) однородных элементов, а затем выборки из каждого слоя.

Просто говоря: я беру популяцию, делю её на группы похожих объектов, и выбираю из каждой группы.

Пример в реальной жизни

Популяция: 1000 пользователей
Возраст:
- 18-25: 100 человек (10%)
- 26-35: 300 человек (30%)
- 36-45: 400 человек (40%)
- 46+: 200 человек (20%)

Что вы хотите: выборка из 100 пользователей

БЕЗ стратификации (random sampling):
- 18-25: ~10 человек (случайно)
- 26-35: ~30 человек
- 36-45: ~40 человек
- 46+: ~20 человек
→ Может быть, что-то пойдёт не так

СО стратификацией:
- 18-25: 10 человек (10% from 10%)
- 26-35: 30 человек (30% from 30%)
- 36-45: 40 человек (40% from 40%)
- 46+: 20 человек (20% from 20%)
→ Выборка отражает структуру популяции

Зачем нужна стратификация?

1. Уменьшить дисперсию (variance)

import numpy as np
from sklearn.model_selection import train_test_split

# Данные: высокая дисперсия между возрастными группами
data = np.array([18, 20, 21, 25, 26, 32, 35, 40, 45, 50])  # 10 человек
target = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])  # 5 молодых, 5 старых

# Обычный split (может быть дисбалансиро вано)
X_train, X_test = train_test_split(data, test_size=0.3, random_state=42)

print(f"Train: молодых={sum(target[i] for i in range(len(data)) if data[i] in X_train)}")
print(f"Test: молодых={sum(target[i] for i in range(len(data)) if data[i] in X_test)}")
# Может быть случайное дисбалансирование

# Стратифицированный split
X_train_strat, X_test_strat = train_test_split(
    data,
    test_size=0.3,
    stratify=target  ← КЛЮЧЕВОЙ параметр
)

print(f"\nТо со стратификацией:")
print(f"Train: 70% молодых, 70% старых")
print(f"Test: 30% молодых, 30% старых")

2. Более репрезентативная выборка

Проблема: если популяция имеет редкий класс (1%),
обычный random sampling может вообще его не выбрать.

Решение: стратифицированный sampling гарантирует
что редкий класс будет представлен пропорционально.

Виды стратификации

1. Пропорциональная стратификация (Proportional Allocation)

import pandas as pd
from sklearn.model_selection import train_test_split

# Датасет: неравномерный
df = pd.DataFrame({
    'customer_id': range(1, 1001),
    'segment': ['Premium']*100 + ['Standard']*300 + ['Economy']*600
})

# Стратификация по сегментам
train, test = train_test_split(
    df,
    test_size=0.2,
    stratify=df['segment']
)

print(train['segment'].value_counts())
# Premium:  80 (10% от всех, как и в оригинале)
# Standard: 240 (30% от всех)
# Economy:  480 (60% от всех)

print(test['segment'].value_counts())
# Premium:  20 (10% от всех)
# Standard: 60 (30% от всех)
# Economy:  120 (60% от всех)

# Пропорции соответствуют оригиналу

2. Оптимальная стратификация (Optimal Allocation)

# Используется когда разные страты имеют разные costs или дисперсии

# Пример: опрос
# Premium клиенты: стоят дорого опрашивать, но дают стабильный результат
# Economy клиенты: дёшево опрашивать, но очень разнородные

# Вместо пропорционального (100 Premium, 900 Economy)
# Делаю оптимальное (150 Premium, 850 Economy)
# → меньше вариантность результата при той же стоимости

3. Неравномерная стратификация (Disproportional Allocation)

# Используется для редких классов

from imblearn.over_sampling import SMOTE
from sklearn.model_selection import train_test_split

# Датасет: 99% нормальные, 1% fraud
df = pd.DataFrame({
    'amount': np.random.normal(100, 50, 10000),
    'is_fraud': [0]*9900 + [1]*100
})

# Проблема: если я возьму 20% выборку, fraud может быть 0-1 пример
# Решение: oversampling fraud до 50%, потом обучаю, потом калибрую

smote = SMOTE(sampling_strategy=0.5)  # Сделай fraud 50% от normal
X_balanced, y_balanced = smote.fit_resample(df[['amount']], df['is_fraud'])

print(f"До SMOTE: {y_balanced.sum()} fraud из {len(y_balanced)}")

Практические примеры

Пример 1: A/B тест с контролем за демографией

users = pd.DataFrame({
    'user_id': range(1, 1001),
    'age_group': ['18-25']*200 + ['26-35']*300 + ['36-45']*300 + ['46+']*200,
    'platform': ['Mobile']*400 + ['Web']*600
})

# Стратифицированное распределение в контрольную и тестовую группы
from sklearn.model_selection import train_test_split

# Стратификую по двум признакам
users['strata'] = users['age_group'] + '_' + users['platform']

control, test = train_test_split(
    users,
    test_size=0.5,
    stratify=users['strata']
)

print(f"Control age distribution:\n{control['age_group'].value_counts(normalize=True)}")
print(f"\nTest age distribution:\n{test['age_group'].value_counts(normalize=True)}")
# Оба имеют одинаковое распределение

Пример 2: Sampling для экономии данных

import pandas as pd

# Огромный датасет с классами
df = pd.DataFrame({
    'id': range(100000),
    'category': ['A']*50000 + ['B']*30000 + ['C']*15000 + ['D']*5000,
    'value': np.random.rand(100000)
})

# Нужна выборка 5000 строк, но сохранить распределение по категориям

from sklearn.model_selection import StratifiedShuffleSplit

split = StratifiedShuffleSplit(n_splits=1, test_size=0.05, random_state=42)

for train_idx, test_idx in split.split(df, df['category']):
    sample = df.iloc[test_idx]

print(f"Original categories:\n{df['category'].value_counts(normalize=True)}")
print(f"\nSample categories:\n{sample['category'].value_counts(normalize=True)}")
# Пропорции совпадают

Пример 3: Cross-validation со стратификацией

from sklearn.model_selection import StratifiedKFold
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score

# Датасет с дисбалансом
X = df[['feature1', 'feature2', 'feature3']]
y = df['target']  # 95% нули, 5% единицы

# Обычный KFold может создать fold где только нули
kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

scores = []
for train_idx, test_idx in kf.split(X, y):
    X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
    y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]
    
    # X_train и X_test имеют одинаковые пропорции y
    
    model = RandomForestClassifier()
    model.fit(X_train, y_train)
    score = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
    scores.append(score)

print(f"CV scores: {scores}")
print(f"Mean: {np.mean(scores):.3f}")

Пример 4: Стратификация для редких событий

# Чёрн лебедь проблема: редкое событие (0.1%)
df = pd.DataFrame({
    'user_id': range(10000),
    'rare_event': [1]*10 + [0]*9990  # только 10 из 10000
})

# Обычная выборка из 1000
random_sample = df.sample(n=1000)
print(f"Random sample rare_events: {random_sample['rare_event'].sum()}")  # ~1

# Стратифицированная выборка
from sklearn.model_selection import train_test_split

train, test = train_test_split(
    df,
    test_size=0.1,
    stratify=df['rare_event']
)
print(f"Stratified sample rare_events: {test['rare_event'].sum()}")  # 1 (гарантированно)

Стратификация в pandas

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame({
    'id': range(100),
    'group': ['A']*60 + ['B']*40,
    'value': np.random.rand(100)
})

# Способ 1: train_test_split
train, test = train_test_split(
    df,
    test_size=0.2,
    stratify=df['group']
)

# Способ 2: groupby (более контролируемо)
samples = []
for group in df['group'].unique():
    group_data = df[df['group'] == group]
    sample = group_data.sample(frac=0.2)  # 20% из каждой группы
    samples.append(sample)

test = pd.concat(samples)

Вывод

Стратификация — разделение данных на слои и выборка из каждого.

Когда использовать:

1. Дисбалансированные классы (редкие события)
2. A/B тесты (контроль за конфаундерами)
3. Cross-validation при дисбалансе
4. Sampling большого датасета
5. Когда нужна репрезентативная выборка

Типы:

• Пропорциональная (сохраняет пропорции)
• Оптимальная (минимизирует дисперсию)
• Неравномерная (oversampling редких классов)

Код:

from sklearn.model_selection import train_test_split, StratifiedKFold

train, test = train_test_split(df, stratify=df['column_name'])
kf = StratifiedKFold(n_splits=5)

Результат: Более надёжные результаты анализа и моделей, особенно на дисбалансированных данных.