python - 有没有办法矢量化这个函数或提高它的效率

Question

此循环旨在以 1:4 的比例将 df2 中的主题与 df1 中的主题进行匹配。这里的关键是随机选择主题，同时避免冗余。任何主题都不应匹配两次。 df1 有几千个主题，而 df2 有超过一百万个。 df1 中的每个科目都将与 df2 中的四个科目匹配，不匹配的将被排除在外。有没有人有提高效率的想法？一种还可以节省 RAM 的方法将是理想的。谢谢。

for x in range(4): # 1:4 matching
    for index, row in df1.iterrows():
        temp = df2.loc[(df2['matched'] != 1) & (df2['race_ethnicity'] == row['race_ethnicity']) & (df2['age'] == row['age']) & ((df2['date1'] > row['date2']) | (df2['date1'].isna()))]
        a = temp.sample()
        a['matched_subject'] = row['subject_id']
        a['matched'] = '1'
        a['possible_matches'] = len(temp)

它可以简化为这个，但我更愿意继续使用 'possible_matches' 行进行诊断。

for x in range(4): # 4 because 1:4 matching
    for index, row in df1.iterrows():
        a = df2.loc[(df2['matched_subject']=='') & (df2['race_ethnicity'] == row['race_ethnicity']) & (df2['age'] == row['age']) & ((df2['date1'] > row['date2']) | (df2['date1'] == 0))].sample()
        a['matched_subject'] = row['pid']

说明:所有行在两个 DataFrame 中都是唯一的，代表将在生存分析中进行比较的主题列表。 Date1 是结果变量事件的日期时间，存在于两个 DataFrame 中的一小部分主题中。日期 2 是自变量事件的日期时间，适用于所有 df1 科目，没有 df2 科目。示例输入包括:

subject_id(数字)

Race_ethnicity(str，3 个类别)

年龄(数字)

匹配(二进制，表示 df2 主题是否与 df1 中的主题匹配)

Matched_subject(数字，匹配主题的subject_id)

date1(日期时间。生存模型的结果变量，存在于两个数据框中的某些而不是其他)

date2(自变量事件的日期时间。df1 中的每个人都有一个 date2，df2 中没有人有 date2。df1 主题的 date2 是他们生存分析的索引日期，也用作没有日期 2 的匹配主题的索引日期变量)

我们希望为 df1 中的每个主题匹配 df2 中的四个主题。 df1 中的 Date2 将是 df2 中匹配主题的索引日期，因此我们有一个条件来确保 df2 中的 date1(结果事件)不会在 df1 中的 date2(索引事件)之前发生
这是与 df1 匹配的两个 df2 主题的示例。还有一个无与伦比的 df1 和 df2 主题。真正的 df2 足够大，可以匹配所有 df1 的科目，而且非常多。
df1:


主题ID
种族_种族
年龄
日期1
日期2
匹配
匹配主题
可能的匹配


3a3r796e
非西类牙裔白人
55
(可以存在于 df1 或 df2 中。如果存在，必须在 df1 中的日期 2 之后)
2012-01-01
1
3a3r796e 匹配基于 df1，因此这些是相同的
仅对 df2 重要。对分析不重要，只是一个诊断值

第1234章
非西类牙裔黑人
58
2017-01-01
2016-01-01
0



df2:


主题ID
种族_种族
年龄
日期1
日期2
匹配
匹配主题
可能的匹配


5c69a756
非西类牙裔白人
55
2015-01-01
(根据定义，不能出现在 df2 中)
1
3a3r796e
571

7as89f75
非西类牙裔白人
55

1
3a3r796e
571

6376asef
西类牙裔
42
2010-01-01

0

Answer 1

首先让我们定义一些辅助函数:

def generate_data(df1_len, df2_len, seed=42):
    """Generate random data to help test different algorithms"""

    np.random.seed(seed)
    d2 = np.random.randint(0, 3000, size=df1_len)
    df1 = pd.DataFrame({
        'subject_id': np.arange(df1_len),
        'race_ethnicity': np.random.choice(list('ABC'), df1_len),
        'age': np.random.randint(18, 100, df1_len),
        'date2': pd.Timestamp('2000-01-01') + pd.to_timedelta(d2, unit='D')
    })

    d1 = np.random.randint(0, 3000, size=int(df2_len * np.random.rand()))
    d1 = np.hstack([d2, np.repeat(np.nan, df2_len - len(d1))])
    df2 = pd.DataFrame({
        'subject_id': np.arange(df2_len),
        'race_ethnicity': np.random.choice(list('ABC'), df2_len),
        'age': np.random.randint(18, 100, df2_len),
        'date1': pd.Timestamp('2000-01-01') + pd.to_timedelta(d1, unit='D')
    })

    return df1, df2


