I am experimenting with different subsets of channels in my EEG experiments, so how can I highlight those subsets in an MNE figure? Either by circling them or coloring them differently from the rest. I have the following code to plot a figure of the channels used in my experiment:

import mne
import matplotlib.pyplot as plt

# Montage, ordered channel names
CH_NAMES = ['Fp1', 'Fz', 'F3', 'F7', 'FT9', 'FC5', 'FC1', 'C3', 'T7', 'TP9', 'CP5', 'CP1', 'Pz', 'P3', 'P7', 'O1', 'Oz', 'O2', 'P4', 'P8', 'TP10', 'CP6', 'CP2', 'C4', 'T8', 'FT10', 'FC6', 'FC2', 'F4', 'F8', 'Fp2', 'AF7', 'AF3', 'AFz', 'F1', 'F5', 'FT7', 'FC3', 'C1', 'C5', 'TP7', 'CP3', 'P1', 'P5', 'PO7', 'PO3', 'POz', 'PO4', 'PO8', 'P6', 'P2', 'CPz', 'CP4', 'TP8', 'C6', 'C2', 'FC4', 'FT8', 'F6', 'AF8', 'AF4', 'F2', 'FCz']

# Create a standard montage for the 10-20 system
montage = mne.channels.make_standard_montage('standard_1020')

# Filter the montage to only include channels in ch_names
ch_pos = {ch: pos for ch, pos in montage.get_positions()['ch_pos'].items() if ch in CH_NAMES}
montage_subset = mne.channels.make_dig_montage(ch_pos=ch_pos, coord_frame='head')

# Create an info object
info = mne.create_info(ch_names=CH_NAMES, sfreq=1000, ch_types='eeg')
info.set_montage(montage_subset)

# Plot the montage
fig = mne.viz.plot_montage(montage_subset, show_names=True)
plt.show()

Which outputs the following figure:

I have tried to use ChatGPT to solve it, but most of the solutions I can make it produce is where it create a separate plot containing only the channels from the subset in a box (so not overlayed on a head topographic). The closest it came to what I wanted was the below figure with specific channels colored blue, but it greatly limits the visibility of the channels in the experiment.

The following code (together with the initial code) produces the above image:

# Highlight specific channels by adding circles around them
fig, ax = plt.subplots()
ax.set_aspect('equal')
ax.set_title('EEG Montage with Highlighted Channels')
for ch_name, (x, y, z) in montage_subset.get_positions()['ch_pos'].items():
    ax.plot(x, y, 'o', markersize=20, markeredgecolor='black', markerfacecolor='blue' if ch_name in highlight_channels else 'gray', markeredgewidth=1)
    ax.text(x, y, ch_name, color='red' if ch_name in highlight_channels else 'black', fontsize=10, ha='center', va='center')

ax.set_xlim(-0.12, 0.12)
ax.set_ylim(-0.125, 0.1)
plt.xlabel('x')
plt.ylabel('y')
plt.grid(True)
plt.show()