Plot

jarvais.utils.plot

plot_corr(corr, size, output_dir, file_name='correlation_matrix.png', title='Correlation Matrix')

Plots a lower-triangle heatmap of the correlation matrix and saves it as an image file.

Parameters:

Name	Type	Description	Default
corr	DataFrame	Correlation matrix to visualize.	required
size	float	Size of the heatmap figure.	required
output_dir	Path	Directory to save the output image.	required
file_name	str	Name of the saved image file. Defaults to 'correlation_matrix.png'.	'correlation_matrix.png'

Example

import pandas as pd
from pathlib import Path

# Sample data
df = pd.DataFrame({
    'A': [1, 2, 3, 4, 5],
    'B': [5, 4, 3, 2, 1],
    'C': [2, 3, 4, 5, 6]
})

# Compute Spearman correlation
corr_matrix = df.corr(method='spearman')

# Plot and save the heatmap
plot_corr(corr=corr_matrix, size=6, output_dir=Path('./output'))

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_corr(
        corr: pd.DataFrame,
        size: float,
        output_dir: Path,
        file_name: str = 'correlation_matrix.png',
        title: str = "Correlation Matrix"
    ) -> None:
    """
    Plots a lower-triangle heatmap of the correlation matrix and saves it as an image file.

    Args:
        corr (pd.DataFrame): Correlation matrix to visualize.
        size (float): Size of the heatmap figure.
        output_dir (Path): Directory to save the output image.
        file_name (str): Name of the saved image file. Defaults to 'correlation_matrix.png'.

    Example:
        ```python
        import pandas as pd
        from pathlib import Path

        # Sample data
        df = pd.DataFrame({
            'A': [1, 2, 3, 4, 5],
            'B': [5, 4, 3, 2, 1],
            'C': [2, 3, 4, 5, 6]
        })

        # Compute Spearman correlation
        corr_matrix = df.corr(method='spearman')

        # Plot and save the heatmap
        plot_corr(corr=corr_matrix, size=6, output_dir=Path('./output'))
        ```
    """
    fig, ax = plt.subplots(1, 1, figsize=(size*1.2, size))
    mask = np.triu(np.ones_like(corr, dtype=bool)) # Keep only lower triangle
    np.fill_diagonal(mask, False)
    sns.heatmap(corr, mask=mask, annot=True, cmap='coolwarm', vmin=-1, vmax=1, linewidth=.5, fmt="1.2f", ax=ax)
    plt.title(title)
    plt.tight_layout()

    figure_path = output_dir / file_name
    fig.savefig(figure_path)
    plt.close()

plot_frequency_table(data, columns, output_dir, save_to_json=False)

Generates and saves heatmap visualizations for frequency tables of all column pair combinations.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input dataset containing the columns to analyze.	required
columns	list	List of column names to create frequency tables for.	required
output_dir	Path	Directory to save the generated heatmaps.	required
save_to_json	bool	Flag to indicate whether to save frequency table data as JSON files.	False

Source code in src/jarvais/utils/plot.py

@config_plot(plot_type='ft')
def plot_frequency_table(
        data: pd.DataFrame,
        columns: list,
        output_dir: Path,
        save_to_json: bool = False
    ) -> None:
    """
    Generates and saves heatmap visualizations for frequency tables of all column pair combinations.

    Args:
        data (pd.DataFrame): Input dataset containing the columns to analyze.
        columns (list): List of column names to create frequency tables for.
        output_dir (Path): Directory to save the generated heatmaps.
        save_to_json (bool): Flag to indicate whether to save frequency table data as JSON files.
    """
    frequency_dir = Path(output_dir) / 'frequency_tables'
    frequency_dir.mkdir(parents=True, exist_ok=True)

    for column_1, column_2 in combinations(columns, 2):
        heatmap_data = pd.crosstab(data[column_1], data[column_2])
        plt.figure(figsize=(8, 6))
        sns.heatmap(heatmap_data, annot=True, cmap='coolwarm', fmt='d', linewidth=.5)
        plt.title(f'Frequency Table for {column_1} and {column_2}')
        plt.xlabel(column_2)
        plt.ylabel(column_1)
        plt.savefig(frequency_dir / f'{column_1}_vs_{column_2}.png')
        plt.close()

        if save_to_json:
            heatmap_data.to_json(frequency_dir / f'{column_1}_vs_{column_2}.json')

