Cluster multiple add stats

[RobinL]

import duckdb
import numpy as np
import pandas as pd

from splink import DuckDBAPI, Linker, splink_datasets
from splink.internals.clustering import (
    cluster_pairwise_predictions_at_multiple_thresholds,
    cluster_pairwise_predictions_at_threshold,
)
from splink.internals.misc import bayes_factor_to_prob, match_weight_to_bayes_factor

# Load the dataset and set up the DuckDB API
df = splink_datasets.historical_50k
df = df.reset_index()

df = df.drop(columns=["unique_id"])
df["unique_id"] = df.index
df = df.drop(columns=["index"])
df["index"] = df.index.astype("int32")

db_api = DuckDBAPI(":default:")

# Create the first_name_surname_concat column
df["first_name_surname_concat"] = df["first_name"] + " " + df["surname"]

# Initialize the Linker
linker = Linker(df, "50k_model.json", db_api=db_api)

# Generate predictions
df_edges = linker.inference.predict(threshold_match_probability=0.01)


# Define thresholds
thresholds_mw = [i / 2 for i in range(-20, 20, 1)]

threshold_probs = [
    bayes_factor_to_prob(match_weight_to_bayes_factor(mw)) for mw in thresholds_mw
]


# logging.getLogger("splink").setLevel(1)

# Cluster at multiple thresholds
all_clusters = cluster_pairwise_predictions_at_multiple_thresholds(
    df,
    df_edges,
    node_id_column_name="unique_id",
    db_api=db_api,
    match_probability_thresholds=threshold_probs,
    output_cluster_summary_stats=True,
)
dc_df = all_clusters.as_pandas_dataframe()


# dc is:
# ┌───────────┬──────────────┬──────────────────┬────────────────────┐
# │ threshold │ num_clusters │ max_cluster_size │  avg_cluster_size  │
# │   float   │    int64     │      int64       │       double       │
# ├───────────┼──────────────┼──────────────────┼────────────────────┤
# │       0.4 │            2 │                5 │                3.5 │
# │       0.5 │            3 │                4 │ 2.3333333333333335 │
# │       0.6 │            4 │                3 │               1.75 │
# │      0.99 │            7 │                1 │                1.0 │
# └───────────┴──────────────┴──────────────────┴────────────────────┘

import altair as alt

base_chart = (
    alt.Chart(dc_df)
    .encode(x=alt.X("threshold:Q", title="Clustering Threshold"))
    .properties(width=400, height=150)
)

num_clusters = (
    base_chart.mark_line()
    .encode(y=alt.Y("num_clusters:Q", title="Number of Clusters"))
    .properties(title="Number of Clusters vs Clustering Threshold")
)

max_cluster_size = (
    base_chart.mark_line()
    .encode(y=alt.Y("max_cluster_size:Q", title="Max Cluster Size"))
    .properties(title="Maximum Cluster Size vs Clustering Threshold")
)

combined_chart = alt.vconcat(num_clusters, max_cluster_size).resolve_scale(
    y="independent"
)

combined_chart