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 4aad47dc..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"
     )
@@ -328,6 +337,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 +348,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,11 +384,16 @@ 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',
                 'vis_format',
+                'vis_print',
                 'vis_z_n_int',
                 'vis_pose_class',
             ]:
@@ -385,7 +402,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..b879d1af 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.
 
@@ -140,7 +146,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:
@@ -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).
@@ -202,11 +209,19 @@ 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)], ()))
-            plot_classes_distribution(train_y, "train", vis_format=vis_format)
+            train_y = _labels(trains.dataset)
+            val_y = _labels(vals.dataset)
+            plot_classes_distribution(
+                train_y,
+                "train",
+                vis_format=vis_format,
+                vis_print=vis_print,
+            )
             plot_classes_distribution(
-                val_y, "validation", vis_format=vis_format
+                val_y,
+                "validation",
+                vis_format=vis_format,
+                vis_print=vis_print,
             )
 
         train_size = len(trains.dataset)
@@ -251,9 +266,12 @@ 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", 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.
 
@@ -309,12 +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,
+            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]
@@ -326,8 +371,22 @@ def get_affinity_matrix(
         return None
 
 
-class AffinityDiskDataset(caked.dataloader.DiskDataset):
+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"""
 
@@ -357,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/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,
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 b8689726..aade1513 100644
--- a/avae/utils.py
+++ b/avae/utils.py
@@ -1,15 +1,6 @@
-import copy
-import logging
 import os.path
-import typing
 
-import matplotlib.pyplot as plt
-import mrcfile
 import numpy as np
-import numpy.typing as npt
-import sklearn.linear_model
-import sklearn.metrics
-import sklearn.metrics.pairwise
 import torch
 
 
@@ -25,239 +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_imshow_png(
-    fname: str,
-    array: npt.NDArray,
-    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,
-) -> None:
-    if not display:
-        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)
-
-        if writer:
-            writer.add_figure(figname, fig, epoch)
-    else:
-        plt.imshow(array, cmap=cmap, vmin=min, vmax=max)  # channels last
-        plt.show()
-
-    plt.close()
-
-
-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.
-    """
-    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())
-
-    decoded_grid = np.reshape(
-        np.array(decoded_grid), (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.
-    """
-    # get same label order as affinity matrix
-    cosine_sim_matrix = sklearn.metrics.pairwise.cosine_similarity(
-        latent_space
-    )
-
-    cosine_sim_mat = np.zeros((num_classes, num_classes))
-
-    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
-
-    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 <u_a u_a^T, u_b u_b^T>_F, so the mean squared similarity is
+    #   E_ij[s^2] = <M_i, M_j>_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 92fd2a09..3613f546 100644
--- a/avae/vis.py
+++ b/avae/vis.py
@@ -1,13 +1,13 @@
-import copy
+import functools
 import logging
 import os.path
 import pathlib
 import random
+import re
 import typing
 import warnings
 
 import altair
-import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
@@ -20,162 +20,853 @@
 import torchvision
 
 from . import utils_learning
-from .utils import (
+from .utils_vis 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
+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 _encoder(i: PIL.Image) -> str:
-    """Encode PIL Image as base64 buffer.
-    Parameters
-    ----------
-    i: PIL.Image
-        Image to be encoded.
 
-    Returns
-    -------
-    str
-        Encoded image.
+def _class_plot_scale(
+    num_classes: int, fig_size: float | None = None
+) -> tuple[float, int, int]:
+    """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.
     """
-    import base64
-    import io
+    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))
+    title_size = max(16, font_size + 6, round(1.5 * resolved_fig_size))
+    return resolved_fig_size, font_size, title_size
 
-    with io.BytesIO() as buffer:
-        i.thumbnail((110, 110))
-        i.save(buffer, "PNG")
-        data = base64.encodebytes(buffer.getvalue()).decode("utf-8")
 