plot_pairplot(data, continuous_columns, output_dir, target_variable=None, n_keep=10)

Generates a pair plot of the specified continuous columns in the dataset.

Parameters:

Name	Type	Description	Default
data	DataFrame	The input DataFrame containing the data to be visualized.	required
continuous_columns	list	A list of column names corresponding to continuous variables.	required
output_dir	Path	Directory where the resulting plot will be saved.	required
target_variable	str	The target variable to use as a hue for coloring the pair plot. Defaults to None.	None
n_keep	int	The maximum number of continuous columns to include in the plot. If exceeded, the most correlated columns are selected. Defaults to 10.	10

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_pairplot(
        data: pd.DataFrame,
        continuous_columns: list,
        output_dir: Path,
        target_variable: str = None,
        n_keep: int = 10
    ) -> None:
    """
    Generates a pair plot of the specified continuous columns in the dataset.

    Args:
        data (pd.DataFrame): The input DataFrame containing the data to be visualized.
        continuous_columns (list): A list of column names corresponding to continuous variables.
        output_dir (Path): Directory where the resulting plot will be saved.
        target_variable (str, optional): The target variable to use as a hue for coloring the pair plot. Defaults to None.
        n_keep (int, optional): The maximum number of continuous columns to include in the plot. 
            If exceeded, the most correlated columns are selected. Defaults to 10.
    """
    if len(continuous_columns) > n_keep:
        spearman_corr = data[continuous_columns].corr(method="spearman") 
        corr_pairs = spearman_corr.abs().unstack().sort_values(
            kind="quicksort",
            ascending=False
        ).drop_duplicates()
        top_10_pairs = corr_pairs[corr_pairs < 1].nlargest(5)
        columns_to_plot = list({index for pair in top_10_pairs.index for index in pair})
    else:
        columns_to_plot = continuous_columns.copy()

    hue = target_variable
    if target_variable is not None:
        columns_to_plot += [target_variable]

    sns.set_theme(style="darkgrid")
    g = sns.pairplot(data[columns_to_plot], hue=hue)
    g.figure.suptitle("Pair Plot", y=1.08)

    figure_path = output_dir / 'pairplot.png'
    plt.savefig(figure_path)
    plt.close()

plot_umap(data, continuous_columns, output_dir)

Generates a 2D UMAP projection of the specified continuous columns and saves the resulting scatter plot.

Parameters:

Name	Type	Description	Default
data	DataFrame	The input DataFrame containing the data to be visualized.	required
continuous_columns	list	A list of column names corresponding to continuous variables to be included in the UMAP projection.	required
output_dir	Path	Directory where the resulting plot will be saved.	required

Returns:

Type	Description
ndarray	np.ndarray: A 2D NumPy array of the UMAP-transformed data.

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_umap(
        data: pd.DataFrame,
        continuous_columns: list,
        output_dir: Path,
    ) -> np.ndarray:
    """
    Generates a 2D UMAP projection of the specified continuous columns and saves the resulting scatter plot.

    Args:
        data (pd.DataFrame): The input DataFrame containing the data to be visualized.
        continuous_columns (list): A list of column names corresponding to continuous variables 
            to be included in the UMAP projection.
        output_dir (Path): Directory where the resulting plot will be saved.

    Returns:
        np.ndarray: A 2D NumPy array of the UMAP-transformed data.
    """
    umap_data = UMAP(n_components=2).fit_transform(data[continuous_columns])

    plt.figure(figsize=(8, 8))
    sns.scatterplot(x=umap_data[:,0], y=umap_data[:,1], alpha=.7)
    plt.title('UMAP of Continuous Variables')
    plt.savefig(output_dir / 'umap_continuous_data.png')
    plt.close()

    return umap_data

plot_kaplan_meier_by_category(data_x, data_y, categorical_columns, output_dir)

Plots Kaplan-Meier survival curves for each category in the specified categorical columns.

Parameters:

