Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
174 changes: 96 additions & 78 deletions src/multicalibration/methods.py

Large diffs are not rendered by default.

66 changes: 36 additions & 30 deletions src/multicalibration/plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -685,43 +685,43 @@ def plot_segment_calibration_errors(


def plot_learning_curve(
mcboost_model: methods.MCBoost, show_all: bool = False
mcgrad_model: methods.MCGrad, show_all: bool = False
) -> go.Figure:
"""
Plots a learning curve for a MCBoost model.
Plots a learning curve for an MCGrad model.

:param mcboost_model: An MCBoost model object.
:param show_all: Whether to show all metrics in the learning curve. If False, only the metric specified in MCBoost's early_stopping_score_func is shown.
:param mcgrad_model: An MCGrad model object.
:param show_all: Whether to show all metrics in the learning curve. If False, only the metric specified in the model's early_stopping_score_func is shown.
:returns: A Plotly Figure object representing the learning curve.
"""
if not mcboost_model.early_stopping:
if not mcgrad_model.early_stopping:
raise ValueError(
"Learning curve can only be plotted for models that have been trained with EARLY_STOPPING=True."
"Learning curve can only be plotted for models that have been trained with early_stopping=True."
)

performance_metrics = mcboost_model._performance_metrics
performance_metrics = mcgrad_model._performance_metrics
extra_evaluation_due_to_early_stopping = (
1
if (
mcboost_model.early_stopping
and len(mcboost_model.mr) < mcboost_model.NUM_ROUNDS
mcgrad_model.early_stopping
and len(mcgrad_model.mr) < mcgrad_model.num_rounds
)
else 0
)
# Calculate the total number of rounds (including the initial round)
tot_num_rounds = min(
1
+ len(mcboost_model.mr)
+ len(mcgrad_model.mr)
+ extra_evaluation_due_to_early_stopping
+ mcboost_model.PATIENCE,
1 + mcboost_model.NUM_ROUNDS,
+ mcgrad_model.patience,
1 + mcgrad_model.num_rounds,
)
x_vals = np.arange(0, tot_num_rounds)
metric_names = [mcboost_model.early_stopping_score_func.name]
metric_names = [mcgrad_model.early_stopping_score_func.name]
for metric_name in performance_metrics.keys():
if (
"valid" in metric_name
and mcboost_model.early_stopping_score_func.name not in metric_name
and mcgrad_model.early_stopping_score_func.name not in metric_name
and show_all
):
metric_names.append(metric_name.split("performance_")[-1])
Expand Down Expand Up @@ -758,8 +758,8 @@ def plot_learning_curve(
row=i + 1,
col=1,
)
if mcboost_model.save_training_performance:
train_performance = mcboost_model._performance_metrics[
if mcgrad_model.save_training_performance:
train_performance = mcgrad_model._performance_metrics[
f"avg_train_performance_{metric_name}"
]
# Plot the training performance (training set)
Expand Down Expand Up @@ -787,7 +787,7 @@ def plot_learning_curve(
)
# Add vertical line for the selected iteration
fig.add_vline(
x=len(mcboost_model.mr),
x=len(mcgrad_model.mr),
line_dash="dash",
line_color="black",
opacity=0.5,
Expand All @@ -800,7 +800,7 @@ def plot_learning_curve(
text="Selected round by early stopping",
xref="x",
yref="y",
x=len(mcboost_model.mr) - 0.05,
x=len(mcgrad_model.mr) - 0.05,
y=max_perf_for_annotation * 1.075,
xanchor="center",
yanchor="middle",
Expand All @@ -815,9 +815,9 @@ def plot_learning_curve(
col=1,
)
if "mce_sigma_scale" in metric_name:
if max_perf_for_annotation >= mcboost_model.MCE_STRONG_EVIDENCE_THRESHOLD:
if max_perf_for_annotation >= mcgrad_model.MCE_STRONG_EVIDENCE_THRESHOLD:
fig.add_hline(
y=mcboost_model.MCE_STRONG_EVIDENCE_THRESHOLD,
y=mcgrad_model.MCE_STRONG_EVIDENCE_THRESHOLD,
line_dash="dash",
line_color="darkgreen",
opacity=1,
Expand All @@ -830,9 +830,9 @@ def plot_learning_curve(
annotation_textangle=90,
)

if max_perf_for_annotation >= mcboost_model.MCE_STAT_SIGN_THRESHOLD:
if max_perf_for_annotation >= mcgrad_model.MCE_STAT_SIGN_THRESHOLD:
fig.add_hline(
y=mcboost_model.MCE_STAT_SIGN_THRESHOLD,
y=mcgrad_model.MCE_STAT_SIGN_THRESHOLD,
line_dash="dash",
line_color="darkorange",
opacity=0.7,
Expand All @@ -846,11 +846,11 @@ def plot_learning_curve(
)

if (
test_performance[len(mcboost_model.mr)]
>= mcboost_model.MCE_STRONG_EVIDENCE_THRESHOLD
test_performance[len(mcgrad_model.mr)]
>= mcgrad_model.MCE_STRONG_EVIDENCE_THRESHOLD
):
fig.add_annotation(
text="<b>WARNING: MCBoost run failed to remove strong evidence of multicalibration!</b>",
text=f"<b>WARNING: {mcgrad_model.__class__.__name__} run failed to remove strong evidence of multicalibration!</b>",
xref="paper",
yref="paper",
x=tot_num_rounds - 1,
Expand Down Expand Up @@ -889,23 +889,29 @@ def plot_learning_curve(
row=i + 1,
col=1,
)
# Update x-axis labels including "without MCBoost" for the 0th iteration
# Update x-axis labels including "without MCGrad" for the 0th iteration
if len(x_vals) <= 10:
fig.update_xaxes(
title_text="MCBoost round" if i == len(metric_names) - 1 else "",
title_text=f"{mcgrad_model.__class__.__name__} round"
if i == len(metric_names) - 1
else "",
tickmode="array",
tickvals=x_vals,
ticktext=["without<br>MCBoost"] + [str(int(val)) for val in x_vals[1:]],
ticktext=[f"without<br>{mcgrad_model.__class__.__name__}"]
+ [str(int(val)) for val in x_vals[1:]],
row=i + 1,
col=1,
)
else:
x_vals = np.arange(0, len(x_vals), np.ceil(len(x_vals) / 5))
fig.update_xaxes(
title_text="MCBoost round" if i == len(metric_names) - 1 else "",
title_text=f"{mcgrad_model.__class__.__name__} round"
if i == len(metric_names) - 1
else "",
tickmode="array",
tickvals=x_vals,
ticktext=["without<br>MCBoost"] + [str(int(val)) for val in x_vals[1:]],
ticktext=[f"without<br>{mcgrad_model.__class__.__name__}"]
+ [str(int(val)) for val in x_vals[1:]],
row=i + 1,
col=1,
)
Expand Down
30 changes: 16 additions & 14 deletions src/multicalibration/tuning.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,8 @@ def _suppress_logger(logger: logging.Logger) -> Generator[None, None, None]:
logger.setLevel(previous_level)


def tune_mcboost_params(
model: methods.MCBoost,
def tune_mcgrad_params(
model: methods.MCGrad,
df_train: pd.DataFrame,
prediction_column_name: str,
label_column_name: str,
Expand All @@ -135,12 +135,11 @@ def tune_mcboost_params(
parameter_configurations: list[ParameterConfig] | None = None,
pass_df_val_into_tuning: bool = False,
pass_df_val_into_final_fit: bool = False,
) -> tuple[methods.MCBoost | None, pd.DataFrame]:
# Make a description for this function
) -> tuple[methods.MCGrad | None, pd.DataFrame]:
"""
Tune the hyperparameters of an MCBoost model using Ax.
Tune the hyperparameters of an MCGrad model using Ax.

:param model: The MCBoost model to be tuned. It could be a fitted model or an unfitted model.
:param model: The MCGrad model to be tuned. It could be a fitted model or an unfitted model.
:param df_train: The training data: 80% of the data is used for training the model, and the remaining 20% is used for validation.
:param prediction_column_name: The name of the prediction column in the data.
:param label_column_name: The name of the label column in the data.
Expand Down Expand Up @@ -219,7 +218,7 @@ def _train_evaluate(parameterization: dict[str, Any]) -> float:

ax_client = AxClient()
ax_client.create_experiment(
name="mcboost_lightgbm_autotuning" + str(uuid.uuid4())[:8],
name="lightgbm_autotuning" + str(uuid.uuid4())[:8],
parameters=[config.to_dict() for config in parameter_configurations],
objectives={"normalized_entropy": ObjectiveProperties(minimize=True)},
# If num_initialization_trials is None, the number of warm starting trials is automatically determined
Expand All @@ -233,25 +232,25 @@ def _train_evaluate(parameterization: dict[str, Any]) -> float:
},
)

# Construct a set of parameters for the first trial which contains the MCBoost defaults for every parameter that is tuned. If a default is not available
# Construct a set of parameters for the first trial which contains the defaults for every parameter that is tuned. If a default is not available
# use the Lightgbm default
initial_trial_parameters = {}
mcboost_defaults = methods.MCBoost.DEFAULT_HYPERPARAMS["lightgbm_params"]
mcgrad_defaults = methods.MCGrad.DEFAULT_HYPERPARAMS["lightgbm_params"]
for config in parameter_configurations:
if config.name in mcboost_defaults:
initial_trial_parameters[config.name] = mcboost_defaults[config.name]
if config.name in mcgrad_defaults:
initial_trial_parameters[config.name] = mcgrad_defaults[config.name]
else:
initial_trial_parameters[config.name] = ORIGINAL_LIGHTGBM_PARAMS[
config.name
]

logger.info(
f"Adding initial configuration from MCBoost defaults to trials: {initial_trial_parameters}"
f"Adding initial configuration from defaults to trials: {initial_trial_parameters}"
)

with _suppress_logger(methods.logger):
# Attach and complete the initial trial with default hyperparameters. Note that we're only using the defaults for the parameters that are being tuned.
# That is, this configuration does not necessarily correspond to the out-of-the-box defaults for MCBoost.
# That is, this configuration does not necessarily correspond to the out-of-the-box defaults.
_, initial_trial_index = ax_client.attach_trial(
parameters=initial_trial_parameters
)
Expand All @@ -273,7 +272,7 @@ def _train_evaluate(parameterization: dict[str, Any]) -> float:
best_params = best_params[0]

logger.info(f"Best parameters: {best_params}")
logger.info("Fitting MCBoost model with best parameters")
logger.info("Fitting model with best parameters")

with _suppress_logger(methods.logger):
model._set_lightgbm_params(best_params)
Expand All @@ -294,3 +293,6 @@ def _train_evaluate(parameterization: dict[str, Any]) -> float:
)

return model, trial_results


# @oss-disable[end= ]: tune_mcboost_params = tune_mcgrad_params
Loading