一. 只计算AUC值

直接计算auc值比较简单,，直接用pyspark.ml.evaluation即可：

evaluator = BinaryClassificationEvaluator(rawPredictionCol="probability", labelCol="label")
areaUnderPR = evaluator.evaluate(dataSetWithProba, {evaluator.metricName: "PR"}) # 0.6474016172819991

二. 获取ROC曲线

1. 如果只需要获取train的ROC曲线点，直接用summnary即可

trainingSummary = LrModel.summary
print('ROC曲线绘图点：')
trainingSummary.roc.show()
print("AUC:"trainingSummary.areaUnderROC) 0.646777751909773

2. 如果需要该模型对其他数据集预测结果的ROC曲线

方法一：利用sklearn

注意：需要转换数据格式

# 数据格式转换
def get_prob_target(dataSetWithProba):
        return dataSetWithProba.select('label','probability').rdd.map(lambda row: (float(row['probability'][1]), float(row['label']))).collect()
def get_roc_curve(dataSetWithProba):
        dataSetWithProba=get_prob_target(dataSetWithProba)
        y_label=[i[1] for i in dataSetWithProba]
        y_pred=[i[0] for i in dataSetWithProba]
        return y_label,y_pred
        
from sklearn import metrics
fpr, tpr, thersholds = roc_curve(y_label, y_pred) # len=10262
auc=metrics.auc(fpr,tpr) # 0.6474016172819987
print('ROC曲线绘图点：')
print("AUC: ",auc)

方法二：利用pyspark.mllib.evaluation

Scala中BinaryClassificationMetrics函数提供了提取ROC曲线的方式，但pyspark中没有提供，因此需要先从Scala模块中借用

from pyspark.mllib.evaluation import BinaryClassificationMetrics

# Scala version implements .roc() and .pr()
# Python: https://spark.apache.org/docs/latest/api/python/_modules/pyspark/mllib/common.html
# Scala: https://spark.apache.org/docs/latest/api/java/org/apache/spark/mllib/evaluation/BinaryClassificationMetrics.html
class CurveMetrics(BinaryClassificationMetrics):
    def __init__(self, *args):
        super(CurveMetrics, self).__init__(*args)

    def _to_list(self, rdd):
        points = []
        # Note this collect could be inefficient for large datasets
        # considering there may be one probability per datapoint (at most)
        # The Scala version takes a numBins parameter,
        # but it doesn't seem possible to pass this from Python to Java
        for row in rdd.collect():
            # Results are returned as type scala.Tuple2,
            # which doesn't appear to have a py4j mapping
            points += [(float(row._1()), float(row._2()))]
        return points

    def get_curve(self, method):
        rdd = getattr(self._java_model, method)().toJavaRDD()
        return self._to_list(rdd)

preds = sdf.select(label,'prediction').rdd.map(lambda row: (float(row['prediction']), float(row[label]))) 

# Returns as a list (false positive rate, true positive rate)
points = CurveMetrics(preds).get_curve('roc')
fpr = [x[0] for x in points] # len=33964,相当于原始预测值去重后数量+2
tpr = [x[1] for x in points]
auc_roc = BinaryClassificationMetrics(preds).areaUnderROC # 0.6474016172819987 

plt.figure()
plt.title(title)
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.plot(fpr, tpr)

result['figure'] = {'title': 'ROC曲线',
                    'AUROC': auc_roc,
                    'x': [1-f for f in fpr],
                    'y': tpr,
                    'xlabel': '特异度',
                    'ylabel': '灵敏度'}