Name	Type	Description	Default
data_x	DataFrame	Dataset containing the categorical columns to group by.	required
data_y	DataFrame	Dataset containing 'time' and 'event' columns for survival analysis.	required
categorical_columns	list	List of categorical column names to generate survival curves for.	required
output_dir	Path	Directory to save the Kaplan-Meier survival curve plots.	required

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_kaplan_meier_by_category(
        data_x: pd.DataFrame,
        data_y: pd.DataFrame,
        categorical_columns: list,
        output_dir: Path
    ) -> None:
    """
    Plots Kaplan-Meier survival curves for each category in the specified categorical columns.

    Args:
        data_x (pd.DataFrame): Dataset containing the categorical columns to group by.
        data_y (pd.DataFrame): Dataset containing 'time' and 'event' columns for survival analysis.
        categorical_columns (list): List of categorical column names to generate survival curves for.
        output_dir (Path): Directory to save the Kaplan-Meier survival curve plots.
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    for cat_col in categorical_columns:
        plt.figure(figsize=(10, 6))
        plt.title(f"Kaplan-Meier Survival Curve by {cat_col}")

        unique_categories = data_x[cat_col].unique()

        # Plot survival curves for each category
        for category in unique_categories:
            mask_category = data_x[cat_col] == category
            try: # To catch when there are not enough samples for category
                time_category, survival_prob_category, conf_int = kaplan_meier_estimator(
                    data_y["event"][mask_category].astype(bool),
                    data_y["time"][mask_category],
                    conf_type="log-log",
                )

                plt.step(
                    time_category,
                    survival_prob_category,
                    where="post",
                    label=f"{cat_col} = {category}"
                )
                plt.fill_between(
                    time_category,
                    conf_int[0],
                    conf_int[1],
                    alpha=0.25,
                    step="post"
                )
            except Exception as _:
                pass

        results_multivariate = multivariate_logrank_test(
            data_y['time'], 
            data_x[cat_col], 
            data_y['event']
        )
        multivariate_p_value = results_multivariate.p_value

        plt.text(0.6, 0.1, f"Multivariate log-rank p-value: {multivariate_p_value:.4e}",
                 fontsize=10, transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.8))
        plt.ylim(0, 1)
        plt.ylabel(r"Estimated Probability of Survival $\hat{S}(t)$")
        plt.xlabel("Time $t$")
        plt.legend(loc="best")
        plt.grid(alpha=0.3)
        plt.savefig(output_dir / f'kaplan_meier_{cat_col}.png')
        plt.close()

plot_feature_importance(df, output_dir, model_name='')

Plots feature importance with standard deviation and p-value significance.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame containing the feature importance data. Look at example for required format.	required
output_dir	Path	Directory to save the feature importance plot.	required
model_name	str	Optional name of the model, included in the plot title.	''

Example

import pandas as pd
from pathlib import Path

df = pd.DataFrame({
    'importance': [0.25, 0.18, 0.12, 0.10],
    'stddev': [0.03, 0.02, 0.01, 0.015],
    'p_value': [0.03, 0.07, 0.01, 0.2]
}, index=['Feature A', 'Feature B', 'Feature C', 'Feature D'])

plot_feature_importance(df, Path('./output'))

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_feature_importance(df: pd.DataFrame, output_dir: Path, model_name: str=''):
    """
    Plots feature importance with standard deviation and p-value significance.

    Args:
        df (pd.DataFrame): DataFrame containing the feature importance data. 
            Look at example for required format.
        output_dir (Path): Directory to save the feature importance plot.
        model_name (str): Optional name of the model, included in the plot title.

    Example:
        ```python
        import pandas as pd
        from pathlib import Path

        df = pd.DataFrame({
            'importance': [0.25, 0.18, 0.12, 0.10],
            'stddev': [0.03, 0.02, 0.01, 0.015],
            'p_value': [0.03, 0.07, 0.01, 0.2]
        }, index=['Feature A', 'Feature B', 'Feature C', 'Feature D'])

        plot_feature_importance(df, Path('./output'))
        ```
    """
    fig, ax = plt.subplots(figsize=(20, 12), dpi=72)

    bars = ax.bar(df.index, df['importance'], yerr=df['stddev'], capsize=5, color='skyblue', edgecolor='black')

    if 'p_value' in df.columns:
        for bar, p_value in zip(bars, df['p_value']):
            height = bar.get_height()
            significance = '*' if p_value < 0.05 else ''
            ax.text(bar.get_x() + bar.get_width() / 2.0, height + 0.002, significance,
                    ha='center', va='bottom', fontsize=10, color='red')

    ax.set_xlabel('Feature', fontsize=14)
    ax.set_ylabel('Importance', fontsize=14)
    ax.set_title(f'Feature Importance with Standard Deviation and p-value Significance ({model_name})', fontsize=16)
    ax.axhline(0, color='grey', linewidth=0.8)

    ax.set_xticks(np.arange(len(df.index.values)))
    ax.set_xticklabels(df.index.values, rotation=60, ha='right', fontsize=10)

    ax.grid(axis='y', linestyle='--', alpha=0.7)

    significance_patch = plt.Line2D([0], [0], color='red', marker='*', linestyle='None', markersize=10, label='p < 0.05')
    ax.legend(handles=[significance_patch], loc='upper right', fontsize=12)

    plt.tight_layout()
    fig.savefig(output_dir / 'feature_importance.png')
    plt.close()

plot_shap_values(predictor, X_train, X_test, output_dir, max_display=10)

Generates and saves SHAP value visualizations, including a heatmap and a bar plot, for a autogluon tabular model.

Parameters:

Name	Type	Description	Default
predictor	TabularPredictor	The trained tabular predictor model for which SHAP values are calculated.	required
X_train	DataFrame	Training dataset used to create the SHAP background data.	required
X_test	DataFrame	Test dataset used to evaluate and compute SHAP values.	required
output_dir	Path	Directory to save the SHAP value visualizations.	required
max_display	int	Maximum number of features to display in the visualizations. Defaults to 10.	10

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_shap_values(
        predictor: TabularPredictor,
        X_train: pd.DataFrame,
        X_test: pd.DataFrame,
        output_dir: Path,
        max_display: int = 10,
    ) -> None:
    """
    Generates and saves SHAP value visualizations, including a heatmap and a bar plot, for a autogluon tabular model.

    Args:
        predictor (TabularPredictor): The trained tabular predictor model for which SHAP values are calculated.
        X_train (pd.DataFrame): Training dataset used to create the SHAP background data.
        X_test (pd.DataFrame): Test dataset used to evaluate and compute SHAP values.
        output_dir (Path): Directory to save the SHAP value visualizations.
        max_display (int): Maximum number of features to display in the visualizations. Defaults to 10.
    """
    predictor = ModelWrapper(predictor, X_train.columns)
    background_data = shap.sample(X_train, 100)
    shap_exp = shap.KernelExplainer(predictor.predict_proba, background_data)

    # sample 100 samples from test set to evaluate with shap values
    test_data = shap.sample(X_test, 100)

    # Compute SHAP values for the test set
    shap_values = shap_exp(test_data)

    sns.set_theme(style="darkgrid")
    fig, ax = plt.subplots(figsize=(20, 12), dpi=72)
    shap.plots.heatmap(shap_values[...,1], max_display=max_display, show=False, ax=ax)
    fig.savefig(output_dir / 'shap_heatmap.png')
    plt.close()

    fig, ax = plt.subplots(figsize=(20, 12), dpi=72)
    shap.plots.bar(shap_values[...,1], max_display=max_display, show=False, ax=ax)
    fig.savefig(output_dir / 'shap_barplot.png')
    plt.close()

plot_violin_of_bootstrapped_metrics(trainer, X_test, y_test, X_val, y_val, X_train, y_train, output_dir)

Generates violin plots for bootstrapped model performance metrics across train, validation, and test datasets.

Parameters:

Name	Type	Description	Default
trainer	TrainerSupervised	The trained model predictor for evaluating performance metrics.	required
X_test	Series	Test features dataset.	required
y_test	Series	Test target values.	required
X_val	Series	Validation features dataset.	required
y_val	Series	Validation target values.	required
X_train	Series	Training features dataset.	required
y_train	Series	Training target values.	required
output_dir	Path	Directory to save the generated violin plots.	required

Source code in src/jarvais/utils/plot.py

def plot_violin_of_bootstrapped_metrics(
        trainer,
        X_test: pd.Series,
        y_test: pd.Series,
        X_val: pd.Series,
        y_val: pd.Series,
        X_train: pd.Series,
        y_train: pd.Series,
        output_dir: Path
    ) -> None:
    """
    Generates violin plots for bootstrapped model performance metrics across train, validation, and test datasets.

    Args:
        trainer (TrainerSupervised): The trained model predictor for evaluating performance metrics.
        X_test (pd.Series): Test features dataset.
        y_test (pd.Series): Test target values.
        X_val (pd.Series): Validation features dataset.
        y_val (pd.Series): Validation target values.
        X_train (pd.Series): Training features dataset.
        y_train (pd.Series): Training target values.
        output_dir (Path): Directory to save the generated violin plots.
    """
    # Define metrics based on the problem type
    if trainer.task == 'regression':
        metrics = [('R Squared', r2_score), ('Root Mean Squared Error', root_mean_squared_error)]
    elif trainer.task == 'binary':
        metrics = [('AUROC', roc_auc_score), ('AUPRC', auprc)]
    elif trainer.task == 'survival':
        metrics = [('Concordance Index', ci_wrapper)]

    # Prepare lists for DataFrame
    results = []

    # Loop through models and metrics to compute bootstrapped values
    for model_name in trainer.model_names():
        y_pred_test = trainer.infer(X_test, model=model_name)
        y_pred_val = trainer.infer(X_val, model=model_name)
        y_pred_train = trainer.infer(X_train, model=model_name)

        for metric_name, metric_func in metrics:
            test_values = bootstrap_metric(y_test.to_numpy(), y_pred_test, metric_func)
            results.extend([(model_name, metric_name, 'Test', value) for value in test_values])

            val_values = bootstrap_metric(y_val.to_numpy(), y_pred_val, metric_func)
            results.extend([(model_name, metric_name, 'Validation', value) for value in val_values])

            train_values = bootstrap_metric(y_train.to_numpy(), y_pred_train, metric_func)
            results.extend([(model_name, metric_name, 'Train', value) for value in train_values])

    # Create a results DataFrame
    result_df = pd.DataFrame(results, columns=['model', 'metric', 'data_split', 'value'])

     # Sort models by median metric value within each combination of metric and data_split
    model_order_per_split = {}
    for split in ['Test', 'Validation', 'Train']:
        split_order = (
            result_df[result_df['data_split'] == split]
            .groupby(['metric', 'model'])['value']
            .median()
            .reset_index()
            .sort_values(by=['metric', 'value'], ascending=[True, False])
            .groupby('metric')['model']
            .apply(list)
            .to_dict()
        )
        model_order_per_split[split] = split_order

    # Function to create violin plots for a specific data split
    def create_violin_plot(data_split, save_path):
        sns.set_theme(style="darkgrid")
        subset = result_df[result_df['data_split'] == data_split]
        g = sns.FacetGrid(
            subset,
            col="metric",
            margin_titles=True,
            height=4,
            aspect=1.5,
            sharex=False,
        )

        # Create violin plots with sorted models
        def violin_plot(data, **kwargs):
            metric = data.iloc[0]['metric']
            order = model_order_per_split[data_split].get(metric, None)
            sns.violinplot(data=data, x="value", y="model", linewidth=1, order=order, **kwargs)

        g.map_dataframe(violin_plot)

        # Adjust the titles and axis labels
        g.set_titles(col_template="{col_name}")
        g.set_axis_labels("", "Model")

        # Add overall title and adjust layout
        g.figure.suptitle(f"Model Performance of {data_split} Data (Bootstrapped)", fontsize=16)
        g.tight_layout(w_pad=0.5, h_pad=1)

        # Save the plot
        g.savefig(save_path, dpi=500)
        plt.close()

    # Generate and save plots for each data split
    create_violin_plot('Test', output_dir / 'test_metrics_bootstrap.png')
    create_violin_plot('Validation', output_dir / 'validation_metrics_bootstrap.png')
    create_violin_plot('Train', output_dir / 'train_metrics_bootstrap.png')

plot_regression_diagnostics(y_true, y_pred, output_dir)

Generates diagnostic plots for evaluating a regression model.

Plots

• True vs. Predicted values plot.
• Residuals plot.
• Histogram of residuals.

Parameters:

Name	Type	Description	Default
y_true	ndarray	Array of true target values.	required
y_pred	ndarray	Array of predicted values from the regression model.	required
output_dir	Path	Directory to save the diagnostic plots.	required

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_regression_diagnostics(
        y_true: np.ndarray, 
        y_pred: np.ndarray, 
        output_dir: Path
    ) -> None:
    """
    Generates diagnostic plots for evaluating a regression model.

    Plots:
        - True vs. Predicted values plot.
        - Residuals plot.
        - Histogram of residuals.

    Args:
        y_true (np.ndarray): Array of true target values.
        y_pred (np.ndarray): Array of predicted values from the regression model.
        output_dir (Path): Directory to save the diagnostic plots.
    """
    residuals = y_true - y_pred

    # Regression Line
    plt.figure(figsize=(10, 6))
    sns.set_theme(style="darkgrid")
    sns.scatterplot(x=y_true, y=y_pred, alpha=0.5)
    sns.lineplot(x=y_true, y=y_true, color='red')  # Perfect prediction line
    plt.xlabel('True Values')
    plt.ylabel('Predictions')
    plt.title('True vs Predicted Values')
    plt.savefig(output_dir / 'true_vs_predicted.png')
    plt.close()

    # Residuals
    plt.figure(figsize=(10, 6))
    sns.set_theme(style="darkgrid")
    sns.scatterplot(x=y_pred, y=residuals, alpha=0.5)
    plt.axhline(0, color='red', linestyle='--')
    plt.xlabel('Fitted Values')
    plt.ylabel('Residuals')
    plt.title('Residual Plot')
    plt.savefig(output_dir / 'residual_plot.png')
    plt.close()

    # Residual Histogram
    plt.figure(figsize=(10, 6))
    sns.set_theme(style="darkgrid")
    sns.histplot(residuals, kde=True, bins=30)
    plt.xlabel('Residuals')
    plt.title('Histogram of Residuals')
    plt.savefig(output_dir / 'residual_hist.png')
    plt.close()

plot_classification_diagnostics(y_test, y_test_pred, y_val, y_val_pred, y_train, y_train_pred, output_dir)

Generates diagnostic plots for evaluating a classification model.

Plots

• Epic model evaluation plots (ROC Curve, Precision-Recall Curve, Calibration Curve, Sensitivity/Flag Curve).
• Confusion Matrix.

Parameters:

Name	Type	Description	Default
y_test	DataFrame	True labels for the test dataset.	required
y_test_pred	DataFrame	Predicted probabilities for the test dataset.	required
y_val	DataFrame	True labels for the validation dataset.	required
y_val_pred	DataFrame	Predicted probabilities for the validation dataset.	required
y_train	DataFrame	True labels for the training dataset.	required
y_train_pred	DataFrame	Predicted probabilities for the training dataset.	required
output_dir	Path	Directory to save the diagnostic plots.	required

Source code in src/jarvais/utils/plot.py

@config_plot()
def plot_classification_diagnostics(
        y_test: pd.DataFrame,
        y_test_pred: pd.DataFrame,
        y_val: pd.DataFrame,
        y_val_pred: pd.DataFrame,
        y_train: pd.DataFrame,
        y_train_pred: pd.DataFrame,
        output_dir: Path
    ) -> None:
    """
    Generates diagnostic plots for evaluating a classification model.

    Plots:
        - Epic model evaluation plots (ROC Curve, Precision-Recall Curve, Calibration Curve, Sensitivity/Flag Curve).
        - Confusion Matrix.

    Args:
        y_test (pd.DataFrame): True labels for the test dataset.
        y_test_pred (pd.DataFrame): Predicted probabilities for the test dataset.
        y_val (pd.DataFrame): True labels for the validation dataset.
        y_val_pred (pd.DataFrame): Predicted probabilities for the validation dataset.
        y_train (pd.DataFrame): True labels for the training dataset.
        y_train_pred (pd.DataFrame): Predicted probabilities for the training dataset.
        output_dir (Path): Directory to save the diagnostic plots.
    """
    plot_epic_copy(
        y_test.to_numpy(),
        y_test_pred.to_numpy(),
        y_val.to_numpy(),
        y_val_pred.to_numpy(),
        y_train.to_numpy(),
        y_train_pred.to_numpy() ,
        output_dir
    )

    conf_matrix = confusion_matrix(y_test, y_test_pred.apply(lambda x: 1 if x >= 0.5 else 0))

    plt.figure(figsize=(8, 6))
    sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.savefig(output_dir / 'confusion_matrix.png')
    plt.close()