-    return f"{data}"
+@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 _decoder(i: str) -> PIL.Image:
-    """Decode base64 buffer as PIL 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)
+
+    for label, values, color, linestyle in series:
+        ax.plot(
+            range(1, epochs + 1),
+            values,
+            c=color,
+            linestyle=linestyle,
+            label=label,
+            linewidth=linewidth,
+        )
+
+    if log_scale:
+        ax.set_yscale("log")
+    ax.set_ylabel(ylabel)
+    ax.set_xlabel("Epochs")
+    if legend:
+        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(
+    data: npt.NDArray,
+    class_labels: typing.Sequence | npt.NDArray,
+    title: str,
+    path: str,
+    vis_format: str = "png",
+    vis_print: bool = False,
+    *,
+    cmap,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    values_format: str | 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,
+    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.
+    is_vector = vis_format in VECTOR_FORMATS
+
+    # 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",
+        cmap=cmap,
+        vmin=vmin,
+        vmax=vmax,
+    )
+
+    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):
+                ax.text(
+                    j,
+                    i,
+                    f"{data[i, j]:{values_format}}",
+                    ha="center",
+                    va="center",
+                    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)
+    ax.set_yticks(np.arange(num_classes), labels=class_labels)
+    ax.set_ylim((num_classes - 0.5, -0.5))
+
+    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")
+
+    ax.tick_params(axis="x", rotation=90, labelsize=font_size)
+    ax.tick_params(axis="y", labelsize=font_size)
+
+    colorbar = fig.colorbar(im, cax=cax)
+    colorbar.ax.tick_params(labelsize=font_size)
+
+    if xlabel:
+        ax.set_xlabel(xlabel, fontsize=font_size)
+    if ylabel:
+        ax.set_ylabel(ylabel, fontsize=font_size)
+
+    _set_title(ax, title, title_font_size)
+    save_figure(fig, path, vis_format, vis_print, display)
+
+
+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,
+        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])
+
+    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(
+    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,
+        )
+
+
+@_styled
+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",
+    vis_print: bool = False,
+) -> None:
+    """Visualise loss over epochs.
+
     Parameters
     ----------
-    i: str
-        Encoded image.
-
-    Returns
-    -------
-    PIL.Image
-        Decoded image.
+    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.
     """
-    import base64
-    import io
+    logging.info(
+        "################################################################",
+    )
+    logging.info("Visualising loss ...\n")
 
-    return PIL.Image.open(io.BytesIO(base64.b64decode(i)))
+    train_loss = np.transpose(np.asarray(train_loss))
+    if val_loss is not None:
+        val_loss = np.transpose(np.asarray(val_loss))
 
+    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
+
+    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],
+            )
+        )
+    else:
+        title = "Loss"
+
+    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
+    )
+
+    # 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
+    )
 
-def format(im: PIL.Image, data_dim: int) -> list[str]:
-    """Format PIL Image as Pandas compatible Altair image display.
+
+@_styled
+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
 
     Parameters
     ----------
-    im: PIL.Image
-        Image to be formatted.
-    data_dim: int
-        Dimension of the data.
+    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,
+    )
 
-    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.
+
+@_styled
+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 = "",
+    epoch: int = 0,
+    vis_format: str = "png",
+    vis_print: bool = False,
+):
+    """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.
     """
-    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()
+    logging.info(
+        "################################################################",
+    )
+
+    logging.info("Visualising accuracy: confusion and F1 scores ...\n")
+
+    classes_list = np.unique(np.concatenate((y_train, ypred_train)))
+    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))
+    )
+
+    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),
+    ):
+        _confusion_figure(
+            y, ypred, labels, f"plots/{stem}", epoch, vis_format, vis_print
         )
-    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()
+
+
+@_styled
+def f1_plot(
+    y_train: npt.NDArray,
+    ypred_train: npt.NDArray,
+    y_val: npt.NDArray,
+    ypred_val: npt.NDArray,
+    mode: str = "",
+    epoch: int = 0,
+    vis_format: str = "png",
+    vis_print: bool = False,
+):
+    """Plot F1 values using classes observed in the training labels.
+
+    Evaluation classes unseen during training are excluded from the scores.
+
+    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
+        Epoch number.
+    """
+    logging.info(
+        "################################################################",
+    )
+    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, font_size, title_font_size = _class_plot_scale(len(classes_list))
+
+    if np.setdiff1d(classes_list_eval, classes_list).size > 0:
+        logging.info(
+            f"Class {np.setdiff1d(classes_list_eval, classes_list)} will not be used to compute F1 values as it was unseen in training data."
         )
