Merge pull request #1153 from repagh/zoom-rlocus

murrayrm · web-flow · commit 34c6d5922cf4 · 2025-06-21T11:30:58.000-07:00
recalculate loci for sisotool and rlocus on axis scaling
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -124,6 +124,17 @@ def plot(self, *args, **kwargs):
         """
         return pole_zero_plot(self, *args, **kwargs)
 
+    def replot(self, cplt: ControlPlot):
+        """Update the pole/zero loci of an existing plot.
+
+        Parameters
+        ----------
+        cplt: ControlPlot
+            Graphics handles of the existing plot.
+        """
+        pole_zero_replot(self, cplt)
+
+
 
 # Pole/zero map
 def pole_zero_map(sysdata):
@@ -513,6 +524,35 @@ def _click_dispatcher(event):
     return ControlPlot(out, ax, fig, legend=legend)
 
 
+def pole_zero_replot(pzmap_responses, cplt):
+    """Update the loci of a plot after zooming/panning.
+
+    Parameters
+    ----------
+    pzmap_responses : PoleZeroMap list
+        Responses to update.
+    cplt : ControlPlot
+        Collection of plot handles.
+    """
+
+    for idx, response in enumerate(pzmap_responses):
+
+        # remove the old data
+        for l in cplt.lines[idx, 2]:
+            l.set_data([], [])
+
+        # update the line data
+        if response.loci is not None:
+
+            for il, locus in enumerate(response.loci.transpose()):
+                try:
+                    cplt.lines[idx,2][il].set_data(real(locus), imag(locus))
+                except IndexError:
+                    # not expected, but more lines apparently needed
+                    cplt.lines[idx,2].append(cplt.ax[0,0].plot(
+                        real(locus), imag(locus)))
+
+
 # Utility function to find gain corresponding to a click event
 def _find_root_locus_gain(event, sys, ax):
     # Get the current axis limits to set various thresholds
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -33,7 +33,7 @@
 
 
 # Root locus map
-def root_locus_map(sysdata, gains=None):
+def root_locus_map(sysdata, gains=None, xlim=None, ylim=None):
     """Compute the root locus map for an LTI system.
 
     Calculate the root locus by finding the roots of 1 + k * G(s) where G
@@ -46,6 +46,10 @@ def root_locus_map(sysdata, gains=None):
     gains : array_like, optional
         Gains to use in computing plot of closed-loop poles.  If not given,
         gains are chosen to include the main features of the root locus map.
+    xlim : tuple or list, optional
+        Set limits of x axis (see `matplotlib.axes.Axes.set_xlim`).
+    ylim : tuple or list, optional
+        Set limits of y axis (see `matplotlib.axes.Axes.set_ylim`).
 
     Returns
     -------
@@ -75,7 +79,7 @@ def root_locus_map(sysdata, gains=None):
         nump, denp = _systopoly1d(sys[0, 0])
 
         if gains is None:
-            kvect, root_array, _, _ = _default_gains(nump, denp, None, None)
+            kvect, root_array, _, _ = _default_gains(nump, denp, xlim, ylim)
         else:
             kvect = np.atleast_1d(gains)
             root_array = _RLFindRoots(nump, denp, kvect)
