From a8c2362529049ef20fff74e3867a8550ab98ce45 Mon Sep 17 00:00:00 2001 From: "Anne Sekscinska (they/them)" Date: Mon, 7 Sep 2026 10:42:08 +0100 Subject: [PATCH 01/12] Fix figure issues --- avae/vis.py | 62 +++++++++++++++++++++++++++++------------------------ 1 file changed, 34 insertions(+), 28 deletions(-) diff --git a/avae/vis.py b/avae/vis.py index cd7ed54c..91d2d8ab 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -638,6 +638,7 @@ def confidence_plot(x, y, s, suffix=None): if len(np.unique(y)) % 2 != 0: rows += 1 fig, ax = plt.subplots(len(np.unique(y)), sharex=True, sharey=True) + ax = np.atleast_1d(ax) for c, cl in enumerate(np.unique(y)): mu_cl = np.take(x, np.where(np.array(y) == cl)[0], axis=0) var_cl = np.take(s, np.where(np.array(y) == cl)[0], axis=0) @@ -667,6 +668,7 @@ def confidence_plot(x, y, s, suffix=None): leg = fig.legend( handles, labels, bbox_to_anchor=(1.06, 0.9) ) # , loc="upper left") + plt.tight_layout() fig.savefig(name, bbox_extra_artists=(leg,), bbox_inches="tight") plt.close() @@ -713,6 +715,10 @@ def accuracy_plot( else: classes_list = np.unique(np.concatenate((y_train, ypred_train))) + #to avoid figure being too small with too little classses + fig_size = max(6, int(len(classes_list)) / 2) + font_size = max(8, int(len(classes_list) / 3) + 3) + # Compute confusion matrix cm = confusion_matrix(y_train, ypred_train, labels=classes_list) @@ -736,10 +742,10 @@ def accuracy_plot( ) with plt.rc_context( - {"font.weight": "bold", "font.size": int(len(classes_list) / 3) + 3} + {"font.weight": "bold", "font.size": font_size} ): fig, ax = plt.subplots( - figsize=(int(len(classes_list)) / 2, int(len(classes_list)) / 2) + figsize=(fig_size, fig_size) ) disp.plot(cmap=plt.cm.Blues, ax=ax, xticks_rotation=90) @@ -747,7 +753,7 @@ def accuracy_plot( plt.tight_layout() plt.title( "Average accuracy at epoch {}: {:.3f}%".format( - epoch, np.mean(avg_accuracy) * 100 + epoch, np.mean(avg_accuracy) ), fontsize=10, ) @@ -755,6 +761,7 @@ def accuracy_plot( if not os.path.exists("plots"): os.mkdir("plots") + plt.tight_layout() plt.savefig( f"plots/confusion_train{mode}.{settings.VIS_FORMAT}", dpi=300 ) @@ -771,7 +778,7 @@ def accuracy_plot( plt.close() fig, ax = plt.subplots( - figsize=(int(len(classes_list)) / 2, int(len(classes_list)) / 2) + figsize=(fig_size, fig_size) ) dispn.plot( @@ -787,7 +794,7 @@ def accuracy_plot( "Average accuracy at epoch {}: {:.3f}%".format( epoch, np.mean(avg_accuracy) * 100 ), - fontsize=12, + fontsize=font_size+2, ) plt.xlabel("Predicted label (%)") @@ -848,13 +855,12 @@ def accuracy_plot( with plt.rc_context( { "font.weight": "bold", - "font.size": int(len(ordered_class_eval) / 3) + 3, + "font.size": font_size, } ): fig, ax = plt.subplots( figsize=( - int(len(ordered_class_eval)) / 2, - int(len(ordered_class_eval)) / 2, + fig_size,fig_size ) ) disp_eval.plot(cmap=plt.cm.Blues, ax=ax, xticks_rotation=90) @@ -863,13 +869,13 @@ def accuracy_plot( "Average accuracy at epoch {}: {:.1f}%".format( epoch, np.mean(avg_accuracy_eval) * 100 ), - fontsize=12, + fontsize=font_size+2, ) plt.savefig(figure_name + f".{settings.VIS_FORMAT}", dpi=300) plt.close() fig, ax = plt.subplots( - figsize=(int(len(classes_list)) / 2, int(len(classes_list)) / 2) + figsize=(fig_size,fig_size) ) dispn_eval.plot( @@ -935,6 +941,8 @@ def f1_plot( classes_list = np.unique(np.concatenate((y_train, ypred_train))) classes_list_eval = np.unique(np.concatenate((y_val, ypred_val))) + fig_size = max(6, int(len(classes_list)) / 2) + font_size = max(8, int(len(classes_list) / 3) + 3) if np.setdiff1d(classes_list_eval, classes_list).size > 0: logging.info( @@ -985,13 +993,12 @@ def f1_plot( with plt.rc_context( { "font.weight": "bold", - "font.size": int(len(classes_list) / 3) + 3, + "font.size": font_size, } ): fig, ax = plt.subplots( figsize=( - int(len(classes_list)) / 2, - int(len(classes_list)) / 2, + fig_size,fig_size ) ) plt.plot(classes_list, train_f1_score, label="train", marker="o") @@ -1000,6 +1007,7 @@ def f1_plot( plt.legend(loc="lower left") plt.title("F1 Score at epoch {}".format(epoch)) plt.ylabel("F1 Score") + plt.tight_layout() plt.savefig(f"plots/f1{mode}.{settings.VIS_FORMAT}", dpi=150) if writer: @@ -1852,21 +1860,17 @@ def plot_affinity_matrix( if len(vis_format) > 0: settings.VIS_FORMAT = vis_format - if fig_size is None: - with plt.rc_context( - {"font.weight": "bold", "font.size": int(len(all_classes) / 3) + 3} - ): - fig, ax = plt.subplots( - figsize=(int(len(all_classes)) / 2, int(len(all_classes)) / 2) - ) - else: - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) + fig_size = fig_size if fig_size is not None else max(6, int(len(all_classes)) / 2) + font_size = max(8, int(len(all_classes) / 3) + 3) + # Create the figure and gridspec + fig = plt.figure(figsize=(fig_size, fig_size)) gs = gridspec.GridSpec(1, 2, width_ratios=[9, 0.4]) # Plot the data on the left grid - ax = plt.subplot(gs[0]) - ax.set_title("Affinity Matrix", fontsize=16) + ax = fig.add_subplot(gs[0]) + ax2 = fig.add_subplot(gs[1]) + ax.set_title("Affinity Matrix", fontsize=font_size+2) im = ax.imshow(lookup, vmin=-1, vmax=1, cmap=plt.get_cmap("RdBu")) @@ -1996,6 +2000,7 @@ def latent_space_similarity_plot( if len(vis_format) > 0: settings.VIS_FORMAT = vis_format + if len(classes_order) == 0: unique_classes = np.unique(class_labels) else: @@ -2011,6 +2016,8 @@ def latent_space_similarity_plot( unique_classes = classes_order num_classes = len(unique_classes) + fig_size = fig_size if fig_size is not None else max(6,int(num_classes / 2)) + font_size = max(8, int(num_classes/ 3) + 3) cosine_sim = latent_space_similarity_mat( latent_space, class_labels, @@ -2024,18 +2031,17 @@ def latent_space_similarity_plot( with plt.rc_context( { "font.weight": "bold", - "font.size": int(len(unique_classes) / 3) + 3, + "font.size": font_size, } ): fig, ax = plt.subplots( figsize=( - int(len(unique_classes)) / 2, - int(len(unique_classes)) / 2, + fig_size, fig_size ) ) else: fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - fig.tight_layout(pad=3) + fig.tight_layout() plt.imshow(cosine_sim, cmap="RdBu", vmin=-1, vmax=1) plt.colorbar(label="Average Cosine Similarity") plt.xticks(ticks=np.arange(num_classes), labels=unique_classes) From 6df56498ba10ae144a58332b5e2548c01d249dc2 Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 16 Sep 2026 16:00:04 +0000 Subject: [PATCH 02/12] Improve visualization output consistency Refine plot sizing, typography, colorbars, image output, and pipeline test behavior. --- avae/config.py | 15 +- avae/data.py | 2 +- avae/utils.py | 14 +- avae/vis.py | 288 ++++++++++++++++++++---------- tests/test_config.py | 7 + tests/test_train_eval_pipeline.py | 1 - 6 files changed, 220 insertions(+), 107 deletions(-) diff --git a/avae/config.py b/avae/config.py index 4aad47dc..32cf7a55 100644 --- a/avae/config.py +++ b/avae/config.py @@ -328,6 +328,8 @@ def load_config_params( # if no config file is provided, start from default and update with command line arguments params = {} + command_line_params = set() + # check for command line input values and overwrite config file values # we're using sys args because click has populated defaults from confing if sys_args is not None: @@ -337,6 +339,7 @@ def load_config_params( continue name = arg[2:].split("=")[0] if name in local_args.keys(): + command_line_params.add(name) # overwrite config file value with command line argument value # but only if its on system args (click has defaults) if name in params.keys(): @@ -372,7 +375,11 @@ def load_config_params( if type(val) == pathlib.Path: # turn relative paths to absolute params[key] = str(val.absolute()) - if 'vis' in key and params['vis_all'] is not None: + if ( + 'vis' in key + and params['vis_all'] is not None + and key not in command_line_params + ): # set visualisation to vis_all if it is not set, except for vis_z_n_int and vis_pose_class if key in [ 'vis_all', @@ -385,7 +392,11 @@ def load_config_params( logging.warning( f"Visualisation parameter 'vis_all' is overriding {key} to {params['vis_all']}" ) - if 'freq' in key and params['freq_all'] is not None: + if ( + 'freq' in key + and params['freq_all'] is not None + and key not in command_line_params + ): # set frequency to freq_all if it is not set if key in ['freq_all']: continue diff --git a/avae/data.py b/avae/data.py index 701b84b5..63217036 100644 --- a/avae/data.py +++ b/avae/data.py @@ -140,7 +140,7 @@ def load_data( logging.info("\n") logging.info("############################################### DATA") - logging.info(f"Loading data...") + logging.info("Loading data...\n") # read the class list, if not provided all classes in the dataset will be used as default if classes is not None: diff --git a/avae/utils.py b/avae/utils.py index b8689726..184262c3 100644 --- a/avae/utils.py +++ b/avae/utils.py @@ -98,10 +98,16 @@ def save_imshow_png( if not os.path.exists("plots"): os.mkdir("plots") - fig, _ = plt.subplots(figsize=(10, 10)) - plt.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last - - plt.savefig("plots/" + fname) + height, width = array.shape[:2] + if width >= height: + image_figure_size = (10, 10 * height / width) + else: + image_figure_size = (10 * width / height, 10) + fig, ax = plt.subplots(figsize=image_figure_size) + ax.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last + ax.axis("off") + fig.subplots_adjust(left=0, right=1, bottom=0, top=1) + fig.savefig("plots/" + fname, bbox_inches="tight", pad_inches=0) if writer: writer.add_figure(figname, fig, epoch) diff --git a/avae/vis.py b/avae/vis.py index bcd029b7..d0e20c1e 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -7,7 +7,6 @@ import warnings import altair -import matplotlib.gridspec as gridspec import matplotlib.pyplot as plt import numpy as np import numpy.typing as npt @@ -18,6 +17,7 @@ import sklearn.metrics.pairwise import torch import torchvision +from mpl_toolkits.axes_grid1 import make_axes_locatable from . import utils_learning from .utils import ( @@ -31,6 +31,18 @@ ) +def _class_plot_scale( + num_classes: int, fig_size: int | None = None +) -> tuple[float, int, int]: + """Return consistent figure and font sizes for class-based plots.""" + font_size = max(8, int(num_classes / 3) + 3) + return ( + fig_size if fig_size is not None else max(6, int(num_classes) / 2), + font_size, + max(font_size + 4, int(num_classes / 2) + 8), + ) + + def _encoder(i: PIL.Image) -> str: """Encode PIL Image as base64 buffer. Parameters @@ -248,6 +260,7 @@ def latent_embed_plot_tsne( ) n_classes = len(classes) + fig_size, font_size, title_font_size = _class_plot_scale(n_classes) if n_classes < 3: # If the number of classes are not moe than 3 the size of the figure would be too @@ -259,6 +272,7 @@ def latent_embed_plot_tsne( fig, ax = plt.subplots( figsize=(int(n_classes / 2) + 4, int(n_classes / 2) + 2) ) + marker_size = max(marker_size, int(marker_size * max(1, n_classes / 5))) # When the number of classes is less than 3 the image becomes two small colours = colour_per_class(classes) @@ -276,12 +290,16 @@ def latent_embed_plot_tsne( label=mol[:4], facecolor=color, edgecolor=color, - alpha=0.2, + alpha=0.5, ) - ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left", fontsize=16) - plt.xlabel("TSNE-1") - plt.ylabel("TSNE-2") + ax.legend( + bbox_to_anchor=(1.05, 1), + loc="upper left", + fontsize=font_size, + ) + ax.set_xlabel("TSNE-1", fontsize=font_size) + ax.set_ylabel("TSNE-2", fontsize=font_size) if xs.shape[-1] == 1: @@ -298,16 +316,21 @@ def latent_embed_plot_tsne( label=mol[:4], linewidth=l_w, ) - plt.legend( - prop={"size": 10}, + ax.legend( bbox_to_anchor=(1.05, 1), loc="upper left", - fontsize=16, + fontsize=font_size, ) - plt.xlabel("dim 1") - plt.ylabel("freq") + ax.set_xlabel("dim 1", fontsize=font_size) + ax.set_ylabel("freq", fontsize=font_size) - plt.tight_layout() + ax.set_title( + f"t-SNE Embedding at epoch: {epoch + 1}", + fontsize=title_font_size, + fontweight="normal", + ) + ax.tick_params(axis="both", labelsize=font_size) + fig.tight_layout() if not display: if not os.path.exists("plots"): @@ -507,12 +530,14 @@ def confidence_plot(x, y, s, suffix=None, vis_format="png"): ) cmap = plt.get_cmap("jet") cols = [cmap(i) for i in np.linspace(0, 1, len(x[0]))] - rows = len(np.unique(y)) // 2 - if len(np.unique(y)) % 2 != 0: + classes = np.unique(y) + _, font_size, title_font_size = _class_plot_scale(len(classes)) + rows = len(classes) // 2 + if len(classes) % 2 != 0: rows += 1 - fig, ax = plt.subplots(len(np.unique(y)), sharex=True, sharey=True) + fig, ax = plt.subplots(len(classes), sharex=True, sharey=True) ax = np.atleast_1d(ax) - for c, cl in enumerate(np.unique(y)): + for c, cl in enumerate(classes): mu_cl = np.take(x, np.where(np.array(y) == cl)[0], axis=0) var_cl = np.take(s, np.where(np.array(y) == cl)[0], axis=0) std_cl = np.exp(0.5 * var_cl) @@ -533,7 +558,7 @@ def confidence_plot(x, y, s, suffix=None, vis_format="png"): color=cols[i], label="lat" + str(i + 1), ) - ax[c].set_title(cl) + ax[c].set_title(cl, fontsize=title_font_size, fontweight="normal") name = f"plots/confidence.{vis_format}" if suffix is not None: name = name[:-4] + "_" + suffix + name[-4:] @@ -584,9 +609,8 @@ def accuracy_plot( classes_list = np.unique(np.concatenate((y_train, ypred_train))) - # to avoid figure being too small with too little classses - fig_size = max(6, int(len(classes_list)) / 2) - font_size = max(8, int(len(classes_list) / 3) + 3) + fig_size, font_size, title_font_size = _class_plot_scale(len(classes_list)) + value_font_size = min(font_size, 10) # Compute confusion matrix cm = sklearn.metrics.confusion_matrix( @@ -635,21 +659,27 @@ def accuracy_plot( ax=ax, xticks_rotation=90, values_format="d", + colorbar=False, ) + divider = make_axes_locatable(ax) + colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) + fig.colorbar(disp.im_, cax=colorbar_ax) + colorbar_ax.tick_params(labelsize=font_size) + plt.setp(disp.text_.ravel(), fontsize=value_font_size) - plt.tight_layout() - plt.title( + ax.set_title( "Average accuracy at epoch {}: {:.3f}%".format( epoch + 1, - np.mean(avg_accuracy[supported_train_classes]), + np.mean(avg_accuracy[supported_train_classes]) * 100, ), - fontsize=10, + fontsize=title_font_size, + fontweight="normal", ) if not os.path.exists("plots"): os.mkdir("plots") - plt.tight_layout() + fig.tight_layout() plt.savefig(f"plots/confusion_train{mode}.{vis_format}", dpi=300) # Save confusion matrix DataFrame to CSV @@ -671,19 +701,26 @@ def accuracy_plot( xticks_rotation=90, values_format=".0f", im_kw={"vmin": 0, "vmax": 100}, + colorbar=False, ) + divider = make_axes_locatable(ax) + colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) + fig.colorbar(dispn.im_, cax=colorbar_ax) + colorbar_ax.tick_params(labelsize=font_size) + plt.setp(dispn.text_.ravel(), fontsize=value_font_size) - plt.tight_layout() - plt.title( + ax.set_title( "Average accuracy at epoch {}: {:.3f}%".format( epoch + 1, np.mean(avg_accuracy[supported_train_classes]) * 100, ), - fontsize=font_size + 2, + fontsize=title_font_size, + fontweight="normal", ) - plt.xlabel("Predicted label (%)") - plt.ylabel("True label (%)") + ax.set_xlabel("Predicted label (%)") + ax.set_ylabel("True label (%)") + fig.tight_layout() plt.savefig(f"plots/confusion_train{mode}_norm.{vis_format}", dpi=300) if writer: @@ -769,15 +806,22 @@ def accuracy_plot( ax=ax, xticks_rotation=90, values_format="d", + colorbar=False, ) - plt.tight_layout() - plt.title( + divider = make_axes_locatable(ax) + colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) + fig.colorbar(disp_eval.im_, cax=colorbar_ax) + colorbar_ax.tick_params(labelsize=font_size) + plt.setp(disp_eval.text_.ravel(), fontsize=value_font_size) + ax.set_title( "Average accuracy at epoch {}: {:.1f}%".format( epoch + 1, np.mean(avg_accuracy_eval[supported_eval_classes]) * 100, ), - fontsize=font_size + 2, + fontsize=title_font_size, + fontweight="normal", ) + fig.tight_layout() plt.savefig(figure_name + f".{vis_format}", dpi=300) plt.close() @@ -789,18 +833,25 @@ def accuracy_plot( xticks_rotation=90, values_format=".0f", im_kw={"vmin": 0, "vmax": 100}, + colorbar=False, ) + divider = make_axes_locatable(ax) + colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) + fig.colorbar(dispn_eval.im_, cax=colorbar_ax) + colorbar_ax.tick_params(labelsize=font_size) + plt.setp(dispn_eval.text_.ravel(), fontsize=value_font_size) - plt.tight_layout() - plt.title( + ax.set_title( "Average accuracy at epoch {}: {:.1f}% ".format( epoch + 1, np.mean(avg_accuracy_eval[supported_eval_classes]) * 100, ), - fontsize=10, + fontsize=title_font_size, + fontweight="normal", ) - plt.xlabel("Predicted label (%)") - plt.ylabel("True label (%)") + ax.set_xlabel("Predicted label (%)") + ax.set_ylabel("True label (%)") + fig.tight_layout() plt.savefig(f"{figure_name}_norm.{vis_format}", dpi=300) plt.close() @@ -842,8 +893,7 @@ def f1_plot( logging.info("Visualising F1 scores ...\n") classes_list = np.unique(np.concatenate((y_train, ypred_train))) classes_list_eval = np.unique(np.concatenate((y_val, ypred_val))) - fig_size = max(6, int(len(classes_list)) / 2) - font_size = max(8, int(len(classes_list) / 3) + 3) + fig_size, font_size, title_font_size = _class_plot_scale(len(classes_list)) if np.setdiff1d(classes_list_eval, classes_list).size > 0: logging.info( @@ -906,11 +956,28 @@ def f1_plot( } ): fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - plt.plot(classes_list, train_f1_score, label="train", marker="o") - plt.plot(classes_list, valid_f1_score, label=label, marker="o") + line_width = max(1.5, font_size / 6) + plt.plot( + classes_list, + train_f1_score, + label="train", + marker="o", + linewidth=line_width, + ) + plt.plot( + classes_list, + valid_f1_score, + label=label, + marker="o", + linewidth=line_width, + ) plt.xticks(rotation=45) plt.legend(loc="lower left") - plt.title("F1 Score at epoch {}".format(epoch + 1)) + plt.title( + "F1 Score at epoch {}".format(epoch + 1), + fontsize=title_font_size, + fontweight="normal", + ) plt.ylabel("F1 Score") plt.tight_layout() plt.savefig(f"plots/f1{mode}.{vis_format}", dpi=150) @@ -970,6 +1037,9 @@ def loss_plot( "VAL KL divergence loss x BETA", "VAL Affinity loss x GAMMA", ] + title_font_size = 12 + label_font_size = 10 + tick_font_size = 9 plt.clf() plt.ticklabel_format(useOffset=False) @@ -1006,17 +1076,19 @@ def loss_plot( return plt.title( "bs: %d, d: %d, ch: %d, lat: %d, lr: %.3f, beta: %.1f, " - "gamma: %.1f" % (p[0], p[1], p[2], p[3], p[4], p[5], p[6]) + "gamma: %.1f" % (p[0], p[1], p[2], p[3], p[4], p[5], p[6]), + fontsize=title_font_size, + fontweight="normal", ) else: - plt.title("Loss") + plt.title("Loss", fontsize=title_font_size, fontweight="normal") plt.yscale("log") - plt.ylabel("Loss", fontsize=16) - plt.xlabel("Epochs", fontsize=16) - plt.xticks(fontsize=16) - plt.yticks(fontsize=16) - plt.legend() + plt.ylabel("Loss", fontsize=label_font_size) + plt.xlabel("Epochs", fontsize=label_font_size) + plt.xticks(fontsize=tick_font_size) + plt.yticks(fontsize=tick_font_size) + plt.legend(fontsize=tick_font_size) plt.tight_layout() if not os.path.exists("plots"): @@ -1041,11 +1113,11 @@ def loss_plot( label=vlabs[0], ) plt.yscale("log") - plt.ylabel("Loss", fontsize=16) - plt.xlabel("Epochs", fontsize=16) - plt.xticks(fontsize=16) - plt.yticks(fontsize=16) - plt.legend() + plt.ylabel("Loss", fontsize=label_font_size) + plt.xlabel("Epochs", fontsize=label_font_size) + plt.xticks(fontsize=tick_font_size) + plt.yticks(fontsize=tick_font_size) + plt.legend(fontsize=tick_font_size) plt.tight_layout() plt.savefig(f"plots/loss_total.{vis_format}", dpi=300) plt.close() @@ -1775,19 +1847,16 @@ def plot_affinity_matrix( ) logging.info("Visualising affinity matrix ...\n") - fig_size = ( - fig_size if fig_size is not None else max(6, int(len(all_classes) / 2)) + resolved_fig_size, font_size, title_font_size = _class_plot_scale( + len(all_classes), fig_size ) - font_size = max(8, int(len(all_classes) / 3) + 3) - # Create the figure and gridspec - fig = plt.figure(figsize=(fig_size, fig_size)) - gs = gridspec.GridSpec(1, 2, width_ratios=[9, 0.4]) - - # Plot the data on the left grid - ax = fig.add_subplot(gs[0]) - ax2 = fig.add_subplot(gs[1]) - ax.set_title("Affinity Matrix", fontsize=font_size + 2) + fig, ax = plt.subplots(figsize=(resolved_fig_size, resolved_fig_size)) + ax.set_title( + "Affinity Matrix", + fontsize=title_font_size, + fontweight="normal", + ) im = ax.imshow(lookup, vmin=-1, vmax=1, cmap=plt.get_cmap("RdBu")) @@ -1796,21 +1865,19 @@ def plot_affinity_matrix( ax.set_yticks(np.arange(0, len(all_classes))) ax.set_yticklabels(all_classes) - # Colour the label for selected classes + # Colour classes absent from the selected dataset. for i, c in enumerate(all_classes): - if c in selected_classes: + if c not in selected_classes: ax.get_xticklabels()[i].set_color("red") ax.get_yticklabels()[i].set_color("red") - ax.tick_params(axis="x", rotation=90, labelsize=16) - ax.tick_params(axis="y", labelsize=16) - - # Create an empty plot on the right grid - ax2 = plt.subplot(gs[1]) + ax.tick_params(axis="x", rotation=90, labelsize=font_size) + ax.tick_params(axis="y", labelsize=font_size) - plt.colorbar(im, cax=ax2) - # Set the height of the color bar to match the height of the plot - # cb.ax.set_position([0.96, 0.01, 0.01, 0.01]) + divider = make_axes_locatable(ax) + colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) + fig.colorbar(im, cax=colorbar_ax) + colorbar_ax.tick_params(labelsize=font_size) fig.tight_layout() if not os.path.exists("plots"): os.mkdir("plots") @@ -1836,17 +1903,22 @@ def plot_classes_distribution( ) logging.info("Visualising classes distribution " + category + "...\n") - fig, ax = plt.subplots(figsize=(9, 9)) labels, counts = np.unique(data, return_counts=True) + fig_size, font_size, title_font_size = _class_plot_scale(len(labels)) + fig, ax = plt.subplots(figsize=(fig_size, fig_size)) ticks = range(len(counts)) - plt.bar(ticks, counts, align="center", color="blue", alpha=0.5) - plt.xticks(ticks, labels) - plt.title("Classes Distribution", fontsize=16) - plt.xlabel("Class", fontsize=16) - plt.ylabel("Number of Entries", fontsize=16) - plt.xticks(fontsize=16, rotation=90) - plt.yticks(fontsize=16) - plt.tight_layout() + ax.bar(ticks, counts, align="center", color="blue", alpha=0.5) + ax.set_xticks(ticks, labels) + ax.set_title( + "Classes Distribution", + fontsize=title_font_size, + fontweight="normal", + ) + ax.set_xlabel("Class", fontsize=font_size) + ax.set_ylabel("Number of Entries", fontsize=font_size) + ax.tick_params(axis="x", labelsize=font_size, rotation=90) + ax.tick_params(axis="y", labelsize=font_size) + fig.tight_layout() if not os.path.exists("plots"): os.mkdir("plots") plt.savefig(f"plots/classes_distribution_{category}.{vis_format}", dpi=300) @@ -1877,7 +1949,7 @@ def plot_cyc_variable( plt.tight_layout() if not os.path.exists("plots"): os.mkdir("plots") - plt.savefig(f"plots/{variable_name}_array.{vis_format}", dpi=300) + plt.savefig(f"plots/hyperaparam_{variable_name}.{vis_format}", dpi=300) plt.close() @@ -1934,10 +2006,9 @@ def latent_space_similarity_plot( unique_classes = classes_order num_classes = len(unique_classes) - fig_size = ( - fig_size if fig_size is not None else max(6, int(num_classes / 2)) + resolved_fig_size, font_size, title_font_size = _class_plot_scale( + num_classes, fig_size ) - font_size = max(8, int(num_classes / 3) + 3) cosine_sim = latent_space_similarity_mat( latent_space, class_labels, @@ -1954,24 +2025,43 @@ def latent_space_similarity_plot( "font.size": font_size, } ): - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) + fig, ax = plt.subplots( + figsize=(resolved_fig_size, resolved_fig_size) + ) else: - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - fig.tight_layout() - plt.imshow(cosine_sim, cmap="RdBu", vmin=-1, vmax=1) - plt.colorbar(label="Average Cosine Similarity") - plt.xticks(ticks=np.arange(num_classes), labels=unique_classes) - plt.yticks(ticks=np.arange(num_classes), labels=unique_classes) - plt.title(f"Average Cosine Similarity Matrix at epoch :{epoch}") - plt.xlabel("Class Labels") - plt.ylabel("Class Labels") + fig, ax = plt.subplots(figsize=(resolved_fig_size, resolved_fig_size)) + im = ax.imshow(cosine_sim, cmap="RdBu", vmin=-1, vmax=1) + divider = make_axes_locatable(ax) + colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) + fig.colorbar(im, cax=colorbar_ax) + colorbar_ax.tick_params(labelsize=font_size) + ax.set_xticks(np.arange(num_classes), labels=unique_classes) + ax.set_yticks(np.arange(num_classes), labels=unique_classes) + ax.set_title( + f"Average Cosine Similarity Matrix at epoch: {epoch + 1}", + fontsize=title_font_size, + fontweight="normal", + ) + ax.set_xlabel("Class Labels", fontsize=font_size) + ax.set_ylabel("Class Labels", fontsize=font_size) ax.tick_params(axis="x", rotation=90, labelsize=font_size) ax.tick_params(axis="y", labelsize=font_size) + classes_in_data = set(np.asarray(class_labels).astype(str)) + for class_name, x_label, y_label in zip( + unique_classes, ax.get_xticklabels(), ax.get_yticklabels() + ): + if str(class_name) not in classes_in_data: + x_label.set_color("red") + y_label.set_color("red") + fig.tight_layout() if not display: if not os.path.exists("plots"): os.mkdir("plots") - plt.savefig(f"plots/similarity_mean{mode}.{vis_format}", dpi=dpi) + plt.savefig( + f"plots/similarity_mean{mode}.{vis_format}", + dpi=dpi, + ) plt.close() else: plt.show() diff --git a/tests/test_config.py b/tests/test_config.py index 44d6af95..de228d03 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -121,6 +121,13 @@ def test_validate_config(self): data_config_only = load_config_params(config_file=self.config) self.assertEqual(data_config_only.epochs, 1000) + data_cli_frequency = load_config_params( + config_file=self.config, + local_args={"freq_all": 5, "freq_acc": 3}, + ) + self.assertEqual(data_cli_frequency.freq_acc, 3) + self.assertEqual(data_cli_frequency.freq_rec, 5) + data_local_data_only = load_config_params(local_args=self.data_local) self.assertEqual(data_local_data_only.epochs, 150) self.assertEqual( diff --git a/tests/test_train_eval_pipeline.py b/tests/test_train_eval_pipeline.py index 95642edc..bf5b2646 100644 --- a/tests/test_train_eval_pipeline.py +++ b/tests/test_train_eval_pipeline.py @@ -23,7 +23,6 @@ def pipeline_params(datapath): "no_val_drop": True, "affinity": os.path.join(datapath, "affinity_fsc_10.csv"), "classes": os.path.join(datapath, "classes.csv"), - "vis_dynamic": True, "epochs": 1, "channels": 3, "depth": 4, From ebe058ccc8ab13f87e1e5ad3909401b460807fbc Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 16 Sep 2026 16:08:17 +0000 Subject: [PATCH 03/12] Fix loss and hyperparameter figure sizing --- avae/vis.py | 73 ++++++++++++++++++++++++++--------------------------- 1 file changed, 36 insertions(+), 37 deletions(-) diff --git a/avae/vis.py b/avae/vis.py index d0e20c1e..b4083d89 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -1037,19 +1037,19 @@ def loss_plot( "VAL KL divergence loss x BETA", "VAL Affinity loss x GAMMA", ] - title_font_size = 12 - label_font_size = 10 - tick_font_size = 9 + title_font_size = 16 + label_font_size = 16 + tick_font_size = 14 - plt.clf() - plt.ticklabel_format(useOffset=False) + fig, ax = plt.subplots(figsize=(10, 7)) + ax.ticklabel_format(useOffset=False) train_loss[-2] = train_loss[-2] * beta train_loss[-1] = train_loss[-1] * gamma for i, loss in enumerate(train_loss): s = "-" - plt.plot( + ax.plot( range(1, epochs + 1), loss, c=cols[i], linestyle=s, label=labs[i] ) if val_loss is not None: @@ -1057,7 +1057,7 @@ def loss_plot( val_loss[-1] = val_loss[-1] * gamma for i, loss in enumerate(val_loss): s = "--" - plt.plot( + ax.plot( range(1, epochs + 1), loss, c=cols[i], @@ -1074,30 +1074,30 @@ def loss_plot( "Exiting.\n", ) return - plt.title( + ax.set_title( "bs: %d, d: %d, ch: %d, lat: %d, lr: %.3f, beta: %.1f, " "gamma: %.1f" % (p[0], p[1], p[2], p[3], p[4], p[5], p[6]), fontsize=title_font_size, fontweight="normal", ) else: - plt.title("Loss", fontsize=title_font_size, fontweight="normal") + ax.set_title("Loss", fontsize=title_font_size, fontweight="normal") - plt.yscale("log") - plt.ylabel("Loss", fontsize=label_font_size) - plt.xlabel("Epochs", fontsize=label_font_size) - plt.xticks(fontsize=tick_font_size) - plt.yticks(fontsize=tick_font_size) - plt.legend(fontsize=tick_font_size) + ax.set_yscale("log") + ax.set_ylabel("Loss", fontsize=label_font_size) + ax.set_xlabel("Epochs", fontsize=label_font_size) + ax.tick_params(axis="both", labelsize=tick_font_size) + ax.legend(fontsize=tick_font_size) - plt.tight_layout() + fig.tight_layout() if not os.path.exists("plots"): os.mkdir("plots") - plt.savefig(f"plots/loss.{vis_format}", dpi=300) - plt.close() + fig.savefig(f"plots/loss.{vis_format}", dpi=300) + plt.close(fig) # plotting only the total loss as it sometimes is a few order of magnitude higher than KLD and affinity losses - plt.plot( + fig, ax = plt.subplots(figsize=(10, 7)) + ax.plot( range(1, epochs + 1), train_loss[0], c=cols[0], @@ -1105,22 +1105,21 @@ def loss_plot( label=labs[0], ) if val_loss is not None: - plt.plot( + ax.plot( range(1, epochs + 1), val_loss[0], c=cols[0], linestyle="--", label=vlabs[0], ) - plt.yscale("log") - plt.ylabel("Loss", fontsize=label_font_size) - plt.xlabel("Epochs", fontsize=label_font_size) - plt.xticks(fontsize=tick_font_size) - plt.yticks(fontsize=tick_font_size) - plt.legend(fontsize=tick_font_size) - plt.tight_layout() - plt.savefig(f"plots/loss_total.{vis_format}", dpi=300) - plt.close() + ax.set_yscale("log") + ax.set_ylabel("Loss", fontsize=label_font_size) + ax.set_xlabel("Epochs", fontsize=label_font_size) + ax.tick_params(axis="both", labelsize=tick_font_size) + ax.legend(fontsize=tick_font_size) + fig.tight_layout() + fig.savefig(f"plots/loss_total.{vis_format}", dpi=300) + plt.close(fig) def recon_plot( @@ -1941,16 +1940,16 @@ def plot_cyc_variable( "################################################################", ) logging.info(f"Visualising {variable_name} ...\n") - plt.plot(array, linewidth=3) - plt.ylabel(rf"$\{variable_name}$", fontsize=16) - plt.xlabel("Epochs", fontsize=16) - plt.xticks(fontsize=16) - plt.yticks(fontsize=16) - plt.tight_layout() + fig, ax = plt.subplots(figsize=(10, 7)) + ax.plot(array, linewidth=3) + ax.set_ylabel(rf"$\{variable_name}$", fontsize=16) + ax.set_xlabel("Epochs", fontsize=16) + ax.tick_params(axis="both", labelsize=14) + fig.tight_layout() if not os.path.exists("plots"): os.mkdir("plots") - plt.savefig(f"plots/hyperaparam_{variable_name}.{vis_format}", dpi=300) - plt.close() + fig.savefig(f"plots/hyperaparam_{variable_name}.{vis_format}", dpi=300) + plt.close(fig) def latent_space_similarity_plot( From e78e5444da7975af1ada203d6784d01cad789270 Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:43:29 +0000 Subject: [PATCH 04/12] Fix GaussianSplatDecoder default axis device The default axis tensor was created on the CPU, while the splat renderer runs on `device`, which breaks GPU runs. --- avae/decoders/differentiable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/avae/decoders/differentiable.py b/avae/decoders/differentiable.py index 1c21767a..4cb25015 100644 --- a/avae/decoders/differentiable.py +++ b/avae/decoders/differentiable.py @@ -250,7 +250,7 @@ def configure_renderer( `shape` of the output. """ self._shape = shape - self._default_axis = default_axis.as_tensor() + self._default_axis = default_axis.as_tensor().to(device) self._splatter = GaussianSplatRenderer( shape, device=device, From 6630a2076a7a134fc9b6c3f0d72256475ac79e1c Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:43:29 +0000 Subject: [PATCH 05/12] Vectorise latent space similarity matrix Compute class-averaged cosine similarity (mean and std) from per-class sums of unit vectors, without building the N x N pairwise matrix. Memory no longer grows quadratically with the number of samples. A test checks the results against the full pairwise computation, including classes absent from the data. --- avae/utils.py | 67 +++++++++++++++++++++++++++------------------ tests/test_utils.py | 29 ++++++++++++++++++++ 2 files changed, 69 insertions(+), 27 deletions(-) diff --git a/avae/utils.py b/avae/utils.py index 184262c3..10e982bb 100644 --- a/avae/utils.py +++ b/avae/utils.py @@ -9,7 +9,6 @@ import numpy.typing as npt import sklearn.linear_model import sklearn.metrics -import sklearn.metrics.pairwise import torch @@ -232,34 +231,48 @@ def latent_space_similarity_mat( display: bool When this variable is set to true, the function only dispalys the plot and doesnt save it. """ - # get same label order as affinity matrix - cosine_sim_matrix = sklearn.metrics.pairwise.cosine_similarity( - latent_space - ) + # The cosine similarity of two vectors is the dot product of their unit + # length versions: cos(a, b) = u_a . u_b, with u = x / |x|. Averaging it + # over all pairs of two classes therefore reduces to dot products of + # per-class sums of unit vectors, so the N x N similarity matrix never + # needs to be built: + # mean_ij = (S_i . S_j) / (n_i n_j), S_c = sum of u over class c + # For the std, the squared similarity (u_a . u_b)^2 equals the Frobenius + # product _F, so the mean squared similarity is + # E_ij[s^2] = _F / (n_i n_j), M_c = sum of u u^T in class c + # and std_ij = sqrt(E_ij[s^2] - mean_ij^2). + # Self-pairs are included on the diagonal, as in the full pairwise version. + latent_space = np.asarray(latent_space, dtype=np.float64) + class_labels = np.asarray(class_labels) + # unit length latent vectors; zero vectors stay zero (similarity 0) + norms = np.linalg.norm(latent_space, axis=1, keepdims=True) + unit = latent_space / np.where(norms == 0, 1.0, norms) + + latent_dims = unit.shape[1] + counts = np.zeros(num_classes) + sums = np.zeros((num_classes, latent_dims)) + outer_sums = np.zeros((num_classes, latent_dims * latent_dims)) + for i in range(num_classes): + class_unit = unit[class_labels == unique_classes[i]] + counts[i] = len(class_unit) + sums[i] = class_unit.sum(axis=0) + if plot_mode == "std": + outer_sums[i] = (class_unit.T @ class_unit).ravel() - cosine_sim_mat = np.zeros((num_classes, num_classes)) + pair_counts = np.outer(counts, counts) + has_pairs = pair_counts > 0 + safe_pair_counts = np.where(has_pairs, pair_counts, 1.0) + mean_sim = (sums @ sums.T) / safe_pair_counts - for i in range(num_classes): - for j in range(i, num_classes): - class_i_indices = np.where(class_labels == unique_classes[i])[0] - class_j_indices = np.where(class_labels == unique_classes[j])[0] - cosine_sims = cosine_sim_matrix[class_i_indices][ - :, class_j_indices - ] - if cosine_sims.size == 0: - cosine_sim_mat[i, j] = 0.0 - cosine_sim_mat[j, i] = 0.0 - continue - if plot_mode == "mean": - cosine_sim_mat[i, j] = np.mean(cosine_sims) - cosine_sim_mat[j, i] = cosine_sim_mat[ - i, j - ] # symmetrical matrix - if plot_mode == "std": - cosine_sim_mat[i, j] = np.std(cosine_sims) - cosine_sim_mat[j, i] = cosine_sim_mat[ - i, j - ] # symmetrical matrix + cosine_sim_mat = np.zeros((num_classes, num_classes)) + if plot_mode == "mean": + cosine_sim_mat = mean_sim + elif plot_mode == "std": + second_moment = (outer_sums @ outer_sums.T) / safe_pair_counts + cosine_sim_mat = np.sqrt( + np.clip(second_moment - mean_sim**2, 0, None) + ) + cosine_sim_mat[~has_pairs] = 0.0 return cosine_sim_mat diff --git a/tests/test_utils.py b/tests/test_utils.py index 84546251..28f0328c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2,8 +2,10 @@ import numpy as np import pytest +import sklearn.metrics.pairwise from avae.base import dims_after_pooling +from avae.utils import latent_space_similarity_mat from avae.utils_learning import combine_accuracy_data @@ -73,3 +75,30 @@ def test_combine_accuracy_data_rejects_missing_rank(): pytest.raises(RuntimeError, match="was not gathered from every rank"), ): combine_accuracy_data(*local_data, True, 2) + + +@pytest.mark.parametrize("plot_mode", ["mean", "std"]) +def test_latent_space_similarity_mat_matches_pairwise(plot_mode): + """Vectorised class similarity equals averaging the full pairwise matrix, + including classes missing from the data.""" + rng = np.random.default_rng(0) + latent_space = rng.normal(size=(60, 5)) + latent_space[0] = 0.0 + class_labels = rng.choice(["a", "b", "c"], size=60) + unique_classes = ["c", "a", "missing", "b"] + + pairwise = sklearn.metrics.pairwise.cosine_similarity(latent_space) + reduce = np.mean if plot_mode == "mean" else np.std + expected = np.zeros((4, 4)) + for i, class_i in enumerate(unique_classes): + for j, class_j in enumerate(unique_classes): + block = pairwise[class_labels == class_i][ + :, class_labels == class_j + ] + if block.size: + expected[i, j] = reduce(block) + + result = latent_space_similarity_mat( + latent_space, class_labels, unique_classes, 4, plot_mode=plot_mode + ) + np.testing.assert_allclose(result, expected, atol=1e-7) From d755abaf675133cfe58b6b8a1b746c64d9c0fafd Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 06/12] Remove TensorBoard, add vis_print and figure helpers - Drop the tensorboard option, SummaryWriter and every writer argument to plot functions, and remove the tensorboard dependency. - Add a vis_print option to render figures at 300 dpi for print, and document svg as a vis_format. - Rename accuracy_plot to confusion_plot and log "Overall accuracy". - Extract the shared _loss_curve_figure, _matrix_figure and _save_grid_figure helpers. Vector formats are embedded at native resolution. --- README.md | 6 +- avae/config.py | 14 +- avae/data.py | 41 +- avae/evaluate.py | 14 +- avae/train.py | 58 +- avae/utils.py | 15 +- avae/vis.py | 2330 +++++++++++++------------ pyproject.toml | 1 - requirements.txt | 1 - tests/test_train_eval_pipeline.py | 1 - tests/test_vis.py | 8 +- tutorials/mnist_data/mnist_config.yml | 1 - 12 files changed, 1288 insertions(+), 1202 deletions(-) diff --git a/README.md b/README.md index fada2623..bd6dd6dd 100644 --- a/README.md +++ b/README.md @@ -216,8 +216,8 @@ Options: per run). -similarity, --vis_sim Visualise train-val model similarity matrix. -va, --vis_all Visualise all above. - -vf, --vis_format TEXT The format of saved images. Options: png , - pdf + -vf, --vis_format TEXT The format of saved images. Options: png, + pdf, svg -fev, --freq_eval INTEGER Frequency at which to evaluate test set. -fs, --freq_sta INTEGER Frequency at which to save state -fac, --freq_acc INTEGER Frequency at which to visualise confusion @@ -246,8 +246,6 @@ Options: and the maximum to one -res, --rescale INTEGER Rescale images to given value (tuple, one value per dim). - -tb, --tensorboard Log metrics and figures to tensorboard - during training -st, --strategy TEXT Define the strategy for distributed training. Options are: 'ddp', 'deepspeed' or 'fsdp diff --git a/avae/config.py b/avae/config.py index 32cf7a55..946e2cd4 100644 --- a/avae/config.py +++ b/avae/config.py @@ -177,9 +177,18 @@ class AffinityConfig(pydantic.BaseModel): #### Visualisation parameters # setup params vis_format: None | str = pydantic.Field( - "png", description="The format of saved images. Options: png , pdf" + "png", + description="The format of saved images. Options: png, pdf, svg " + "(vector formats stay sharp at any size and are fastest for " + "large matrices)", + ) + vis_print: bool = pydantic.Field( + False, + description="Render matrix visualisations (confusion, affinity, " + "similarity) at print quality (300dpi) instead of the faster " + "default resolution. The dpi setting does not apply to vector " + "vis_format (pdf, svg)", ) - tensorboard: bool = pydantic.Field(False, description="Use tensorboard") vis_his: bool = pydantic.Field( False, description="Visualise class distribution histogram" ) @@ -384,6 +393,7 @@ def load_config_params( if key in [ 'vis_all', 'vis_format', + 'vis_print', 'vis_z_n_int', 'vis_pose_class', ]: diff --git a/avae/data.py b/avae/data.py index 63217036..e43944ea 100644 --- a/avae/data.py +++ b/avae/data.py @@ -34,6 +34,7 @@ def load_data( vis_his: bool = False, vis_aff: bool = False, vis_format: str = "png", + vis_print: bool = False, ) -> tuple[torch.utils.data.DataLoader, int]: ... @@ -57,6 +58,7 @@ def load_data( vis_his: bool = False, vis_aff: bool = False, vis_format: str = "png", + vis_print: bool = False, ) -> tuple[ torch.utils.data.DataLoader, torch.utils.data.DataLoader, @@ -85,13 +87,17 @@ def load_data( vis_his: bool = False, vis_aff: bool = False, vis_format: str = "png", -) -> tuple[ - torch.utils.data.DataLoader, - torch.utils.data.DataLoader, - torch.utils.data.DataLoader, - pd.DataFrame, - int, -] | tuple[torch.utils.data.DataLoader, int]: + vis_print: bool = False, +) -> ( + tuple[ + torch.utils.data.DataLoader, + torch.utils.data.DataLoader, + torch.utils.data.DataLoader, + pd.DataFrame, + int, + ] + | tuple[torch.utils.data.DataLoader, int] +): """This function a wrapper around the DiskDataLoader class from the caked library. It loads data from a given path, selects a subset of classes if requested, splits it into train / val and test in batch sets, and loads an affinity matrix. The function is overloaded to return different types of data depending on the value of the eval parameter. If eval is True, the function returns only the test data and the dimension of the data. If eval is False, the function returns train, validation, and test data, the affinity matrix, and the dimension of the data. @@ -188,6 +194,7 @@ def load_data( classes_list, vis_aff=vis_aff and fabric.global_rank == 0, vis_format=vis_format, + vis_print=vis_print, ) # assign the affinity matrix to the dataset (small modification from the caked DiskDataset, which only returns data and labels, and we need the affinity matrix indexes for training). @@ -204,9 +211,17 @@ def load_data( if vis_his and fabric.global_rank == 0: train_y = list(sum([y[1] for _, y in enumerate(trains)], ())) val_y = list(sum([y[1] for _, y in enumerate(vals)], ())) - plot_classes_distribution(train_y, "train", vis_format=vis_format) plot_classes_distribution( - val_y, "validation", vis_format=vis_format + train_y, + "train", + vis_format=vis_format, + vis_print=vis_print, + ) + plot_classes_distribution( + val_y, + "validation", + vis_format=vis_format, + vis_print=vis_print, ) train_size = len(trains.dataset) @@ -253,7 +268,10 @@ def load_data( if vis_his and fabric.global_rank == 0: eval_y = list(sum([y[1] for _, y in enumerate(tests)], ())) plot_classes_distribution( - eval_y, "evaluation", vis_format=vis_format + eval_y, + "evaluation", + vis_format=vis_format, + vis_print=vis_print, ) tests = fabric.setup_dataloaders(tests) @@ -278,6 +296,7 @@ def get_affinity_matrix( classes: list = [], vis_aff: bool = False, vis_format: str = "png", + vis_print: bool = False, ) -> pd.DataFrame: """Loads affinity matrix from a given path, subsets it given selected classes and returns it as a pandas DataFrame. @@ -314,6 +333,7 @@ def get_affinity_matrix( all_classes=affinity.columns.tolist(), selected_classes=classes, vis_format=vis_format, + vis_print=vis_print, ) # subset affinity matrix with only the relevant classes @@ -327,7 +347,6 @@ def get_affinity_matrix( class AffinityDiskDataset(caked.dataloader.DiskDataset): - """Modified version of the caked DiskDataset to include the affinity matrix and data metadata that is needed for the affinity pipeline""" diff --git a/avae/evaluate.py b/avae/evaluate.py index 222f5954..8381cff6 100644 --- a/avae/evaluate.py +++ b/avae/evaluate.py @@ -192,6 +192,7 @@ def evaluate(params): data_dim, mode="evl", vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise latent disentanglement @@ -204,6 +205,7 @@ def evaluate(params): poses=p_test, mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise pose disentanglement @@ -216,6 +218,7 @@ def evaluate(params): device, mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) if pose_dims != 0 and params.vis_pose_class: @@ -229,6 +232,7 @@ def evaluate(params): device, mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise interpolations if params.vis_int: @@ -241,6 +245,7 @@ def evaluate(params): poses=p_test, mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise embeddings @@ -251,6 +256,7 @@ def evaluate(params): classes_list, "_eval", vis_format=params.vis_format, + vis_print=params.vis_print, embedding=latent_embedding[eval_mask], ) @@ -260,6 +266,7 @@ def evaluate(params): np.array(y_test), mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) # ############################# Predict ############################# @@ -288,6 +295,7 @@ def evaluate(params): classes_list, "_train_eval_comparison", vis_format=params.vis_format, + vis_print=params.vis_print, embedding=latent_embedding, ) @@ -306,10 +314,10 @@ def evaluate(params): classifier=params.classifier, ) logging.info( - "------------------->>> Accuracy: Train: %f | Val : %f | Val with unseen labels: %f\n" + "------------------->>> Overall accuracy: Train: %f | Val : %f | Val with unseen labels: %f\n" % (train_acc, val_acc_selected, val_acc) ) - vis.accuracy_plot( + vis.confusion_plot( np.array(latents_training_id), ypred_train, y_test, @@ -317,6 +325,7 @@ def evaluate(params): params.classes, mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) vis.f1_plot( np.array(latents_training_id), @@ -325,6 +334,7 @@ def evaluate(params): ypred_val, mode="_eval", vis_format=params.vis_format, + vis_print=params.vis_print, ) logging.info("Saving meta files with evaluation data.") diff --git a/avae/train.py b/avae/train.py index e42164a2..c2e3d75a 100644 --- a/avae/train.py +++ b/avae/train.py @@ -5,7 +5,6 @@ import numpy as np import pandas as pd import torch -import torch.utils.tensorboard from . import utils_learning, vis from .cyc_annealing import setup_annealing @@ -36,12 +35,6 @@ def train(params): lt.pytorch.seed_everything(42) - # ############################### LOGGING ################################# - - writer = ( - torch.utils.tensorboard.SummaryWriter() if params.tensorboard else None - ) - # ############################### GPU SETUP ################################ fabric = setup_gpus( @@ -83,8 +76,18 @@ def train(params): cycle_load=params.gamma_load, ) if rank_zero and params.vis_cyc: - vis.plot_cyc_variable(beta_arr, "beta", vis_format=params.vis_format) - vis.plot_cyc_variable(gamma_arr, "gamma", vis_format=params.vis_format) + vis.plot_cyc_variable( + beta_arr, + "beta", + vis_format=params.vis_format, + vis_print=params.vis_print, + ) + vis.plot_cyc_variable( + gamma_arr, + "gamma", + vis_format=params.vis_format, + vis_print=params.vis_print, + ) # ############################### DATA ############################### trains, vals, tests, affinity_matrix, data_dim = load_data( @@ -104,6 +107,7 @@ def train(params): vis_his=params.vis_his, vis_aff=params.vis_aff, vis_format=params.vis_format, + vis_print=params.vis_print, fabric=fabric, ) @@ -300,12 +304,6 @@ def train(params): ) ) - if writer: - for i, loss_name in enumerate( - ["Loss", "Recon loss", "KLdiv loss", "Affin loss"] - ): - writer.add_scalar(loss_name, v_history[-1][i], epoch) - # ########################## TEST ##################################### if params.freq_eval != 0 and (epoch + 1) % params.freq_eval == 0: for batch_number, (t, ys, aff, meta_data) in enumerate( @@ -469,17 +467,17 @@ def train(params): ) logging.info( - "------------------->>> Accuracy: Train: %f | Val: %f\n" + "------------------->>> Overall accuracy: Train: %f | Val: %f\n" % (train_acc, val_acc), ) - vis.accuracy_plot( + vis.confusion_plot( y_train_all, ypred_train, y_val_all, ypred_val, epoch=epoch, - writer=writer, vis_format=params.vis_format, + vis_print=params.vis_print, ) vis.f1_plot( @@ -488,8 +486,8 @@ def train(params): y_val_all, ypred_val, epoch=epoch, - writer=writer, vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise loss @@ -511,6 +509,7 @@ def train(params): v_history, p=p, vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise reconstructions - last batch @@ -521,9 +520,8 @@ def train(params): y_train, data_dim, mode="trn", - epoch=epoch, - writer=writer, vis_format=params.vis_format, + vis_print=params.vis_print, ) vis.recon_plot( v, @@ -531,9 +529,8 @@ def train(params): y_val, data_dim, mode="val", - epoch=epoch, - writer=writer, vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise mean and logvar similarity matrix @@ -545,6 +542,7 @@ def train(params): epoch=epoch, affinity_matrix=params.affinity, vis_format=params.vis_format, + vis_print=params.vis_print, ) vis.latent_space_similarity_plot( z_val, @@ -553,6 +551,7 @@ def train(params): epoch=epoch, affinity_matrix=params.affinity, vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise embeddings @@ -561,8 +560,8 @@ def train(params): embedding_xs, embedding_ys, epoch=epoch, - writer=writer, vis_format=params.vis_format, + vis_print=params.vis_print, embedding=latent_embedding, ) @@ -571,9 +570,9 @@ def train(params): pose_xs, embedding_ys, epoch=epoch, - writer=writer, mode="pose", vis_format=params.vis_format, + vis_print=params.vis_print, embedding=pose_embedding, ) @@ -600,6 +599,7 @@ def train(params): device, poses=p_train if pose else None, vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise pose disentanglement @@ -616,6 +616,7 @@ def train(params): vae, device, vis_format=params.vis_format, + vis_print=params.vis_print, ) if params.vis_pose_class is not None: @@ -628,6 +629,7 @@ def train(params): vae, device, vis_format=params.vis_format, + vis_print=params.vis_print, ) # visualise interpolations @@ -657,6 +659,7 @@ def train(params): params.vis_z_n_int, poses=ps, vis_format=params.vis_format, + vis_print=params.vis_print, ) vis.interpolations_plot( @@ -667,6 +670,7 @@ def train(params): device, poses=ps, # do we need val and test here? vis_format=params.vis_format, + vis_print=params.vis_print, ) # ########################## SAVE STATE ############################### if params.freq_sta != 0 and (epoch + 1) % params.freq_sta == 0: @@ -720,7 +724,3 @@ def train(params): logging.info(f"Saved meta file : {filename} for evaluation \n") fabric.barrier() - - if writer: - writer.flush() - writer.close() diff --git a/avae/utils.py b/avae/utils.py index 10e982bb..c8c1fbcb 100644 --- a/avae/utils.py +++ b/avae/utils.py @@ -1,7 +1,6 @@ import copy import logging import os.path -import typing import matplotlib.pyplot as plt import mrcfile @@ -88,10 +87,8 @@ def save_imshow_png( cmap: str | None = None, min: float | None = None, max: float | None = None, - writer: typing.Any = None, - figname: str | None = None, - epoch: int = 0, display: bool = False, + vis_print: bool = False, ) -> None: if not display: if not os.path.exists("plots"): @@ -106,10 +103,12 @@ def save_imshow_png( ax.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last ax.axis("off") fig.subplots_adjust(left=0, right=1, bottom=0, top=1) - fig.savefig("plots/" + fname, bbox_inches="tight", pad_inches=0) - - if writer: - writer.add_figure(figname, fig, epoch) + fig.savefig( + "plots/" + fname, + **({"dpi": 300} if vis_print else {}), + bbox_inches="tight", + pad_inches=0, + ) else: plt.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last plt.show() diff --git a/avae/vis.py b/avae/vis.py index b4083d89..23158435 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -17,7 +17,6 @@ import sklearn.metrics.pairwise import torch import torchvision -from mpl_toolkits.axes_grid1 import make_axes_locatable from . import utils_learning from .utils import ( @@ -30,703 +29,457 @@ save_mrc_file, ) +MAX_CLASS_FIG_SIZE = 20 +VECTOR_FORMATS = ("pdf", "svg", "eps") + def _class_plot_scale( num_classes: int, fig_size: int | None = None ) -> tuple[float, int, int]: - """Return consistent figure and font sizes for class-based plots.""" - font_size = max(8, int(num_classes / 3) + 3) - return ( - fig_size if fig_size is not None else max(6, int(num_classes) / 2), - font_size, - max(font_size + 4, int(num_classes / 2) + 8), - ) + """Return consistent figure and font sizes for class-based plots. + The default figure grows by a quarter inch per class, capped at + MAX_CLASS_FIG_SIZE inches since rendering cost grows with its area. + Fonts are sized to one class row so tick labels do not overlap; save + large matrices in a vector format (pdf/svg) to zoom in on them. + """ + resolved_fig_size = ( + fig_size + if fig_size is not None + else min(max(6, num_classes / 4), MAX_CLASS_FIG_SIZE) + ) + row_height_pt = 72 * resolved_fig_size / max(num_classes, 1) + font_size = int(np.clip(0.7 * row_height_pt, 4, 14)) + return resolved_fig_size, font_size, max(14, font_size + 6) -def _encoder(i: PIL.Image) -> str: - """Encode PIL Image as base64 buffer. - Parameters - ---------- - i: PIL.Image - Image to be encoded. - Returns - ------- - str - Encoded image. +def _loss_curve_figure( + epochs: int, + series: typing.Sequence[tuple[str | None, npt.ArrayLike, str | None, str]], + path: str, + vis_format: str = "png", + vis_print: bool = False, + title: str | None = None, + ylabel: str = "Loss", + log_scale: bool = True, + linewidth: float | None = None, + legend: bool = True, +) -> None: + """Render and save one line-vs-epoch figure. `series` is a list of + (label, values, colour, linestyle) tuples, one per line drawn.""" + fig, ax = plt.subplots(figsize=(10, 7)) + ax.ticklabel_format(useOffset=False) - """ - import base64 - import io + for label, values, color, linestyle in series: + ax.plot( + range(1, epochs + 1), + values, + c=color, + linestyle=linestyle, + label=label, + linewidth=linewidth, + ) - with io.BytesIO() as buffer: - i.thumbnail((110, 110)) - i.save(buffer, "PNG") - data = base64.encodebytes(buffer.getvalue()).decode("utf-8") + if title: + ax.set_title(title, fontsize=16, fontweight="normal") - return f"{data}" + if log_scale: + ax.set_yscale("log") + ax.set_ylabel(ylabel, fontsize=16) + ax.set_xlabel("Epochs", fontsize=16) + ax.tick_params(axis="both", labelsize=14) + if legend: + ax.legend(fontsize=14) + fig.tight_layout() + if not os.path.exists("plots"): + os.mkdir("plots") + fig.savefig(f"{path}.{vis_format}", **({"dpi": 300} if vis_print else {})) + plt.close(fig) -def _decoder(i: str) -> PIL.Image: - """Decode base64 buffer as PIL Image. - Parameters - ---------- - i: str - Encoded image. - Returns - ------- - PIL.Image - Decoded image. +def _matrix_figure( + fig_size: float, + font_size: int, + title_font_size: int, + title: str, + path: str, + vis_format: str = "png", + vis_print: bool = False, + print_dpi: int = 300, + *, + data: npt.NDArray, + class_labels: typing.Sequence, + cmap, + vmin: float | None = None, + vmax: float | None = None, + values_format: str | None = None, + value_font_size: int | None = None, + highlight: npt.NDArray | None = None, + xlabel: str | None = None, + ylabel: str | None = None, + display: bool = False, +) -> None: + """Render and save (or display) one + class x class matrix figure: image, optional per-cell values (colour + picked for contrast against the cell, the way sklearn's + ConfusionMatrixDisplay does), optional red-highlighted tick labels, + colorbar, and title. """ - import base64 - import io + data = np.asarray(data) + num_classes = len(class_labels) + # Vector formats embed the matrix at its native class x class resolution + # ("none" skips resampling), so they stay sharp at any size and dpi has + # no effect on them. + is_vector = vis_format in VECTOR_FORMATS - return PIL.Image.open(io.BytesIO(base64.b64decode(i))) + fig, ax = plt.subplots(figsize=(fig_size, fig_size)) + im = ax.imshow( + data, + interpolation="none" if is_vector else "nearest", + cmap=cmap, + vmin=vmin, + vmax=vmax, + ) + if values_format is not None: + color_min, color_max = im.cmap(0.0), im.cmap(1.0) + thresh = (data.max() + data.min()) / 2.0 + for i in range(num_classes): + for j in range(num_classes): + color = color_max if data[i, j] < thresh else color_min + ax.text( + j, + i, + f"{data[i, j]:{values_format}}", + ha="center", + va="center", + color=color, + fontsize=value_font_size, + ) -def format(im: PIL.Image, data_dim: int) -> list[str]: - """Format PIL Image as Pandas compatible Altair image display. + ax.set_xticks(np.arange(num_classes), labels=class_labels) + ax.set_yticks(np.arange(num_classes), labels=class_labels) + ax.set_ylim((num_classes - 0.5, -0.5)) - Parameters - ---------- - im: PIL.Image - Image to be formatted. - data_dim: int - Dimension of the data. + if highlight is not None: + for is_highlighted, x_label, y_label in zip( + highlight, ax.get_xticklabels(), ax.get_yticklabels() + ): + if is_highlighted: + x_label.set_color("red") + y_label.set_color("red") - Returns - ------- - list[str] - Formatted images compatible with Altair image display. Batch input - returns one encoded item per sample, while single-sample input returns - a one-element list. - """ - if len(im.shape) == 5 and data_dim == 3: - batch = True - im = np.copy( - im[:, :, im.shape[2] // 2, :, :] - .squeeze(dim=1) - .cpu() - .detach() - .numpy() - ) - elif len(im.shape) == 4 and data_dim == 3: - batch = False - im = np.copy( - im[:, im.shape[1] // 2, :, :].squeeze(dim=0).cpu().detach().numpy() - ) - elif len(im.shape) == 4 and data_dim == 2: - batch = True - im = np.copy(im.squeeze(dim=1).cpu().detach().numpy()) - # .astype(np.uint8) - elif len(im.shape) == 3 and data_dim == 2: - batch = False - im = np.copy(im.squeeze(dim=0).cpu().detach().numpy()) - # .astype(np.uint8) - else: - logging.warning( - "\n\nWARNING: Wrong data format, please pass either a single " - "unsqueezed tensor or a batch to image formatter. Exiting.\n", - ) - return [] - im *= 255 - im = im.astype(np.uint8) - if batch: - # if batch we are adding a reconstruction to an input that already - # exists in the DF, '&' is a separator. - return ["&" + _encoder(PIL.Image.fromarray(i)) for i in im] - else: - return [_encoder(PIL.Image.fromarray(im))] + ax.tick_params(axis="x", rotation=90, labelsize=font_size) + ax.tick_params(axis="y", labelsize=font_size) + colorbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + colorbar.ax.tick_params(labelsize=font_size) -def merge(im: str) -> str | None: - """Merge 2 base64 buffers as PIL Images and encode back to base64 - buffers. + ax.set_title(title, fontsize=title_font_size, fontweight="normal") + if xlabel: + ax.set_xlabel(xlabel, fontsize=font_size) + if ylabel: + ax.set_ylabel(ylabel, fontsize=font_size) - Parameters - ---------- - im: str - Input PIL Images to be merged. + fig.tight_layout() - Returns - ------- - str - Merged image. - """ - if im.startswith("data:image/"): - return im + if display: + plt.show() + return - i = im.split("&") - if len(i) != 2: - logging.warning( - "\n\nWARNING: Image format corrupt. Number of images in meta_df: {}. " - "Exiting. \n".format(len(i)), - ) - return None + if not os.path.exists("plots"): + os.mkdir("plots") + fig.savefig( + path, + **({"dpi": print_dpi} if vis_print and not is_vector else {}), + **( + {"pil_kwargs": {"compress_level": 3}} + if vis_format == "png" + else {} + ), + ) - im1 = _decoder(i[0]) - im2 = _decoder(i[1]) + plt.close(fig) - new_image = PIL.Image.new("L", (im2.size[0] + im1.size[0], im2.size[1])) - new_image.paste(im1, (0, 0)) - new_image.paste(im2, (im1.size[0], 0)) - data = _encoder(new_image) - return f"data:image/png;base64,{data}" +def _save_grid_figure( + rows: int, + cols: int, + dsize: tuple, + images: npt.NDArray, + name: str, + vis_format: str = "png", + padding: int = 0, + display: bool = False, + vis_print: bool = False, +) -> None: + """Arrange `images` into a grid and save it: as an .mrc file for 3D + data (`dsize` has 3 dims), or a raster image for 2D data.""" + grid_for_napari = create_grid_for_plotting(rows, cols, dsize, padding) + grid_for_napari = fill_grid_for_plottting( + rows, cols, grid_for_napari, dsize, images, padding + ) + data_dim = len(dsize) + if data_dim == 3: + save_mrc_file(f"{name}.mrc", grid_for_napari) + elif data_dim == 2: + save_imshow_png( + f"{name}.{vis_format}", + grid_for_napari, + display=display, + vis_print=vis_print, + ) -def latent_embed_plot_tsne( - xs: npt.NDArray, - ys: npt.NDArray, - classes: list | None = None, - mode: str = "", - epoch: int = 0, - writer: typing.Any = None, - perplexity: int = 40, - marker_size: int = 24, - l_w: int = 2, - display: bool = False, + +def loss_plot( + epochs: int, + beta: float, + gamma: float, + train_loss: list[float], + val_loss: list[float] | None = None, + p: list | None = None, vis_format: str = "png", - embedding: npt.NDArray | None = None, + vis_print: bool = False, ) -> None: - """Plot static TSNE embedding. + """Visualise loss over epochs. Parameters ---------- - xs: np.ndarray - Array of latent vectors. - ys: np.ndarray - Array of labels. - classes: list - List of classes. - mode: str - Added data mode to the name of the saved figure (e.g train, valid, eval). - epoch: int - Current epoch - writer: SummaryWriter - Tensorboard summary writer - perplexity: int - display: bool - When this variable is set to true, the save_imshow_png function only dispalys the plot - and does not save a png image. This is to allow the use of this function in jupyter notebook - post calculation. This features is not yet implemented for save_mrc_file. - vis_format: str - format of the image saved - embedding: np.ndarray | None - Precomputed embedding coordinates. If omitted, they are calculated - from ``xs``. + epochs: int + Number of epochs. + beta: list[float] + List of beta values. + gamma: list[float] + List of gamma values. + train_loss: list + Training loss over epochs. + val_loss: list + Validation loss over epochs. + p: list + List of 7 hyperparameters: batch size, depth, " + "channel init, latent dimension, learning rate, beta, gamma. """ logging.info( "################################################################", ) - if not mode: - logging.info("Visualising static TSNE embedding...\n") - else: - logging.info("Visualising static TSNE embedding " + mode + "...\n") + logging.info("Visualising loss ...\n") - xs = np.asarray(xs) - ys = np.asarray(ys) + train_loss = np.transpose(np.asarray(train_loss)) + if val_loss is not None: + val_loss = np.transpose(np.asarray(val_loss)) - if xs.ndim != 2: - raise ValueError("Embedding only accepts 2D arrays.") - if xs.shape[-1] == 1: - logging.warning( - "Data contains 1 dimension, cannot create embedding," - " plotting histogram instead...\n" - ) - if xs.shape[-1] == 2: - logging.warning( - "Data already contains 2 dimensions, cannot create" - " embedding, plotting scatter of original data...\n" - ) + cols = ["blue", "red", "green", "orange"] + labs = [ + "Total loss", + "Reconstruction loss", + "KL divergence loss x BETA", + "Affinity loss x GAMMA", + ] + vlabs = [ + "VAL Total loss", + "VAL Reconstruction loss", + "VAL KL divergence loss x BETA", + "VAL Affinity loss x GAMMA", + ] + train_loss[-2] = train_loss[-2] * beta + train_loss[-1] = train_loss[-1] * gamma + if val_loss is not None: + val_loss[-2] = val_loss[-2] * beta + val_loss[-1] = val_loss[-1] * gamma - lats = ( - utils_learning.tsne_embedding(xs, perplexity=perplexity) - if embedding is None - else np.asarray(embedding) - ) - if len(lats) != len(xs): - raise ValueError( - "Embedding and latent vectors must have equal lengths." - ) - - if classes is None: - classes = sorted(list(np.unique(ys))) - else: - if np.setdiff1d(ys, classes).size > 0: - classes = list( - np.concatenate((classes, np.setdiff1d(ys, classes))) + if p is not None: + if len(p) != 7: + logging.warning( + "\n\nWARNING: Function vis.loss_plot is expecting 'p' parameter " + "to be a list of 7 hyperparameters: batch size, depth, " + "channel init, latent dimension, learning rate, beta, gamma. " + "Exiting.\n", + ) + return + title = ( + "bs: %d, d: %d, ch: %d, lat: %d, lr: %.3f, beta: %.1f, " + "gamma: %.1f" + % ( + p[0], + p[1], + p[2], + p[3], + p[4], + p[5], + p[6], ) - - n_classes = len(classes) - fig_size, font_size, title_font_size = _class_plot_scale(n_classes) - - if n_classes < 3: - # If the number of classes are not moe than 3 the size of the figure would be too - # small and matplotlib would through a singularity error - fig, ax = plt.subplots( - figsize=(int(n_classes / 2) + 7, int(n_classes / 2) + 5) ) else: - fig, ax = plt.subplots( - figsize=(int(n_classes / 2) + 4, int(n_classes / 2) + 2) - ) - marker_size = max(marker_size, int(marker_size * max(1, n_classes / 5))) - # When the number of classes is less than 3 the image becomes two small - colours = colour_per_class(classes) - - if xs.shape[-1] != 1: + title = "Loss" - for mol_id, mol in enumerate(set(ys.tolist())): - idx = np.where(np.array(ys.tolist()) == mol)[0] - - color = colours[classes.index(mol)] - - plt.scatter( - lats[idx, 0], - lats[idx, 1], - s=marker_size, - label=mol[:4], - facecolor=color, - edgecolor=color, - alpha=0.5, - ) + full_series = [(labs[i], train_loss[i], cols[i], "-") for i in range(4)] + if val_loss is not None: + full_series += [ + (vlabs[i], val_loss[i], cols[i], "--") for i in range(4) + ] + _loss_curve_figure( + epochs, full_series, "plots/loss", vis_format, vis_print, title=title + ) - ax.legend( - bbox_to_anchor=(1.05, 1), - loc="upper left", - fontsize=font_size, - ) - ax.set_xlabel("TSNE-1", fontsize=font_size) - ax.set_ylabel("TSNE-2", fontsize=font_size) + # plotting only the total loss as it sometimes is a few order of magnitude higher than KLD and affinity losses + total_series = [(labs[0], train_loss[0], cols[0], "-")] + if val_loss is not None: + total_series.append((vlabs[0], val_loss[0], cols[0], "--")) + _loss_curve_figure( + epochs, total_series, "plots/loss_total", vis_format, vis_print + ) - if xs.shape[-1] == 1: - for mol_id, mol in enumerate(set(ys.tolist())): - idx = np.where(np.array(ys.tolist()) == mol)[0] - cols = colours[classes.index(mol)] - plt.hist( - lats[idx], - 100, - color=cols, - histtype="step", - stacked=True, - fill=False, - label=mol[:4], - linewidth=l_w, - ) - ax.legend( - bbox_to_anchor=(1.05, 1), - loc="upper left", - fontsize=font_size, - ) - ax.set_xlabel("dim 1", fontsize=font_size) - ax.set_ylabel("freq", fontsize=font_size) +def plot_cyc_variable( + array: list, + variable_name: str, + vis_format: str = "png", + vis_print: bool = False, +): + """Plot evolution of variable from the cyclical training - ax.set_title( - f"t-SNE Embedding at epoch: {epoch + 1}", - fontsize=title_font_size, - fontweight="normal", + Parameters + ---------- + array : list + List of values for the variable + variable_name : str + Name of the variable + """ + logging.info( + "################################################################", + ) + logging.info(f"Visualising {variable_name} ...\n") + _loss_curve_figure( + len(array), + [(None, array, None, "-")], + f"plots/hyperaparam_{variable_name}", + vis_format, + vis_print, + ylabel=rf"$\{variable_name}$", + log_scale=False, + linewidth=3, + legend=False, ) - ax.tick_params(axis="both", labelsize=font_size) - fig.tight_layout() - - if not display: - if not os.path.exists("plots"): - os.mkdir("plots") - plt.savefig(f"plots/embedding_TSNE{mode}.{vis_format}") - else: - plt.show() - if writer: - writer.add_figure("TSNE embedding", fig, epoch) - - plt.close() -def dyn_latentembed_plot( - df: pd.DataFrame, - epoch: int, +def confusion_plot( + y_train: npt.NDArray, + ypred_train: npt.NDArray, + y_val: npt.NDArray, + ypred_val: npt.NDArray, + classes: str | None = None, mode: str = "", - embedding: npt.NDArray | None = None, + epoch: int = 0, + vis_format: str = "png", + vis_print: bool = False, ): - """Plot dynamic TSNE embedding. + """Plot confusion matrix . Parameters ---------- - df: pd.DataFrame - Dataframe containing the latent vectors. + y_train: np.array + Training labels. + ypred_train: np.array + Predicted training labels. + y_val: np.array + Validation labels (unseen data). + ypred_val: np.array + Predicted validation labels (unseen data). + mode: str + Added data mode to the name of the saved figures (e.g train, valid, eval). epoch: int Current epoch. - mode: str - Added data mode to the name of the saved figure (e.g train, valid, eval). - embedding: np.ndarray | None - Precomputed embedding coordinates. If omitted, they are calculated - from the latent columns in ``df``. """ logging.info( "################################################################", ) - logging.info("Visualising dynamic embedding...\n") - latentspace = df[[col for col in df if col.startswith("lat")]].to_numpy() - lat_emb = np.asarray( - utils_learning.tsne_embedding(latentspace) - if embedding is None - else embedding - ) - if len(lat_emb) != len(df): - raise ValueError("Embedding and metadata must have equal lengths.") - if lat_emb.shape[1] == 1: - lat_emb = np.column_stack([lat_emb[:, 0], np.zeros(len(lat_emb))]) - titlex = "t-SNE-1" - titley = "t-SNE-2" - df["emb-x"], df["emb-y"] = np.array(lat_emb)[:, 0], np.array(lat_emb)[:, 1] + logging.info("Visualising accuracy: confusion and F1 scores ...\n") - # create column select options for radio buttons - # add no std to radio select options - change value here for marker size! - if "std-off" not in df.columns: - df.insert(loc=0, column="std-off", value=np.zeros(len(lat_emb)) + 0.5) - # add certainty averaged across all latent dimensions as an extra option - std_dim_cols = [ - col - for col in df.columns - if col.startswith("std-") and col not in ("std-off", "std-avg") - ] - if std_dim_cols and "std-avg" not in df.columns: - df["std-avg"] = df[std_dim_cols].mean(axis=1) - opts = [col for col in df.columns if col.startswith("std")] + classes_list = np.unique(np.concatenate((y_train, ypred_train))) - # create radio buttons and bind to a folded column select - bind_checkbox = altair.binding_radio( - options=opts, - labels=[ - "off" - if i == "std-off" - else "avg" - if i == "std-avg" - else str(int(i.split("-")[-1]) + 1) - for i in opts - ], - name="Certainty of prediction per dimension:", - ) - column_select = altair.selection_point( - fields=["column"], - bind=bind_checkbox, - value=[{"column": "std-off"}], - name="certainty", + fig_size, font_size, title_font_size = _class_plot_scale(len(classes_list)) + value_font_size = min(font_size, 10) + + # Compute confusion matrix + cm = sklearn.metrics.confusion_matrix( + y_train, ypred_train, labels=classes_list ) - # mode, class and in-chart selections (also work with shft+click for multi) - selection = altair.selection_point(fields=["id"], on="mouseover") - selection_mode = altair.selection_point(fields=["mode"]) - selection_class = altair.selection_point(fields=["id"]) + # Convert confusion matrix to a DataFrame to be saved as csv file + cm_df = pd.DataFrame(cm) - # in-chart color condition - color = altair.condition( - (selection & selection_mode & selection_class), - altair.Color("id:N", legend=None), - altair.value("lightgray"), + train_support = cm.sum(axis=1) + supported_train_classes = train_support > 0 + avg_accuracy = np.divide( + cm.diagonal(), + train_support, + out=np.zeros_like(train_support, dtype=float), + where=supported_train_classes, ) - # tooltip disla on-mouseover - tooltip = ["id", "meta", "mode", "image"] # .append(opts) - - # main scatter plot - scatter = ( - altair.Chart(df, title="shift+click for multi-select") - .mark_point(size=100, opacity=0.5, filled=True) - .transform_fold(fold=opts, as_=["column", "value"]) - .transform_filter(column_select) - .encode( - altair.X("emb-x", title=titlex), - altair.Y("emb-y", title=titley), - altair.Shape( - "mode", - scale=altair.Scale( - range=["circle", "square", "triangle", "diamond"] - ), - legend=None, - ), - altair.Tooltip(tooltip), - color=color, - size=altair.Size("value:Q", legend=None, aggregate="mean").scale( - type="log" - ), + # Normalize confusion matrix + cmn = ( + np.divide( + cm.astype(float), + train_support[:, np.newaxis], + out=np.zeros_like(cm, dtype=float), + where=supported_train_classes[:, np.newaxis], ) - .interactive() - .properties(width=800, height=500) - .add_params(selection) - .add_params(selection_class) - .add_params(selection_mode) - .add_params(column_select) - ) - - # interactive class legend - legend_class = ( - altair.Chart(df, title="Class") - .mark_point(size=100, opacity=0.5, filled=True) - .encode( - y=altair.Y("id:N", axis=altair.Axis(title=None, orient="right")), - color=altair.condition( - selection_class, - altair.Color("id:N"), - altair.value("lightgrey"), - ), - ) - .add_params(selection_class) - ) - - # interactive mode legend - legend_mode = ( - altair.Chart(df, title="Mode") - .mark_point(size=100, opacity=0.5, filled=True) - .encode( - y=altair.Y("mode:N", axis=altair.Axis(title=None, orient="right")), - shape=altair.Shape( - "mode", - scale=altair.Scale( - range=["circle", "square", "triangle", "diamond"] - ), - legend=None, - ), - color=altair.condition( - selection_mode, - altair.value("steelblue"), - altair.value("lightgrey"), - ), - ) - .add_params(selection_mode) - ) - - # organise charts in window and configure fonts - # class/mode selections are shared across the scatter and legend views; - # Altair merges the duplicates and warns benignly, so silence that message. - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message="Automatically deduplicated selection parameter.*", - ) - chart = ( - (scatter | altair.vconcat(legend_class, legend_mode)) - .configure_axis(labelFontSize=20, titleFontSize=20) - .configure_legend(labelFontSize=20, titleFontSize=20) - .configure_title(fontSize=20) - ) - - # save charts and latent embedding - if not os.path.exists("latents"): - os.mkdir("latents") - # save latentspace and ids - chart.save(f"latents/latent_epoch_{epoch + 1}_{mode}.html") - - -def confidence_plot(x, y, s, suffix=None, vis_format="png"): - logging.info( - "################################################################", - ) - logging.info( - "Visualising class-average confidence metrics " + suffix + "...\n", - ) - cmap = plt.get_cmap("jet") - cols = [cmap(i) for i in np.linspace(0, 1, len(x[0]))] - classes = np.unique(y) - _, font_size, title_font_size = _class_plot_scale(len(classes)) - rows = len(classes) // 2 - if len(classes) % 2 != 0: - rows += 1 - fig, ax = plt.subplots(len(classes), sharex=True, sharey=True) - ax = np.atleast_1d(ax) - for c, cl in enumerate(classes): - mu_cl = np.take(x, np.where(np.array(y) == cl)[0], axis=0) - var_cl = np.take(s, np.where(np.array(y) == cl)[0], axis=0) - std_cl = np.exp(0.5 * var_cl) - mu_cl = np.mean(mu_cl, axis=0) - std_cl = np.mean(std_cl, axis=0) - - min_mu = np.min(mu_cl) - max_mu = np.max(mu_cl) - max_sig = np.max(std_cl) - step = (2 * 4 * max_sig) / 100 - - xs = np.arange(min_mu - (4 * max_sig), max_mu + (4 * max_sig), step) - - for i in range(len(mu_cl)): - ax[c].plot( - xs, - scipy.stats.norm.pdf(xs, mu_cl[i], std_cl[i]), - color=cols[i], - label="lat" + str(i + 1), - ) - ax[c].set_title(cl, fontsize=title_font_size, fontweight="normal") - name = f"plots/confidence.{vis_format}" - if suffix is not None: - name = name[:-4] + "_" + suffix + name[-4:] - handles, labels = ax[-1].get_legend_handles_labels() - leg = fig.legend( - handles, labels, bbox_to_anchor=(1.06, 0.9) - ) # , loc="upper left") - plt.tight_layout() - fig.savefig(name, bbox_extra_artists=(leg,), bbox_inches="tight") - plt.close() - - -def accuracy_plot( - y_train: npt.NDArray, - ypred_train: npt.NDArray, - y_val: npt.NDArray, - ypred_val: npt.NDArray, - classes: str | None = None, - mode: str = "", - epoch: int = 0, - writer: typing.Any = None, - vis_format: str = "png", -): - """Plot confusion matrix . - - Parameters - ---------- - y_train: np.array - Training labels. - ypred_train: np.array - Predicted training labels. - y_val: np.array - Validation labels (unseen data). - ypred_val: np.array - Predicted validation labels (unseen data). - mode: str - Added data mode to the name of the saved figures (e.g train, valid, eval). - epoch: int - Current epoch. - writer: SummaryWriter - Tensorboard summary writer - """ - logging.info( - "################################################################", - ) - - logging.info("Visualising accuracy: confusion and F1 scores ...\n") - - classes_list = np.unique(np.concatenate((y_train, ypred_train))) - - fig_size, font_size, title_font_size = _class_plot_scale(len(classes_list)) - value_font_size = min(font_size, 10) - - # Compute confusion matrix - cm = sklearn.metrics.confusion_matrix( - y_train, ypred_train, labels=classes_list - ) - - # Convert confusion matrix to a DataFrame to be saved as csv file - cm_df = pd.DataFrame(cm) - - disp = sklearn.metrics.ConfusionMatrixDisplay( - confusion_matrix=cm, display_labels=classes_list - ) - - train_support = cm.sum(axis=1) - supported_train_classes = train_support > 0 - avg_accuracy = np.divide( - cm.diagonal(), - train_support, - out=np.zeros_like(train_support, dtype=float), - where=supported_train_classes, - ) - - # Normalize confusion matrix - cmn = ( - np.divide( - cm.astype(float), - train_support[:, np.newaxis], - out=np.zeros_like(cm, dtype=float), - where=supported_train_classes[:, np.newaxis], - ) - * 100 + * 100 ) # Convert normalised confusion matrix to a DataFrame to be saved as csv file cmn_df = pd.DataFrame(cmn) - dispn = sklearn.metrics.ConfusionMatrixDisplay( - confusion_matrix=cmn, display_labels=classes_list - ) - with plt.rc_context({"font.weight": "bold", "font.size": font_size}): - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) + train_title = "Balanced accuracy at epoch {}: {:.1f}%".format( + epoch + 1, np.mean(avg_accuracy[supported_train_classes]) * 100 + ) - disp.plot( + _matrix_figure( + fig_size, + font_size, + title_font_size, + train_title, + f"plots/confusion_train{mode}.{vis_format}", + vis_format, + vis_print, + data=cm, + class_labels=classes_list, cmap=plt.cm.Blues, - ax=ax, - xticks_rotation=90, - values_format="d", - colorbar=False, - ) - divider = make_axes_locatable(ax) - colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) - fig.colorbar(disp.im_, cax=colorbar_ax) - colorbar_ax.tick_params(labelsize=font_size) - plt.setp(disp.text_.ravel(), fontsize=value_font_size) - - ax.set_title( - "Average accuracy at epoch {}: {:.3f}%".format( - epoch + 1, - np.mean(avg_accuracy[supported_train_classes]) * 100, - ), - fontsize=title_font_size, - fontweight="normal", + values_format="d" if vis_print else None, + value_font_size=value_font_size, + xlabel="Predicted label", + ylabel="True label", ) - if not os.path.exists("plots"): - os.mkdir("plots") - - fig.tight_layout() - plt.savefig(f"plots/confusion_train{mode}.{vis_format}", dpi=300) - - # Save confusion matrix DataFrame to CSV - cm_df.to_csv(f'plots/confusion_train{mode}.csv', index=False) - - # Save normalised confusion matrix DataFrame to CSV - cmn_df.to_csv(f'plots/confusion_train{mode}_norm.csv', index=False) - - if writer: - writer.add_figure("Accuracy", fig, epoch) - - plt.close() - - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - - dispn.plot( + _matrix_figure( + fig_size, + font_size, + title_font_size, + train_title, + f"plots/confusion_train{mode}_norm.{vis_format}", + vis_format, + vis_print, + data=cmn, + class_labels=classes_list, cmap=plt.cm.Blues, - ax=ax, - xticks_rotation=90, - values_format=".0f", - im_kw={"vmin": 0, "vmax": 100}, - colorbar=False, - ) - divider = make_axes_locatable(ax) - colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) - fig.colorbar(dispn.im_, cax=colorbar_ax) - colorbar_ax.tick_params(labelsize=font_size) - plt.setp(dispn.text_.ravel(), fontsize=value_font_size) - - ax.set_title( - "Average accuracy at epoch {}: {:.3f}%".format( - epoch + 1, - np.mean(avg_accuracy[supported_train_classes]) * 100, - ), - fontsize=title_font_size, - fontweight="normal", + vmin=0, + vmax=100, + values_format=".0f" if vis_print else None, + value_font_size=value_font_size, + xlabel="Predicted label (%)", + ylabel="True label (%)", ) - ax.set_xlabel("Predicted label (%)") - ax.set_ylabel("True label (%)") - fig.tight_layout() - plt.savefig(f"plots/confusion_train{mode}_norm.{vis_format}", dpi=300) - - if writer: - writer.add_figure("Accuracy (Norm)", fig, epoch) + # Save confusion matrix DataFrame to CSV + cm_df.to_csv(f"plots/confusion_train{mode}.csv", index=False) - plt.close() + # Save normalised confusion matrix DataFrame to CSV + cmn_df.to_csv(f"plots/confusion_train{mode}_norm.csv", index=False) classes_list_eval = np.unique(np.concatenate((y_val, ypred_val))) @@ -744,9 +497,6 @@ def accuracy_plot( # Convert confusion matrix to a DataFrame to be saved as csv file cm_eval_df = pd.DataFrame(cm_eval) - disp_eval = sklearn.metrics.ConfusionMatrixDisplay( - confusion_matrix=cm_eval, display_labels=ordered_class_eval - ) eval_support = cm_eval.sum(axis=1) supported_eval_classes = eval_support > 0 if not np.all(supported_eval_classes): @@ -778,14 +528,10 @@ def accuracy_plot( cmn_eval_df = pd.DataFrame(cmn_eval) # Save confusion matrix DataFrame to CSV - cm_eval_df.to_csv(f'plots/confusion_valid{mode}.csv', index=False) + cm_eval_df.to_csv(f"plots/confusion_valid{mode}.csv", index=False) # Save normalised confusion matrix DataFrame to CSV - cmn_eval_df.to_csv(f'plots/confusion_valid{mode}_norm.csv', index=False) - - dispn_eval = sklearn.metrics.ConfusionMatrixDisplay( - confusion_matrix=cmn_eval, display_labels=ordered_class_eval - ) + cmn_eval_df.to_csv(f"plots/confusion_valid{mode}_norm.csv", index=False) if mode == "_eval": figure_name = "plots/confusion_eval" @@ -800,60 +546,45 @@ def accuracy_plot( "font.size": font_size, } ): - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - disp_eval.plot( - cmap=plt.cm.Blues, - ax=ax, - xticks_rotation=90, - values_format="d", - colorbar=False, - ) - divider = make_axes_locatable(ax) - colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) - fig.colorbar(disp_eval.im_, cax=colorbar_ax) - colorbar_ax.tick_params(labelsize=font_size) - plt.setp(disp_eval.text_.ravel(), fontsize=value_font_size) - ax.set_title( - "Average accuracy at epoch {}: {:.1f}%".format( - epoch + 1, - np.mean(avg_accuracy_eval[supported_eval_classes]) * 100, - ), - fontsize=title_font_size, - fontweight="normal", + eval_title = "Balanced accuracy at epoch {}: {:.1f}%".format( + epoch + 1, np.mean(avg_accuracy_eval[supported_eval_classes]) * 100 ) - fig.tight_layout() - plt.savefig(figure_name + f".{vis_format}", dpi=300) - plt.close() - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - - dispn_eval.plot( + _matrix_figure( + fig_size, + font_size, + title_font_size, + eval_title, + figure_name + f".{vis_format}", + vis_format, + vis_print, + data=cm_eval, + class_labels=ordered_class_eval, cmap=plt.cm.Blues, - ax=ax, - xticks_rotation=90, - values_format=".0f", - im_kw={"vmin": 0, "vmax": 100}, - colorbar=False, + values_format="d" if vis_print else None, + value_font_size=value_font_size, + xlabel="Predicted label", + ylabel="True label", ) - divider = make_axes_locatable(ax) - colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) - fig.colorbar(dispn_eval.im_, cax=colorbar_ax) - colorbar_ax.tick_params(labelsize=font_size) - plt.setp(dispn_eval.text_.ravel(), fontsize=value_font_size) - - ax.set_title( - "Average accuracy at epoch {}: {:.1f}% ".format( - epoch + 1, - np.mean(avg_accuracy_eval[supported_eval_classes]) * 100, - ), - fontsize=title_font_size, - fontweight="normal", + + _matrix_figure( + fig_size, + font_size, + title_font_size, + eval_title, + f"{figure_name}_norm.{vis_format}", + vis_format, + vis_print, + data=cmn_eval, + class_labels=ordered_class_eval, + cmap=plt.cm.Blues, + vmin=0, + vmax=100, + values_format=".0f" if vis_print else None, + value_font_size=value_font_size, + xlabel="Predicted label (%)", + ylabel="True label (%)", ) - ax.set_xlabel("Predicted label (%)") - ax.set_ylabel("True label (%)") - fig.tight_layout() - plt.savefig(f"{figure_name}_norm.{vis_format}", dpi=300) - plt.close() def f1_plot( @@ -863,8 +594,8 @@ def f1_plot( ypred_val: npt.NDArray, mode: str = "", epoch: int = 0, - writer: typing.Any = None, vis_format: str = "png", + vis_print: bool = False, ): """Plot F1 values using classes observed in the training labels. @@ -884,8 +615,6 @@ def f1_plot( Added data mode to the name of the saved figures (e.g train, valid, eval). epoch: int Epoch number. - writer: SummaryWriter - Tensorboard summary writer """ logging.info( "################################################################", @@ -980,281 +709,658 @@ def f1_plot( ) plt.ylabel("F1 Score") plt.tight_layout() - plt.savefig(f"plots/f1{mode}.{vis_format}", dpi=150) - - if writer: - writer.add_figure("F1 score", fig, epoch) + plt.savefig( + f"plots/f1{mode}.{vis_format}", + **({"dpi": 300} if vis_print else {}), + ) plt.close() -def loss_plot( - epochs: int, - beta: float, - gamma: float, - train_loss: list[float], - val_loss: list[float] | None = None, - p: list | None = None, +def recon_plot( + img: torch.Tensor, + rec: torch.Tensor, + label: list, + data_dim: int, + mode: str = "trn", + display: bool = False, vis_format: str = "png", + vis_print: bool = False, ) -> None: - """Visualise loss over epochs. + """Visualise reconstructions. Parameters ---------- - epochs: int - Number of epochs. - beta: list[float] - List of beta values. - gamma: list[float] - List of gamma values. - train_loss: list - Training loss over epochs. - val_loss: list - Validation loss over epochs. - p: list - List of 7 hyperparameters: batch size, depth, " - "channel init, latent dimension, learning rate, beta, gamma. - """ - logging.info( - "################################################################", - ) - logging.info("Visualising loss ...\n") + img: torch.Tensor + Input images. + rec: torch.Tensor + Reconstructed images. + label: list + List of labels for reconstructed images. + data_dim: int + Dimensionality of the data. + mode: str + Type of image in the training set: trn or val. + display: bool + When this variable is set to true, the save_imshow_png function only dispalys the plot + and does not save a png image. This is to allow the use of this function in jupyter notebook + post calculation. This features is not yet implemented for save_mrc_file. + """ + logging.info( + "################################################################", + ) + logging.info("Visualising reconstructions " + mode + "...\n") - train_loss = np.transpose(np.asarray(train_loss)) - if val_loss is not None: - val_loss = np.transpose(np.asarray(val_loss)) + fname_in = f"recon_{mode}_in.{vis_format}" + fname_out = f"recon_{mode}_out.{vis_format}" - cols = ["blue", "red", "green", "orange"] - labs = [ - "Total loss", - "Reconstruction loss", - "KL divergence loss x BETA", - "Affinity loss x GAMMA", - ] - vlabs = [ - "VAL Total loss", - "VAL Reconstruction loss", - "VAL KL divergence loss x BETA", - "VAL Affinity loss x GAMMA", - ] - title_font_size = 16 - label_font_size = 16 - tick_font_size = 14 + if data_dim == 3: + img_2d = img[:, :, img.shape[2] // 2, :, :] + rec_2d = rec[:, :, rec.shape[2] // 2, :, :] + elif data_dim == 2: + img_2d = img + rec_2d = rec - fig, ax = plt.subplots(figsize=(10, 7)) - ax.ticklabel_format(useOffset=False) + img_2d = torchvision.utils.make_grid(img_2d.cpu(), 10, 2).numpy() + rec_2d = torchvision.utils.make_grid(rec_2d.detach().cpu(), 10, 2).numpy() - train_loss[-2] = train_loss[-2] * beta - train_loss[-1] = train_loss[-1] * gamma + save_imshow_png( + fname_in, + np.transpose(img_2d, (1, 2, 0)), + display=display, + vis_print=vis_print, + ) + save_imshow_png( + fname_out, + np.transpose(rec_2d, (1, 2, 0)), + display=display, + vis_print=vis_print, + ) - for i, loss in enumerate(train_loss): - s = "-" - ax.plot( - range(1, epochs + 1), loss, c=cols[i], linestyle=s, label=labs[i] + if data_dim == 3: + rec = rec.detach().cpu().numpy() + img = img.detach().cpu().numpy() + labels = np.asarray(label)[-len(img) :] + dsize = rec.shape[-data_dim:] + + # The number of reconstruction and input images to be displayed in the .mrc output file + number_of_random_samples = 10 + number_of_columns = 3 + padding = 0 + + if len(labels) < number_of_random_samples * number_of_columns: + # In there are not enough images for the stack, do one column only + number_of_random_samples = min( + number_of_random_samples, len(labels) + ) + number_of_columns = 1 + + # define the dimensions for the napari grid + grid_for_napari = np.zeros( + ( + number_of_random_samples * dsize[0], + 2 * dsize[1] * number_of_columns + padding * number_of_columns, + dsize[2], + ), + dtype=np.float32, ) - if val_loss is not None: - val_loss[-2] = val_loss[-2] * beta - val_loss[-1] = val_loss[-1] * gamma - for i, loss in enumerate(val_loss): - s = "--" - ax.plot( - range(1, epochs + 1), - loss, - c=cols[i], - linestyle=s, - label=vlabs[i], + + reconstruction_index: list[dict[str, int | str]] = [] + for k in range(number_of_columns): + # select 10 images at random + selected_indices = np.random.choice( + len(img), size=number_of_random_samples, replace=False ) + selected_img = img[selected_indices] + selected_rec = rec[selected_indices] - if p is not None: - if len(p) != 7: - logging.warning( - "\n\nWARNING: Function vis.loss_plot is expecting 'p' parameter " - "to be a list of 7 hyperparameters: batch size, depth, " - "channel init, latent dimension, learning rate, beta, gamma. " - "Exiting.\n", + # stack the images together with their reconstruction + rec_img = np.hstack((selected_img, selected_rec)) + + # Create and save the mrc file with single transversals + for j in range(2): + for i in range(number_of_random_samples): + grid_for_napari[ + i * dsize[0] : (i + 1) * dsize[0], + j * dsize[1] + + (dsize[1] * 2 + padding) * k : (j + 1) * dsize[1] + + (dsize[1] * 2 + padding) * k, + :, + ] = rec_img[i, j, :, :, :] + + reconstruction_index.extend( + { + "grid_column": k, + "grid_row": i, + "batch_index": int(source_index), + "label": labels[source_index], + } + for i, source_index in enumerate(selected_indices) ) - return - ax.set_title( - "bs: %d, d: %d, ch: %d, lat: %d, lr: %.3f, beta: %.1f, " - "gamma: %.1f" % (p[0], p[1], p[2], p[3], p[4], p[5], p[6]), - fontsize=title_font_size, - fontweight="normal", + + save_mrc_file(f"recon_{mode}.mrc", grid_for_napari) + pd.DataFrame(reconstruction_index).to_csv( + f"plots/recon_{mode}.csv", index=False + ) + + +def latent_embed_plot_tsne( + xs: npt.NDArray, + ys: npt.NDArray, + classes: list | None = None, + mode: str = "", + epoch: int = 0, + perplexity: int = 40, + marker_size: int = 24, + l_w: int = 2, + display: bool = False, + vis_format: str = "png", + vis_print: bool = False, + embedding: npt.NDArray | None = None, +) -> None: + """Plot static TSNE embedding. + + Parameters + ---------- + xs: np.ndarray + Array of latent vectors. + ys: np.ndarray + Array of labels. + classes: list + List of classes. + mode: str + Added data mode to the name of the saved figure (e.g train, valid, eval). + epoch: int + Current epoch + perplexity: int + display: bool + When this variable is set to true, the save_imshow_png function only dispalys the plot + and does not save a png image. This is to allow the use of this function in jupyter notebook + post calculation. This features is not yet implemented for save_mrc_file. + vis_format: str + format of the image saved + embedding: np.ndarray | None + Precomputed embedding coordinates. If omitted, they are calculated + from ``xs``. + """ + logging.info( + "################################################################", + ) + if not mode: + logging.info("Visualising static TSNE embedding...\n") + else: + logging.info("Visualising static TSNE embedding " + mode + "...\n") + + xs = np.asarray(xs) + ys = np.asarray(ys) + + if xs.ndim != 2: + raise ValueError("Embedding only accepts 2D arrays.") + if xs.shape[-1] == 1: + logging.warning( + "Data contains 1 dimension, cannot create embedding," + " plotting histogram instead...\n" + ) + if xs.shape[-1] == 2: + logging.warning( + "Data already contains 2 dimensions, cannot create" + " embedding, plotting scatter of original data...\n" + ) + + lats = ( + utils_learning.tsne_embedding(xs, perplexity=perplexity) + if embedding is None + else np.asarray(embedding) + ) + if len(lats) != len(xs): + raise ValueError( + "Embedding and latent vectors must have equal lengths." + ) + + if classes is None: + classes = sorted(list(np.unique(ys))) + else: + if np.setdiff1d(ys, classes).size > 0: + classes = list( + np.concatenate((classes, np.setdiff1d(ys, classes))) + ) + + n_classes = len(classes) + fig_size, font_size, title_font_size = _class_plot_scale(n_classes) + + if n_classes < 3: + # If the number of classes are not moe than 3 the size of the figure would be too + # small and matplotlib would through a singularity error + fig, ax = plt.subplots( + figsize=(int(n_classes / 2) + 7, int(n_classes / 2) + 5) + ) + else: + fig, ax = plt.subplots( + figsize=(int(n_classes / 2) + 4, int(n_classes / 2) + 2) + ) + marker_size = max(marker_size, int(marker_size * max(1, n_classes / 5))) + # When the number of classes is less than 3 the image becomes two small + colours = colour_per_class(classes) + + if xs.shape[-1] != 1: + + for mol_id, mol in enumerate(set(ys.tolist())): + idx = np.where(np.array(ys.tolist()) == mol)[0] + + color = colours[classes.index(mol)] + + plt.scatter( + lats[idx, 0], + lats[idx, 1], + s=marker_size, + label=mol[:4], + facecolor=color, + edgecolor=color, + alpha=0.5, + ) + + ax.legend( + bbox_to_anchor=(1.05, 1), + loc="upper left", + fontsize=font_size, + ) + ax.set_xlabel("TSNE-1", fontsize=font_size) + ax.set_ylabel("TSNE-2", fontsize=font_size) + + if xs.shape[-1] == 1: + + for mol_id, mol in enumerate(set(ys.tolist())): + idx = np.where(np.array(ys.tolist()) == mol)[0] + cols = colours[classes.index(mol)] + plt.hist( + lats[idx], + 100, + color=cols, + histtype="step", + stacked=True, + fill=False, + label=mol[:4], + linewidth=l_w, + ) + ax.legend( + bbox_to_anchor=(1.05, 1), + loc="upper left", + fontsize=font_size, + ) + ax.set_xlabel("dim 1", fontsize=font_size) + ax.set_ylabel("freq", fontsize=font_size) + + ax.set_title( + f"t-SNE Embedding at epoch: {epoch + 1}", + fontsize=title_font_size, + fontweight="normal", + ) + ax.tick_params(axis="both", labelsize=font_size) + fig.tight_layout() + + if not display: + if not os.path.exists("plots"): + os.mkdir("plots") + plt.savefig( + f"plots/embedding_TSNE{mode}.{vis_format}", + **({"dpi": 300} if vis_print else {}), + ) + else: + plt.show() + + plt.close() + + +def dyn_latentembed_plot( + df: pd.DataFrame, + epoch: int, + mode: str = "", + embedding: npt.NDArray | None = None, +): + """Plot dynamic TSNE embedding. + + Parameters + ---------- + df: pd.DataFrame + Dataframe containing the latent vectors. + epoch: int + Current epoch. + mode: str + Added data mode to the name of the saved figure (e.g train, valid, eval). + embedding: np.ndarray | None + Precomputed embedding coordinates. If omitted, they are calculated + from the latent columns in ``df``. + """ + logging.info( + "################################################################", + ) + logging.info("Visualising dynamic embedding...\n") + + latentspace = df[[col for col in df if col.startswith("lat")]].to_numpy() + lat_emb = np.asarray( + utils_learning.tsne_embedding(latentspace) + if embedding is None + else embedding + ) + if len(lat_emb) != len(df): + raise ValueError("Embedding and metadata must have equal lengths.") + if lat_emb.shape[1] == 1: + lat_emb = np.column_stack([lat_emb[:, 0], np.zeros(len(lat_emb))]) + titlex = "t-SNE-1" + titley = "t-SNE-2" + df["emb-x"], df["emb-y"] = np.array(lat_emb)[:, 0], np.array(lat_emb)[:, 1] + + # create column select options for radio buttons + # add no std to radio select options - change value here for marker size! + if "std-off" not in df.columns: + df.insert(loc=0, column="std-off", value=np.zeros(len(lat_emb)) + 0.5) + # add certainty averaged across all latent dimensions as an extra option + std_dim_cols = [ + col + for col in df.columns + if col.startswith("std-") and col not in ("std-off", "std-avg") + ] + if std_dim_cols and "std-avg" not in df.columns: + df["std-avg"] = df[std_dim_cols].mean(axis=1) + opts = [col for col in df.columns if col.startswith("std")] + + # create radio buttons and bind to a folded column select + bind_checkbox = altair.binding_radio( + options=opts, + labels=[ + "off" + if i == "std-off" + else "avg" + if i == "std-avg" + else str(int(i.split("-")[-1]) + 1) + for i in opts + ], + name="Certainty of prediction per dimension:", + ) + column_select = altair.selection_point( + fields=["column"], + bind=bind_checkbox, + value=[{"column": "std-off"}], + name="certainty", + ) + + # mode, class and in-chart selections (also work with shft+click for multi) + selection = altair.selection_point(fields=["id"], on="mouseover") + selection_mode = altair.selection_point(fields=["mode"]) + selection_class = altair.selection_point(fields=["id"]) + + # in-chart color condition + color = altair.condition( + (selection & selection_mode & selection_class), + altair.Color("id:N", legend=None), + altair.value("lightgray"), + ) + + # tooltip disla on-mouseover + tooltip = ["id", "meta", "mode", "image"] # .append(opts) + + # main scatter plot + scatter = ( + altair.Chart(df, title="shift+click for multi-select") + .mark_point(size=100, opacity=0.5, filled=True) + .transform_fold(fold=opts, as_=["column", "value"]) + .transform_filter(column_select) + .encode( + altair.X("emb-x", title=titlex), + altair.Y("emb-y", title=titley), + altair.Shape( + "mode", + scale=altair.Scale( + range=["circle", "square", "triangle", "diamond"] + ), + legend=None, + ), + altair.Tooltip(tooltip), + color=color, + size=altair.Size("value:Q", legend=None, aggregate="mean").scale( + type="log" + ), + ) + .interactive() + .properties(width=800, height=500) + .add_params(selection) + .add_params(selection_class) + .add_params(selection_mode) + .add_params(column_select) + ) + + # interactive class legend + legend_class = ( + altair.Chart(df, title="Class") + .mark_point(size=100, opacity=0.5, filled=True) + .encode( + y=altair.Y("id:N", axis=altair.Axis(title=None, orient="right")), + color=altair.condition( + selection_class, + altair.Color("id:N"), + altair.value("lightgrey"), + ), + ) + .add_params(selection_class) + ) + + # interactive mode legend + legend_mode = ( + altair.Chart(df, title="Mode") + .mark_point(size=100, opacity=0.5, filled=True) + .encode( + y=altair.Y("mode:N", axis=altair.Axis(title=None, orient="right")), + shape=altair.Shape( + "mode", + scale=altair.Scale( + range=["circle", "square", "triangle", "diamond"] + ), + legend=None, + ), + color=altair.condition( + selection_mode, + altair.value("steelblue"), + altair.value("lightgrey"), + ), + ) + .add_params(selection_mode) + ) + + # organise charts in window and configure fonts + # class/mode selections are shared across the scatter and legend views; + # Altair merges the duplicates and warns benignly, so silence that message. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="Automatically deduplicated selection parameter.*", + ) + chart = ( + (scatter | altair.vconcat(legend_class, legend_mode)) + .configure_axis(labelFontSize=20, titleFontSize=20) + .configure_legend(labelFontSize=20, titleFontSize=20) + .configure_title(fontSize=20) ) - else: - ax.set_title("Loss", fontsize=title_font_size, fontweight="normal") - ax.set_yscale("log") - ax.set_ylabel("Loss", fontsize=label_font_size) - ax.set_xlabel("Epochs", fontsize=label_font_size) - ax.tick_params(axis="both", labelsize=tick_font_size) - ax.legend(fontsize=tick_font_size) + # save charts and latent embedding + if not os.path.exists("latents"): + os.mkdir("latents") + # save latentspace and ids + chart.save(f"latents/latent_epoch_{epoch + 1}_{mode}.html") - fig.tight_layout() - if not os.path.exists("plots"): - os.mkdir("plots") - fig.savefig(f"plots/loss.{vis_format}", dpi=300) - plt.close(fig) - # plotting only the total loss as it sometimes is a few order of magnitude higher than KLD and affinity losses - fig, ax = plt.subplots(figsize=(10, 7)) - ax.plot( - range(1, epochs + 1), - train_loss[0], - c=cols[0], - linestyle="-", - label=labs[0], +def confidence_plot(x, y, s, suffix=None, vis_format="png", vis_print=False): + logging.info( + "################################################################", ) - if val_loss is not None: - ax.plot( - range(1, epochs + 1), - val_loss[0], - c=cols[0], - linestyle="--", - label=vlabs[0], - ) - ax.set_yscale("log") - ax.set_ylabel("Loss", fontsize=label_font_size) - ax.set_xlabel("Epochs", fontsize=label_font_size) - ax.tick_params(axis="both", labelsize=tick_font_size) - ax.legend(fontsize=tick_font_size) - fig.tight_layout() - fig.savefig(f"plots/loss_total.{vis_format}", dpi=300) - plt.close(fig) + logging.info( + "Visualising class-average confidence metrics " + suffix + "...\n", + ) + cmap = plt.get_cmap("jet") + cols = [cmap(i) for i in np.linspace(0, 1, len(x[0]))] + classes = np.unique(y) + _, font_size, title_font_size = _class_plot_scale(len(classes)) + rows = len(classes) // 2 + if len(classes) % 2 != 0: + rows += 1 + fig, ax = plt.subplots(len(classes), sharex=True, sharey=True) + ax = np.atleast_1d(ax) + for c, cl in enumerate(classes): + mu_cl = np.take(x, np.where(np.array(y) == cl)[0], axis=0) + var_cl = np.take(s, np.where(np.array(y) == cl)[0], axis=0) + std_cl = np.exp(0.5 * var_cl) + mu_cl = np.mean(mu_cl, axis=0) + std_cl = np.mean(std_cl, axis=0) + + min_mu = np.min(mu_cl) + max_mu = np.max(mu_cl) + max_sig = np.max(std_cl) + step = (2 * 4 * max_sig) / 100 + xs = np.arange(min_mu - (4 * max_sig), max_mu + (4 * max_sig), step) -def recon_plot( - img: torch.Tensor, - rec: torch.Tensor, - label: list, - data_dim: int, - mode: str = "trn", + for i in range(len(mu_cl)): + ax[c].plot( + xs, + scipy.stats.norm.pdf(xs, mu_cl[i], std_cl[i]), + color=cols[i], + label="lat" + str(i + 1), + ) + ax[c].set_title(cl, fontsize=title_font_size, fontweight="normal") + name = f"plots/confidence.{vis_format}" + if suffix is not None: + name = name[:-4] + "_" + suffix + name[-4:] + handles, labels = ax[-1].get_legend_handles_labels() + leg = fig.legend( + handles, labels, bbox_to_anchor=(1.06, 0.9) + ) # , loc="upper left") + plt.tight_layout() + fig.savefig( + name, + **({"dpi": 300} if vis_print else {}), + bbox_extra_artists=(leg,), + bbox_inches="tight", + ) + plt.close() + + +def latent_space_similarity_plot( + latent_space: npt.NDArray, + class_labels: npt.NDArray, + mode: str = "", epoch: int = 0, - writer: typing.Any = None, + affinity_matrix: pathlib.Path | None = None, + plot_mode: str = "mean", display: bool = False, + font_size: int = 16, + fig_size: int | None = None, + dpi: int = 300, vis_format: str = "png", + vis_print: bool = False, ) -> None: - """Visualise reconstructions. - + """ + This function calculates the similarity (affinity) between classes in the latent space and builds a matrix. Parameters ---------- - img: torch.Tensor - Input images. - rec: torch.Tensor - Reconstructed images. - label: list - List of labels for reconstructed images. - data_dim: int - Dimensionality of the data. + latent_space: np.ndarray + The latent space + class_labels: np.array + The labels of the latent space mode: str - Type of image in the training set: trn or val. + Mode of the calculation (train, test, val) epoch: int - Current epoch - writer: SummaryWriter - Tensorboard summary writer - display: bool - When this variable is set to true, the save_imshow_png function only dispalys the plot - and does not save a png image. This is to allow the use of this function in jupyter notebook - post calculation. This features is not yet implemented for save_mrc_file. + Epoch number for title + affinity_matrix: pathlib.Path | None + File path to affinity matrix + vis_format: str + format of the image saved """ logging.info( "################################################################", ) - logging.info("Visualising reconstructions " + mode + "...\n") - - fname_in = f"recon_{mode}_in.{vis_format}" - fname_out = f"recon_{mode}_out.{vis_format}" - - if data_dim == 3: - img_2d = img[:, :, img.shape[2] // 2, :, :] - rec_2d = rec[:, :, rec.shape[2] // 2, :, :] - elif data_dim == 2: - img_2d = img - rec_2d = rec - - img_2d = torchvision.utils.make_grid(img_2d.cpu(), 10, 2).numpy() - rec_2d = torchvision.utils.make_grid(rec_2d.detach().cpu(), 10, 2).numpy() - - save_imshow_png( - fname_in, - np.transpose(img_2d, (1, 2, 0)), - writer=writer, - figname="Recon. (In)", - epoch=epoch, - display=display, - ) - save_imshow_png( - fname_out, - np.transpose(rec_2d, (1, 2, 0)), - writer=writer, - figname="Recon. (Out)", - epoch=epoch, - display=display, - ) - - if data_dim == 3: - rec = rec.detach().cpu().numpy() - img = img.detach().cpu().numpy() - labels = np.asarray(label)[-len(img) :] - dsize = rec.shape[-data_dim:] - - # The number of reconstruction and input images to be displayed in the .mrc output file - number_of_random_samples = 10 - number_of_columns = 3 - padding = 0 - - if len(labels) < number_of_random_samples * number_of_columns: - # In there are not enough images for the stack, do one column only - number_of_random_samples = min( - number_of_random_samples, len(labels) - ) - number_of_columns = 1 + logging.info("Visualising the latent space similarity matrix ...\n") - # define the dimensions for the napari grid - grid_for_napari = np.zeros( - ( - number_of_random_samples * dsize[0], - 2 * dsize[1] * number_of_columns + padding * number_of_columns, - dsize[2], - ), - dtype=np.float32, + if affinity_matrix is None: + unique_classes = np.unique(class_labels) + else: + classes_order = ( + pd.read_csv(affinity_matrix, header=0).columns.astype(str).tolist() ) - - reconstruction_index: list[dict[str, int | str]] = [] - for k in range(number_of_columns): - # select 10 images at random - selected_indices = np.random.choice( - len(img), size=number_of_random_samples, replace=False + unique_classes_in_data = np.unique(class_labels) + if np.setdiff1d(unique_classes_in_data, classes_order).size > 0: + unique_classes = np.concatenate( + ( + classes_order, + np.setdiff1d(unique_classes_in_data, classes_order), + ) ) - selected_img = img[selected_indices] - selected_rec = rec[selected_indices] + else: + unique_classes = classes_order - # stack the images together with their reconstruction - rec_img = np.hstack((selected_img, selected_rec)) + num_classes = len(unique_classes) + resolved_fig_size, font_size, title_font_size = _class_plot_scale( + num_classes, fig_size + ) + cosine_sim = latent_space_similarity_mat( + latent_space, + class_labels, + unique_classes, + num_classes, + plot_mode=plot_mode, + ) - # Create and save the mrc file with single transversals - for j in range(2): - for i in range(number_of_random_samples): - grid_for_napari[ - i * dsize[0] : (i + 1) * dsize[0], - j * dsize[1] - + (dsize[1] * 2 + padding) * k : (j + 1) * dsize[1] - + (dsize[1] * 2 + padding) * k, - :, - ] = rec_img[i, j, :, :, :] + # Visualize average cosine similarity matrix + classes_in_data = set(np.asarray(class_labels).astype(str)) + highlight = np.array( + [str(c) not in classes_in_data for c in unique_classes] + ) + title = f"Average Cosine Similarity Matrix at epoch: {epoch + 1}" + path = f"plots/similarity_mean{mode}.{vis_format}" - reconstruction_index.extend( - { - "grid_column": k, - "grid_row": i, - "batch_index": int(source_index), - "label": labels[source_index], - } - for i, source_index in enumerate(selected_indices) + if fig_size is None: + with plt.rc_context( + { + "font.weight": "bold", + "font.size": font_size, + } + ): + _matrix_figure( + resolved_fig_size, + font_size, + title_font_size, + title, + path, + vis_format, + vis_print, + print_dpi=dpi, + data=cosine_sim, + class_labels=unique_classes, + cmap="RdBu", + vmin=-1, + vmax=1, + highlight=highlight, + xlabel="Class Labels", + ylabel="Class Labels", + display=display, ) - - save_mrc_file(f"recon_{mode}.mrc", grid_for_napari) - pd.DataFrame(reconstruction_index).to_csv( - f"plots/recon_{mode}.csv", index=False + else: + _matrix_figure( + resolved_fig_size, + font_size, + title_font_size, + title, + path, + vis_format, + vis_print, + print_dpi=dpi, + data=cosine_sim, + class_labels=unique_classes, + cmap="RdBu", + vmin=-1, + vmax=1, + highlight=highlight, + xlabel="Class Labels", + ylabel="Class Labels", + display=display, ) @@ -1268,6 +1374,7 @@ def latent_4enc_interpolate_plot( poses: list, display: bool = False, vis_format: str = "png", + vis_print: bool = False, ) -> None: """Visualise the interpolation of latent space between 4 randomly selected encodings. The number of plots and the number of interpolation steps is modifyable. @@ -1362,30 +1469,18 @@ def latent_4enc_interpolate_plot( np.array(decoded_grid), (num_steps, num_steps, *dsize) ) - grid_for_napari = create_grid_for_plotting( - num_steps, num_steps, dsize, padding - ) - grid_for_napari = fill_grid_for_plottting( + _save_grid_figure( num_steps, num_steps, - grid_for_napari, dsize, decoded_grid, + f"latent_interpolate_{num_fig}", + vis_format, padding, + display=display, + vis_print=vis_print, ) - if data_dim == 3: - save_mrc_file( - f"latent_interpolate_{num_fig}.mrc", - grid_for_napari, - ) - elif data_dim == 2: - save_imshow_png( - f"latent_interpolate_{num_fig}.{vis_format}", - grid_for_napari, - display=display, - ) - def latent_disentamglement_plot( dsize: tuple, @@ -1395,6 +1490,7 @@ def latent_disentamglement_plot( poses: list | None = None, mode: str = "trn", vis_format: str = "png", + vis_print: bool = False, ) -> None: """Visualise latent content disentanglement. @@ -1414,8 +1510,6 @@ def latent_disentamglement_plot( List of pose vectors. mode: str Mode of evaluation (trn: Training, vld: Validation, eval: Evaluation ) - writer: SummaryWriter - Tensorboard summary writer """ logging.info( "################################################################" @@ -1469,26 +1563,17 @@ def latent_disentamglement_plot( recon_images, (latent_dims, number_of_samples, *dsize) ) - grid_for_napari = create_grid_for_plotting( - latent_dims, number_of_samples, dsize, padding - ) - grid_for_napari = fill_grid_for_plottting( + _save_grid_figure( latent_dims, number_of_samples, - grid_for_napari, dsize, recon_images, + f"disentanglement-latent_{mode}", + vis_format, padding, + vis_print=vis_print, ) - if data_dim == 3: - save_mrc_file(f"disentanglement-latent_{mode}.mrc", grid_for_napari) - elif data_dim == 2: - save_imshow_png( - f"disentanglement-latent_{mode}.{vis_format}", - grid_for_napari, - ) - def pose_class_disentanglement_plot( dsize: tuple, @@ -1503,6 +1588,7 @@ def pose_class_disentanglement_plot( specific_enc: npt.NDArray = None, display: bool = False, vis_format: str = "png", + vis_print: bool = False, ): """Visualise Pose interpolation per class. This function creates a pose interpolatoion @@ -1579,27 +1665,18 @@ def pose_class_disentanglement_plot( device, ) - grid_for_napari = create_grid_for_plotting( - pos_dims, number_of_samples, dsize, padding - ) - grid_for_napari = fill_grid_for_plottting( + _save_grid_figure( pos_dims, number_of_samples, - grid_for_napari, dsize, decoded_grid, + f"pose_interpolate_{mode}_{i}", + vis_format, padding, + display=display, + vis_print=vis_print, ) - if data_dim == 3: - save_mrc_file(f"pose_interpolate_{mode}_{i}.mrc", grid_for_napari) - elif data_dim == 2: - save_imshow_png( - f"pose_interpolate_{mode}_{i}.{vis_format}", - grid_for_napari, - display=display, - ) - def pose_disentanglement_plot( dsize: tuple, @@ -1611,6 +1688,7 @@ def pose_disentanglement_plot( mode: str = "trn", display: bool = False, vis_format: str = "png", + vis_print: bool = False, ): """Visualise pose disentanglement. @@ -1668,26 +1746,18 @@ def pose_disentanglement_plot( device, ) - grid_for_napari = create_grid_for_plotting( - pos_dims, number_of_samples, dsize, padding - ) - - # Create and save the mrc file with single transversals - grid_for_napari = fill_grid_for_plottting( - pos_dims, number_of_samples, grid_for_napari, dsize, recon, padding + _save_grid_figure( + pos_dims, + number_of_samples, + dsize, + recon, + f"disentanglement-pose_{mode}_{label}", + vis_format, + padding, + display=display, + vis_print=vis_print, ) - if data_dim == 3: - save_mrc_file( - f"disentanglement-pose_{mode}_{label}.mrc", grid_for_napari - ) - elif data_dim == 2: - save_imshow_png( - f"disentanglement-pose_{mode}_{label}.{vis_format}", - grid_for_napari, - display=display, - ) - def interpolations_plot( dsize: tuple, @@ -1698,6 +1768,7 @@ def interpolations_plot( poses: list | None = None, mode: str = "trn", vis_format: str = "png", + vis_print: bool = False, ) -> None: """Visualise interpolations. @@ -1781,10 +1852,10 @@ def interpolations_plot( if poses is not None: interpolated_pose = ( - (1 - h) * (1 - v) * class_reps_poses[0] - + h * (1 - v) * class_reps_poses[1] - + (1 - h) * v * class_reps_poses[2] - + h * v * class_reps_poses[3] + (1 - h) * (1 - v) * class_rep_poses[0] + + h * (1 - v) * class_rep_poses[1] + + (1 - h) * v * class_rep_poses[2] + + h * v * class_rep_poses[3] ) interpolated_pose_t = torch.from_numpy( np.asarray(interpolated_pose, dtype=np.float32) @@ -1807,16 +1878,15 @@ def interpolations_plot( np.array(decoded_grid), (grid_size, grid_size, *dsize) ) - grid_for_napari = create_grid_for_plotting(grid_size, grid_size, dsize) - grid_for_napari = fill_grid_for_plottting( - grid_size, grid_size, grid_for_napari, dsize, decoded_grid + _save_grid_figure( + grid_size, + grid_size, + dsize, + decoded_grid, + f"interpolations_{mode}", + vis_format, + vis_print=vis_print, ) - # Create an mrc file with interpolations - - if data_dim == 3: - save_mrc_file(f"interpolations_{mode}.mrc", grid_for_napari) - elif data_dim == 2: - save_imshow_png(f"interpolations_{mode}.{vis_format}", grid_for_napari) def plot_affinity_matrix( @@ -1825,6 +1895,7 @@ def plot_affinity_matrix( selected_classes: list, fig_size: int | None = None, vis_format: str = "png", + vis_print: bool = False, ) -> None: """ This function plots the Affinity matrix and highlights the @@ -1841,6 +1912,7 @@ def plot_affinity_matrix( vis_format: str format of the image saved """ + logging.info("\n") logging.info( "################################################################", ) @@ -1850,42 +1922,30 @@ def plot_affinity_matrix( len(all_classes), fig_size ) - fig, ax = plt.subplots(figsize=(resolved_fig_size, resolved_fig_size)) - ax.set_title( + highlight = np.array([c not in selected_classes for c in all_classes]) + + _matrix_figure( + resolved_fig_size, + font_size, + title_font_size, "Affinity Matrix", - fontsize=title_font_size, - fontweight="normal", + f"plots/affinity_matrix.{vis_format}", + vis_format, + vis_print, + data=lookup, + class_labels=all_classes, + cmap="RdBu", + vmin=-1, + vmax=1, + highlight=highlight, ) - im = ax.imshow(lookup, vmin=-1, vmax=1, cmap=plt.get_cmap("RdBu")) - - ax.set_xticks(np.arange(0, len(all_classes))) - ax.set_xticklabels(all_classes) - ax.set_yticks(np.arange(0, len(all_classes))) - ax.set_yticklabels(all_classes) - - # Colour classes absent from the selected dataset. - for i, c in enumerate(all_classes): - if c not in selected_classes: - ax.get_xticklabels()[i].set_color("red") - ax.get_yticklabels()[i].set_color("red") - - ax.tick_params(axis="x", rotation=90, labelsize=font_size) - ax.tick_params(axis="y", labelsize=font_size) - - divider = make_axes_locatable(ax) - colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) - fig.colorbar(im, cax=colorbar_ax) - colorbar_ax.tick_params(labelsize=font_size) - fig.tight_layout() - if not os.path.exists("plots"): - os.mkdir("plots") - plt.savefig(f"plots/affinity_matrix.{vis_format}", dpi=300) - plt.close() - def plot_classes_distribution( - data: list, category: str, vis_format: str = "png" + data: list, + category: str, + vis_format: str = "png", + vis_print: bool = False, ) -> None: """Plot histogram with classes distribution @@ -1920,147 +1980,141 @@ def plot_classes_distribution( fig.tight_layout() if not os.path.exists("plots"): os.mkdir("plots") - plt.savefig(f"plots/classes_distribution_{category}.{vis_format}", dpi=300) + plt.savefig( + f"plots/classes_distribution_{category}.{vis_format}", + **({"dpi": 300} if vis_print else {}), + ) plt.close() -def plot_cyc_variable( - array: list, variable_name: str, vis_format: str = "png" -): - """Plot evolution of variable from the cyclical training - +def _encoder(i: PIL.Image) -> str: + """Encode PIL Image as base64 buffer. Parameters ---------- - array : list - List of values for the variable - variable_name : str - Name of the variable - """ - logging.info( - "################################################################", - ) - logging.info(f"Visualising {variable_name} ...\n") - fig, ax = plt.subplots(figsize=(10, 7)) - ax.plot(array, linewidth=3) - ax.set_ylabel(rf"$\{variable_name}$", fontsize=16) - ax.set_xlabel("Epochs", fontsize=16) - ax.tick_params(axis="both", labelsize=14) - fig.tight_layout() - if not os.path.exists("plots"): - os.mkdir("plots") - fig.savefig(f"plots/hyperaparam_{variable_name}.{vis_format}", dpi=300) - plt.close(fig) + i: PIL.Image + Image to be encoded. + Returns + ------- + str + Encoded image. -def latent_space_similarity_plot( - latent_space: npt.NDArray, - class_labels: npt.NDArray, - mode: str = "", - epoch: int = 0, - affinity_matrix: pathlib.Path | None = None, - plot_mode: str = "mean", - display: bool = False, - font_size: int = 16, - fig_size: int | None = None, - dpi: int = 300, - vis_format: str = "png", -) -> None: """ - This function calculates the similarity (affinity) between classes in the latent space and builds a matrix. + import base64 + import io + + with io.BytesIO() as buffer: + i.thumbnail((110, 110)) + i.save(buffer, "PNG") + data = base64.encodebytes(buffer.getvalue()).decode("utf-8") + + return f"{data}" + + +def _decoder(i: str) -> PIL.Image: + """Decode base64 buffer as PIL Image. Parameters ---------- - latent_space: np.ndarray - The latent space - class_labels: np.array - The labels of the latent space - mode: str - Mode of the calculation (train, test, val) - epoch: int - Epoch number for title - affinity_matrix: pathlib.Path | None - File path to affinity matrix - vis_format: str - format of the image saved + i: str + Encoded image. + + Returns + ------- + PIL.Image + Decoded image. """ - logging.info( - "################################################################", - ) - logging.info("Visualising the latent space similarity matrix ...\n") + import base64 + import io - if affinity_matrix is None: - unique_classes = np.unique(class_labels) + return PIL.Image.open(io.BytesIO(base64.b64decode(i))) + + +def format(im: PIL.Image, data_dim: int) -> list[str]: + """Format PIL Image as Pandas compatible Altair image display. + + Parameters + ---------- + im: PIL.Image + Image to be formatted. + data_dim: int + Dimension of the data. + + Returns + ------- + list[str] + Formatted images compatible with Altair image display. Batch input + returns one encoded item per sample, while single-sample input returns + a one-element list. + """ + if len(im.shape) == 5 and data_dim == 3: + batch = True + im = np.copy( + im[:, :, im.shape[2] // 2, :, :] + .squeeze(dim=1) + .cpu() + .detach() + .numpy() + ) + elif len(im.shape) == 4 and data_dim == 3: + batch = False + im = np.copy( + im[:, im.shape[1] // 2, :, :].squeeze(dim=0).cpu().detach().numpy() + ) + elif len(im.shape) == 4 and data_dim == 2: + batch = True + im = np.copy(im.squeeze(dim=1).cpu().detach().numpy()) + # .astype(np.uint8) + elif len(im.shape) == 3 and data_dim == 2: + batch = False + im = np.copy(im.squeeze(dim=0).cpu().detach().numpy()) + # .astype(np.uint8) else: - classes_order = ( - pd.read_csv(affinity_matrix, header=0).columns.astype(str).tolist() + logging.warning( + "\n\nWARNING: Wrong data format, please pass either a single " + "unsqueezed tensor or a batch to image formatter. Exiting.\n", ) - unique_classes_in_data = np.unique(class_labels) - if np.setdiff1d(unique_classes_in_data, classes_order).size > 0: - unique_classes = np.concatenate( - ( - classes_order, - np.setdiff1d(unique_classes_in_data, classes_order), - ) - ) - else: - unique_classes = classes_order + return [] + im *= 255 + im = im.astype(np.uint8) + if batch: + # if batch we are adding a reconstruction to an input that already + # exists in the DF, '&' is a separator. + return ["&" + _encoder(PIL.Image.fromarray(i)) for i in im] + else: + return [_encoder(PIL.Image.fromarray(im))] - num_classes = len(unique_classes) - resolved_fig_size, font_size, title_font_size = _class_plot_scale( - num_classes, fig_size - ) - cosine_sim = latent_space_similarity_mat( - latent_space, - class_labels, - unique_classes, - num_classes, - plot_mode=plot_mode, - ) - # Visualize average cosine similarity matrix - if fig_size is None: - with plt.rc_context( - { - "font.weight": "bold", - "font.size": font_size, - } - ): - fig, ax = plt.subplots( - figsize=(resolved_fig_size, resolved_fig_size) - ) - else: - fig, ax = plt.subplots(figsize=(resolved_fig_size, resolved_fig_size)) - im = ax.imshow(cosine_sim, cmap="RdBu", vmin=-1, vmax=1) - divider = make_axes_locatable(ax) - colorbar_ax = divider.append_axes("right", size="5%", pad=0.1) - fig.colorbar(im, cax=colorbar_ax) - colorbar_ax.tick_params(labelsize=font_size) - ax.set_xticks(np.arange(num_classes), labels=unique_classes) - ax.set_yticks(np.arange(num_classes), labels=unique_classes) - ax.set_title( - f"Average Cosine Similarity Matrix at epoch: {epoch + 1}", - fontsize=title_font_size, - fontweight="normal", - ) - ax.set_xlabel("Class Labels", fontsize=font_size) - ax.set_ylabel("Class Labels", fontsize=font_size) - ax.tick_params(axis="x", rotation=90, labelsize=font_size) - ax.tick_params(axis="y", labelsize=font_size) - classes_in_data = set(np.asarray(class_labels).astype(str)) - for class_name, x_label, y_label in zip( - unique_classes, ax.get_xticklabels(), ax.get_yticklabels() - ): - if str(class_name) not in classes_in_data: - x_label.set_color("red") - y_label.set_color("red") - fig.tight_layout() +def merge(im: str) -> str | None: + """Merge 2 base64 buffers as PIL Images and encode back to base64 + buffers. - if not display: - if not os.path.exists("plots"): - os.mkdir("plots") - plt.savefig( - f"plots/similarity_mean{mode}.{vis_format}", - dpi=dpi, + Parameters + ---------- + im: str + Input PIL Images to be merged. + + Returns + ------- + str + Merged image. + """ + if im.startswith("data:image/"): + return im + + i = im.split("&") + if len(i) != 2: + logging.warning( + "\n\nWARNING: Image format corrupt. Number of images in meta_df: {}. " + "Exiting. \n".format(len(i)), ) - plt.close() - else: - plt.show() + return None + + im1 = _decoder(i[0]) + im2 = _decoder(i[1]) + + new_image = PIL.Image.new("L", (im2.size[0] + im1.size[0], im2.size[1])) + new_image.paste(im1, (0, 0)) + new_image.paste(im2, (im1.size[0], 0)) + data = _encoder(new_image) + + return f"data:image/png;base64,{data}" diff --git a/pyproject.toml b/pyproject.toml index 0809177b..27131e3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,6 @@ dependencies = [ "requests", "scikit-image", "scikit-learn", - "tensorboard", "umap-learn", "caked@git+https://github.com/alan-turing-institute/caked/" #to be changed when caked is in pypi. ] diff --git a/requirements.txt b/requirements.txt index ac9aa72b..af299284 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,6 @@ requests scikit-image scikit-learn scipy -tensorboard torch torchvision umap-learn diff --git a/tests/test_train_eval_pipeline.py b/tests/test_train_eval_pipeline.py index bf5b2646..6315bc1b 100644 --- a/tests/test_train_eval_pipeline.py +++ b/tests/test_train_eval_pipeline.py @@ -41,7 +41,6 @@ def pipeline_params(datapath): "normalise": True, "shift_min": True, "rescale": 32, - "tensorboard": False, "classifier": "NN", "opt_method": "adam", "gpu_devices": "0", diff --git a/tests/test_vis.py b/tests/test_vis.py index a18a548b..d5f00587 100644 --- a/tests/test_vis.py +++ b/tests/test_vis.py @@ -75,7 +75,7 @@ def test_pipeline_dispatches_visualisations(self): self.data.vis_z_n_int = "1,2" vis_calls = { - "accuracy_plot": 1, + "confusion_plot": 1, "f1_plot": 1, "recon_plot": 2, "latent_space_similarity_plot": 2, @@ -149,7 +149,7 @@ def test_pipeline_dispatches_visualisations(self): # Accuracy is computed once and reused by both accuracy renderers. self.assertEqual(compute_accuracy.call_count, 1) self.assertEqual( - patched_vis["accuracy_plot"].call_args.kwargs["epoch"], 0 + patched_vis["confusion_plot"].call_args.kwargs["epoch"], 0 ) # Separate latent and pose spaces each require one t-SNE calculation. self.assertEqual(compute_tsne.call_count, 2) @@ -223,7 +223,7 @@ def test_loss_plot_writes_both_views(self): # Loss rendering writes both component and total-loss views. self.assertEqual(plot_files, {"loss.png", "loss_total.png"}) - def test_accuracy_plot_handles_class_absent_from_validation(self): + def test_confusion_plot_handles_class_absent_from_validation(self): original_cwd = os.getcwd() with tempfile.TemporaryDirectory(prefix="avae-accuracy-") as temp_dir: @@ -236,7 +236,7 @@ def test_accuracy_plot_handles_class_absent_from_validation(self): ) as plot_title, ): with np.errstate(divide="raise", invalid="raise"): - vis.accuracy_plot( + vis.confusion_plot( np.array(["a", "a", "b", "b"]), np.array(["a", "a", "b", "b"]), np.array(["a", "a"]), diff --git a/tutorials/mnist_data/mnist_config.yml b/tutorials/mnist_data/mnist_config.yml index 80635eae..8e7fc39d 100644 --- a/tutorials/mnist_data/mnist_config.yml +++ b/tutorials/mnist_data/mnist_config.yml @@ -60,7 +60,6 @@ restart: false #### Visualization parameters vis_format: png -tensorboard: false vis_his: false vis_aff: false vis_cyc: false From 535960ff841e0b987648c6f0486f3b86ccaa85da Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 07/12] Unify figure style, saving and matrix plots Saving and layout - Add save_figure as the single save path. Non-print rasters are capped at 720 px tall (top half of a 1440p screen); vis_print saves at 300 dpi and full size; image grids keep native resolution. - Replace tight_layout with margins estimated from the tick labels the axes will draw, with equal outer padding on all sides. This avoids an extra text-layout pass. Style - Apply one rc style (_STYLE) to every plot through @_styled, and drop the scattered bold rc_context blocks. - Use one title format for every epoch-based figure: "X at epoch N". Titles scale with the figure and shrink only when they would not fit. Matrix plots - Build matrix figures around a fixed-size image, so the matrix keeps its size and the padding is even. - Draw the affinity and similarity matrices through one _similarity_figure helper, so they look identical. Remove the "Class Labels" axis labels and the ignored font_size and dpi parameters. - Draw only the normalised confusion matrices; both CSVs are still written. Train and validation share a _confusion_figure helper. - Write per-cell values only for MAX_ANNOTATED_CLASSES (25) classes or fewer. At 53 classes, N^2 text objects took about 14 s per call. Embeddings - Give t-SNE plots an 8:5 data area, matching the interactive embedding, with the legend outside the axes. Stop truncating legend labels to four characters and stop scaling markers with the class count. --- avae/utils.py | 77 ++- avae/vis.py | 805 ++++++++++++++--------------- tests/test_integration_pipeline.py | 4 +- 3 files changed, 441 insertions(+), 445 deletions(-) diff --git a/avae/utils.py b/avae/utils.py index c8c1fbcb..0226c54f 100644 --- a/avae/utils.py +++ b/avae/utils.py @@ -1,6 +1,7 @@ import copy import logging import os.path +import pathlib import matplotlib.pyplot as plt import mrcfile @@ -10,6 +11,9 @@ import sklearn.metrics import torch +SCREEN_HEIGHT_PX = 720 # top half of a 1440p screen +VECTOR_FORMATS = ("pdf", "svg", "eps") + def as_list(value: object) -> list: if isinstance(value, torch.Tensor): @@ -81,6 +85,40 @@ def fill_grid_for_plottting( return grid +def save_figure( + fig: plt.Figure, + path: str, + vis_format: str, + vis_print: bool = False, + display: bool = False, + cap_height: bool = True, +) -> None: + """Save (or, with `display`, only show) a figure and close it. + + Vector formats take no dpi. `vis_print` saves at 300 dpi and full size. + Otherwise, with `cap_height`, the dpi is lowered so the figure is at most + SCREEN_HEIGHT_PX tall, as large figures cost time and space to rasterise. + """ + if display: + plt.show() + plt.close(fig) + return + + pathlib.Path(path).parent.mkdir(parents=True, exist_ok=True) + kwargs: dict = {} + if vis_format not in VECTOR_FORMATS: + dpi = plt.rcParams["figure.dpi"] + if vis_print: + dpi = 300 + elif cap_height: + dpi = min(dpi, SCREEN_HEIGHT_PX / fig.get_figheight()) + kwargs["dpi"] = dpi + if vis_format == "png": + kwargs["pil_kwargs"] = {"compress_level": 3} + fig.savefig(path, **kwargs) + plt.close(fig) + + def save_imshow_png( fname: str, array: npt.NDArray, @@ -90,30 +128,23 @@ def save_imshow_png( display: bool = False, vis_print: bool = False, ) -> None: - if not display: - if not os.path.exists("plots"): - os.mkdir("plots") - - height, width = array.shape[:2] - if width >= height: - image_figure_size = (10, 10 * height / width) - else: - image_figure_size = (10 * width / height, 10) - fig, ax = plt.subplots(figsize=image_figure_size) - ax.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last - ax.axis("off") - fig.subplots_adjust(left=0, right=1, bottom=0, top=1) - fig.savefig( - "plots/" + fname, - **({"dpi": 300} if vis_print else {}), - bbox_inches="tight", - pad_inches=0, - ) + height, width = array.shape[:2] + if width >= height: + image_figure_size = (10, 10 * height / width) else: - plt.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last - plt.show() - - plt.close() + image_figure_size = (10 * width / height, 10) + fig, ax = plt.subplots(figsize=image_figure_size) + ax.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last + ax.axis("off") + fig.subplots_adjust(left=0, right=1, bottom=0, top=1) + save_figure( + fig, + "plots/" + fname, + fname.rsplit(".", 1)[-1], + vis_print, + display, + cap_height=False, + ) def save_mrc_file(fname: str, array: npt.NDArray) -> None: diff --git a/avae/vis.py b/avae/vis.py index 23158435..c8cceb4e 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -1,8 +1,10 @@ import copy +import functools import logging import os.path import pathlib import random +import re import typing import warnings @@ -20,21 +22,117 @@ from . import utils_learning from .utils import ( + VECTOR_FORMATS, colour_per_class, create_grid_for_plotting, fill_grid_for_plottting, latent_space_similarity_mat, pose_interpolation, + save_figure, save_imshow_png, save_mrc_file, ) MAX_CLASS_FIG_SIZE = 20 -VECTOR_FORMATS = ("pdf", "svg", "eps") +MAX_ANNOTATED_CLASSES = 25 +_PAD = 0.28 # outer padding in inches, equal on all four sides +_STYLE: dict[str, typing.Any] = { + "font.weight": "normal", + "axes.titleweight": "normal", + "axes.labelweight": "normal", + "font.size": 14, + "axes.titlesize": 18, + "axes.labelsize": 16, + "xtick.labelsize": 14, + "ytick.labelsize": 14, + "legend.fontsize": 14, +} + + +def _styled(func): + """Run a plotting function under the shared matplotlib style.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + with plt.rc_context(_STYLE): + return func(*args, **kwargs) + + return wrapper + + +def _epoch_title(text: str, epoch: int, detail: str | None = None) -> str: + title = f"{text} at epoch {epoch + 1}" + return f"{title}: {detail}" if detail else title + + +def _ytick_labels(ax) -> list[str]: + """The y tick labels the axis will show, available before rendering.""" + axis = ax.yaxis + labels = axis.get_major_formatter().format_ticks(axis.get_majorticklocs()) + # Strip mathtext markup (log axes) so only the visible glyphs are counted + return [re.sub(r"\\mathdefault|[${}\\^]", "", label) for label in labels] + + +def _set_title(ax, text: str, size: float) -> None: + """Set the title, shrinking its font only if it would be wider than the + room left either side of the axes centre. Call after _fit_margins.""" + centre = (ax.get_position().x0 + ax.get_position().x1) / 2 + width_in = ax.figure.get_size_inches()[0] + available = 2 * (min(centre, 1 - centre) * width_in - _PAD) + width = 0.6 * len(text) * size / 72 + ax.set_title(text, fontsize=size * min(1, available / width)) + + +def _margins( + font_size: float, + *, + title_size: float | None = None, + xlabels: typing.Sequence | npt.NDArray = (), + ylabels: typing.Sequence = (), + x_rotation: int = 0, + xlabel: bool = False, + ylabel: bool = False, + right: float = 0.0, +) -> tuple[float, float, float, float]: + """Estimate the (left, right, bottom, top) margins in inches from text + extents. `right` is extra room for a colorbar's labels or an outside + legend.""" + pad = _PAD + char = 0.65 * font_size / 72 # generous average glyph width in inches + line = 1.5 * font_size / 72 + + def longest(labels): + return max((len(str(label)) for label in labels), default=0) + + tick_x = longest(xlabels) * char + if x_rotation == 90: + bottom = tick_x + elif x_rotation: + bottom = tick_x * np.sin(np.radians(x_rotation)) + line + else: + bottom = line + left = longest(ylabels) * char if len(ylabels) else 4 * char + bottom += pad + 0.1 + (line if xlabel else 0) + left += pad + 0.1 + (line if ylabel else 0) + top = pad + (1.5 * title_size / 72 if title_size else 0) + return left, pad + right, bottom, top + + +def _fit_margins(fig, font_size: float, **kwargs) -> None: + """Set the subplot margins from estimated text extents, avoiding the + extra layout pass of tight_layout (see _margins for the arguments).""" + width, height = fig.get_size_inches() + left, right, bottom, top = _margins(font_size, **kwargs) + fig.subplots_adjust( + left=left / width, + right=1 - right / width, + bottom=bottom / height, + top=1 - top / height, + ) def _class_plot_scale( - num_classes: int, fig_size: int | None = None + num_classes: int, fig_size: float | None = None ) -> tuple[float, int, int]: """Return consistent figure and font sizes for class-based plots. @@ -50,7 +148,8 @@ def _class_plot_scale( ) row_height_pt = 72 * resolved_fig_size / max(num_classes, 1) font_size = int(np.clip(0.7 * row_height_pt, 4, 14)) - return resolved_fig_size, font_size, max(14, font_size + 6) + title_size = max(16, font_size + 6, round(1.5 * resolved_fig_size)) + return resolved_fig_size, font_size, title_size def _loss_curve_figure( @@ -80,60 +179,86 @@ def _loss_curve_figure( linewidth=linewidth, ) - if title: - ax.set_title(title, fontsize=16, fontweight="normal") - if log_scale: ax.set_yscale("log") - ax.set_ylabel(ylabel, fontsize=16) - ax.set_xlabel("Epochs", fontsize=16) - ax.tick_params(axis="both", labelsize=14) + ax.set_ylabel(ylabel) + ax.set_xlabel("Epochs") if legend: - ax.legend(fontsize=14) - - fig.tight_layout() - if not os.path.exists("plots"): - os.mkdir("plots") - fig.savefig(f"{path}.{vis_format}", **({"dpi": 300} if vis_print else {})) - plt.close(fig) + ax.legend() + + _fit_margins( + fig, + 14, + title_size=_STYLE["axes.titlesize"] if title else None, + ylabels=_ytick_labels(ax), + xlabel=True, + ylabel=True, + ) + if title: + _set_title(ax, title, _STYLE["axes.titlesize"]) + save_figure(fig, f"{path}.{vis_format}", vis_format, vis_print) def _matrix_figure( - fig_size: float, - font_size: int, - title_font_size: int, + data: npt.NDArray, + class_labels: typing.Sequence | npt.NDArray, title: str, path: str, vis_format: str = "png", vis_print: bool = False, - print_dpi: int = 300, *, - data: npt.NDArray, - class_labels: typing.Sequence, cmap, vmin: float | None = None, vmax: float | None = None, values_format: str | None = None, - value_font_size: int | None = None, highlight: npt.NDArray | None = None, xlabel: str | None = None, ylabel: str | None = None, + fig_size: float | None = None, display: bool = False, ) -> None: - """Render and save (or display) one - class x class matrix figure: image, optional per-cell values (colour - picked for contrast against the cell, the way sklearn's - ConfusionMatrixDisplay does), optional red-highlighted tick labels, - colorbar, and title. + """Render and save (or display) one class x class matrix figure: image, + per-cell values when `values_format` is given and there are at most + MAX_ANNOTATED_CLASSES classes (colour picked for contrast against the + cell, as sklearn's ConfusionMatrixDisplay does), optional red-highlighted + tick labels, colorbar, and title. Sizes are derived from the class count. """ data = np.asarray(data) num_classes = len(class_labels) + fig_size, font_size, title_font_size = _class_plot_scale( + num_classes, fig_size + ) # Vector formats embed the matrix at its native class x class resolution - # ("none" skips resampling), so they stay sharp at any size and dpi has - # no effect on them. + # ("none" skips resampling), so they stay sharp at any size. is_vector = vis_format in VECTOR_FORMATS - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) + # Size the figure around a fig_size x fig_size image so the margins are + # even on all sides whatever the label lengths. + cbar_gap, cbar_width = 0.15, max(0.15, 0.04 * fig_size) + left, right, bottom, top = _margins( + font_size, + title_size=title_font_size, + xlabels=class_labels, + ylabels=class_labels, + x_rotation=90, + xlabel=bool(xlabel), + ylabel=bool(ylabel), + right=cbar_gap + cbar_width + 0.1 + 5 * 0.65 * font_size / 72, + ) + top += 0.5 * font_size / 72 # room for the top colorbar tick label + width, height = left + fig_size + right, bottom + fig_size + top + fig = plt.figure(figsize=(width, height)) + ax = fig.add_axes( + (left / width, bottom / height, fig_size / width, fig_size / height) + ) + cax = fig.add_axes( + ( + (left + fig_size + cbar_gap) / width, + bottom / height, + cbar_width / width, + fig_size / height, + ) + ) im = ax.imshow( data, interpolation="none" if is_vector else "nearest", @@ -142,20 +267,19 @@ class x class matrix figure: image, optional per-cell values (colour vmax=vmax, ) - if values_format is not None: + if values_format is not None and num_classes <= MAX_ANNOTATED_CLASSES: color_min, color_max = im.cmap(0.0), im.cmap(1.0) thresh = (data.max() + data.min()) / 2.0 for i in range(num_classes): for j in range(num_classes): - color = color_max if data[i, j] < thresh else color_min ax.text( j, i, f"{data[i, j]:{values_format}}", ha="center", va="center", - color=color, - fontsize=value_font_size, + color=color_max if data[i, j] < thresh else color_min, + fontsize=min(font_size, 10), ) ax.set_xticks(np.arange(num_classes), labels=class_labels) @@ -173,34 +297,98 @@ class x class matrix figure: image, optional per-cell values (colour ax.tick_params(axis="x", rotation=90, labelsize=font_size) ax.tick_params(axis="y", labelsize=font_size) - colorbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + colorbar = fig.colorbar(im, cax=cax) colorbar.ax.tick_params(labelsize=font_size) - ax.set_title(title, fontsize=title_font_size, fontweight="normal") if xlabel: ax.set_xlabel(xlabel, fontsize=font_size) if ylabel: ax.set_ylabel(ylabel, fontsize=font_size) - fig.tight_layout() + _set_title(ax, title, title_font_size) + save_figure(fig, path, vis_format, vis_print, display) - if display: - plt.show() - return - if not os.path.exists("plots"): - os.mkdir("plots") - fig.savefig( +def _similarity_figure( + data: npt.NDArray, + class_labels: typing.Sequence | npt.NDArray, + highlight: npt.NDArray, + title: str, + path: str, + vis_format: str, + vis_print: bool, + fig_size: float | None = None, + display: bool = False, +) -> None: + """Class x class matrix in [-1, 1] (affinity or cosine similarity) with + the `highlight`ed classes' tick labels in red.""" + _matrix_figure( + data, + class_labels, + title, path, - **({"dpi": print_dpi} if vis_print and not is_vector else {}), - **( - {"pil_kwargs": {"compress_level": 3}} - if vis_format == "png" - else {} - ), + vis_format, + vis_print, + cmap="RdBu", + vmin=-1, + vmax=1, + highlight=highlight, + fig_size=fig_size, + display=display, + ) + + +def _confusion_figure( + y: npt.NDArray, + ypred: npt.NDArray, + labels: npt.NDArray, + stem: str, + epoch: int, + vis_format: str, + vis_print: bool, +) -> None: + """Write the count and normalised (% of true class) confusion matrices + as `{stem}.csv` / `{stem}_norm.csv` and draw `{stem}_norm.{vis_format}` + titled with the balanced accuracy over classes that have samples.""" + cm = sklearn.metrics.confusion_matrix(y, ypred, labels=labels) + support = cm.sum(axis=1) + supported = support > 0 + if not np.all(supported): + logging.warning( + "Confusion matrix has no samples for classes %s. " + "Their normalised rows will be set to zero and excluded from " + "average per-class accuracy.", + labels[~supported], + ) + cmn = ( + np.divide( + cm.astype(float), + support[:, np.newaxis], + out=np.zeros(cm.shape), + where=supported[:, np.newaxis], + ) + * 100 ) + acc = np.mean(cmn.diagonal()[supported]) - plt.close(fig) + os.makedirs(os.path.dirname(stem), exist_ok=True) + pd.DataFrame(cm).to_csv(f"{stem}.csv", index=False) + pd.DataFrame(cmn).to_csv(f"{stem}_norm.csv", index=False) + + _matrix_figure( + cmn, + labels, + _epoch_title("Balanced accuracy", epoch, f"{acc:.1f}%"), + f"{stem}_norm.{vis_format}", + vis_format, + vis_print, + cmap=plt.cm.Blues, + vmin=0, + vmax=100, + values_format=".0f", + xlabel="Predicted label (%)", + ylabel="True label (%)", + ) def _save_grid_figure( @@ -233,6 +421,7 @@ def _save_grid_figure( ) +@_styled def loss_plot( epochs: int, beta: float, @@ -332,6 +521,7 @@ def loss_plot( ) +@_styled def plot_cyc_variable( array: list, variable_name: str, @@ -364,6 +554,7 @@ def plot_cyc_variable( ) +@_styled def confusion_plot( y_train: npt.NDArray, ypred_train: npt.NDArray, @@ -399,194 +590,23 @@ def confusion_plot( logging.info("Visualising accuracy: confusion and F1 scores ...\n") classes_list = np.unique(np.concatenate((y_train, ypred_train))) - - fig_size, font_size, title_font_size = _class_plot_scale(len(classes_list)) - value_font_size = min(font_size, 10) - - # Compute confusion matrix - cm = sklearn.metrics.confusion_matrix( - y_train, ypred_train, labels=classes_list - ) - - # Convert confusion matrix to a DataFrame to be saved as csv file - cm_df = pd.DataFrame(cm) - - train_support = cm.sum(axis=1) - supported_train_classes = train_support > 0 - avg_accuracy = np.divide( - cm.diagonal(), - train_support, - out=np.zeros_like(train_support, dtype=float), - where=supported_train_classes, - ) - - # Normalize confusion matrix - cmn = ( - np.divide( - cm.astype(float), - train_support[:, np.newaxis], - out=np.zeros_like(cm, dtype=float), - where=supported_train_classes[:, np.newaxis], - ) - * 100 + classes_eval = np.unique(np.concatenate((y_val, ypred_val))) + # validation keeps the training class order, with unseen classes appended + ordered_class_eval = np.concatenate( + (classes_list, np.setdiff1d(classes_eval, classes_list)) ) - # Convert normalised confusion matrix to a DataFrame to be saved as csv file - cmn_df = pd.DataFrame(cmn) - - with plt.rc_context({"font.weight": "bold", "font.size": font_size}): - train_title = "Balanced accuracy at epoch {}: {:.1f}%".format( - epoch + 1, np.mean(avg_accuracy[supported_train_classes]) * 100 - ) - - _matrix_figure( - fig_size, - font_size, - title_font_size, - train_title, - f"plots/confusion_train{mode}.{vis_format}", - vis_format, - vis_print, - data=cm, - class_labels=classes_list, - cmap=plt.cm.Blues, - values_format="d" if vis_print else None, - value_font_size=value_font_size, - xlabel="Predicted label", - ylabel="True label", - ) - - _matrix_figure( - fig_size, - font_size, - title_font_size, - train_title, - f"plots/confusion_train{mode}_norm.{vis_format}", - vis_format, - vis_print, - data=cmn, - class_labels=classes_list, - cmap=plt.cm.Blues, - vmin=0, - vmax=100, - values_format=".0f" if vis_print else None, - value_font_size=value_font_size, - xlabel="Predicted label (%)", - ylabel="True label (%)", - ) - - # Save confusion matrix DataFrame to CSV - cm_df.to_csv(f"plots/confusion_train{mode}.csv", index=False) - - # Save normalised confusion matrix DataFrame to CSV - cmn_df.to_csv(f"plots/confusion_train{mode}_norm.csv", index=False) - - classes_list_eval = np.unique(np.concatenate((y_val, ypred_val))) - - if np.setdiff1d(classes_list_eval, classes_list).size > 0: - ordered_class_eval = np.concatenate( - (classes_list, np.setdiff1d(classes_list_eval, classes_list)) - ) - else: - ordered_class_eval = classes_list - - cm_eval = sklearn.metrics.confusion_matrix( - y_val, ypred_val, labels=ordered_class_eval - ) - - # Convert confusion matrix to a DataFrame to be saved as csv file - cm_eval_df = pd.DataFrame(cm_eval) - - eval_support = cm_eval.sum(axis=1) - supported_eval_classes = eval_support > 0 - if not np.all(supported_eval_classes): - logging.warning( - "Validation confusion matrix has no samples for classes %s. " - "Their normalised rows will be set to zero and excluded from " - "average per-class accuracy.", - ordered_class_eval[~supported_eval_classes], - ) - avg_accuracy_eval = np.divide( - cm_eval.diagonal(), - eval_support, - out=np.zeros_like(eval_support, dtype=float), - where=supported_eval_classes, - ) - - # Normalise the validation confusion matrix - cmn_eval = ( - np.divide( - cm_eval.astype(float), - eval_support[:, np.newaxis], - out=np.zeros_like(cm_eval, dtype=float), - where=supported_eval_classes[:, np.newaxis], - ) - * 100 - ) - - # Convert confusion matrix to a DataFrame to be saved as csv file - cmn_eval_df = pd.DataFrame(cmn_eval) - - # Save confusion matrix DataFrame to CSV - cm_eval_df.to_csv(f"plots/confusion_valid{mode}.csv", index=False) - - # Save normalised confusion matrix DataFrame to CSV - cmn_eval_df.to_csv(f"plots/confusion_valid{mode}_norm.csv", index=False) - - if mode == "_eval": - figure_name = "plots/confusion_eval" - elif mode == "": - figure_name = "plots/confusion_valid" - else: - figure_name = f"plots/confusion_{mode}" - - with plt.rc_context( - { - "font.weight": "bold", - "font.size": font_size, - } + valid_name = {"": "valid", "_eval": "eval"}.get(mode, mode) + for stem, y, ypred, labels in ( + (f"confusion_train{mode}", y_train, ypred_train, classes_list), + (f"confusion_{valid_name}", y_val, ypred_val, ordered_class_eval), ): - eval_title = "Balanced accuracy at epoch {}: {:.1f}%".format( - epoch + 1, np.mean(avg_accuracy_eval[supported_eval_classes]) * 100 - ) - - _matrix_figure( - fig_size, - font_size, - title_font_size, - eval_title, - figure_name + f".{vis_format}", - vis_format, - vis_print, - data=cm_eval, - class_labels=ordered_class_eval, - cmap=plt.cm.Blues, - values_format="d" if vis_print else None, - value_font_size=value_font_size, - xlabel="Predicted label", - ylabel="True label", - ) - - _matrix_figure( - fig_size, - font_size, - title_font_size, - eval_title, - f"{figure_name}_norm.{vis_format}", - vis_format, - vis_print, - data=cmn_eval, - class_labels=ordered_class_eval, - cmap=plt.cm.Blues, - vmin=0, - vmax=100, - values_format=".0f" if vis_print else None, - value_font_size=value_font_size, - xlabel="Predicted label (%)", - ylabel="True label (%)", + _confusion_figure( + y, ypred, labels, f"plots/{stem}", epoch, vis_format, vis_print ) +@_styled def f1_plot( y_train: npt.NDArray, ypred_train: npt.NDArray, @@ -678,45 +698,30 @@ def f1_plot( else: valid_df.to_csv(f1_valid_file, index=False) - with plt.rc_context( - { - "font.weight": "bold", - "font.size": font_size, - } - ): - fig, ax = plt.subplots(figsize=(fig_size, fig_size)) - line_width = max(1.5, font_size / 6) - plt.plot( - classes_list, - train_f1_score, - label="train", - marker="o", - linewidth=line_width, - ) - plt.plot( - classes_list, - valid_f1_score, - label=label, - marker="o", - linewidth=line_width, - ) - plt.xticks(rotation=45) - plt.legend(loc="lower left") - plt.title( - "F1 Score at epoch {}".format(epoch + 1), - fontsize=title_font_size, - fontweight="normal", - ) - plt.ylabel("F1 Score") - plt.tight_layout() - plt.savefig( - f"plots/f1{mode}.{vis_format}", - **({"dpi": 300} if vis_print else {}), + fig, ax = plt.subplots(figsize=(fig_size, fig_size)) + line_width = max(1.5, font_size / 6) + for name, scores in (("train", train_f1_score), (label, valid_f1_score)): + ax.plot( + classes_list, scores, label=name, marker="o", linewidth=line_width ) - - plt.close() + ax.tick_params(axis="x", rotation=45, labelsize=font_size) + ax.tick_params(axis="y", labelsize=font_size) + ax.legend(loc="lower left", fontsize=font_size) + ax.set_ylabel("F1 Score", fontsize=font_size) + _fit_margins( + fig, + font_size, + title_size=title_font_size, + xlabels=classes_list, + ylabels=_ytick_labels(ax), + x_rotation=45, + ylabel=True, + ) + _set_title(ax, _epoch_title("F1 Score", epoch), title_font_size) + save_figure(fig, f"plots/f1{mode}.{vis_format}", vis_format, vis_print) +@_styled def recon_plot( img: torch.Tensor, rec: torch.Tensor, @@ -844,6 +849,7 @@ def recon_plot( ) +@_styled def latent_embed_plot_tsne( xs: npt.NDArray, ys: npt.NDArray, @@ -928,18 +934,11 @@ def latent_embed_plot_tsne( n_classes = len(classes) fig_size, font_size, title_font_size = _class_plot_scale(n_classes) - if n_classes < 3: - # If the number of classes are not moe than 3 the size of the figure would be too - # small and matplotlib would through a singularity error - fig, ax = plt.subplots( - figsize=(int(n_classes / 2) + 7, int(n_classes / 2) + 5) - ) - else: - fig, ax = plt.subplots( - figsize=(int(n_classes / 2) + 4, int(n_classes / 2) + 2) - ) - marker_size = max(marker_size, int(marker_size * max(1, n_classes / 5))) - # When the number of classes is less than 3 the image becomes two small + # 8:5 axes, matching the interactive embedding; the figure is built around + # it (in inches) once the legend and tick labels are known. + axes_width, axes_height = fig_size, fig_size * 5 / 8 + fig = plt.figure() + ax = fig.add_axes((0.1, 0.1, 0.8, 0.8)) colours = colour_per_class(classes) if xs.shape[-1] != 1: @@ -949,21 +948,16 @@ def latent_embed_plot_tsne( color = colours[classes.index(mol)] - plt.scatter( + ax.scatter( lats[idx, 0], lats[idx, 1], s=marker_size, - label=mol[:4], + label=mol, facecolor=color, edgecolor=color, alpha=0.5, ) - ax.legend( - bbox_to_anchor=(1.05, 1), - loc="upper left", - fontsize=font_size, - ) ax.set_xlabel("TSNE-1", fontsize=font_size) ax.set_ylabel("TSNE-2", fontsize=font_size) @@ -972,43 +966,56 @@ def latent_embed_plot_tsne( for mol_id, mol in enumerate(set(ys.tolist())): idx = np.where(np.array(ys.tolist()) == mol)[0] cols = colours[classes.index(mol)] - plt.hist( + ax.hist( lats[idx], 100, color=cols, histtype="step", stacked=True, fill=False, - label=mol[:4], + label=mol, linewidth=l_w, ) - ax.legend( - bbox_to_anchor=(1.05, 1), - loc="upper left", - fontsize=font_size, - ) ax.set_xlabel("dim 1", fontsize=font_size) ax.set_ylabel("freq", fontsize=font_size) - ax.set_title( - f"t-SNE Embedding at epoch: {epoch + 1}", - fontsize=title_font_size, - fontweight="normal", + # Legend outside the axes, in as many columns as fit the axes height + ncol = int(np.ceil(n_classes * 1.6 * font_size / 72 / axes_height)) + ax.legend( + bbox_to_anchor=(1.02, 1), + loc="upper left", + fontsize=font_size, + ncol=ncol, ) + legend_width = ncol * 10 * font_size / 72 + 0.3 ax.tick_params(axis="both", labelsize=font_size) - fig.tight_layout() - - if not display: - if not os.path.exists("plots"): - os.mkdir("plots") - plt.savefig( - f"plots/embedding_TSNE{mode}.{vis_format}", - **({"dpi": 300} if vis_print else {}), + left, right, bottom, top = _margins( + font_size, + title_size=title_font_size, + ylabels=_ytick_labels(ax), + xlabel=True, + ylabel=True, + right=legend_width, + ) + width = left + axes_width + right + height = bottom + axes_height + top + fig.set_size_inches(width, height) + ax.set_position( + ( + left / width, + bottom / height, + axes_width / width, + axes_height / height, ) - else: - plt.show() - - plt.close() + ) + _set_title(ax, _epoch_title("t-SNE Embedding", epoch), title_font_size) + save_figure( + fig, + f"plots/embedding_TSNE{mode}.{vis_format}", + vis_format, + vis_print, + display, + ) def dyn_latentembed_plot( @@ -1068,11 +1075,13 @@ def dyn_latentembed_plot( bind_checkbox = altair.binding_radio( options=opts, labels=[ - "off" - if i == "std-off" - else "avg" - if i == "std-avg" - else str(int(i.split("-")[-1]) + 1) + ( + "off" + if i == "std-off" + else ( + "avg" if i == "std-avg" else str(int(i.split("-")[-1]) + 1) + ) + ) for i in opts ], name="Certainty of prediction per dimension:", @@ -1188,6 +1197,7 @@ def dyn_latentembed_plot( chart.save(f"latents/latent_epoch_{epoch + 1}_{mode}.html") +@_styled def confidence_plot(x, y, s, suffix=None, vis_format="png", vis_print=False): logging.info( "################################################################", @@ -1198,11 +1208,10 @@ def confidence_plot(x, y, s, suffix=None, vis_format="png", vis_print=False): cmap = plt.get_cmap("jet") cols = [cmap(i) for i in np.linspace(0, 1, len(x[0]))] classes = np.unique(y) - _, font_size, title_font_size = _class_plot_scale(len(classes)) - rows = len(classes) // 2 - if len(classes) % 2 != 0: - rows += 1 - fig, ax = plt.subplots(len(classes), sharex=True, sharey=True) + fig_size, font_size, title_font_size = _class_plot_scale(len(classes)) + fig, ax = plt.subplots( + len(classes), sharex=True, sharey=True, figsize=(fig_size, fig_size) + ) ax = np.atleast_1d(ax) for c, cl in enumerate(classes): mu_cl = np.take(x, np.where(np.array(y) == cl)[0], axis=0) @@ -1225,24 +1234,33 @@ def confidence_plot(x, y, s, suffix=None, vis_format="png", vis_print=False): color=cols[i], label="lat" + str(i + 1), ) - ax[c].set_title(cl, fontsize=title_font_size, fontweight="normal") + ax[c].tick_params(labelsize=font_size) name = f"plots/confidence.{vis_format}" if suffix is not None: name = name[:-4] + "_" + suffix + name[-4:] handles, labels = ax[-1].get_legend_handles_labels() - leg = fig.legend( - handles, labels, bbox_to_anchor=(1.06, 0.9) - ) # , loc="upper left") - plt.tight_layout() - fig.savefig( - name, - **({"dpi": 300} if vis_print else {}), - bbox_extra_artists=(leg,), - bbox_inches="tight", + ncol = int(np.ceil(len(labels) * 1.6 * font_size / 72 / fig_size)) + fig.legend( + handles, + labels, + loc="center right", + fontsize=font_size, + ncol=ncol, ) - plt.close() + _fit_margins( + fig, + font_size, + title_size=font_size, + ylabels=_ytick_labels(ax[0]), + right=ncol * 8 * font_size / 72 + 0.3, + ) + fig.subplots_adjust(hspace=0.8) + for axis, cl in zip(ax, classes): + _set_title(axis, str(cl), font_size) + save_figure(fig, name, vis_format, vis_print) +@_styled def latent_space_similarity_plot( latent_space: npt.NDArray, class_labels: npt.NDArray, @@ -1251,9 +1269,7 @@ def latent_space_similarity_plot( affinity_matrix: pathlib.Path | None = None, plot_mode: str = "mean", display: bool = False, - font_size: int = 16, fig_size: int | None = None, - dpi: int = 300, vis_format: str = "png", vis_print: bool = False, ) -> None: @@ -1280,7 +1296,7 @@ def latent_space_similarity_plot( logging.info("Visualising the latent space similarity matrix ...\n") if affinity_matrix is None: - unique_classes = np.unique(class_labels) + unique_classes: npt.NDArray | list = np.unique(class_labels) else: classes_order = ( pd.read_csv(affinity_matrix, header=0).columns.astype(str).tolist() @@ -1296,72 +1312,26 @@ def latent_space_similarity_plot( else: unique_classes = classes_order - num_classes = len(unique_classes) - resolved_fig_size, font_size, title_font_size = _class_plot_scale( - num_classes, fig_size - ) cosine_sim = latent_space_similarity_mat( latent_space, class_labels, unique_classes, - num_classes, + len(unique_classes), plot_mode=plot_mode, ) - # Visualize average cosine similarity matrix classes_in_data = set(np.asarray(class_labels).astype(str)) - highlight = np.array( - [str(c) not in classes_in_data for c in unique_classes] - ) - title = f"Average Cosine Similarity Matrix at epoch: {epoch + 1}" - path = f"plots/similarity_mean{mode}.{vis_format}" - - if fig_size is None: - with plt.rc_context( - { - "font.weight": "bold", - "font.size": font_size, - } - ): - _matrix_figure( - resolved_fig_size, - font_size, - title_font_size, - title, - path, - vis_format, - vis_print, - print_dpi=dpi, - data=cosine_sim, - class_labels=unique_classes, - cmap="RdBu", - vmin=-1, - vmax=1, - highlight=highlight, - xlabel="Class Labels", - ylabel="Class Labels", - display=display, - ) - else: - _matrix_figure( - resolved_fig_size, - font_size, - title_font_size, - title, - path, - vis_format, - vis_print, - print_dpi=dpi, - data=cosine_sim, - class_labels=unique_classes, - cmap="RdBu", - vmin=-1, - vmax=1, - highlight=highlight, - xlabel="Class Labels", - ylabel="Class Labels", - display=display, - ) + _similarity_figure( + cosine_sim, + unique_classes, + np.array([str(c) not in classes_in_data for c in unique_classes]), + _epoch_title("Average Cosine Similarity Matrix", epoch), + f"plots/similarity_mean{mode}.{vis_format}", + vis_format, + vis_print, + fig_size, + display, + ) def latent_4enc_interpolate_plot( @@ -1889,6 +1859,7 @@ def interpolations_plot( ) +@_styled def plot_affinity_matrix( lookup: pd.DataFrame, all_classes: list, @@ -1918,29 +1889,19 @@ def plot_affinity_matrix( ) logging.info("Visualising affinity matrix ...\n") - resolved_fig_size, font_size, title_font_size = _class_plot_scale( - len(all_classes), fig_size - ) - - highlight = np.array([c not in selected_classes for c in all_classes]) - - _matrix_figure( - resolved_fig_size, - font_size, - title_font_size, + _similarity_figure( + lookup, + all_classes, + np.array([c not in selected_classes for c in all_classes]), "Affinity Matrix", f"plots/affinity_matrix.{vis_format}", vis_format, vis_print, - data=lookup, - class_labels=all_classes, - cmap="RdBu", - vmin=-1, - vmax=1, - highlight=highlight, + fig_size, ) +@_styled def plot_classes_distribution( data: list, category: str, @@ -1968,23 +1929,27 @@ def plot_classes_distribution( ticks = range(len(counts)) ax.bar(ticks, counts, align="center", color="blue", alpha=0.5) ax.set_xticks(ticks, labels) - ax.set_title( - "Classes Distribution", - fontsize=title_font_size, - fontweight="normal", - ) ax.set_xlabel("Class", fontsize=font_size) ax.set_ylabel("Number of Entries", fontsize=font_size) ax.tick_params(axis="x", labelsize=font_size, rotation=90) ax.tick_params(axis="y", labelsize=font_size) - fig.tight_layout() - if not os.path.exists("plots"): - os.mkdir("plots") - plt.savefig( + _fit_margins( + fig, + font_size, + title_size=title_font_size, + xlabels=labels, + ylabels=_ytick_labels(ax), + x_rotation=90, + xlabel=True, + ylabel=True, + ) + _set_title(ax, "Classes Distribution", title_font_size) + save_figure( + fig, f"plots/classes_distribution_{category}.{vis_format}", - **({"dpi": 300} if vis_print else {}), + vis_format, + vis_print, ) - plt.close() def _encoder(i: PIL.Image) -> str: diff --git a/tests/test_integration_pipeline.py b/tests/test_integration_pipeline.py index 956b8ca3..36989417 100644 --- a/tests/test_integration_pipeline.py +++ b/tests/test_integration_pipeline.py @@ -31,10 +31,10 @@ def test_cnn_mrc_artifacts(self): ) = helper_train_eval(data) self.assertEqual(n_dir_train, 3) - self.assertEqual(n_plots_train, 34) + self.assertEqual(n_plots_train, 32) self.assertEqual(n_latent_train, 1) self.assertEqual(n_states_train, 2) - self.assertEqual(n_plots_eval, 56) + self.assertEqual(n_plots_eval, 52) self.assertEqual(n_latent_eval, 2) self.assertEqual(n_states_eval, 3) From 1e0c87e7718b4b46817b9af36dc814d109417464 Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 08/12] Batch decoder calls in interpolation plots Decode each latent/pose grid in a single decoder call instead of one call per cell (21-36 GPU round trips per plot). Outputs match the previous loops to within 1e-7. interpolations_plot now interpolates poses from the four drawn class representatives rather than the first four classes. --- avae/utils.py | 37 +++++------ avae/vis.py | 169 ++++++++++++++++++++++---------------------------- 2 files changed, 92 insertions(+), 114 deletions(-) diff --git a/avae/utils.py b/avae/utils.py index 0226c54f..6aa41021 100644 --- a/avae/utils.py +++ b/avae/utils.py @@ -1,4 +1,3 @@ -import copy import logging import os.path import pathlib @@ -190,7 +189,6 @@ def pose_interpolation( vae: torch.nn.Module, device: torch.device, ) -> npt.NDArray: - """This function: 1- interpolates within each pose channels for the number_of_samples requested. @@ -216,24 +214,27 @@ def pose_interpolation( device: torch.device Device to run the model on. """ - decoded_grid = [] - # Generate vectors representing single transversals along each lat_dim - for p_dim in range(pos_dims): - for grid_spot in range(number_of_samples): - means = copy.deepcopy(pose_mean) - means[p_dim] += pose_std[p_dim] * (-1.2 + 0.4 * grid_spot) - - pos = torch.from_numpy(np.array(means)).unsqueeze(0).to(device) - lat = torch.from_numpy(np.array(enc)).unsqueeze(0).to(device) - - # Decode interpolated vectors - with torch.no_grad(): - decoded_img = vae.decoder(lat, pos) - - decoded_grid.append(decoded_img.cpu().squeeze().numpy()) + # Single transversals along each pose channel, row-major: p_dim, grid_spot + steps = -1.2 + 0.4 * np.arange(number_of_samples) + pose_mean = np.asarray(pose_mean) + means = np.tile(pose_mean, (pos_dims, number_of_samples, 1)) + means[ + np.arange(pos_dims)[:, None], + np.arange(number_of_samples)[None, :], + np.arange(pos_dims)[:, None], + ] += (np.asarray(pose_std)[:, None] * steps[None, :]).astype(means.dtype) + means = means.reshape(-1, means.shape[-1]) + + n = means.shape[0] + pos = torch.from_numpy(means).to(device) + lat = torch.from_numpy(np.array(enc)).unsqueeze(0).expand(n, -1).to(device) + + # Decode all interpolated vectors in a single batched call + with torch.no_grad(): + decoded_img = vae.decoder(lat, pos) decoded_grid = np.reshape( - np.array(decoded_grid), (pos_dims, number_of_samples, *dsize) + decoded_img.cpu().numpy(), (pos_dims, number_of_samples, *dsize) ) return decoded_grid diff --git a/avae/vis.py b/avae/vis.py index c8cceb4e..99f36cbb 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -1,4 +1,3 @@ -import copy import functools import logging import os.path @@ -1384,7 +1383,6 @@ def latent_4enc_interpolate_plot( ) padding = 0 - data_dim = len(dsize) classes = np.unique(np.asarray(ys)) latent_dim = xs.shape[1] @@ -1409,34 +1407,34 @@ def latent_4enc_interpolate_plot( ) enc.append(lat) - enc = np.asarray(enc) - alpha_values = torch.linspace(0, 1, num_steps) - beta_values = torch.linspace(0, 1, num_steps) - decoded_grid = [] - - for h in alpha_values: - for v in beta_values: - - # bilinear interpolation in the latent space - interpolated_z = ( - (1 - h) * (1 - v) * enc[0] - + h * (1 - v) * enc[1] - + (1 - h) * v * enc[2] - + h * v * enc[3] - ) - - # Decode the interpolated encoding to generate an image - with torch.no_grad(): - decoded_images = vae.decoder( - interpolated_z.view(-1, latent_dim).to(device=device), - (torch.zeros(1, poses[0].shape[0]) + pose_mean).to( - device=device - ), - ) - decoded_grid.append(decoded_images.cpu().squeeze().numpy()) + enc_t = torch.as_tensor( + np.asarray(enc, dtype=np.float32).reshape(4, latent_dim) + ) + h = torch.linspace(0, 1, num_steps)[:, None, None] + v = torch.linspace(0, 1, num_steps)[None, :, None] + + # bilinear interpolation in the latent space, all cells at once + z_batch = ( + (1 - h) * (1 - v) * enc_t[0] + + h * (1 - v) * enc_t[1] + + (1 - h) * v * enc_t[2] + + h * v * enc_t[3] + ).reshape(-1, latent_dim) + z_batch = z_batch.to(device=device, dtype=torch.float32) + + pose_batch = None + if poses is not None: + pose_batch = ( + torch.zeros(num_steps * num_steps, poses[0].shape[0]) + + pose_mean + ).to(device=device, dtype=torch.float32) + + # Decode all interpolated encodings in a single call + with torch.no_grad(): + decoded_images = vae.decoder(z_batch, pose_batch) decoded_grid = np.reshape( - np.array(decoded_grid), (num_steps, num_steps, *dsize) + decoded_images.cpu().numpy(), (num_steps, num_steps, *dsize) ) _save_grid_figure( @@ -1487,7 +1485,6 @@ def latent_disentamglement_plot( logging.info("Visualising latent content disentanglement ...\n") number_of_samples = 7 padding = 0 - data_dim = len(dsize) latents = np.asarray(lats) if poses is not None: poses_space = np.asarray(poses) @@ -1500,37 +1497,28 @@ def latent_disentamglement_plot( pos_means = np.mean(poses_space, axis=0) pos_dims = poses_space.shape[-1] - recon_images = [] - - # Generate vectors representing single transversals along each latent_dims - for l_dim in range(latent_dims): - for grid_spot in range(number_of_samples): - means = copy.deepcopy(lat_means) - # every 0.4 interval from -1.2 to 1.2 sigma - means[l_dim] += lat_stds[l_dim] * (-1.2 + 0.4 * grid_spot) + # Single traversals along each latent dim, every 0.4 sigma from -1.2 to 1.2 + steps = -1.2 + 0.4 * np.arange(number_of_samples) + z_np = np.tile(lat_means, (latent_dims, number_of_samples, 1)) + idx = np.arange(latent_dims) + z_np[idx, :, idx] += lat_stds[:, None] * steps[None, :] + z_batch = torch.from_numpy( + z_np.reshape(-1, latent_dims).astype(np.float32) + ).to(device) - # Decode the current vector - with torch.no_grad(): - current_lat_grid = torch.from_numpy(np.array([means])).to( - device - ) - - if poses is not None: - current_pos_grid = torch.from_numpy( - np.array([pos_means]) - ).to(device) - current_recon = vae.decoder( - current_lat_grid, current_pos_grid - ) - else: - current_recon = vae.decoder(current_lat_grid, None) + pose_batch = None + if poses is not None: + pose_batch = torch.from_numpy( + np.tile(pos_means, (latent_dims * number_of_samples, 1)).astype( + np.float32 + ) + ).to(device) - recon_images.append(current_recon.cpu().squeeze().numpy()) + with torch.no_grad(): + recon = vae.decoder(z_batch, pose_batch) - # Combine the individual decoded images into a single array - recon_images = np.array(recon_images) recon_images = np.reshape( - recon_images, (latent_dims, number_of_samples, *dsize) + recon.cpu().numpy(), (latent_dims, number_of_samples, *dsize) ) _save_grid_figure( @@ -1763,7 +1751,6 @@ def interpolations_plot( "################################################################", ) logging.info("Visualising interpolations ...\n") - data_dim = len(dsize) lats = np.asarray(lats) classes = np.asarray(classes) @@ -1805,47 +1792,37 @@ def interpolations_plot( grid_size = 6 alpha_values = np.linspace(0.0, 1.0, grid_size, dtype=np.float32) beta_values = np.linspace(0.0, 1.0, grid_size, dtype=np.float32) - decoded_grid = [] - for i, h in enumerate(alpha_values): - for j, v in enumerate(beta_values): - - # bilinear interpolation in the latent space - interpolated_z = ( - (1 - h) * (1 - v) * class_rep_lats[0] - + h * (1 - v) * class_rep_lats[1] - + (1 - h) * v * class_rep_lats[2] - + h * v * class_rep_lats[3] - ) - interpolated_z_t = torch.from_numpy( - np.asarray(interpolated_z, dtype=np.float32) - ).view(-1, latent_dim) - - if poses is not None: - interpolated_pose = ( - (1 - h) * (1 - v) * class_rep_poses[0] - + h * (1 - v) * class_rep_poses[1] - + (1 - h) * v * class_rep_poses[2] - + h * v * class_rep_poses[3] - ) - interpolated_pose_t = torch.from_numpy( - np.asarray(interpolated_pose, dtype=np.float32) - ).view(-1, poses_dim) - with torch.no_grad(): - if poses is not None: - decoded_images = vae.decoder( - interpolated_z_t.to(device=device), - interpolated_pose_t.to(device=device), - ) - else: - decoded_images = vae.decoder( - interpolated_z_t.to(device=device), - None, - ) - - decoded_grid.append(decoded_images.cpu().squeeze().numpy()) + h = alpha_values[:, None, None] + v = beta_values[None, :, None] + w = ( + (1 - h) * (1 - v), + h * (1 - v), + (1 - h) * v, + h * v, + ) + + # bilinear interpolation in the latent space, all grid cells at once + reps_z = np.asarray(class_rep_lats, dtype=np.float32) + z_batch = sum(w[k] * reps_z[k] for k in range(4)).reshape(-1, latent_dim) + z_batch = torch.from_numpy(np.asarray(z_batch, dtype=np.float32)).to( + device=device + ) + + pose_batch = None + if poses is not None: + reps_p = np.asarray(class_rep_poses, dtype=np.float32) + p_batch = sum(w[k] * reps_p[k] for k in range(4)).reshape( + -1, poses_dim + ) + pose_batch = torch.from_numpy( + np.asarray(p_batch, dtype=np.float32) + ).to(device=device) + + with torch.no_grad(): + decoded_images = vae.decoder(z_batch, pose_batch) decoded_grid = np.reshape( - np.array(decoded_grid), (grid_size, grid_size, *dsize) + decoded_images.cpu().numpy(), (grid_size, grid_size, *dsize) ) _save_grid_figure( From 9e054a6441fd0ed76d2bc9d09c4040be44ac6337 Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 09/12] Skip redrawing an up-to-date affinity matrix The affinity matrix is fixed input, so skip plotting it again when the existing plot is newer than the CSV and was made with the same CSV path, vis_print setting and selected classes, which are recorded in plots/affinity_matrix_classes.txt. --- avae/data.py | 37 +++++++++++++++---- tests/test_data.py | 58 +++++++++++++++++++++++++++++- tests/test_integration_pipeline.py | 4 +-- 3 files changed, 90 insertions(+), 9 deletions(-) diff --git a/avae/data.py b/avae/data.py index e43944ea..2093cd09 100644 --- a/avae/data.py +++ b/avae/data.py @@ -328,13 +328,38 @@ def get_affinity_matrix( ) ) if vis_aff: - plot_affinity_matrix( - lookup=affinity, - all_classes=affinity.columns.tolist(), - selected_classes=classes, - vis_format=vis_format, - vis_print=vis_print, + plot_path = f"plots/affinity_matrix.{vis_format}" + # selected classes are highlighted in the plot, so they are part + # of what makes an existing plot up to date + classes_path = "plots/affinity_matrix_classes.txt" + selected = "\n".join( + [ + str(pathlib.Path(affinity_path).resolve()), + f"vis_print={vis_print}", + *map(str, classes), + ] ) + if ( + os.path.exists(plot_path) + and os.path.getmtime(plot_path) + > os.path.getmtime(affinity_path) + and os.path.exists(classes_path) + and pathlib.Path(classes_path).read_text() == selected + ): + logging.info( + "Affinity matrix plot is up to date, skipping: " + f"{plot_path}" + ) + else: + plot_affinity_matrix( + lookup=affinity, + all_classes=affinity.columns.tolist(), + selected_classes=classes, + vis_format=vis_format, + vis_print=vis_print, + ) + os.makedirs("plots", exist_ok=True) + pathlib.Path(classes_path).write_text(selected) # subset affinity matrix with only the relevant classes index = [affinity.columns.get_loc(f"{columns}") for columns in classes] diff --git a/tests/test_data.py b/tests/test_data.py index a9fae402..c65d391e 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -7,7 +7,7 @@ import lightning as lt import numpy as np -from avae.data import load_data +from avae.data import get_affinity_matrix, load_data from tests import testdata_mrc @@ -160,3 +160,59 @@ def test_nonzero_rank_does_not_plot_data_visualisations(self): plot_affinity.assert_not_called() plot_distribution.assert_not_called() + + +class AffinityPlotCacheTest(unittest.TestCase): + def setUp(self) -> None: + self._orig_dir = os.getcwd() + self.test_dir = tempfile.mkdtemp(prefix="avae_") + os.chdir(self.test_dir) + self.csv = os.path.join(self.test_dir, "affinity.csv") + with open(self.csv, "w") as f: + f.write("a,b\n1,0\n0,1\n") + self.plot = os.path.join("plots", "affinity_matrix.png") + + def tearDown(self): + os.chdir(self._orig_dir) + shutil.rmtree(self.test_dir, ignore_errors=True) + + def _run(self, classes=("a", "b"), vis_print=False): + def fake_plot(**kwargs): + os.makedirs("plots", exist_ok=True) + with open(self.plot, "w") as f: + f.write("x") + + with mock.patch( + "avae.data.plot_affinity_matrix", side_effect=fake_plot + ) as m: + get_affinity_matrix( + self.csv, list(classes), vis_aff=True, vis_print=vis_print + ) + return m + + def test_replots_when_vis_print_changes(self): + self._run() + t = os.path.getmtime(self.csv) + os.utime(self.plot, (t + 10, t + 10)) + self.assertEqual(self._run(vis_print=True).call_count, 1) + + def test_replots_when_selected_classes_change(self): + self._run() + t = os.path.getmtime(self.csv) + os.utime(self.plot, (t + 10, t + 10)) + self.assertEqual(self._run(classes=["a"]).call_count, 1) + + def test_plots_when_missing(self): + self.assertEqual(self._run().call_count, 1) + + def test_skips_when_plot_newer(self): + self._run() + t = os.path.getmtime(self.csv) + os.utime(self.plot, (t + 10, t + 10)) + self.assertEqual(self._run().call_count, 0) + + def test_replots_when_csv_newer(self): + self._run() + t = os.path.getmtime(self.plot) + os.utime(self.csv, (t + 10, t + 10)) + self.assertEqual(self._run().call_count, 1) diff --git a/tests/test_integration_pipeline.py b/tests/test_integration_pipeline.py index 36989417..fc21d9b0 100644 --- a/tests/test_integration_pipeline.py +++ b/tests/test_integration_pipeline.py @@ -31,10 +31,10 @@ def test_cnn_mrc_artifacts(self): ) = helper_train_eval(data) self.assertEqual(n_dir_train, 3) - self.assertEqual(n_plots_train, 32) + self.assertEqual(n_plots_train, 33) self.assertEqual(n_latent_train, 1) self.assertEqual(n_states_train, 2) - self.assertEqual(n_plots_eval, 52) + self.assertEqual(n_plots_eval, 53) self.assertEqual(n_latent_eval, 2) self.assertEqual(n_states_eval, 3) From ae600fd3f78670a623242040b17322811f58a509 Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 10/12] Read class labels from file names for distribution plots The class-distribution histograms collected labels by iterating every train, validation and test batch. That read, transformed and PNG-encoded every image just to get its label, and sum(..., ()) made the collection quadratic in the number of batches. Labels come from the file name, so read them from the dataset paths instead. Histograms now count every sample in each split, including the last partial training batch. --- avae/data.py | 23 +++++++++++++++++++---- tests/test_data.py | 20 +++++++++++++++++++- 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/avae/data.py b/avae/data.py index 2093cd09..b879d1af 100644 --- a/avae/data.py +++ b/avae/data.py @@ -209,8 +209,8 @@ def load_data( # Plot the complete splits once, before Fabric shards the loaders. if vis_his and fabric.global_rank == 0: - train_y = list(sum([y[1] for _, y in enumerate(trains)], ())) - val_y = list(sum([y[1] for _, y in enumerate(vals)], ())) + train_y = _labels(trains.dataset) + val_y = _labels(vals.dataset) plot_classes_distribution( train_y, "train", @@ -266,7 +266,7 @@ def load_data( tests = test_loader.get_loader(batch_size=batch) if vis_his and fabric.global_rank == 0: - eval_y = list(sum([y[1] for _, y in enumerate(tests)], ())) + eval_y = _labels(test_loader.dataset) plot_classes_distribution( eval_y, "evaluation", @@ -371,6 +371,21 @@ def get_affinity_matrix( return None +def _labels(dataset, indices=None) -> list: + """Get labels from file names (no data loading) for a dataset with `paths` + or a torch Subset of one. `indices` selects positions within `dataset`.""" + if isinstance(dataset, torch.utils.data.Subset): + indices = ( + list(dataset.indices) + if indices is None + else [dataset.indices[i] for i in indices] + ) + dataset = dataset.dataset + if indices is None: + indices = range(len(dataset.paths)) + return [pathlib.Path(dataset.paths[i]).name.split("_")[0] for i in indices] + + class AffinityDiskDataset(caked.dataloader.DiskDataset): """Modified version of the caked DiskDataset to include the affinity matrix and data metadata that is needed for the affinity pipeline""" @@ -401,7 +416,7 @@ def __getitem__(self, index): # get file basename filename = pathlib.Path(self.paths[index]).name # ground truth - y = pathlib.Path(filename).name.split("_")[0] + y = _labels(self, [index])[0] # similarity column / vector if self.affinity is not None: diff --git a/tests/test_data.py b/tests/test_data.py index c65d391e..dba8acb1 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -6,8 +6,9 @@ import lightning as lt import numpy as np +import torch.utils.data -from avae.data import get_affinity_matrix, load_data +from avae.data import _labels, get_affinity_matrix, load_data from tests import testdata_mrc @@ -216,3 +217,20 @@ def test_replots_when_csv_newer(self): t = os.path.getmtime(self.plot) os.utime(self.csv, (t + 10, t + 10)) self.assertEqual(self._run().call_count, 1) + + +class LabelsTest(unittest.TestCase): + def setUp(self) -> None: + self.dataset = mock.MagicMock() + self.dataset.paths = ["/a/x_1.mrc", "/a/y_2.mrc", "/a/x_3.mrc"] + + def test_plain_dataset(self): + self.assertEqual(_labels(self.dataset), ["x", "y", "x"]) + self.dataset.read.assert_not_called() + self.dataset.__getitem__.assert_not_called() + + def test_subset(self): + subset = torch.utils.data.Subset(self.dataset, [2, 1]) + self.assertEqual(_labels(subset), ["x", "y"]) + self.dataset.read.assert_not_called() + self.dataset.__getitem__.assert_not_called() From 341ed3d2b6366912eaaa20ca269714be3741121d Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 11/12] Speed up t-SNE/similarity plots and guard edge cases - Build t-SNE class masks once instead of rebuilding the label array for every class. - Read the affinity CSV header once per path with a cached _affinity_classes helper, placed with the other private helpers. - Warn and return instead of crashing in pose_class_disentanglement_plot (no poses) and latent_4enc_interpolate_plot (fewer than 4 classes), and accept suffix=None in confidence_plot. - Remove unused locals. --- avae/vis.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/avae/vis.py b/avae/vis.py index 99f36cbb..715af6d2 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -151,6 +151,12 @@ def _class_plot_scale( return resolved_fig_size, font_size, title_size +@functools.lru_cache(maxsize=None) +def _affinity_classes(path: str) -> tuple: + """Read only the header of the affinity CSV (cached per path).""" + return tuple(pd.read_csv(path, nrows=0).columns.astype(str).tolist()) + + def _loss_curve_figure( epochs: int, series: typing.Sequence[tuple[str | None, npt.ArrayLike, str | None, str]], @@ -942,8 +948,8 @@ def latent_embed_plot_tsne( if xs.shape[-1] != 1: - for mol_id, mol in enumerate(set(ys.tolist())): - idx = np.where(np.array(ys.tolist()) == mol)[0] + for mol in np.unique(ys): + idx = ys == mol color = colours[classes.index(mol)] @@ -962,8 +968,8 @@ def latent_embed_plot_tsne( if xs.shape[-1] == 1: - for mol_id, mol in enumerate(set(ys.tolist())): - idx = np.where(np.array(ys.tolist()) == mol)[0] + for mol in np.unique(ys): + idx = ys == mol cols = colours[classes.index(mol)] ax.hist( lats[idx], @@ -1202,7 +1208,9 @@ def confidence_plot(x, y, s, suffix=None, vis_format="png", vis_print=False): "################################################################", ) logging.info( - "Visualising class-average confidence metrics " + suffix + "...\n", + "Visualising class-average confidence metrics" + + (f" {suffix}" if suffix is not None else "") + + "...\n", ) cmap = plt.get_cmap("jet") cols = [cmap(i) for i in np.linspace(0, 1, len(x[0]))] @@ -1297,9 +1305,7 @@ def latent_space_similarity_plot( if affinity_matrix is None: unique_classes: npt.NDArray | list = np.unique(class_labels) else: - classes_order = ( - pd.read_csv(affinity_matrix, header=0).columns.astype(str).tolist() - ) + classes_order = list(_affinity_classes(str(affinity_matrix))) unique_classes_in_data = np.unique(class_labels) if np.setdiff1d(unique_classes_in_data, classes_order).size > 0: unique_classes = np.concatenate( @@ -1384,6 +1390,12 @@ def latent_4enc_interpolate_plot( padding = 0 classes = np.unique(np.asarray(ys)) + if len(classes) <= 3: + logging.warning( + "\n\nWARNING: Interpolation plot needs at least 4 distinct classes, " + "cannot visualise interpolations. Exiting.\n", + ) + return latent_dim = xs.shape[1] # Number of plots (each have 4 random corners) @@ -1548,7 +1560,6 @@ def pose_class_disentanglement_plot( vis_format: str = "png", vis_print: bool = False, ): - """Visualise Pose interpolation per class. This function creates a pose interpolatoion plot for all classes listed in pose_vis_class. @@ -1590,9 +1601,9 @@ def pose_class_disentanglement_plot( logging.warning( "Pose interpolation cannot be done if pose dimension is not specified" ) + return padding = 0 - data_dim = len(dsize) x = np.asarray(x) poses_space = np.asarray(poses) @@ -1600,7 +1611,6 @@ def pose_class_disentanglement_plot( pose_vis_class_list = pose_vis_class.replace(" ", "").split(",") for i in pose_vis_class_list: - decoded_grid = [] class_x = np.take(x, np.where(np.array(y) == i)[0], axis=0) class_x_indx = np.random.choice(class_x.shape[0]) enc = class_x[class_x_indx, :] @@ -1683,7 +1693,6 @@ def pose_disentanglement_plot( number_of_samples = 7 padding = 0 - data_dim = len(dsize) latents = np.asarray(lats) poses_space = np.asarray(poses) From ccb3f28de3b6c92c64681e76bbd480c3dc28cdfa Mon Sep 17 00:00:00 2001 From: jolaem Date: Wed, 7 Oct 2026 15:45:38 +0000 Subject: [PATCH 12/12] Move visualisation helpers into utils_vis Move the plotting, grid, mrc, colour, pose-interpolation and similarity-matrix helpers out of utils into a new utils_vis module, leaving as_list and latest_file in utils. A pure move with no behaviour changes. --- avae/utils.py | 292 ----------------------------- avae/utils_vis.py | 293 ++++++++++++++++++++++++++++++ avae/vis.py | 2 +- tests/test_utils.py | 2 +- tools/plotting_post_process.ipynb | 2 +- 5 files changed, 296 insertions(+), 295 deletions(-) create mode 100644 avae/utils_vis.py diff --git a/avae/utils.py b/avae/utils.py index 6aa41021..aade1513 100644 --- a/avae/utils.py +++ b/avae/utils.py @@ -1,18 +1,8 @@ -import logging import os.path -import pathlib -import matplotlib.pyplot as plt -import mrcfile import numpy as np -import numpy.typing as npt -import sklearn.linear_model -import sklearn.metrics import torch -SCREEN_HEIGHT_PX = 720 # top half of a 1440p screen -VECTOR_FORMATS = ("pdf", "svg", "eps") - def as_list(value: object) -> list: if isinstance(value, torch.Tensor): @@ -26,288 +16,6 @@ def as_list(value: object) -> list: return [value] -def create_grid_for_plotting( - rows: int, columns: int, dsize: tuple, padding: int = 0 -) -> npt.NDArray: - - # define the dimensions for the napari grid - - if len(dsize) == 3: - grid_for_napari = np.zeros( - ( - rows * dsize[0], - dsize[1] * columns + padding * columns, - dsize[2], - ), - dtype=np.float32, - ) - - elif len(dsize) == 2: - grid_for_napari = np.zeros( - ( - rows * dsize[0], - dsize[1] * columns + padding * columns, - ), - dtype=np.float32, - ) - - return grid_for_napari - - -def fill_grid_for_plottting( - rows: int, - columns: int, - grid: npt.NDArray, - dsize: tuple, - array: npt.NDArray, - padding: int = 0, -) -> npt.NDArray: - - if len(dsize) == 3: - for j in range(columns): - for i in range(rows): - grid[ - i * dsize[0] : (i + 1) * dsize[0], - j * (dsize[1] + padding) : (j + 1) * dsize[1] - + padding * j, - :, - ] = array[i, j, :, :, :] - - elif len(dsize) == 2: - for j in range(columns): - for i in range(rows): - grid[ - i * dsize[0] : (i + 1) * dsize[0], - j * (dsize[1] + padding) : (j + 1) * dsize[1] - + padding * j, - ] = array[i, j, :, :] - return grid - - -def save_figure( - fig: plt.Figure, - path: str, - vis_format: str, - vis_print: bool = False, - display: bool = False, - cap_height: bool = True, -) -> None: - """Save (or, with `display`, only show) a figure and close it. - - Vector formats take no dpi. `vis_print` saves at 300 dpi and full size. - Otherwise, with `cap_height`, the dpi is lowered so the figure is at most - SCREEN_HEIGHT_PX tall, as large figures cost time and space to rasterise. - """ - if display: - plt.show() - plt.close(fig) - return - - pathlib.Path(path).parent.mkdir(parents=True, exist_ok=True) - kwargs: dict = {} - if vis_format not in VECTOR_FORMATS: - dpi = plt.rcParams["figure.dpi"] - if vis_print: - dpi = 300 - elif cap_height: - dpi = min(dpi, SCREEN_HEIGHT_PX / fig.get_figheight()) - kwargs["dpi"] = dpi - if vis_format == "png": - kwargs["pil_kwargs"] = {"compress_level": 3} - fig.savefig(path, **kwargs) - plt.close(fig) - - -def save_imshow_png( - fname: str, - array: npt.NDArray, - cmap: str | None = None, - min: float | None = None, - max: float | None = None, - display: bool = False, - vis_print: bool = False, -) -> None: - height, width = array.shape[:2] - if width >= height: - image_figure_size = (10, 10 * height / width) - else: - image_figure_size = (10 * width / height, 10) - fig, ax = plt.subplots(figsize=image_figure_size) - ax.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last - ax.axis("off") - fig.subplots_adjust(left=0, right=1, bottom=0, top=1) - save_figure( - fig, - "plots/" + fname, - fname.rsplit(".", 1)[-1], - vis_print, - display, - cap_height=False, - ) - - -def save_mrc_file(fname: str, array: npt.NDArray) -> None: - if not os.path.exists("plots"): - os.mkdir("plots") - with mrcfile.new("plots/" + fname, overwrite=True) as mrc: - mrc.set_data(array) - - -def colour_per_class(classes: list) -> list: - # Define the number of colors you want - num_colors = len(classes) - - # Choose colormaps for combining - cmap_1 = plt.get_cmap("tab20") - cmap_2 = plt.get_cmap("Accent") - cmap_3 = plt.get_cmap("Pastel1") - cmap_4 = plt.get_cmap("Set1") - - # Combine the four colormaps - combined_cmap = [cmap_1(i % 20) for i in range(20)] - combined_cmap.extend([cmap_2(i % 8) for i in range(8)]) - combined_cmap.extend([cmap_3(i % 8) for i in range(8)]) - combined_cmap.extend([cmap_4(i % 8) for i in range(8)]) - - # Create the colormap object - custom_cmap = plt.cm.colors.ListedColormap( - combined_cmap, name="custom_cmap" - ) - - # Generate a list of colors based on the modulo operation of i with respect to the number of colors in the combined colormap - colours = [custom_cmap(i % len(combined_cmap)) for i in range(num_colors)] - return colours - - -def pose_interpolation( - enc: npt.NDArray, - pos_dims: int, - pose_mean: npt.NDArray, - pose_std: npt.NDArray, - dsize: tuple, - number_of_samples: int, - vae: torch.nn.Module, - device: torch.device, -) -> npt.NDArray: - """This function: - 1- interpolates within each pose channels - for the number_of_samples requested. - 2- returns all decoded images based on the input latent - and the interpolated pose - - Parameters - ---------- - enc: numpy array - the latent encoding. - pos_dims: int - the pose channel dimension - pose_mean: numpy array - mean of each pose channel. - pose_std: numpy array - standard deviation of each pose channel. - dsize: torch.size - the dimension of the data. Example [32,32,32] - number_of_samples: int - number of samples to interpolate for. - vae: torch.nn.Module - Affinity vae model. - device: torch.device - Device to run the model on. - """ - # Single transversals along each pose channel, row-major: p_dim, grid_spot - steps = -1.2 + 0.4 * np.arange(number_of_samples) - pose_mean = np.asarray(pose_mean) - means = np.tile(pose_mean, (pos_dims, number_of_samples, 1)) - means[ - np.arange(pos_dims)[:, None], - np.arange(number_of_samples)[None, :], - np.arange(pos_dims)[:, None], - ] += (np.asarray(pose_std)[:, None] * steps[None, :]).astype(means.dtype) - means = means.reshape(-1, means.shape[-1]) - - n = means.shape[0] - pos = torch.from_numpy(means).to(device) - lat = torch.from_numpy(np.array(enc)).unsqueeze(0).expand(n, -1).to(device) - - # Decode all interpolated vectors in a single batched call - with torch.no_grad(): - decoded_img = vae.decoder(lat, pos) - - decoded_grid = np.reshape( - decoded_img.cpu().numpy(), (pos_dims, number_of_samples, *dsize) - ) - - return decoded_grid - - -def latent_space_similarity_mat( - latent_space: npt.NDArray, - class_labels: npt.NDArray, - unique_classes: list, - num_classes: int, - plot_mode: str = "", -) -> npt.NDArray: - """ - This function calculates the similarity (affinity) between classes in the latent space and builds a matrix. - Parameters - ---------- - latent_space: np.ndarray - The latent space - class_labels: np.array - The labels of the latent space - mode: str - Mode of the calculation (train, test, val) - epoch: int - Epoch number for title - display: bool - When this variable is set to true, the function only dispalys the plot and doesnt save it. - """ - # The cosine similarity of two vectors is the dot product of their unit - # length versions: cos(a, b) = u_a . u_b, with u = x / |x|. Averaging it - # over all pairs of two classes therefore reduces to dot products of - # per-class sums of unit vectors, so the N x N similarity matrix never - # needs to be built: - # mean_ij = (S_i . S_j) / (n_i n_j), S_c = sum of u over class c - # For the std, the squared similarity (u_a . u_b)^2 equals the Frobenius - # product _F, so the mean squared similarity is - # E_ij[s^2] = _F / (n_i n_j), M_c = sum of u u^T in class c - # and std_ij = sqrt(E_ij[s^2] - mean_ij^2). - # Self-pairs are included on the diagonal, as in the full pairwise version. - latent_space = np.asarray(latent_space, dtype=np.float64) - class_labels = np.asarray(class_labels) - # unit length latent vectors; zero vectors stay zero (similarity 0) - norms = np.linalg.norm(latent_space, axis=1, keepdims=True) - unit = latent_space / np.where(norms == 0, 1.0, norms) - - latent_dims = unit.shape[1] - counts = np.zeros(num_classes) - sums = np.zeros((num_classes, latent_dims)) - outer_sums = np.zeros((num_classes, latent_dims * latent_dims)) - for i in range(num_classes): - class_unit = unit[class_labels == unique_classes[i]] - counts[i] = len(class_unit) - sums[i] = class_unit.sum(axis=0) - if plot_mode == "std": - outer_sums[i] = (class_unit.T @ class_unit).ravel() - - pair_counts = np.outer(counts, counts) - has_pairs = pair_counts > 0 - safe_pair_counts = np.where(has_pairs, pair_counts, 1.0) - mean_sim = (sums @ sums.T) / safe_pair_counts - - cosine_sim_mat = np.zeros((num_classes, num_classes)) - if plot_mode == "mean": - cosine_sim_mat = mean_sim - elif plot_mode == "std": - second_moment = (outer_sums @ outer_sums.T) / safe_pair_counts - cosine_sim_mat = np.sqrt( - np.clip(second_moment - mean_sim**2, 0, None) - ) - cosine_sim_mat[~has_pairs] = 0.0 - - return cosine_sim_mat - - def latest_file(path: str, extension: str) -> str: most_recent_file = "" diff --git a/avae/utils_vis.py b/avae/utils_vis.py new file mode 100644 index 00000000..3d910208 --- /dev/null +++ b/avae/utils_vis.py @@ -0,0 +1,293 @@ +import os.path +import pathlib + +import matplotlib.pyplot as plt +import mrcfile +import numpy as np +import numpy.typing as npt +import torch + +SCREEN_HEIGHT_PX = 720 # top half of a 1440p screen +VECTOR_FORMATS = ("pdf", "svg", "eps") + + +def create_grid_for_plotting( + rows: int, columns: int, dsize: tuple, padding: int = 0 +) -> npt.NDArray: + + # define the dimensions for the napari grid + + if len(dsize) == 3: + grid_for_napari = np.zeros( + ( + rows * dsize[0], + dsize[1] * columns + padding * columns, + dsize[2], + ), + dtype=np.float32, + ) + + elif len(dsize) == 2: + grid_for_napari = np.zeros( + ( + rows * dsize[0], + dsize[1] * columns + padding * columns, + ), + dtype=np.float32, + ) + + return grid_for_napari + + +def fill_grid_for_plottting( + rows: int, + columns: int, + grid: npt.NDArray, + dsize: tuple, + array: npt.NDArray, + padding: int = 0, +) -> npt.NDArray: + + if len(dsize) == 3: + for j in range(columns): + for i in range(rows): + grid[ + i * dsize[0] : (i + 1) * dsize[0], + j * (dsize[1] + padding) : (j + 1) * dsize[1] + + padding * j, + :, + ] = array[i, j, :, :, :] + + elif len(dsize) == 2: + for j in range(columns): + for i in range(rows): + grid[ + i * dsize[0] : (i + 1) * dsize[0], + j * (dsize[1] + padding) : (j + 1) * dsize[1] + + padding * j, + ] = array[i, j, :, :] + return grid + + +def save_figure( + fig: plt.Figure, + path: str, + vis_format: str, + vis_print: bool = False, + display: bool = False, + cap_height: bool = True, +) -> None: + """Save (or, with `display`, only show) a figure and close it. + + Vector formats take no dpi. `vis_print` saves at 300 dpi and full size. + Otherwise, with `cap_height`, the dpi is lowered so the figure is at most + SCREEN_HEIGHT_PX tall, as large figures cost time and space to rasterise. + """ + if display: + plt.show() + plt.close(fig) + return + + pathlib.Path(path).parent.mkdir(parents=True, exist_ok=True) + kwargs: dict = {} + if vis_format not in VECTOR_FORMATS: + dpi = plt.rcParams["figure.dpi"] + if vis_print: + dpi = 300 + elif cap_height: + dpi = min(dpi, SCREEN_HEIGHT_PX / fig.get_figheight()) + kwargs["dpi"] = dpi + if vis_format == "png": + kwargs["pil_kwargs"] = {"compress_level": 3} + fig.savefig(path, **kwargs) + plt.close(fig) + + +def save_imshow_png( + fname: str, + array: npt.NDArray, + cmap: str | None = None, + min: float | None = None, + max: float | None = None, + display: bool = False, + vis_print: bool = False, +) -> None: + height, width = array.shape[:2] + if width >= height: + image_figure_size = (10, 10 * height / width) + else: + image_figure_size = (10 * width / height, 10) + fig, ax = plt.subplots(figsize=image_figure_size) + ax.imshow(array, cmap=cmap, vmin=min, vmax=max) # channels last + ax.axis("off") + fig.subplots_adjust(left=0, right=1, bottom=0, top=1) + save_figure( + fig, + "plots/" + fname, + fname.rsplit(".", 1)[-1], + vis_print, + display, + cap_height=False, + ) + + +def save_mrc_file(fname: str, array: npt.NDArray) -> None: + if not os.path.exists("plots"): + os.mkdir("plots") + with mrcfile.new("plots/" + fname, overwrite=True) as mrc: + mrc.set_data(array) + + +def colour_per_class(classes: list) -> list: + # Define the number of colors you want + num_colors = len(classes) + + # Choose colormaps for combining + cmap_1 = plt.get_cmap("tab20") + cmap_2 = plt.get_cmap("Accent") + cmap_3 = plt.get_cmap("Pastel1") + cmap_4 = plt.get_cmap("Set1") + + # Combine the four colormaps + combined_cmap = [cmap_1(i % 20) for i in range(20)] + combined_cmap.extend([cmap_2(i % 8) for i in range(8)]) + combined_cmap.extend([cmap_3(i % 8) for i in range(8)]) + combined_cmap.extend([cmap_4(i % 8) for i in range(8)]) + + # Create the colormap object + custom_cmap = plt.cm.colors.ListedColormap( + combined_cmap, name="custom_cmap" + ) + + # Generate a list of colors based on the modulo operation of i with respect to the number of colors in the combined colormap + colours = [custom_cmap(i % len(combined_cmap)) for i in range(num_colors)] + return colours + + +def pose_interpolation( + enc: npt.NDArray, + pos_dims: int, + pose_mean: npt.NDArray, + pose_std: npt.NDArray, + dsize: tuple, + number_of_samples: int, + vae: torch.nn.Module, + device: torch.device, +) -> npt.NDArray: + """This function: + 1- interpolates within each pose channels + for the number_of_samples requested. + 2- returns all decoded images based on the input latent + and the interpolated pose + + Parameters + ---------- + enc: numpy array + the latent encoding. + pos_dims: int + the pose channel dimension + pose_mean: numpy array + mean of each pose channel. + pose_std: numpy array + standard deviation of each pose channel. + dsize: torch.size + the dimension of the data. Example [32,32,32] + number_of_samples: int + number of samples to interpolate for. + vae: torch.nn.Module + Affinity vae model. + device: torch.device + Device to run the model on. + """ + # Single transversals along each pose channel, row-major: p_dim, grid_spot + steps = -1.2 + 0.4 * np.arange(number_of_samples) + pose_mean = np.asarray(pose_mean) + means = np.tile(pose_mean, (pos_dims, number_of_samples, 1)) + means[ + np.arange(pos_dims)[:, None], + np.arange(number_of_samples)[None, :], + np.arange(pos_dims)[:, None], + ] += (np.asarray(pose_std)[:, None] * steps[None, :]).astype(means.dtype) + means = means.reshape(-1, means.shape[-1]) + + n = means.shape[0] + pos = torch.from_numpy(means).to(device) + lat = torch.from_numpy(np.array(enc)).unsqueeze(0).expand(n, -1).to(device) + + # Decode all interpolated vectors in a single batched call + with torch.no_grad(): + decoded_img = vae.decoder(lat, pos) + + decoded_grid = np.reshape( + decoded_img.cpu().numpy(), (pos_dims, number_of_samples, *dsize) + ) + + return decoded_grid + + +def latent_space_similarity_mat( + latent_space: npt.NDArray, + class_labels: npt.NDArray, + unique_classes: list, + num_classes: int, + plot_mode: str = "", +) -> npt.NDArray: + """ + This function calculates the similarity (affinity) between classes in the latent space and builds a matrix. + Parameters + ---------- + latent_space: np.ndarray + The latent space + class_labels: np.array + The labels of the latent space + mode: str + Mode of the calculation (train, test, val) + epoch: int + Epoch number for title + display: bool + When this variable is set to true, the function only dispalys the plot and doesnt save it. + """ + # The cosine similarity of two vectors is the dot product of their unit + # length versions: cos(a, b) = u_a . u_b, with u = x / |x|. Averaging it + # over all pairs of two classes therefore reduces to dot products of + # per-class sums of unit vectors, so the N x N similarity matrix never + # needs to be built: + # mean_ij = (S_i . S_j) / (n_i n_j), S_c = sum of u over class c + # For the std, the squared similarity (u_a . u_b)^2 equals the Frobenius + # product _F, so the mean squared similarity is + # E_ij[s^2] = _F / (n_i n_j), M_c = sum of u u^T in class c + # and std_ij = sqrt(E_ij[s^2] - mean_ij^2). + # Self-pairs are included on the diagonal, as in the full pairwise version. + latent_space = np.asarray(latent_space, dtype=np.float64) + class_labels = np.asarray(class_labels) + # unit length latent vectors; zero vectors stay zero (similarity 0) + norms = np.linalg.norm(latent_space, axis=1, keepdims=True) + unit = latent_space / np.where(norms == 0, 1.0, norms) + + latent_dims = unit.shape[1] + counts = np.zeros(num_classes) + sums = np.zeros((num_classes, latent_dims)) + outer_sums = np.zeros((num_classes, latent_dims * latent_dims)) + for i in range(num_classes): + class_unit = unit[class_labels == unique_classes[i]] + counts[i] = len(class_unit) + sums[i] = class_unit.sum(axis=0) + if plot_mode == "std": + outer_sums[i] = (class_unit.T @ class_unit).ravel() + + pair_counts = np.outer(counts, counts) + has_pairs = pair_counts > 0 + safe_pair_counts = np.where(has_pairs, pair_counts, 1.0) + mean_sim = (sums @ sums.T) / safe_pair_counts + + cosine_sim_mat = np.zeros((num_classes, num_classes)) + if plot_mode == "mean": + cosine_sim_mat = mean_sim + elif plot_mode == "std": + second_moment = (outer_sums @ outer_sums.T) / safe_pair_counts + cosine_sim_mat = np.sqrt( + np.clip(second_moment - mean_sim**2, 0, None) + ) + cosine_sim_mat[~has_pairs] = 0.0 + + return cosine_sim_mat diff --git a/avae/vis.py b/avae/vis.py index 715af6d2..3613f546 100644 --- a/avae/vis.py +++ b/avae/vis.py @@ -20,7 +20,7 @@ import torchvision from . import utils_learning -from .utils import ( +from .utils_vis import ( VECTOR_FORMATS, colour_per_class, create_grid_for_plotting, diff --git a/tests/test_utils.py b/tests/test_utils.py index 28f0328c..d6c36b3a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -5,8 +5,8 @@ import sklearn.metrics.pairwise from avae.base import dims_after_pooling -from avae.utils import latent_space_similarity_mat from avae.utils_learning import combine_accuracy_data +from avae.utils_vis import latent_space_similarity_mat @pytest.mark.parametrize( diff --git a/tools/plotting_post_process.ipynb b/tools/plotting_post_process.ipynb index fafe5318..a5c2e86f 100644 --- a/tools/plotting_post_process.ipynb +++ b/tools/plotting_post_process.ipynb @@ -23,7 +23,7 @@ "import matplotlib.pyplot as plt\n", "from ipywidgets import interact, IntSlider, fixed, Dropdown\n", "from avae.vis import latent_embed_plot_tsne\n", - "from avae.utils import colour_per_class\n", + "from avae.utils_vis import colour_per_class\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "%matplotlib inline"