-    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)
+
+        index = np.argwhere(np.isin(y_val, classes_list)).ravel()
+
+        y_val = np.array(y_val)[index].tolist()
+        ypred_val = np.array(ypred_val)[index].tolist()
+
+    train_f1_score = sklearn.metrics.f1_score(
+        y_train,
+        ypred_train,
+        average=None,
+        labels=classes_list,
+        zero_division=0,
+    ).tolist()
+    valid_f1_score = sklearn.metrics.f1_score(
+        y_val,
+        ypred_val,
+        average=None,
+        labels=classes_list,
+        zero_division=0,
+    ).tolist()
+
+    if mode == "_eval":
+        label = "eval"
     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]
+        label = "valid"
+
+    f1_train_file = f"plots/f1_train{mode}.csv"
+    f1_valid_file = f"plots/f1_{label}.csv"
+
+    train_df = pd.DataFrame([train_f1_score], columns=classes_list)
+    valid_df = pd.DataFrame([valid_f1_score], columns=classes_list)
+
+    train_df["epoch"] = epoch + 1
+    valid_df["epoch"] = epoch + 1
+
+    train_df["f1_avg"] = np.mean(train_f1_score)
+    valid_df["f1_avg"] = np.mean(valid_f1_score)
+
+    if os.path.exists(f1_train_file) and epoch != 0:
+        f1_train = pd.read_csv(f1_train_file)
+        pd.concat([f1_train, train_df]).to_csv(f1_train_file, index=False)
     else:
-        return [_encoder(PIL.Image.fromarray(im))]
+        train_df.to_csv(f1_train_file, index=False)
 
+    if os.path.exists(f1_valid_file) and epoch != 0:
+        f1_valid = pd.read_csv(f1_valid_file)
+        pd.concat([f1_valid, valid_df]).to_csv(f1_valid_file, index=False)
+    else:
+        valid_df.to_csv(f1_valid_file, index=False)
 
-def merge(im: str) -> str | None:
-    """Merge 2 base64 buffers as PIL Images and encode back to base64
-    buffers.
+    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
+        )
+    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,
+    label: list,
+    data_dim: int,
+    mode: str = "trn",
+    display: bool = False,
+    vis_format: str = "png",
+    vis_print: bool = False,
+) -> None:
+    """Visualise reconstructions.
 
     Parameters
     ----------
-    im: str
-        Input PIL Images to be merged.
-
-    Returns
-    -------
-    str
-        Merged image.
+    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.
     """
-    if im.startswith("data:image/"):
-        return im
+    logging.info(
+        "################################################################",
+    )
+    logging.info("Visualising reconstructions " + mode + "...\n")
 
-    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)),
+    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)),
+        display=display,
+        vis_print=vis_print,
+    )
+    save_imshow_png(
+        fname_out,
+        np.transpose(rec_2d, (1, 2, 0)),
+        display=display,
+        vis_print=vis_print,
+    )
+
+    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,
         )
-        return None
 
-    im1 = _decoder(i[0])
-    im2 = _decoder(i[1])
+        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]
 
-    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)
+            # stack the images together with their reconstruction
+            rec_img = np.hstack((selected_img, selected_rec))
 
-    return f"data:image/png;base64,{data}"
+            # 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)
+            )
 
+        save_mrc_file(f"recon_{mode}.mrc", grid_for_napari)
+        pd.DataFrame(reconstruction_index).to_csv(
+            f"plots/recon_{mode}.csv", index=False
+        )
+
+
+@_styled
 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,
     vis_format: str = "png",
+    vis_print: bool = False,
     embedding: npt.NDArray | None = None,
 ) -> None:
     """Plot static TSNE embedding.
