[ENH] Support for Multi-Color Line Legends in Plots

[nikosarcevic]

Need

This enhancement proposes adding native support for multi-color lines in the legend of line (and other styles!) plots in Matplotlib. This feature will allow users to represent data series with gradients or multiple colors more intuitively within the plot legends.

Proposed solution

Hello, awesome matplotlib people!

Why This Enhancement is Needed

• Improved Data Visualization: Multi-color lines in legends can significantly enhance the readability and interpretability of plots, especially in cases where data series represent a range of values or a transition over time.
• Consistency: Currently, users need to implement custom handlers for multi-color lines, which can be cumbersome and inconsistent. Native support will standardize the implementation.
• User Demand: As data visualizations become more complex and integral to analysis, the demand for more advanced and intuitive plotting features increases. This feature aligns with modern data visualization needs.

Benefits

• Enhanced Clarity: Users can visually match plot lines with legend entries more easily when colors in the legend reflect the actual colors used in the plot.
• Time-Saving: Eliminates the need for users to write and maintain custom legend handlers for multi-color lines.
• Professional Presentation: Provides a more polished and professional look to data visualizations, which is crucial for presentations and publications.

Proposed Implementation

Add support in the Legend class for a new handler that can process and display multi-color lines.
This handler will generate a line segment with a gradient or series of colors, consistent with the corresponding plot line.

My current implementation:

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import LineCollection
from matplotlib.legend_handler import HandlerLineCollection
from matplotlib.colors import Normalize
import cmasher as cmr

class HandlerColorLineCollection(HandlerLineCollection):
    def __init__(self, cmap, **kwargs):
        self.cmap = cmap
        super().__init__(**kwargs)
    
    def create_artists(self, legend, artist, xdescent, ydescent, width, height, fontsize, trans):
        x = np.linspace(0, width, self.get_numpoints(legend) + 1)
        y = np.zeros(self.get_numpoints(legend) + 1) + height / 2. - ydescent
        points = np.array([x, y]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = LineCollection(segments, cmap=self.cmap, transform=trans)
        lc.set_array(x)
        lc.set_linewidth(artist.get_linewidth())
        return [lc]

def add_normalized_line_collection(ax, cmap, linewidth=3, linestyle='-'):
    norm = Normalize(vmin=0., vmax=1.)
    t = np.linspace(0, 1, 100)  # Smooth gradient
    lc = LineCollection([np.column_stack([t, t * 0])], cmap=cmap, norm=norm, linewidth=linewidth, linestyle=linestyle)
    lc.set_array(np.linspace(0., 1, len(t)))  # Ensure this spans 0 to 1 for correct normalization
    ax.add_collection(lc)  # Add the LineCollection to the axis
    return lc

# Sample data
x_data = np.logspace(1, 3, 50)
y_data = np.random.rand(3, len(x_data))

fig, ax = plt.subplots(figsize=(8, 6))

colors = cmr.take_cmap_colors('cmr.rainforest', 3, cmap_range=(0.15, 0.85), return_fmt='hex')
for i in range(y_data.shape[0]):
    ax.plot(x_data, y_data[i], color=colors[i], lw=3)
    
    ax.set_xlabel('$x$', fontsize=18)
    ax.set_ylabel('$y$', fontsize=18)

# Create color lines for legend
color_line = add_normalized_line_collection(ax, cmap="cmr.rainforest", linewidth=4)

# Existing legend handles and labels
handles, labels = ax.get_legend_handles_labels()
handles.append(color_line)
labels.append("Colormap-based Line")

ax.legend(handles, labels, handler_map={color_line: HandlerColorLineCollection(cmap="cmr.rainforest", numpoints=30)}, loc="upper right", frameon=True, fontsize=15)

plt.show()

Proposed usage

import matplotlib.pyplot as plt
import numpy as np
import cmasher as cmr

# Sample data
x_data = np.logspace(1, 3, 50)
y_data = np.random.rand(3, len(x_data))

fig, ax = plt.subplots(figsize=(8, 6))

colors = cmr.take_cmap_colors('cmr.rainforest', 3, cmap_range=(0.15, 0.85), return_fmt='hex')
for i in range(y_data.shape[0]):
    ax.plot(x_data, y_data[i], color=colors[i], lw=3)
    
    ax.set_xlabel('$x$', fontsize=18)
    ax.set_ylabel('$y$', fontsize=18)

# New API for colormap-based lines in legend
ax.legend(use_colormap=True, loc="upper right", frameon=True, fontsize=15, new_arg='cmr.rainforest', )

plt.show()

Please note that it would be beneficial to allow passing either a colormap or a list of colors (in HEX) to this new argument (new_arg) in the legend. Ideally, the legend object should be able to internally determine the colors used in the plot lines, making the process seamless and intuitive for the user.

I am also including a version from a colleague.

Many thanks!
Niko