ucl-pond / pySuStaIn

Subtype and Stage Inference (SuStaIn) algorithm with an example using simulated data.
MIT License
135 stars 64 forks source link

Enabling PVD Plot Legends #49

Open katrinaCode opened 8 months ago

katrinaCode commented 8 months ago

Hi all,

I just wanted to share some code here in case it may help anyone else: I've added a legend to the PVDs which shows the Z_val associated with each colour. It assumes that the Z_vals are the same for all biomarkers when creating the legend labels. I'm aware that there's an open issue on PVD colourbars and am not offering a solution to that problem, but simply a fix in the meantime to show clearly which value maps to which colour (I myself found this quite confusing to parse out when using more than 3 Z_vals as there was no documentation showing the mappings). It would be well paired with explanations I've seen used in papers about what the intensity, etc. of each colour indicates.

This requires an extra import from Matplotlib in ZscoreSustain: import matplotlib.patches as mpatches.

To implement it, add the following to plot_positional_var in ZscoreSustain between lines 660 (ax.set_title(title_i, fontsize=title_font_size)) and 661 (# Tighten up the figure). Further customization (i.e. legend size, location, etc) is possible through use of the ax.legend arguments.

                ax.set_xlabel(stage_label, fontsize=stage_font_size+2)
                ax.set_title(title_i, fontsize=title_font_size)

                # Add a legend 
                # adding an extra dim to colour mat for RGB reasons
                legend_colour_mat = np.array([[[1, 0, 0], [1, 0, 1], [0, 0, 1], [0.5, 0, 1], [0, 1, 1], [0, 1, 0.5]]])[:N_z]

                patches = [ mpatches.Patch(color=legend_colour_mat[0][i], label=f"Z_val = {zvalues[i]}") for i in range(len(zvalues)) ]

                # put those patches as legend-handles into the legend
                ax.legend(handles=patches, loc = "best" )

            # Tighten up the figure
            #plt.tight_layout()
            fig.tight_layout()


this is an example of the result using the Sustain Workshop data:

Figure 2024-04-03 130249


I also found it helpful to visualize the general mapping of Z_vals to colours, especially if legends aren't used:

Figure 2024-04-03 153502

which I generated using:

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
legend_colour_mat = np.array([[[1, 0, 0], [1, 0, 1], [0, 0, 1], [0.5, 0, 1], [0, 1, 1], [0, 1, 0.5]]])
ax = plt.subplot()
ax.imshow(legend_colour_mat)
ax.set_yticks([])
ax.set_xticks(range(6), ["1st Z_val", "2nd Z_val", "3rd Z_val", "4th Z_val", "5th Z_val", "6th Z_val"])
plt.show()


I hope this will be helpful to some! Thanks 😊

noxtoby commented 5 months ago

Thanks @katrinaCode! I haven't tried this out, so will leave this issue open for now.

Cheers