@@ -192,8 +883,6 @@ def latent_embed_plot_tsne(
         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
@@ -248,77 +937,90 @@ 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)
-        )
-    # 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:
 
-        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)]
 
-            plt.scatter(
+            ax.scatter(
                 lats[idx, 0],
                 lats[idx, 1],
                 s=marker_size,
-                label=mol[:4],
+                label=mol,
                 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.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]
+        for mol in np.unique(ys):
+            idx = ys == mol
             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,
             )
-        plt.legend(
-            prop={"size": 10},
-            bbox_to_anchor=(1.05, 1),
-            loc="upper left",
-            fontsize=16,
-        )
-        plt.xlabel("dim 1")
-        plt.ylabel("freq")
-
-    plt.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()
+        ax.set_xlabel("dim 1", fontsize=font_size)
+        ax.set_ylabel("freq", fontsize=font_size)
+
+    # 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)
+    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,
+        )
+    )
+    _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(
@@ -378,11 +1080,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:",
@@ -495,704 +1199,144 @@ def dyn_latentembed_plot(
     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]))]
-    rows = len(np.unique(y)) // 2
-    if len(np.unique(y)) % 2 != 0:
-        rows += 1
-    fig, ax = plt.subplots(len(np.unique(y)), sharex=True, sharey=True)
-    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)
-        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)
-    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")
-    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)))
-
-    # 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
-    )
-
-    # 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": int(len(classes_list) / 3) + 3}
-    ):
-        fig, ax = plt.subplots(
-            figsize=(int(len(classes_list)) / 2, int(len(classes_list)) / 2)
-        )
-
-        disp.plot(
-            cmap=plt.cm.Blues,
-            ax=ax,
-            xticks_rotation=90,
-            values_format="d",
-        )
-
-        plt.tight_layout()
-        plt.title(
-            "Average accuracy at epoch {}: {:.3f}%".format(
-                epoch + 1,
-                np.mean(avg_accuracy[supported_train_classes]) * 100,
-            ),
-            fontsize=10,
-        )
-
-        if not os.path.exists("plots"):
-            os.mkdir("plots")
-
-        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=(int(len(classes_list)) / 2, int(len(classes_list)) / 2)
-        )
-
-        dispn.plot(
-            cmap=plt.cm.Blues,
-            ax=ax,
-            xticks_rotation=90,
-            values_format=".0f",
-            im_kw={"vmin": 0, "vmax": 100},
-        )
-
-        plt.tight_layout()
-        plt.title(
-            "Average accuracy at epoch {}: {:.3f}%".format(
-                epoch + 1,
-                np.mean(avg_accuracy[supported_train_classes]) * 100,
-            ),
-            fontsize=12,
-        )
-
-        plt.xlabel("Predicted label (%)")
-        plt.ylabel("True label (%)")
-        plt.savefig(f"plots/confusion_train{mode}_norm.{vis_format}", dpi=300)
-
-        if writer:
-            writer.add_figure("Accuracy (Norm)", fig, epoch)
-
-        plt.close()
-
-    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)
-
-    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):
-        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)
-
-    dispn_eval = sklearn.metrics.ConfusionMatrixDisplay(
-        confusion_matrix=cmn_eval, display_labels=ordered_class_eval
-    )
-
-    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": int(len(ordered_class_eval) / 3) + 3,
-        }
-    ):
-        fig, ax = plt.subplots(
-            figsize=(
-                int(len(ordered_class_eval)) / 2,
-                int(len(ordered_class_eval)) / 2,
-            )
-        )
-        disp_eval.plot(
-            cmap=plt.cm.Blues,
-            ax=ax,
-            xticks_rotation=90,
-            values_format="d",
-        )
-        plt.tight_layout()
-        plt.title(
-            "Average accuracy at epoch {}: {:.1f}%".format(
-                epoch + 1,
-                np.mean(avg_accuracy_eval[supported_eval_classes]) * 100,
-            ),
-            fontsize=12,
-        )
-        plt.savefig(figure_name + f".{vis_format}", dpi=300)
-        plt.close()
-
-        fig, ax = plt.subplots(
-            figsize=(int(len(classes_list)) / 2, int(len(classes_list)) / 2)
-        )
-
-        dispn_eval.plot(
-            cmap=plt.cm.Blues,
-            ax=ax,
-            xticks_rotation=90,
-            values_format=".0f",
-            im_kw={"vmin": 0, "vmax": 100},
-        )
-
-        plt.tight_layout()
-        plt.title(
-            "Average accuracy at epoch {}: {:.1f}% ".format(
-                epoch + 1,
-                np.mean(avg_accuracy_eval[supported_eval_classes]) * 100,
-            ),
-            fontsize=10,
-        )
-        plt.xlabel("Predicted label (%)")
-        plt.ylabel("True label (%)")
-        plt.savefig(f"{figure_name}_norm.{vis_format}", dpi=300)
-        plt.close()
-
-
-def f1_plot(
-    y_train: npt.NDArray,
-    ypred_train: npt.NDArray,
-    y_val: npt.NDArray,
-    ypred_val: npt.NDArray,
-    mode: str = "",
-    epoch: int = 0,
-    writer: typing.Any = None,
-    vis_format: str = "png",
-):
-    """Plot F1 values using classes observed in the training labels.
-
-    Evaluation classes unseen during training are excluded from the scores.
-
-    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
-        Epoch number.
-    writer: SummaryWriter
-        Tensorboard summary writer
-    """
-    logging.info(
-        "################################################################",
-    )
-    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)))
-
-    if np.setdiff1d(classes_list_eval, classes_list).size > 0:
-        logging.info(
-            f"Class {np.setdiff1d(classes_list_eval, classes_list)} will not be used to compute F1 values as it was unseen in training data."
-        )
-
-        index = np.argwhere(np.isin(y_val, classes_list)).ravel()
-
-        y_val = np.array(y_val)[index].tolist()
-        ypred_val = np.array(ypred_val)[index].tolist()
-
-    train_f1_score = sklearn.metrics.f1_score(
-        y_train,
-        ypred_train,
-        average=None,
-        labels=classes_list,
-        zero_division=0,
-    ).tolist()
-    valid_f1_score = sklearn.metrics.f1_score(
-        y_val,
-        ypred_val,
-        average=None,
-        labels=classes_list,
-        zero_division=0,
-    ).tolist()
-
-    if mode == "_eval":
-        label = "eval"
-    else:
-        label = "valid"
-
-    f1_train_file = f"plots/f1_train{mode}.csv"
-    f1_valid_file = f"plots/f1_{label}.csv"
-
-    train_df = pd.DataFrame([train_f1_score], columns=classes_list)
-    valid_df = pd.DataFrame([valid_f1_score], columns=classes_list)
-
-    train_df["epoch"] = epoch + 1
-    valid_df["epoch"] = epoch + 1
-
-    train_df["f1_avg"] = np.mean(train_f1_score)
-    valid_df["f1_avg"] = np.mean(valid_f1_score)
-
-    if os.path.exists(f1_train_file) and epoch != 0:
-        f1_train = pd.read_csv(f1_train_file)
-        pd.concat([f1_train, train_df]).to_csv(f1_train_file, index=False)
-    else:
-        train_df.to_csv(f1_train_file, index=False)
-
-    if os.path.exists(f1_valid_file) and epoch != 0:
-        f1_valid = pd.read_csv(f1_valid_file)
-        pd.concat([f1_valid, valid_df]).to_csv(f1_valid_file, index=False)
-    else:
-        valid_df.to_csv(f1_valid_file, index=False)
-
-    with plt.rc_context(
-        {
-            "font.weight": "bold",
-            "font.size": int(len(classes_list) / 3) + 3,
-        }
-    ):
-        fig, ax = plt.subplots(
-            figsize=(
-                int(len(classes_list)) / 2,
-                int(len(classes_list)) / 2,
-            )
-        )
-        plt.plot(classes_list, train_f1_score, label="train", marker="o")
-        plt.plot(classes_list, valid_f1_score, label=label, marker="o")
-        plt.xticks(rotation=45)
-        plt.legend(loc="lower left")
-        plt.title("F1 Score at epoch {}".format(epoch + 1))
-        plt.ylabel("F1 Score")
-        plt.savefig(f"plots/f1{mode}.{vis_format}", dpi=150)
-
-        if writer:
-            writer.add_figure("F1 score", fig, epoch)
-
-        plt.close()
-
+    chart.save(f"latents/latent_epoch_{epoch + 1}_{mode}.html")
 
