Spark Calibration - A python package for calibrating probabilities predicted by ML model involving large training & test datasets as spark dataframes
from .betacal import Betacal
from .metrics import display_classification_calib_metrics
from .visualisation import plot_calibration_curve
src/spark_calibration/__init__.py
Meesho/spark_calibration
import pyspark.sql.functions as F

from pyspark.sql.types import DoubleType

from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import VectorAssembler
from pyspark.sql.types import DoubleType

from pyspark.sql.dataframe import DataFrame


class Betacal:
    def __init__(self, parameters):
        assert parameters == "abm"

        self.parameters = parameters
        self.a = None
        self.b = None
        self.lr_model = None

    def fit(self, df: DataFrame):
        assert (
            "score" in df.columns and "label" in df.columns
        ), "score and label columns should be present in the dataframe"

        if self.parameters == "abm":
            df = df.withColumn("score2", 1 - F.col("score"))
            df = df.withColumn("score", F.log("score")).withColumn(
                "score2", -1 * F.log("score2")
            )
            lr = LogisticRegression()
            featurizer = VectorAssembler(
                inputCols=["score", "score2"], outputCol="features"
            )

            train_data = featurizer.transform(df)["label", "features"]

            lr_fitted = lr.fit(train_data)

            lr_coef = lr_fitted.coefficients

            if lr_coef[0] < 0:
                featurizer = VectorAssembler(inputCols=["score2"], outputCol="features")
                train_data = featurizer.transform(df)["label", "features"]
                lr_fitted = lr.fit(train_data)
                a = 0
                b = lr_fitted.coefficients[0]

            elif lr_coef[1] < 0:
                featurizer = VectorAssembler(inputCols=["score"], outputCol="features")
                train_data = featurizer.transform(df)["label", "features"]
                lr_fitted = lr.fit(train_data)
                b = 0
                a = lr_fitted.coefficients[0]
            else:
                a = lr_coef[0]
                b = lr_coef[1]

            self.a = a
            self.b = b
            self.lr_model = lr_fitted

    def predict(self, df: DataFrame):
        cols = df.columns

        assert "score" in cols, "score column should be present in the dataframe"

        if self.parameters == "abm":

            def pick_value(v):
                return float(v[1])

            pick_value = F.udf(pick_value, DoubleType())

            df = df.withColumn("orig_score", F.col("score"))

            df = df.withColumn("score2", 1 - F.col("score"))

            df = df.withColumn("score", F.log("score")).withColumn(
                "score2", -1 * F.log("score2")
            )

            if self.a == 0:
                featurizer = VectorAssembler(inputCols=["score2"], outputCol="features")

            elif self.b == 0:
                featurizer = VectorAssembler(inputCols=["score"], outputCol="features")

            else:
                featurizer = VectorAssembler(
                    inputCols=["score", "score2"], outputCol="features"
                )

            test_data = featurizer.transform(df)

            df = (
                self.lr_model.transform(test_data)
                .withColumn("prediction", pick_value("probability"))
                .drop("score")
                .withColumnRenamed("orig_score", "score")
            )

            df = df.select(cols + ["prediction"])

            return df

src/spark_calibration/betacal.py
Meesho/spark_calibration
from pyspark.ml.evaluation import BinaryClassificationEvaluator
import pyspark.sql.functions as F

from pyspark.sql.dataframe import DataFrame


def display_classification_calib_metrics(df: DataFrame):
    """Print pre and post calibration metrics for comparison

    Args:
        df: dataframe with score, label and prediction(calibratied score) columns
    """

    assert (
        "score" in df.columns and "label" in df.columns and "prediction" in df.columns
    ), "score and label columns should be present in the dataframe"

    model_bs = df.select(F.avg(F.pow(df["label"] - df["score"], 2))).collect()[0][0]
    model_ll = df.select(
        F.avg(
            -F.col("label") * F.log(F.col("score"))
            - (1 - F.col("label")) * F.log(1 - F.col("score"))
        )
    ).collect()[0][0]

    model_aucpr = BinaryClassificationEvaluator(
        rawPredictionCol="score", metricName="areaUnderPR"
    ).evaluate(df)
    model_roc_auc = BinaryClassificationEvaluator(
        rawPredictionCol="score", metricName="areaUnderROC"
    ).evaluate(df)
    iso_bs = df.select(F.avg(F.pow(df["label"] - df["prediction"], 2))).collect()[0][0]

    print(f"model brier score loss: {model_bs}")
    print(f"calibrated model brier score loss: {iso_bs}")

    print(f"delta: {round((iso_bs/model_bs - 1) * 100, 2)}%")
    iso_ll = df.select(
        F.avg(
            -F.col("label") * F.log(F.col("prediction"))
            - (1 - F.col("label")) * F.log(1 - F.col("prediction"))
        )
    ).collect()[0][0]

    print("")

    print(f"model log loss: {model_ll}")
    print(f"calibrated model log loss: {iso_ll}")
    print(f"delta: {round((iso_ll/model_ll - 1) * 100, 2)}%")
    iso_aucpr = BinaryClassificationEvaluator(
        rawPredictionCol="prediction", metricName="areaUnderPR"
    ).evaluate(df)

    print("")

    print(f"model aucpr: {model_aucpr}")
    print(f"calibrated model aucpr: {iso_aucpr}")
    print(f"delta: {round((iso_aucpr/model_aucpr - 1) * 100, 2)}%")
    iso_roc_auc = BinaryClassificationEvaluator(
        rawPredictionCol="prediction", metricName="areaUnderROC"
    ).evaluate(df)

    print("")

    print(f"model roc_auc: {model_roc_auc}")
    print(f"calibrated model roc_auc: {iso_roc_auc}")
    print(f"delta: {round((iso_roc_auc/model_roc_auc - 1) * 100, 2)}%")

src/spark_calibration/metrics.py
Meesho/spark_calibration
import plotly.graph_objects as go
from sklearn.calibration import calibration_curve
from pyspark.sql.dataframe import DataFrame


def plot_calibration_curve(df: DataFrame):
    assert (
        "score" in df.columns and "label" in df.columns and "prediction" in df.columns
    ), "score and label columns should be present in the dataframe"

    df_p_v = df.select("label", "score", "prediction").toPandas().values

    y_test_true, y_test_pred, y_test_pred_cal = df_p_v[:, 0], df_p_v[:, 1], df_p_v[:, 2]

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=[0, 1],
            y=[0, 1],
            mode="lines",
            name="Perfect Calibration Baseline (y=x)",
            line=dict(dash="dash"),
        )
    )

    prob_true, prob_pred = calibration_curve(
        y_test_true, y_test_pred, n_bins=50, strategy="quantile"
    )
    print(f"number of bins for pre-calibration scores: {prob_true.shape[0]}")

    fig.add_trace(
        go.Scatter(x=prob_pred, y=prob_true, mode="lines+markers", name="Model")
    )

    prob_true, prob_pred = calibration_curve(
        y_test_true, y_test_pred_cal, n_bins=50, strategy="quantile"
    )
    print(f"number of bins for post-calibration scores: {prob_true.shape[0]}")

    fig.add_trace(
        go.Scatter(
            x=prob_pred, y=prob_true, mode="lines+markers", name="Calibrated Model"
        )
    )

    fig.update_layout(
        title=dict(text="Calibration Curve on test data (quantile bins)", x=0.5),
        xaxis_title="Mean Predicted Probability",
        yaxis_title="Actual Fraction of Positives",
    )

    fig.show()

src/spark_calibration/visualisation.py
Meesho/spark_calibration