如何实现含亲属关系样本的NEOFAC_N分层抽样拆分？

问题背景 #

我有一个包含174名受试者的DataFrame（df_behavioral），每行对应一名受试者。部分受试者通过Family_ID标记存在亲属关系。需要将样本拆分为两个规模近似相等的子样本，满足两个核心要求：

• 同家庭的受试者不能在同一个子样本中，即每个子样本内的受试者均无亲属关系
• 按神经质得分（NEOFAC_N）分层，确保两个子样本的得分分布近似一致

目前已有代码实现了无亲属关系的拆分，但还没完成基于NEOFAC_N的分层功能，需要技术帮助。现有代码如下：

import numpy as np
import pandas as pd

# Set seed for reproducibility
np.random.seed(42)

# Generate sample dataset
num_subjects = 174
num_families = 87  # 87 unique families

# Create Family_IDs (maximally two subjects per family)
family_ids = np.repeat(np.arange(num_families), 2)[:num_subjects]
np.random.shuffle(family_ids)

# Generate Neuroticism scores (normally distributed, scale 0-40, integers)
neuroticism_scores = np.clip(np.random.normal(loc=20, scale=5, size=num_subjects), 0, 40).astype(int)

# Create random subject names
subject_ids = [f'Subject_{i}' for i in range(num_subjects)]

# Create dataframe
df_behavioral = pd.DataFrame({
    'Subject_ID': subject_ids,
    'Family_ID': family_ids,
    'NEOFAC_N': neuroticism_scores
})

print(df_behavioral.head())  # Preview dataset

siblings_list = []
 
for family in df_behavioral["Family_ID"].unique():
    siblings_list.append(df_behavioral.query("Family_ID==@family").index.values)
 
print("There are {} unique families in the dataset".format(len(siblings_list)))
 
subj_set1 = [family[0] for family in siblings_list]
subj_set2 = [family[1] for family in siblings_list if len(family)>1]
 
# Make sure subjects are sorted
subj_set1.sort()
subj_set2.sort()
 
print("There are {} unrelated subjects in set 1".format(len(subj_set1)))
print("There are {} unrelated subjects in set 2".format(len(subj_set2)))

# Save subsamples as DataFrames 
sample_1 = df_behavioral.loc[subj_set1].reset_index(drop=True)
sample_2 = df_behavioral.loc[subj_set2].reset_index(drop=True)

解决方案 #

核心思路 #

要同时满足亲属隔离和分层要求，必须以家庭为单位进行分层分配，不能直接对单个受试者分层。具体逻辑：

1. 按Family_ID分组，计算每个家庭的平均NEOFAC_N得分（代表家庭的神经质水平）
2. 将所有家庭按平均得分分箱，构建分层
3. 在每个分层内，随机将家庭近似均分到两个子样本组
4. 对每个家庭的成员，将一名分配到家庭所属的子样本，另一名分配到对立子样本，确保亲属不在同一组

完整代码实现 #

import numpy as np
import pandas as pd

# 固定随机种子保证结果可复现
np.random.seed(42)

# --------------------------
# 模拟数据生成（真实数据可跳过此段）
# --------------------------
num_subjects = 174
num_families = 87

family_ids = np.repeat(np.arange(num_families), 2)[:num_subjects]
np.random.shuffle(family_ids)
neuroticism_scores = np.clip(np.random.normal(loc=20, scale=5, size=num_subjects), 0, 40).astype(int)
subject_ids = [f'Subject_{i}' for i in range(num_subjects)]

df_behavioral = pd.DataFrame({
    'Subject_ID': subject_ids,
    'Family_ID': family_ids,
    'NEOFAC_N': neuroticism_scores
})

# --------------------------
# 样本拆分核心逻辑
# --------------------------
# 1. 按家庭分组，计算平均神经质得分并记录成员索引
family_summary = df_behavioral.groupby('Family_ID').agg(
    avg_neuro=('NEOFAC_N', 'mean'),
    member_indices=('Subject_ID', lambda x: x.index.tolist())
).reset_index()

# 2. 对家庭平均得分进行分箱分层（这里用5分位数，可根据数据调整）
family_summary['neuro_bin'] = pd.qcut(family_summary['avg_neuro'], q=5, labels=False)

# 3. 分层内拆分家庭：每个分层内随机将家庭均分到两组
def split_bin_families(group):
    shuffled = group.sample(frac=1, random_state=42)
    mid = len(shuffled) // 2
    return shuffled.iloc[:mid], shuffled.iloc[mid:]

bin_splits = family_summary.groupby('neuro_bin').apply(split_bin_families).reset_index()

# 收集两组对应的家庭ID
group1_families = []
group2_families = []
for _, row in bin_splits.iterrows():
    group1_families.extend(row[0]['Family_ID'].tolist())
    group2_families.extend(row[1]['Family_ID'].tolist())

# 4. 为每个家庭成员分配子样本：同家庭成员分属不同组
subject_group_map = {}
# 处理第一组家庭
for fam_id in group1_families:
    members = family_summary.loc[family_summary['Family_ID'] == fam_id, 'member_indices'].iloc[0]
    subject_group_map[members[0]] = 1
    if len(members) > 1:
        subject_group_map[members[1]] = 2
# 处理第二组家庭
for fam_id in group2_families:
    members = family_summary.loc[family_summary['Family_ID'] == fam_id, 'member_indices'].iloc[0]
    subject_group_map[members[0]] = 2
    if len(members) > 1:
        subject_group_map[members[1]] = 1

# 映射到原DataFrame并提取子样本
df_behavioral['sample_group'] = df_behavioral.index.map(subject_group_map)
sample_1 = df_behavioral[df_behavioral['sample_group'] == 1].reset_index(drop=True)
sample_2 = df_behavioral[df_behavioral['sample_group'] == 2].reset_index(drop=True)

# --------------------------
# 结果验证
# --------------------------
print(f"样本1规模：{len(sample_1)} | 样本2规模：{len(sample_2)}")
print("\n样本1神经质得分统计：")
print(sample_1['NEOFAC_N'].describe())
print("\n样本2神经质得分统计：")
print(sample_2['NEOFAC_N'].describe())

# 验证同家庭成员是否不在同一子样本
def check_no_relatives(sample):
    family_counts = sample.groupby('Family_ID').size()
    return all(count <= 1 for count in family_counts)

print(f"\n样本1无亲属关系：{check_no_relatives(sample_1)}")
print(f"样本2无亲属关系：{check_no_relatives(sample_2)}")

关键说明 #

• 家庭分层：以家庭平均得分作为分层依据，避免了单个受试者分层导致的亲属同组问题
• 分层内随机分配：保证每个得分区间的家庭在两个子样本中数量近似相等，进而让子样本的得分分布高度一致
• 亲属隔离逻辑：对每个家庭的成员，固定交叉分配到不同子样本，严格满足同家庭成员不同组的要求
• 验证环节：通过统计描述和亲属检查，确保最终结果符合所有要求

内容的提问来源于stack exchange，提问作者Johanna Popp