How to modify scatter-plot figure legend to show different formats for the same types of handles?

Question

I am trying to modify the legend of a figure that contains two overlayed scatter plots. More specifically, I want two legend handles and labels: the first handle will contain multiple points (each colored differently), while the other handle consists of a single point.

As per this related question, I can modify the legend handle to show multiple points, each one being a different color.

As per this similar question, I am aware that I can change the number of points shown by a specified handle. However, this applies the change to all handles in the legend. Can it be applied to one handle only?

My goal is to combine both approaches. Is there a way to do this?

In case it isn't clear, I would like to modify the embedded figure (see below) such that Z vs X handle shows only one-point next to the corresponding legend label, while leaving the Y vs X handle unchanged.

My failed attempt at producing such a figure is below:

To replicate this figure, one can run the code below:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerTuple, HandlerRegularPolyCollection

class ScatterHandler(HandlerRegularPolyCollection):

    def update_prop(self, legend_handle, orig_handle, legend):
        """ """
        legend._set_artist_props(legend_handle)
        legend_handle.set_clip_box(None)
        legend_handle.set_clip_path(None)

    def create_collection(self, orig_handle, sizes, offsets, transOffset):
        """ """
        p = type(orig_handle)([orig_handle.get_paths()[0]], sizes=sizes, offsets=offsets, transOffset=transOffset, cmap=orig_handle.get_cmap(), norm=orig_handle.norm)
        a = orig_handle.get_array()
        if type(a) != type(None):
            p.set_array(np.linspace(a.min(), a.max(), len(offsets)))
        else:
            self._update_prop(p, orig_handle)
        return p

x = np.arange(10)
y = np.sin(x)
z = np.cos(x)

fig, ax = plt.subplots()
hy = ax.scatter(x, y, cmap='plasma', c=y, label='Y vs X')
hz = ax.scatter(x, z, color='k', label='Z vs X')
ax.grid(color='k', linestyle=':', alpha=0.3)
fig.subplots_adjust(bottom=0.2)
handler_map = {type(hz) : ScatterHandler()}
fig.legend(mode='expand', ncol=2, loc='lower center', handler_map=handler_map, scatterpoints=5)

plt.show()
plt.close(fig)

One solution that I do not like is to create two legends - one for Z vs X and one for Y vs X. But, my actual use case involves an optional number of handles (which can exceed two) and I would prefer not having to calculate the optimal width/height of each legend box. How else can this problem be approached?

Answer 1

This is a dirty trick and not an elegant solution, but you can set the sizes of other points for Z-X legend to 0. Just change your last two lines to the following.

leg = fig.legend(mode='expand', ncol=2, loc='lower center', handler_map=handler_map, scatterpoints=5)
# The third dot of the second legend stays the same size, others are set to 0
leg.legendHandles[1].set_sizes([0,0,leg.legendHandles[1].get_sizes()[2],0,0])

The result is as shown.