@@ -205,13 +209,52 @@ def root_locus_plot(
     # Plot the root loci
     cplt = responses.plot(grid=grid, **kwargs)
 
+    # Add a reaction to axis scale changes, if given LTI systems, and
+    # there is no set of pre-defined gains
+    if gains is None:
+        add_loci_recalculate(sysdata, cplt, cplt.axes[0,0])
+
     # Legacy processing: return locations of poles and zeros as a tuple
     if plot is True:
         return responses.loci, responses.gains
 
     return ControlPlot(cplt.lines, cplt.axes, cplt.figure)
 
 
+def add_loci_recalculate(sysdata, cplt, axis):
+    """Add a callback to re-calculate the loci data fitting a zoom action.
+
+    Parameters
+    ----------
+    sysdata: LTI object or list
+        Linear input/output systems (SISO only, for now).
+    cplt: ControlPlot
+        Collection of plot handles.
+    axis: matplotlib.axes.Axis
+        Axis on which callbacks are installed.
+    """
+
+    # if LTI, treat everything as a list of lti
+    if isinstance(sysdata, LTI):
+        sysdata = [sysdata]
+
+    # check that we can actually recalculate the loci
+    if isinstance(sysdata, list) and all(
+       [isinstance(sys, LTI) for sys in sysdata]):
+
+        # callback function for axis change (zoom, pan) events
+        # captures the sysdata object and cplt
+        def _zoom_adapter(_ax):
+            newresp = root_locus_map(sysdata, None,
+                                     _ax.get_xlim(),
+                                     _ax.get_ylim())
+            newresp.replot(cplt)
+
+        # connect the callback to axis changes
+        axis.callbacks.connect('xlim_changed', _zoom_adapter)
+        axis.callbacks.connect('ylim_changed', _zoom_adapter)
+
+
 def _default_gains(num, den, xlim, ylim):
     """Unsupervised gains calculation for root locus plot.
 
@@ -288,7 +331,7 @@ def _default_gains(num, den, xlim, ylim):
         # Root locus is on imaginary axis (rare), use just y distance
         tolerance = y_tolerance
     elif y_tolerance == 0:
-        # Root locus is on imaginary axis (common), use just x distance
+        # Root locus is on real axis (common), use just x distance
         tolerance = x_tolerance
     else:
         tolerance = np.min([x_tolerance, y_tolerance])
diff --git a/control/sisotool.py b/control/sisotool.py
@@ -22,6 +22,7 @@
 from .statesp import ss, summing_junction
 from .timeresp import step_response
 from .xferfcn import tf
+from .rlocus import add_loci_recalculate
 
 _sisotool_defaults = {
     'sisotool.initial_gain': 1
@@ -105,7 +106,7 @@ def sisotool(sys, initial_gain=None, xlim_rlocus=None, ylim_rlocus=None,
     fig = plt.gcf()
     if fig.canvas.manager.get_window_title() != 'Sisotool':
         plt.close(fig)
-        fig,axes = plt.subplots(2, 2)
+        fig, axes = plt.subplots(2, 2)
         fig.canvas.manager.set_window_title('Sisotool')
     else:
         axes = np.array(fig.get_axes()).reshape(2, 2)
@@ -137,15 +138,18 @@ def sisotool(sys, initial_gain=None, xlim_rlocus=None, ylim_rlocus=None,
     #     sys[0, 0], initial_gain=initial_gain, xlim=xlim_rlocus,
     #     ylim=ylim_rlocus, plotstr=plotstr_rlocus, grid=rlocus_grid,
     #     ax=fig.axes[1])
-    ax_rlocus = fig.axes[1]
-    root_locus_map(sys[0, 0]).plot(
+    ax_rlocus = axes[0,1]  # fig.axes[1]
+    cplt = root_locus_map(sys[0, 0]).plot(
         xlim=xlim_rlocus, ylim=ylim_rlocus,
         initial_gain=initial_gain, ax=ax_rlocus)
     if rlocus_grid is False:
         # Need to generate grid manually, since root_locus_plot() won't
         from .grid import nogrid
         nogrid(sys.dt, ax=ax_rlocus)
 
+    # install a zoom callback on the root-locus axis
+    add_loci_recalculate(sys, cplt, ax_rlocus)
+
     # Reset the button release callback so that we can update all plots
     fig.canvas.mpl_connect(
         'button_release_event', partial(
@@ -155,9 +159,8 @@ def sisotool(sys, initial_gain=None, xlim_rlocus=None, ylim_rlocus=None,
 
 def _click_dispatcher(event, sys, ax, bode_plot_params, tvect):
     # Zoom handled by specialized callback in rlocus, only handle gain plot
-    if event.inaxes == ax.axes and \
-       plt.get_current_fig_manager().toolbar.mode not in \
-       {'zoom rect', 'pan/zoom'}:
+    if event.inaxes == ax.axes:
+
         fig = ax.figure
 
         # if a point is clicked on the rootlocus plot visually emphasize it