def verify(df1, df2):
    """Verify that df1 and df2 are matched according to predefined rules"""
    
    tmp = df1.merge(df2, how='left', left_on='subject_id', right_on='matched_subject', suffixes=('_1', "_2"))

    assert (tmp['race_ethnicity_1'] == tmp['race_ethnicity_2']).all(), 'race_ethnicity does not match'
    assert (tmp['age_1'] == tmp['age_2']).all(), 'age does not match'
    assert ((tmp['date1'] > tmp['date2']) | tmp['date1'].isna()).all(), 'date1 must be NaT or grater than date2'
    assert tmp.groupby('matched_subject').size().eq(4).all(), 'Invalid match ratio'
    
    print('All is good')

原来的解决方案
为了清楚起见，请允许我做一些更改。此版本运行于
在我的 Mac 上约 28 秒:

df1, df2 = generate_data(500, 100_000)

df2['matched'] = False
df2['matched_subject'] = None
df2['possible_matches'] = None

for x in range(4): # 1:4 matching
    for index, row in df1.iterrows():
        cond = (
            (df2['matched'] != 1) &
            (df2['race_ethnicity'] == row['race_ethnicity']) &
            (df2['age'] == row['age']) &
            ((df2['date1'] > row['date2']) | df2['date1'].isna())
        )
        temp = df2.loc[cond]
        if temp.empty:
            continue

        idx = temp.sample().index
        df2.loc[idx, 'matched_subject'] = row['subject_id']
        df2.loc[idx, 'matched'] = True
        df2.loc[idx, 'possible_matches'] = len(temp)

改进版
通过取出外循环( for _ in range(4) )，可以提高性能
差不多4次。下面的代码在 7s 内执行:

df1, df2 = generate_data(5000, 1_000_000)

df2['matched'] = False
df2['matched_subject'] = None
df2['possible_matches'] = None

for index, row in df1.iterrows():
    cond = (
        (df2['matched'] != 1) &
        (df2['race_ethnicity'] == row['race_ethnicity']) &
        (df2['age'] == row['age']) &
        ((df2['date1'] > row['date2']) | df2['date1'].isna())
    )
    temp = df2.loc[cond]
    if temp.empty:
        continue

    idx = temp.sample(4).index
    df2.loc[idx, 'matched_subject'] = row['subject_id']
    df2.loc[idx, 'matched'] = True
    df2.loc[idx, 'possible_matches'] = len(temp)

进一步改进的版本
认为一次处理多行比这样做更快
一次一个，我们可以根据具有相似特征的一组行进行循环
而不是循环使用单个行。此代码以 600 毫秒或约 46 倍的速度运行
比原始版本:

df1, df2 = generate_data(500, 100_000)

# Shuffle df2 so the matches will be random
df2 = df2.sample(frac=1)

# A dictionary to hold the result. Its keys are the indexes in df2 and its
# values are the indexes of df1
matches = {}

# We loop by group instead of individual row
grouped1 = df1.groupby(['race_ethnicity', 'age', 'date2'])
grouped2 = df2.groupby(['race_ethnicity', 'age'])

for (race_ethnicity, age, date2), subset1 in grouped1:
    # Get all rows from df2 that have the same `race_ethnicity` and `age`
    subset2 = grouped2.get_group((race_ethnicity, age))

    # pd.Series is slow. Switch to np.array for speed
    index2 = subset2.index.to_numpy()
    date1 = subset2['date1'].to_numpy()

    # Since all rows in subset1 and subset2 have already had the same
    # `race_ethnicity` and `age`, we only need to filter for two things:
    #   1. The relationship between `date1` and `date2`; and
    #   2. That the row in `df2` has NOT been matched before
    cond = (
        (np.isnan(date1) | (date1 > date2))
        & np.isin(index2, list(matches.keys()), invert=True)
    )

    # The match ratio
    index1 = np.repeat(subset1.index.to_numpy(), 4)

    # There is no way to know in advance how many rows in `subset2` will meet
    # the matching criteria:
    #   * Ideally: cond.sum() == len(index1), ie. 4 rows in `subset2` for every
    #     row in `subset1`
    #   * If there are more matches than we need: we will take the first `4 *
    #     len(subset1)` rows
    #   * If there are not enough matches: eg. 6 rows in `subset2` for 2 rows in
    #     `subset1`, some rows in `subset1` will have to accept < 4 matches
    n = min(cond.sum(), len(index1))

    matches.update({
        key: value for key, value in zip(index2[cond][:n], index1[:n])
    })

tmp = pd.DataFrame({
    'index2': matches.keys(),
    'index1': matches.values()
})
df2 = (
    df2.merge(tmp, how='left', left_index=True, right_on='index2')
       .merge(df1['subject_id'].to_frame('matched_subject'), how='left', left_on='index1', right_index=True)
       .drop(columns=['index1', 'index2'])
)

您可以验证解决方案:

verify(df1, df2)
# Output: All is good