-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",
-) -> None:
-    """Visualise loss over epochs.
 
-    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.
-    """
+@_styled
+def confidence_plot(x, y, s, suffix=None, vis_format="png", vis_print=False):
     logging.info(
         "################################################################",
     )
-    logging.info("Visualising loss ...\n")
-
-    train_loss = np.transpose(np.asarray(train_loss))
-    if val_loss is not None:
-        val_loss = np.transpose(np.asarray(val_loss))
-
-    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",
-    ]
-
-    plt.clf()
-    plt.ticklabel_format(useOffset=False)
+    logging.info(
+        "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]))]
+    classes = np.unique(y)
+    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)
+        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)
 
-    train_loss[-2] = train_loss[-2] * beta
-    train_loss[-1] = train_loss[-1] * gamma
+        min_mu = np.min(mu_cl)
+        max_mu = np.max(mu_cl)
+        max_sig = np.max(std_cl)
+        step = (2 * 4 * max_sig) / 100
 
-    for i, loss in enumerate(train_loss):
-        s = "-"
-        plt.plot(
-            range(1, epochs + 1), loss, c=cols[i], linestyle=s, label=labs[i]
-        )
-    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 = "--"
-            plt.plot(
-                range(1, epochs + 1),
-                loss,
-                c=cols[i],
-                linestyle=s,
-                label=vlabs[i],
-            )
+        xs = np.arange(min_mu - (4 * max_sig), max_mu + (4 * max_sig), step)
 
