# Author: Juan Parras & Patricia A. Apellániz # Email: patricia.alonsod@upm.es # Date: 05/08/2025 # Package imports import os import sys import sympy import pickle import numpy as np import pandas as pd import matplotlib.pyplot as plt import matplotlib.patches as mpatches from tqdm import tqdm from kan import ex_round from copy import deepcopy from tueplots import bundles current_dir = os.path.dirname(os.path.abspath(__file__)) parent_dir = os.path.abspath(os.path.join(current_dir, "..")) sys.path.append(parent_dir) from src.data import load_data from src.utils import get_config, create_results_folder from src.models.models import Kan_model from src.models.model_utils import get_metrics from src.representation import plot_binary_explanation_plot, plot_sorted_variances, radar_factory def run_kan_gam(): _ = model.run_model(x_train, x_test, y_train, y_test) predicted_prob_test = model.predict_proba(x_test.to_numpy()) pred_proba_train = model.predict_proba(x_train.to_numpy()) y_pred_test = model.predict(x_test.to_numpy()) y_pred_train = model.predict(x_train.to_numpy()) if predicted_prob_test.shape[1] == 2: binary_data = True n_classes_data = 2 n_logits_data = 1 else: binary_data = False n_classes_data = predicted_prob_test.shape[1] n_logits_data = n_classes_data test_metrics = get_metrics(y_test.to_numpy(), y_pred_test, predicted_prob_test) train_metrics = get_metrics(y_train.to_numpy(), y_pred_train, pred_proba_train) return predicted_prob_test, binary_data, n_classes_data, n_logits_data, test_metrics, train_metrics def get_patient_values(formula, x_in): x = x_in.to_numpy() if isinstance(formula, sympy.Float): # WE have a constant! delta = np.zeros((x.shape[0], x.shape[1] + 1)) delta[:, -1] = float(formula) # There is only a constant term! else: delta = np.zeros( (x.shape[0], x.shape[1] + 1)) # One input per covariate, one extra output for the constant term for i in range(x.shape[0]): # For each patient for fs in formula.args: formula_sum_term = deepcopy(fs) if isinstance(formula_sum_term, sympy.Float): # We have a constant! delta[i, -1] = float(formula_sum_term) else: # Since it is a KAAM, it depends on a single variable assert len(formula_sum_term.free_symbols) == 1 variable_in_the_expresion = list(formula_sum_term.free_symbols)[0] variable_index = x_in.columns.get_loc(str(variable_in_the_expresion)) delta[i, variable_index] += float( formula_sum_term.subs(variable_in_the_expresion, x[i, variable_index])) delta = pd.DataFrame(delta, columns=x_in.columns.tolist() + ['const']) return delta def adjust_polynomial(tr_x, tr_y, test_x, test_y, metrics_train, metrics_test, binary_dataset, n_logits_dataset, results_folder, dataset_name): # lib = ['x', 'x^2', 'x^3', 'x^4', 'x^5'] # Polynomial model, used because it tends to provide a nice symbolic formula model.model(model.dataset['train_input']) # To have activations updated! model.model.auto_symbolic() # Add the lib in here if desired formula = model.model.symbolic_formula()[0] # We have several formulae, one per logit! formula = [ex_round(f, 3) for f in formula] # Number of digits to approximate if binary_dataset: delta_formula = [formula[1] - formula[0]] # Keep a single formula (this is the delta) else: delta_formula = [f for f in formula] # Keep all the formulae, one per logit. IMPORTANT NOTE: We could simplify the formulae, but this may get rid of the additive separability property!! # In the delta_formula, replace x_i by its name for i, col in enumerate(tr_x.columns): for j in range(n_logits_dataset): delta_formula[j] = delta_formula[j].subs(sympy.symbols(f'x_{i + 1}'), sympy.symbols(col)) # Save formula as a text file with open(os.path.join(results_folder, dataset_name, 'formula.txt'), 'w') as f: for i in range(n_logits_dataset): f.write(f"Logit {i}: {delta_formula[i]}\n") f.write(f"Logit {i} Latex: {sympy.latex(delta_formula[i])}\n") # Since the formula may have pruned variables, we keep only the variables that are present in the formula actual_vars = [] for f in delta_formula: actual_vars += [str(s) for s in f.free_symbols] actual_vars = list(set(actual_vars)) # Remove duplicates tr_x = tr_x[actual_vars] test_x = test_x[actual_vars] delta_train, delta_test = [], [] for i in range(n_logits_dataset): d = get_patient_values(delta_formula[i], tr_x) delta_train.append(d) d.to_csv(os.path.join(results_folder, dataset_name, f'delta_train_{i}.csv'), index=False) d = get_patient_values(delta_formula[i], test_x) delta_test.append(d) d.to_csv(os.path.join(results_folder, dataset_name, f'delta_test_{i}.csv'), index=False) if binary_dataset: proba_train_numpy = 1 / (1 + np.exp(-delta_train[0].sum(axis=1).values)) proba_test_numpy = 1 / (1 + np.exp(-delta_test[0].sum(axis=1).values)) values_train_numpy = (proba_train_numpy > 0.5).astype(int) values_test_numpy = (proba_test_numpy > 0.5).astype(int) else: proba_train_numpy = np.array( [np.exp(np.array(d).sum(axis=1)) / np.sum(np.exp(np.array(delta_train).sum(axis=2)), axis=0) for d in delta_train]).T proba_test_numpy = np.array( [np.exp(np.array(d).sum(axis=1)) / np.sum(np.exp(np.array(delta_test).sum(axis=2)), axis=0) for d in delta_test]).T values_train_numpy = np.argmax(proba_train_numpy, axis=1) values_test_numpy = np.argmax(proba_test_numpy, axis=1) metrics_train_numpy = get_metrics(tr_y, values_train_numpy, proba_train_numpy) metrics_test_numpy = get_metrics(test_y, values_test_numpy, proba_test_numpy) print("\n--------Metrics comparison--------") print(f"Train metrics without formula: {metrics_train}") print(f"Test metrics without formula: {metrics_test}") print(f"Train metrics with formula: {metrics_train_numpy}") print(f"Test metrics with formula: {metrics_test_numpy}") # Save metrics results in df metrics_df = pd.DataFrame([metrics_train, metrics_test, metrics_train_numpy, metrics_test_numpy], index=['train', 'test', 'train_formula', 'test_formula']) metrics_df.to_csv(os.path.join(args['results_folder'], dataset_name, 'metrics.csv')) if binary_dataset: with plt.rc_context({**bundles.icml2024(column='half', nrows=1, ncols=1)}): plot_binary_explanation_plot(test_y, proba_test_numpy, ['0', '1'], 0.5, os.path.join(args['results_folder'], dataset_name, 'prob_plot_formula.pdf'), title='Probability of positive class') plt.close() print(f"\nNumber of variables in formula: {len(actual_vars)}\nVariables: {actual_vars}") for i in range(n_logits_dataset): print(f"\nFormula for logit {i}: {delta_formula[i]}") return delta_formula, delta_train, delta_test, proba_train_numpy, proba_test_numpy, tr_x, test_x if __name__ == '__main__': # Get the configuration args = get_config('interpretability') create_results_folder(args['results_folder'], args) for dataset in args['datasets']: # Check if the best model is saved if os.path.exists(os.path.join(args['base_folder'], 'results_performance', dataset, 'kan_gam.pkl')): with open(os.path.join(args['base_folder'], 'results_performance', dataset, 'kan_gam.pkl'), 'rb') as f: metrics = pickle.load(f) print(f"\n\n------------Model for dataset {dataset} found------------\nUsing the following parameters:") for param in metrics: if param not in ['accuracy', 'precision', 'recall', 'f1', 'roc_auc', 'time', 'dataset']: print(f"{param}: {metrics[param]}") print('\n') model = Kan_model(hidden_dim=metrics['hidden_dim'], batch_size=metrics['batch_size'], grid=metrics['grid'], k=metrics['k'], seed=metrics['seed'], lr=metrics['lr'], early_stop=metrics['early_stop'], steps=metrics['steps'], lamb=metrics['lamb'], lamb_entropy=metrics['lamb_entropy'], weight=metrics['weight'], sparse_init=metrics['sparse_init'], mult_kan=metrics['mult_kan']) else: model = Kan_model() print(f"Model for dataset {dataset} not found. Using default parameters.") # Load the data and run model x_train, x_test, y_train, y_test = load_data(dataset, args) pred_prob_test, binary, n_classes, n_logits, test_m, train_m = run_kan_gam() # Rename y_train to have a "class_target" column y_train = pd.DataFrame(y_train.values, columns=['class_target']) y_test = pd.DataFrame(y_test.values, columns=['class_target']) # Represent predictions in histograms if binary: plot_binary_explanation_plot(y_test, pred_prob_test[:, 1], ['0', '1'], 0.5, os.path.join(args['results_folder'], dataset, 'prob_plot.pdf'), title='Probability of positive class') with plt.rc_context({**bundles.icml2024(column='half', nrows=1, ncols=1, usetex=True)}): model.model.plot(folder=os.path.join(args['results_folder'], dataset, 'kan'), in_vars=x_train.columns.tolist(), out_vars=[f'logit_{i}' for i in range(n_classes)], varscale=0.2, scale=1) plt.savefig(os.path.join(args['results_folder'], dataset, 'kan.pdf'), bbox_inches='tight', dpi=300) plt.close() # Adjust a polynomial model (delta_formula, delta_train, delta_test, proba_train_numpy, proba_test_numpy, x_train,x_test) = adjust_polynomial(x_train, y_train, x_test, y_test, train_m, test_m, binary, n_logits, args['results_folder'], dataset) # Plot of the sorted variances in training and testing for each logit (feat imp is the variance of the delta vals) plot_sorted_variances(x_train, x_test, binary, delta_train, delta_test, n_logits, args, dataset) ##### PATIENTS ##### # Create a folder for the patients patients_results_folder = os.path.join(args['results_folder'], dataset, 'patients') os.makedirs(patients_results_folder, exist_ok=True) if binary: # Add a dimension to the proba arrays proba_train_numpy = proba_train_numpy[:, None] proba_test_numpy = proba_test_numpy[:, None] n_dists = args['n_dists'] max_atribs_radar = args['max_atribs_radar'] max_pats_to_save = args['max_pats_to_save'] max_plot_curves = args['max_plot_curves'] for l in range(n_logits): logit = 1 if binary else l print(f"Processing logit {logit}") # Get the most important features for the radar plot, be careful to use only training data! variances = delta_train[l].var(axis=0) all_cols = delta_train[l].columns.tolist() idx_vars = np.argsort(variances.values)[::-1] num_of_zero_var = (variances < 1e-6).sum() idx_vars = idx_vars[:-num_of_zero_var] if delta_train[l].shape[1] > max_atribs_radar: idx_vars = idx_vars[:max_atribs_radar] for i in tqdm(range(min(x_test.shape[0], max_pats_to_save))): actual_label = y_test.iloc[i]['class_target'] current_patient_info = np.concatenate((x_test.iloc[i].values, [proba_test_numpy[i][l], actual_label])) # Find the n_dists closest patients in the training set dists = np.linalg.norm(delta_train[l] - delta_test[l].iloc[i], axis=1) idx_closest = np.argsort(dists)[:n_dists].tolist() pred_prob = (1 / (1 + np.exp(-delta_train[l].iloc[idx_closest].sum(axis=1)))).values real_label = y_train.iloc[idx_closest].values closest_data = x_train.iloc[idx_closest].values closest_data = np.concatenate((closest_data, pred_prob[:, None], real_label), axis=1) closest_data = np.vstack((current_patient_info[None, :], closest_data)) # Add the current patient as the first row new_df = pd.DataFrame(closest_data, columns=x_train.columns.tolist() + ['pred_prob', 'real_label']) # Limit all new_df values to having 3 decimal numbers at most new_df = new_df.map(lambda x: round(x, 3) if isinstance(x, float) else x) new_df.to_csv(os.path.join(patients_results_folder, f'patient_{i}_closest_{n_dists}_logit_{logit}.csv'), index=False) # Prepare the radar plot, show only the attributes with highest variance # Change the order of idx_vars and cols_vars to have the importance in clockwise order in the plot idx_vars = idx_vars[::-1] cols_vars = [all_cols[i] for i in idx_vars.tolist()] n_feats = min(max_atribs_radar, len(cols_vars)) # Number of features to show in the radar plot if n_feats >= 3: # We need at least 3 features to plot a proper radar plot theta = radar_factory(n_feats, frame='polygon') if binary: avg_proba = 1 / (1 + np.exp(-delta_train[l].mean(axis=0).sum())) * np.ones(len(cols_vars)) # Average probability (i.e., "average patient") title = f"Test Patient \n Predicted: {proba_test_numpy[i][l]:.3f} | True: {actual_label} | Average: {avg_proba[0]:.3f}" else: avg_proba = np.exp(delta_train[l].mean(axis=0).sum()) / sum( [np.exp(d.mean(axis=0).sum()) for d in delta_train]) * np.ones(len(cols_vars)) title = f"Test Patient \n Predicted: {proba_test_numpy[i][l]:.3f} | True: {actual_label} | Average: {avg_proba[0]:.3f}" with plt.rc_context({**bundles.icml2024(column='half', ncols=1, nrows=1, usetex=True)}): fig, ax = plt.subplots(subplot_kw=dict(projection='radar')) ax.set_rgrids([0.2, 0.4, 0.6, 0.8]) ax.set_title(title) # Plot the average of all train patients _ = ax.plot(theta, avg_proba, label='Average', color='tab:blue', linewidth=0.5) ax.fill(theta, avg_proba, alpha=0.1, color='tab:blue') # Prepare for individual patient plotting avg_delta = delta_train[l].mean(axis=0).values[None, :] avg_matrix = np.repeat(avg_delta, delta_train[l].shape[1], axis=0) # Plot the closest patients for j in range(n_dists): label = 'Closest patients' if j == 0 else None np.fill_diagonal(avg_matrix, delta_train[l].iloc[idx_closest[j]].values) if binary: pat_proba = 1 / (1 + np.exp(-avg_matrix.sum(axis=1))) else: den_term = np.zeros(delta_train[0].shape[1]) for ll in range(n_logits): den_matrix = np.repeat(delta_train[ll].mean(axis=0).values[None, :], delta_train[ll].shape[1], axis=0) np.fill_diagonal(den_matrix, delta_train[ll].iloc[idx_closest[j]].values) den_term += np.exp(den_matrix.sum(axis=1)) pat_proba = np.exp(avg_matrix.sum(axis=1)) / den_term _ = ax.plot(theta, pat_proba[idx_vars], label=label, color='tab:green', alpha=0.5) ax.fill(theta, pat_proba[idx_vars], alpha=0.1, color='tab:green') # Plot the current patient np.fill_diagonal(avg_matrix, delta_test[l].iloc[i].values) if binary: pat_proba = 1 / (1 + np.exp(-avg_matrix.sum(axis=1))) else: den_term = np.zeros(delta_train[0].shape[1]) for ll in range(n_logits): den_matrix = np.repeat(delta_train[ll].mean(axis=0).values[None, :], delta_train[ll].shape[1], axis=0) np.fill_diagonal(den_matrix, delta_test[ll].iloc[i].values) den_term += np.exp(den_matrix.sum(axis=1)) pat_proba = np.exp(avg_matrix.sum(axis=1)) / den_term _ = ax.plot(theta, pat_proba[idx_vars], label='Test patient', color='tab:red') ax.fill(theta, pat_proba[idx_vars], alpha=0.1, color='tab:red') ax.set_varlabels(cols_vars, fontsize=6) plt.legend(loc='lower center', bbox_to_anchor=(0.5, -0.5), ncol=3) # Note: this can be uncommented, but may clutter the plot plt.savefig(os.path.join(patients_results_folder, f'patient_{i}_radar_logit_{logit}.pdf'), bbox_inches='tight', dpi=600) plt.close() # Curves plot: show only the ones that do matter!! # Revert again the order to have the right plot order idx_vars = idx_vars[::-1] cols_vars = [all_cols[i] for i in idx_vars.tolist()] n_feats = min(len(cols_vars), max_plot_curves) # Number of features to show in the curves plot with plt.rc_context({**bundles.icml2024(column='full', ncols=1, nrows=1, usetex=True)}): if n_feats > 0: # There is something to show if binary: fig, axs = plt.subplots(n_feats + 1, 1, figsize=(3, 5.5)) x_vals = np.arange(delta_train[l].sum(axis=1).min(), delta_train[l].sum(axis=1).max(), 0.01) theor_proba = 1 / (1 + np.exp(-x_vals)) axs[0].plot(x_vals, theor_proba, 'b', alpha=0.2) axs[0].scatter(delta_test[l].sum(axis=1)[i], proba_test_numpy[i][l], color='r') axs[0].set_xlabel('Logit') axs[0].set_ylabel('Probability') axs[0].set_title(f'Patient {i}') else: fig, axs = plt.subplots(n_feats, 1, figsize=(3, 5.5)) for idj, feat_name in enumerate(cols_vars): if idj < n_feats: # Only plot the first n_feats features if binary: j = idj + 1 # The first plot is already used for the theoretical curve else: j = idj # Keep only unique values of x_test[feat_name] idxs = np.unique(x_train[feat_name].values, return_index=True)[1] if n_feats > 1: # Multiple plots axs[j].plot(x_train[feat_name].values[idxs], delta_train[l][feat_name].values[idxs], color='tab:blue') axs[j].scatter(x_train[feat_name].values[idxs], delta_train[l][feat_name].values[idxs], color='tab:blue', s=6, alpha=0.4) axs[j].scatter(x_test[feat_name].values[i], delta_test[l][feat_name].values[i], color='tab:red', s=40, alpha=1) for jj in range(n_dists): axs[j].scatter(x_train.iloc[idx_closest[jj]][feat_name], delta_train[l].iloc[idx_closest[jj]][feat_name], color='tab:green', s=8, alpha=1) axs[j].set_ylabel('Contribution') axs[j].set_xlabel(feat_name) else: # Single plot axs.plot(x_train[feat_name].values[idxs], delta_train[l][feat_name].values[idxs], color='tab:blue') axs.scatter(x_train[feat_name].values[idxs], delta_train[l][feat_name].values[idxs], color='tab:blue', s=6, alpha=0.4) axs.scatter(x_test[feat_name].values[i], delta_test[l][feat_name].values[i], color='tab:red', s=40, alpha=1) for jj in range(n_dists): axs.scatter(x_train.iloc[idx_closest[jj]][feat_name], delta_train[l].iloc[idx_closest[jj]][feat_name], color='tab:green', s=8, alpha=1) axs.set_ylabel('Contribution') axs.set_xlabel(feat_name) red_patch = mpatches.Patch(color='tab:red', label='Test patient') green_patch = mpatches.Patch(color='tab:green', label='Closest patients') blue_patch = mpatches.Patch(color='tab:blue', label='Train patients') plt.tight_layout(rect=[0, 0.05, 1, 1]) # TODO: Adjust legend box location based on dataset!!! if n_feats > 1: axs[0].set_title(f'Patient PDPs') axs[-1].legend(handles=[red_patch, green_patch, blue_patch], loc='lower center', bbox_to_anchor=(0.5, -1.0), ncol=3) else: axs.set_title(f'Patient PDPs') axs.legend(handles=[red_patch, green_patch, blue_patch], loc='lower center', bbox_to_anchor=(0.5, -0.7), ncol=3) plt.rcParams.update(bundles.icml2024(usetex=False)) plt.savefig(os.path.join(patients_results_folder, f'patient_{i}_curves_logit_{logit}.pdf'), bbox_inches='tight', dpi=600) plt.close()