fix bugs in legend generation for pole/zero plots

murrayrm · murrayrm · commit acfadf003488 · 2024-08-08T14:54:49.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -275,25 +275,24 @@ def _process_line_labels(label, ntraces, ninputs=0, noutputs=0):
 
 # Get labels for all lines in an axes
 def _get_line_labels(ax, use_color=True):
-    labels, lines = [], []
+    labels_colors, lines = [], []
     last_color, counter = None, 0       # label unknown systems
     for i, line in enumerate(ax.get_lines()):
         label = line.get_label()
+        color = line.get_color()
         if use_color and label.startswith("Unknown"):
             label = f"Unknown-{counter}"
-            if last_color is None:
-                last_color = line.get_color()
-            elif last_color != line.get_color():
+            if last_color != color:
                 counter += 1
-                last_color = line.get_color()
+            last_color = color
         elif label[0] == '_':
             continue
 
-        if label not in labels:
+        if (label, color) not in labels_colors:
             lines.append(line)
-            labels.append(label)
+            labels_colors.append((label, color))
 
-    return lines, labels
+    return lines, [label for label, color in labels_colors]
 
 
 # Utility function to make legend labels
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -391,12 +391,13 @@ def pole_zero_plot(
                 real(poles), imag(poles), marker='x', linestyle='',
                 markeredgecolor=color, markerfacecolor=color,
                 markersize=marker_size, markeredgewidth=marker_width,
-                label=label)
+                color=color, label=label)
         if len(zeros) > 0:
             out[idx, 1] = ax.plot(
                 real(zeros), imag(zeros), marker='o', linestyle='',
                 markeredgecolor=color, markerfacecolor='none',
-                markersize=marker_size, markeredgewidth=marker_width)
+                markersize=marker_size, markeredgewidth=marker_width,
+                color=color)
 
         # Plot the loci, if present
         if response.loci is not None:
@@ -447,8 +448,8 @@ def pole_zero_plot(
                     markeredgecolor=pole_line.get_markerfacecolor(),
                     markerfacecolor='none', markersize=marker_size,
                     markeredgewidth=marker_width)
-            handle = (pole_line, zero_line)
-            line_tuples.append(handle)
+                handle = (pole_line, zero_line)
+                line_tuples.append(handle)
 
             with plt.rc_context(freqplot_rcParams):
                 legend = ax.legend(