-    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",
+        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),
             )
-            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])
-        )
-    else:
-        plt.title("Loss")
-
-    plt.yscale("log")
-    plt.ylabel("Loss", fontsize=16)
-    plt.xlabel("Epochs", fontsize=16)
-    plt.xticks(fontsize=16)
-    plt.yticks(fontsize=16)
-    plt.legend()
-
-    plt.tight_layout()
-    if not os.path.exists("plots"):
-        os.mkdir("plots")
-    plt.savefig(f"plots/loss.{vis_format}", dpi=300)
-    plt.close()
-
-    # plotting only the total loss as it sometimes is a few order of magnitude higher than KLD and affinity losses
-    plt.plot(
-        range(1, epochs + 1),
-        train_loss[0],
-        c=cols[0],
-        linestyle="-",
-        label=labs[0],
+        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()
+    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,
     )
-    if val_loss is not None:
-        plt.plot(
-            range(1, epochs + 1),
-            val_loss[0],
-            c=cols[0],
-            linestyle="--",
-            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.tight_layout()
-    plt.savefig(f"plots/loss_total.{vis_format}", dpi=300)
-    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)
 
 
-def recon_plot(
-    img: torch.Tensor,
-    rec: torch.Tensor,
-    label: list,
-    data_dim: int,
-    mode: str = "trn",
+@_styled
+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,
+    fig_size: int | None = None,
     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
-
-        # 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,
-        )
+    logging.info("Visualising the latent space similarity matrix ...\n")
 
-        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
+    if affinity_matrix is None:
+        unique_classes: npt.NDArray | list = np.unique(class_labels)
+    else:
+        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(
+                (
+                    classes_order,
+                    np.setdiff1d(unique_classes_in_data, classes_order),
+                )
             )
-            selected_img = img[selected_indices]
-            selected_rec = rec[selected_indices]
-
-            # 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, :, :, :]
+        else:
+            unique_classes = classes_order
 
-            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)
-            )
+    cosine_sim = latent_space_similarity_mat(
+        latent_space,
+        class_labels,
+        unique_classes,
+        len(unique_classes),
+        plot_mode=plot_mode,
+    )
 
