PySpark中从混淆矩阵输出各标签精确率与召回率的可行方案

我已经创建了标签列表labels = [1,2,3,4,5,6,7]，并通过metrics.confusionMatrix().toArray().astype(int)生成混淆矩阵数组，转为Pandas DataFrame展示。我需要输出各标签对应的精确率（precision）与召回率（recall），期望格式如下：

Label       Precision   Recall
------------------------------
1                ?           ?
2                ?           ?
...
7                ?           ?

但使用以下循环代码调用metrics.precision(i)和metrics.recall(i)时：

print('Label       Precision   Recall')
print('------------------------------')
for i, lab in enumerate(labels):
    print(f'{lab:<12}{metrics.precision(i):<12.4f}{metrics.recall(i):.4f}')

出现错误：java.util.NoSuchElementException: key not found: 0.0，需要替代方法实现需求。

错误原因是你传入metrics.precision()和metrics.recall()的是索引值0-6，但你的标签实际是1-7，方法找不到对应标签0，所以报错。可以用两种方法解决：

方法一：直接传入标签值而非索引 #

循环时直接使用标签列表中的值作为参数：

print('Label       Precision   Recall')
print('------------------------------')
for lab in labels:
    print(f'{lab:<12}{metrics.precision(lab):<12.4f}{metrics.recall(lab):.4f}')

方法二：从混淆矩阵手动计算 #

如果metrics方法调用仍有问题，可以直接通过混淆矩阵数组计算精确率和召回率：
假设混淆矩阵已转为Pandas DataFrame cm_df，索引和列名都是labels：

import pandas as pd

# 假设cm是metrics.confusionMatrix().toArray().astype(int)生成的数组
cm_df = pd.DataFrame(cm, index=labels, columns=labels)

print('Label       Precision   Recall')
print('------------------------------')
for lab in labels:
    # 精确率：TP/(TP+FP)，即当前列的TP除以列总和
    precision = cm_df.loc[lab, lab] / cm_df[lab].sum()
    # 召回率：TP/(TP+FN)，即当前行的TP除以行总和
    recall = cm_df.loc[lab, lab] / cm_df.loc[lab].sum()
    print(f'{lab:<12}{precision:<12.4f}{recall:.4f}')

内容的提问来源于stack exchange，提问作者Kayleen Carlson