recalculate loci for sisotool and rlocus on axis scaling

repagh · repagh · commit 7d567d698d4f · 2025-06-05T21:34:45.000+02:00
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, cp):
+    """Update the loci of a plot after zooming/panning
+
+    Parameters
+    ----------
+    pzmap_responses : PoleZeroMap list
+        Responses to update
+    cp : ControlPlot
+        Collection of plot handles
+    """
+
+    for idx, response in enumerate(pzmap_responses):
+
+        # remove the old data
+        for l in cp.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:
+                    cp.lines[idx,2][il].set_data(real(locus), imag(locus))
+                except IndexError:
+                    # not expected, but more lines apparently needed
+                    cp.lines[idx,2].append(cp.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
@@ -75,7 +75,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 +205,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 calback 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 +327,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
@@ -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