-        save_mrc_file(f"recon_{mode}.mrc", grid_for_napari)
-        pd.DataFrame(reconstruction_index).to_csv(
-            f"plots/recon_{mode}.csv", index=False
-        )
+    classes_in_data = set(np.asarray(class_labels).astype(str))
+    _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(
@@ -1205,6 +1349,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.
@@ -1244,8 +1389,13 @@ def latent_4enc_interpolate_plot(
     )
 
     padding = 0
-    data_dim = len(dsize)
     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)
@@ -1269,60 +1419,48 @@ 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)
         )
 
-        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,
@@ -1332,6 +1470,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.
 
@@ -1351,8 +1490,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(
         "################################################################"
@@ -1360,7 +1497,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)
@@ -1373,59 +1509,41 @@ def latent_disentamglement_plot(
         pos_means = np.mean(poses_space, axis=0)
         pos_dims = poses_space.shape[-1]
 
-    recon_images = []
+    # 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)
 
-    # 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)
-
-            # 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)
     )
 
-    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,
@@ -1440,8 +1558,8 @@ 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
     plot for all classes listed in pose_vis_class.
 
@@ -1483,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)
@@ -1493,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, :]
@@ -1516,27 +1633,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,
@@ -1548,6 +1656,7 @@ def pose_disentanglement_plot(
     mode: str = "trn",
     display: bool = False,
     vis_format: str = "png",
+    vis_print: bool = False,
 ):
     """Visualise pose disentanglement.
 
@@ -1584,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)
 
@@ -1605,26 +1713,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,
@@ -1635,6 +1735,7 @@ def interpolations_plot(
     poses: list | None = None,
     mode: str = "trn",
     vis_format: str = "png",
+    vis_print: bool = False,
 ) -> None:
     """Visualise interpolations.
 
@@ -1659,7 +1760,6 @@ def interpolations_plot(
         "################################################################",
     )
     logging.info("Visualising interpolations ...\n")
-    data_dim = len(dsize)
     lats = np.asarray(lats)
     classes = np.asarray(classes)
 
@@ -1701,67 +1801,58 @@ 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_reps_poses[0]
-                    + h * (1 - v) * class_reps_poses[1]
-                    + (1 - h) * v * class_reps_poses[2]
-                    + h * v * class_reps_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,
+    )
 
-    decoded_grid = np.reshape(
-        np.array(decoded_grid), (grid_size, grid_size, *dsize)
+    # 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
     )
 
-    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
+    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(
+        decoded_images.cpu().numpy(), (grid_size, grid_size, *dsize)
     )
-    # 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)
+    _save_grid_figure(
+        grid_size,
+        grid_size,
+        dsize,
+        decoded_grid,
+        f"interpolations_{mode}",
+        vis_format,
+        vis_print=vis_print,
+    )
 
 
+@_styled
 def plot_affinity_matrix(
     lookup: pd.DataFrame,
     all_classes: list,
     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
@@ -1778,58 +1869,30 @@ def plot_affinity_matrix(
     vis_format: str
         format of the image saved
     """
+    logging.info("\n")
     logging.info(
         "################################################################",
     )
     logging.info("Visualising affinity matrix ...\n")
 
-    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))
-    # Create the figure and gridspec
-    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)
-
-    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 the label for selected classes
-    for i, c in enumerate(all_classes):
-        if c 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])
-
-    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])
-    fig.tight_layout()
-    if not os.path.exists("plots"):
-        os.mkdir("plots")
-    plt.savefig(f"plots/affinity_matrix.{vis_format}", dpi=300)
-    plt.close()
+    _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,
+        fig_size,
+    )
 
 
+@_styled
 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
 
@@ -1846,143 +1909,163 @@ 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()
-    if not os.path.exists("plots"):
-        os.mkdir("plots")
-    plt.savefig(f"plots/classes_distribution_{category}.{vis_format}", dpi=300)
-    plt.close()
+    ax.bar(ticks, counts, align="center", color="blue", alpha=0.5)
+    ax.set_xticks(ticks, labels)
+    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)
+    _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}",
+        vis_format,
+        vis_print,
+    )
 
 
-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
+    ----------
+    i: PIL.Image
+        Image to be encoded.
+
+    Returns
+    -------
+    str
+        Encoded image.
+
+    """
+    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
     ----------
-    array : list
-        List of values for the variable
-    variable_name : str
-        Name of the variable
+    i: str
+        Encoded image.
+
+    Returns
+    -------
+    PIL.Image
+        Decoded image.
     """
-    logging.info(
-        "################################################################",
-    )
-    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()
-    if not os.path.exists("plots"):
-        os.mkdir("plots")
-    plt.savefig(f"plots/{variable_name}_array.{vis_format}", dpi=300)
-    plt.close()
+    import base64
+    import io
 
+    return PIL.Image.open(io.BytesIO(base64.b64decode(i)))
 
-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:
+
+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.
     """
-    This function calculates the similarity (affinity) between classes in the latent space and builds a matrix.
+    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))]
+
+
+def merge(im: str) -> str | None:
+    """Merge 2 base64 buffers as PIL Images and encode back to base64
+    buffers.
+
     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
+    im: str
+        Input PIL Images to be merged.
+
+    Returns
+    -------
+    str
+        Merged image.
     """
