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

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 是他们用于生存分析的索引日期，也可以作为其匹配对象的索引日期没有 date2 变量）

我们希望将 df2 中的四个主题与 df1 中的每个主题进行匹配。 df1 中的 Date2 将是 df2 中匹配主题的索引日期，因此我们有一个条件来确保 df2 中的 date1（结果事件）不会在 df1（索引事件）中的 date2 之前发生

以下是与 df1 匹配的两个 df2 主题的示例。还有一个无与伦比的 df1 和 df2 主题。真正的 df2 足够大，可以匹配所有 df1 的主题。

df1:

subject_id	race_ethnicity	age	date1	date2	matched	matched_subject	possible_matches
3a3r796e	Non-Hispanic white	55	(can be present in df1 or df2. if present, must be after date 2 in df1)	2012-01-01	1	3a3r796e matching based on df1, therefore these are the same	only important for df2. not important for analysis, just a diagnostic value
1234abcd	Non-Hispanic black	58	2017-01-01	2016-01-01	0

df2:

subject_id	race_ethnicity	age	date1	date2	matched	matched_subject	possible_matches
5c69a756	Non-Hispanic white	55	2015-01-01	(cannot be present in df2, by definition)	1	3a3r796e	571
7as89f75	Non-Hispanic white	55			1	3a3r796e	571
6376asef	Hispanic	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)

进一步改进的版本

认为一次处理多行比这样做更快 一次一个，我们可以基于 group 具有相似特征的行进行循环 而不是循环单个行。此代码运行时间为 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