scikit-hep / mplhep

Extended histogram plotting on top of matplotlib and HEP collaboration compatible styling
https://mplhep.readthedocs.io
MIT License
189 stars 66 forks source link

feat: add function to merge handles for identical labels in legend #516

Closed cverstege closed 2 months ago

cverstege commented 2 months ago

This function is useful, when using multiple plotting functions and one only wants to use a single label. This function combines all the handles for identical labels into a single handle. I think this comes in handy. e.g. when doing a step plot for a main line and a band around for the uncertainty.

I'm open for feedback.

andrzejnovak commented 2 months ago

Hi @cverstege , this looks ok for me and we already collect some "random" helpers. For posterity, could you add a couple of snippets of how you use it as a comment to this PR?

cverstege commented 2 months ago

Sure. Here's one example:

import mplhep as hep
import numpy as np
import matplotlib.pyplot as plt

if __name__=="__main__":
    s = np.random.normal(size=1000)
    hist, bins = np.histogram(s)
    error = np.sqrt(hist)

    hep.style.use("CMS")
    fig, (ax1, ax2) = plt.subplots(2, 1)
    for ax in (ax1, ax2):
        hep.histplot(hist, bins=bins, yerr=error, histtype="band", label="same label", ax=ax, zorder=-1)
        hep.histplot(hist, bins=bins, histtype="step", color="black", label="same label", ax=ax, zorder=0)
    ax1.legend(loc="upper right")
    handles, labels = ax2.get_legend_handles_labels()
    handles, labels = hep.merge_legend_handles_labels(handles, labels)
    ax2.legend(handles, labels, loc="upper right")
    plt.savefig("test.png")
cverstege commented 2 months ago

And here the plot for reference. The upper one is not using the new function, while the lower one has the combined legend entry. test