-    logging.info(
-        "################################################################",
-    )
-    logging.info("Visualising the latent space similarity matrix ...\n")
+    if im.startswith("data:image/"):
+        return im
 
-    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()
+    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)),
         )
-        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 None
 
-    num_classes = len(unique_classes)
-    cosine_sim = latent_space_similarity_mat(
-        latent_space,
-        class_labels,
-        unique_classes,
-        num_classes,
-        plot_mode=plot_mode,
-    )
+    im1 = _decoder(i[0])
+    im2 = _decoder(i[1])
 
-    # Visualize average cosine similarity matrix
-    if fig_size is None:
-        with plt.rc_context(
-            {
-                "font.weight": "bold",
-                "font.size": int(len(unique_classes) / 3) + 3,
-            }
-        ):
-            fig, ax = plt.subplots(
-                figsize=(
-                    int(len(unique_classes)) / 2,
-                    int(len(unique_classes)) / 2,
-                )
-            )
-    else:
-        fig, ax = plt.subplots(figsize=(fig_size, fig_size))
-    fig.tight_layout(pad=3)
-    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")
-    ax.tick_params(axis="x", rotation=90, labelsize=font_size)
-    ax.tick_params(axis="y", labelsize=font_size)
+    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)
 
-    if not display:
-        if not os.path.exists("plots"):
-            os.mkdir("plots")
-        plt.savefig(f"plots/similarity_mean{mode}.{vis_format}", dpi=dpi)
-        plt.close()
-    else:
-        plt.show()
+    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_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_data.py b/tests/test_data.py
index a9fae402..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 load_data
+from avae.data import _labels, get_affinity_matrix, load_data
 from tests import testdata_mrc
 
 
@@ -160,3 +161,76 @@ 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)
+
+
+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()
diff --git a/tests/test_integration_pipeline.py b/tests/test_integration_pipeline.py
index 956b8ca3..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, 34)
+        self.assertEqual(n_plots_train, 33)
         self.assertEqual(n_latent_train, 1)
         self.assertEqual(n_states_train, 2)
-        self.assertEqual(n_plots_eval, 56)
+        self.assertEqual(n_plots_eval, 53)
         self.assertEqual(n_latent_eval, 2)
         self.assertEqual(n_states_eval, 3)
 
diff --git a/tests/test_train_eval_pipeline.py b/tests/test_train_eval_pipeline.py
index 95642edc..6315bc1b 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,
@@ -42,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_utils.py b/tests/test_utils.py
index 84546251..d6c36b3a 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -2,9 +2,11 @@
 
 import numpy as np
 import pytest
+import sklearn.metrics.pairwise
 
 from avae.base import dims_after_pooling
 from avae.utils_learning import combine_accuracy_data
+from avae.utils_vis import latent_space_similarity_mat
 
 
 @pytest.mark.parametrize(
@@ -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)
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/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"
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