作用于索引的矢量化函数

Question

TL;DR：我知道 .apply() 在 pandas 中很慢。但是，我有一个作用于索引的函数，我无法弄清楚如何对其进行矢量化。我希望这个函数作用于两组参数（分别为 500,000 和 1,500 长）。我想生成一个数据框，其中第一个参数作为行索引，第二个参数作为列名，以及包含该特定行和列的函数输出的单元格。就目前而言，代码看起来需要几天时间才能运行。更多细节和最小可重复示例如下：

---

输入数据：

我有一系列唯一的学生 ID，有 500,000 名学生。我有一个由这些学生 ID 索引的 df (exam_score_df)，其中包含每个学生在数学、语言、历史和科学方面的相应分数。

我还有一系列学校代码（每个学校代码对应一所学校），有 1,500 所学校。我有一个由学校代码索引的 df (school_weight_df)，其中包含学校在数学、语言、历史和科学方面的权重，用于计算学生的分数。每行还包含一个“Y”或“N”索引的“Alternative_Score”，因为有些学校允许您在历史和科学之间取最好的科目分数来计算您的总分。

我写的要矢量化的函数：

def calc_score(student_ID, program_code):
'''
For a given student and program, returns students score for that program.
'''
if school_weight_df.loc[program_code]['Alternative_Score'] == 'N':
    
    return np.dot(np.array(exam_score_df.loc[student_ID][['LANG', 'MAT', 'HIST', 'SCI']]), 
                  np.array(school_weight_df.loc[program_code][['%LANG','%MAT','%HIST','%SCI']]))

elif school_weight_df.loc[program_code]['Alternative_Score'] == 'Y':

    history_score = np.dot(np.array(exam_score_df.loc[student_ID][['LANG', 'MAT', 'HIST']]), 
                           np.array(school_weight_df.loc[program_code][['%LANG','%MAT','%HIST']]))
    
    science_score = np.dot(np.array(exam_score_df.loc[student_ID][['LANG', 'MAT', 'SCI']]), 
                           np.array(school_weight_df.loc[program_code][['%LANG','%MAT','%SCI']]))

    return max(history_score, science_score)

示例 DF：

以下是exam_score_df 和school_weight_df 的示例df：

student_data = [[3, 620, 688, 638, 688], [5, 534, 602, 606, 700], [9, 487, 611, 477, 578]]
exam_score_df = pd.DataFrame(student_data, columns = ['student_ID', 'LANG', 'MAT', 'HIST', 'SCI'])
exam_score_df.set_index('student_ID')

program_data = [[101, 20, 30, 25, 25, 'N'], [102, 40, 10, 50, 50, 'Y']]
school_weight_df = pd.DataFrame(program_data, columns = ['program_code', '%LANG','%MAT','%HIST','%SCI', 'Alternative_Score'])
school_weight_df.set_index('program_code', inplace = True)

这是用于索引以下代码的系列：

series_student_IDs = pd.Series(exam_score_df.index, inplace = True)
series_program_codes = pd.Series(school_weight_df.index, inplace = True)

使用函数创建 DF 的代码：

为了创建每个程序中所有学生分数的 df，我使用了嵌套的 .apply()：

new_df = pd.DataFrame(series_student_IDs.apply(lambda x: series_program_codes.apply(lambda y: calc_score(x, y))))

我已经阅读了几本关于在 pandas 中优化代码的入门书，包括写得很好的Guide by Sofia Heisler。我最关心的问题，也是我不知道如何向量化这段代码的原因，是我的函数需要对索引进行操作。我还有一个次要问题是，即使我进行矢量化，np.dot on large matrices 也存在这个问题，我无论如何都想循环。

感谢大家的帮助！我才编码几个月，所以非常感谢所有有帮助的 cmets。

Answer 1

应用 = 糟糕，双重应用 = 非常糟糕

如果你在函数中使用 Numpy，为什么不一直使用 Numpy？您仍然更喜欢分批方法，因为整个矩阵会占用大量内存。检查以下方法。

在低端 macbook pro 上对 5000 名学生进行了一次迭代，每次迭代花费了 2.05 秒。这意味着对于 500,000 名学生，您可以预期大约 200 秒，这还不错。

我对 100000 名学生和 1500 所学校进行了以下测试，总共花费了大约 30-40 秒。

1. 首先我创建了一个虚拟集数据：考试成绩（100,000 名学生，4 个分数）、学校权重（1500 所学校，4 个权重）和一个 布尔 标志，学校可以选择 Y 或 N , Y==True, N==False

2. 接下来，对于一批 5000 名学生，我只需使用 np.einsum 计算 2 个矩阵之间 4 个科目中每个科目的元素乘积。这给了我(5000,4) * (1500,4) -> (1500,5000,4)。将此视为点积的第一部分（不包括总和）。

3. 我这样做的原因是因为这对于您的条件 N 或 Y 都是必要的步骤。

4. 接下来，FOR N: 我只是简单地根据alt_flag 过滤上面的矩阵，在最后一个轴上减少它（总和）并转置得到(5000, 766)，其中766 是alternative == N 的学校数量

5. FOR Y:，我根据alt_flag 进行过滤，然后我计算前 2 个主题的总和（因为它们很常见）并将它们分别添加到第 3 和第 4 个主题中，取一个最大值并将其作为我的返回最终得分。发布一个转置。这给了我(5000, 734)。

6. 我对每批 5000 个都执行此操作，直到我添加了所有批次，然后简单地 np.vstack 获得决赛桌 (100000, 766) 和 (100000, 734)。

7. 现在我可以简单地将它们堆叠在axis = 0上以获得(100000, 1500)，但如果我想将它们映射到ID（学生，学校），使用pd.DataFrame(data, columns=list_of_schools_alt_Y, index=list_of_student_ids单独进行然后组合它们会更容易.为您阅读最后一步。

8. 最后一步是由您执行，因为我没有完整的数据集。由于索引的顺序是通过批量矢量化保留的，因此您现在可以简单地将 766 个学校 ID 与 N、734 个学校 ID 与 Y 以及 100000 个学生 ID 按它们在主数据集中出现的顺序进行映射。然后只需附加 2 个数据帧即可创建最终（大量）数据帧。

9. 注意：您必须在 for 循环中将 100000 更改为 500000，不要忘记！！

import numpy as np
import pandas as pd
from tqdm import notebook

exam_scores = np.random.randint(300,800,(100000,4))
school_weights = np.random.randint(10,50,(1500,4))
alt_flag = np.random.randint(0,2,(1500,), dtype=bool) #0 for N, 1 for Y
batch = 5000

n_alts = []
y_alts = []

for i in notebook.tqdm(range(0,100000,batch)):
    scores = np.einsum('ij,kj->kij', exam_scores[i:i+batch], school_weights) #(1500,5000,4)

    #Alternative == N
    n_alt = scores[~alt_flag].sum(-1).T   #(766, 5000, 4) -> (5000, 766)

    #Alternative == Y
    lm = scores[alt_flag,:,:2].sum(-1)    #(734, 5000, 2) -> (734, 5000); lang+math
    h = scores[alt_flag,:,2]              #(734, 5000); history
    s = scores[alt_flag,:,3]              #(734, 5000); science
    y_alt = np.maximum(lm+h, lm+s).T      #(5000, 734)
    
    n_alts.append(n_alt)
    y_alts.append(y_alt)

final_n_alts = np.vstack(n_alts)
final_y_alts = np.vstack(y_alts)

print(final_n_alts.shape)
print(final_y_alts.shape)