Skip to content

Commit 1888854

Browse files
murrayrmmurrayrm
committed
update arrow styles for nyquist
1 parent 85faffa commit 1888854

2 files changed

Lines changed: 243 additions & 44 deletions

File tree

‎control/freqplot.py‎

Lines changed: 158 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -157,9 +157,10 @@ def bode_plot(syslist, omega=None,
157157
generate the frequency response for a single system.
158158
159159
2. If a discrete time model is given, the frequency response is plotted
160-
along the upper branch of the unit circle, using the mapping ``z = exp(1j
161-
* omega * dt)`` where `omega` ranges from 0 to `pi/dt` and `dt` is the discrete
162-
timebase. If timebase not specified (``dt=True``), `dt` is set to 1.
160+
along the upper branch of the unit circle, using the mapping ``z =
161+
exp(1j * omega * dt)`` where `omega` ranges from 0 to `pi/dt` and `dt`
162+
is the discrete timebase. If timebase not specified (``dt=True``),
163+
`dt` is set to 1.
163164
164165
Examples
165166
--------
@@ -198,7 +199,8 @@ def bode_plot(syslist, omega=None,
198199
omega_range_given = True if omega is not None else False
199200

200201
if omega is None:
201-
omega_num = config._get_param('freqplot','number_of_samples', omega_num)
202+
omega_num = config._get_param(
203+
'freqplot', 'number_of_samples', omega_num)
202204
if omega_limits is None:
203205
# Select a default range if none is provided
204206
omega = _default_frequency_range(syslist,
@@ -246,11 +248,9 @@ def bode_plot(syslist, omega=None,
246248
# If no axes present, create them from scratch
247249
if ax_mag is None or ax_phase is None:
248250
plt.clf()
249-
ax_mag = plt.subplot(211,
250-
label='control-bode-magnitude')
251-
ax_phase = plt.subplot(212,
252-
label='control-bode-phase',
253-
sharex=ax_mag)
251+
ax_mag = plt.subplot(211, label='control-bode-magnitude')
252+
ax_phase = plt.subplot(
253+
212, label='control-bode-phase', sharex=ax_mag)
254254

255255
mags, phases, omegas, nyquistfrqs = [], [], [], []
256256
for sys in syslist:
@@ -340,9 +340,10 @@ def bode_plot(syslist, omega=None,
340340
np.nan, 0.7*min(mag_plot), 1.3*max(mag_plot)))
341341
mag_plot = np.hstack((mag_plot, mag_nyq_line))
342342
phase_range = max(phase_plot) - min(phase_plot)
343-
phase_nyq_line = np.array((np.nan,
344-
min(phase_plot) - 0.2 * phase_range,
345-
max(phase_plot) + 0.2 * phase_range))
343+
phase_nyq_line = np.array(
344+
(np.nan,
345+
min(phase_plot) - 0.2 * phase_range,
346+
max(phase_plot) + 0.2 * phase_range))
346347
phase_plot = np.hstack((phase_plot, phase_nyq_line))
347348

348349
#
@@ -351,7 +352,7 @@ def bode_plot(syslist, omega=None,
351352

352353
if dB:
353354
ax_mag.semilogx(omega_plot, 20 * np.log10(mag_plot),
354-
*args, **kwargs)
355+
*args, **kwargs)
355356
else:
356357
ax_mag.loglog(omega_plot, mag_plot, *args, **kwargs)
357358

@@ -523,11 +524,13 @@ def gen_zero_centered_series(val_min, val_max, period):
523524
# Default values for module parameter variables
524525
_nyquist_defaults = {
525526
'nyquist.mirror_style': '--',
527+
'nyquist.arrows': 2,
528+
'nyquist.arrow_size': 8,
526529
}
527530

531+
528532
def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
529-
omega_num=None, label_freq=0, arrowhead_length=0.1,
530-
arrowhead_width=0.1, color=None, *args, **kwargs):
533+
omega_num=None, label_freq=0, color=None, *args, **kwargs):
531534
"""
532535
Nyquist plot for a system
533536
@@ -542,32 +545,44 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
542545
If True, plot magnitude
543546
544547
omega : array_like
545-
Set of frequencies to be evaluated in rad/sec.
548+
Set of frequencies to be evaluated, in rad/sec.
546549
547550
omega_limits : array_like of two values
548551
Limits to the range of frequencies. Ignored if omega
549552
is provided, and auto-generated if omitted.
550553
551554
omega_num : int
552-
Number of samples to plot. Defaults to
555+
Number of frequency samples to plot. Defaults to
553556
config.defaults['freqplot.number_of_samples'].
554557
555558
color : string
556-
Used to specify the color of the line and arrowhead
559+
Used to specify the color of the line and arrowhead.
557560
558561
mirror_style : string or False
559562
Linestyle for mirror image of the Nyquist curve. If `False` then
560563
omit completely. Default linestyle ('--') is determined by
561564
config.defaults['nyquist.mirror_style'].
562565
563566
label_freq : int
564-
Label every nth frequency on the plot
565-
566-
arrowhead_width : float
567-
Arrow head width
568-
569-
arrowhead_length : float
570-
Arrow head length
567+
Label every nth frequency on the plot. If not specified, no labels
568+
are generated.
569+
570+
arrows : int or 1D/2D array of floats
571+
Specify the number of arrows to plot on the Nyquist curve. If an
572+
integer is passed. that number of equally spaced arrows will be
573+
plotted on each of the primary segment and the mirror image. If a 1D
574+
array is passed, it should consist of a sorted list of floats between
575+
0 and 1, indicating the location along the curve to plot an arrow. If
576+
a 2D array is passed, the first row will be used to specify arrow
577+
locations for the primary curve and the second row will be used for
578+
the mirror image.
579+
580+
arrow_size : float
581+
Arrowhead width and length (in display coordinates). Default value is
582+
8 and can be set using config.defaults['nyquist.arrow_size'].
583+
584+
arrow_style : matplotlib.patches.ArrowStyle
585+
Define style used for Nyquist curve arrows (overrides `arrow_size`).
571586
572587
*args : :func:`matplotlib.pyplot.plot` positional properties, optional
573588
Additional arguments for `matplotlib` plots (color, linestyle, etc)
@@ -588,7 +603,7 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
588603
589604
Examples
590605
--------
591-
>>> sys = ss("1. -2; 3. -4", "5.; 7", "6. 8", "9.")
606+
>>> sys = ss([[1, -2], [3, -4]], [[5], [7]], [[6, 8]], [[9]])
592607
>>> real, imag, freq = nyquist_plot(sys)
593608
594609
"""
@@ -608,9 +623,21 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
608623
# Map 'labelFreq' keyword to 'label_freq' keyword
609624
label_freq = kwargs.pop('labelFreq')
610625

626+
# Check to see if legacy 'arrow_width' or 'arrow_length' were used
627+
if 'arrow_width' in kwargs or 'arrow_length' in kwargs:
628+
warn("'arrow_width' and 'arrow_length' keywords are deprecated in "
629+
"nyquist_plot; use `arrow_size` instead", FutureWarning)
630+
kwargs['arrow_size'] = \
631+
(kwargs.get('arrow_width', 0) + kwargs.get('arrow_width', 0)) / 2
632+
611633
# Get values for params (and pop from list to allow keyword use in plot)
612634
mirror_style = config._get_param(
613635
'nyquist', 'mirror_style', kwargs, _nyquist_defaults, pop=True)
636+
arrows = config._get_param(
637+
'nyquist', 'arrows', kwargs, _nyquist_defaults, pop=True)
638+
arrow_size = config._get_param(
639+
'nyquist', 'arrow_size', kwargs, _nyquist_defaults, pop=True)
640+
arrow_style = config._get_param('nyquist', 'arrow_style', kwargs, None)
614641

615642
# If argument was a singleton, turn it into a list
616643
if not hasattr(syslist, '__iter__'):
@@ -620,7 +647,8 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
620647
omega_range_given = True if omega is not None else False
621648

622649
if omega is None:
623-
omega_num = config._get_param('freqplot','number_of_samples',omega_num)
650+
omega_num = config._get_param(
651+
'freqplot', 'number_of_samples', omega_num)
624652
if omega_limits is None:
625653
# Select a default range if none is provided
626654
omega = _default_frequency_range(syslist,
@@ -636,6 +664,11 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
636664

637665
xs, ys, omegas = [], [], []
638666
for sys in syslist:
667+
if not sys.issiso():
668+
# TODO: Add MIMO nyquist plots.
669+
raise ControlMIMONotImplemented(
670+
"Nyquist plot currently only supports SISO systems.")
671+
639672
omega_sys = np.asarray(omega)
640673
if sys.isdtime(strict=True):
641674
nyquistfrq = math.pi / sys.dt
@@ -655,28 +688,35 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
655688
omegas.append(omega_sys)
656689

657690
if plot:
658-
if not sys.issiso():
659-
# TODO: Add MIMO nyquist plots.
660-
raise ControlMIMONotImplemented(
661-
"Nyquist plot currently supports SISO systems.")
691+
# Parse the arrows keyword
692+
if isinstance(arrows, int):
693+
N = arrows
694+
# Space arrows out, starting midway along each "region"
695+
arrow_pos = np.linspace(0.5/N, 1 + 0.5/N, N, endpoint=False)
696+
elif isinstance(arrows, (list, np.ndarray)):
697+
arrow_pos = np.sort(np.atleast_1d(arrows))
698+
elif not arrows:
699+
arrow_pos = []
700+
else:
701+
raise ValueError("unknown or unsupported arrow location")
702+
703+
# Set the arrow style
704+
if arrow_style is None:
705+
arrow_style = mpl.patches.ArrowStyle(
706+
'simple', head_width=arrow_size, head_length=arrow_size)
662707

663708
# Plot the primary curve
664709
p = plt.plot(x, y, '-', color=color, *args, **kwargs)
665710
c = p[0].get_color()
666711
ax = plt.gca()
667-
668-
# Plot arrow to indicate Nyquist encirclement orientation
669-
ax.arrow(x[0], y[0], (x[1]-x[0])/2, (y[1]-y[0])/2, fc=c, ec=c,
670-
head_width=arrowhead_width,
671-
head_length=arrowhead_length)
712+
_add_arrows_to_line2D(
713+
ax, p[0], arrow_pos, arrowstyle=arrow_style, dir=1)
672714

673715
# Plot the mirror image
674716
if mirror_style is not False:
675-
plt.plot(x, -y, mirror_style, color=c, *args, **kwargs)
676-
ax.arrow(
677-
x[-1], -y[-1], (x[-1]-x[-2])/2, (y[-1]-y[-2])/2,
678-
fc=c, ec=c, head_width=arrowhead_width,
679-
head_length=arrowhead_length)
717+
p = plt.plot(x, -y, mirror_style, color=c, *args, **kwargs)
718+
_add_arrows_to_line2D(
719+
ax, p[0], arrow_pos, arrowstyle=arrow_style, dir=-1)
680720

681721
# Mark the -1 point
682722
plt.plot([-1], [0], 'r+')
@@ -702,8 +742,8 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
702742
# instead of 1.0, and this would otherwise be
703743
# truncated to 0.
704744
plt.text(xpt, ypt, ' ' +
705-
str(int(np.round(f / 1000 ** pow1000, 0))) + ' ' +
706-
prefix + 'Hz')
745+
str(int(np.round(f / 1000 ** pow1000, 0))) + ' ' +
746+
prefix + 'Hz')
707747

708748
if plot:
709749
ax = plt.gca()
@@ -716,6 +756,80 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
716756
else:
717757
return xs, ys, omegas
718758

759+
760+
# Internal function to add arrows to a curve
761+
def _add_arrows_to_line2D(
762+
axes, line, arrow_locs=[0.2, 0.4, 0.6, 0.8],
763+
arrowstyle='-|>', arrowsize=1, dir=1, transform=None):
764+
"""
765+
Add arrows to a matplotlib.lines.Line2D at selected locations.
766+
767+
Parameters:
768+
-----------
769+
axes: Axes object as returned by axes command (or gca)
770+
line: Line2D object as returned by plot command
771+
arrow_locs: list of locations where to insert arrows, % of total length
772+
arrowstyle: style of the arrow
773+
arrowsize: size of the arrow
774+
transform: a matplotlib transform instance, default to data coordinates
775+
776+
Returns:
777+
--------
778+
arrows: list of arrows
779+
780+
Based on https://stackoverflow.com/questions/26911898/
781+
782+
"""
783+
if not isinstance(line, mpl.lines.Line2D):
784+
raise ValueError("expected a matplotlib.lines.Line2D object")
785+
x, y = line.get_xdata(), line.get_ydata()
786+
787+
arrow_kw = {
788+
"arrowstyle": arrowstyle,
789+
}
790+
791+
color = line.get_color()
792+
use_multicolor_lines = isinstance(color, np.ndarray)
793+
if use_multicolor_lines:
794+
raise NotImplementedError("multicolor lines not supported")
795+
else:
796+
arrow_kw['color'] = color
797+
798+
linewidth = line.get_linewidth()
799+
if isinstance(linewidth, np.ndarray):
800+
raise NotImplementedError("multiwidth lines not supported")
801+
else:
802+
arrow_kw['linewidth'] = linewidth
803+
804+
if transform is None:
805+
transform = axes.transData
806+
807+
# Compute the arc length along the curve
808+
s = np.cumsum(np.sqrt(np.diff(x) ** 2 + np.diff(y) ** 2))
809+
810+
arrows = []
811+
for loc in arrow_locs:
812+
n = np.searchsorted(s, s[-1] * loc)
813+
814+
# Figure out what direction to paint the arrow
815+
if dir == 1:
816+
arrow_tail = (x[n], y[n])
817+
arrow_head = (np.mean(x[n:n + 2]), np.mean(y[n:n + 2]))
818+
elif dir == -1:
819+
# Orient the arrow in the other direction on the segment
820+
arrow_tail = (x[n + 1], y[n + 1])
821+
arrow_head = (np.mean(x[n:n + 2]), np.mean(y[n:n + 2]))
822+
else:
823+
raise ValueError("unknown value for keyword 'dir'")
824+
825+
p = mpl.patches.FancyArrowPatch(
826+
arrow_tail, arrow_head, transform=transform, lw=0,
827+
**arrow_kw)
828+
axes.add_patch(p)
829+
arrows.append(p)
830+
return arrows
831+
832+
719833
#
720834
# Gang of Four plot
721835
#
@@ -840,7 +954,7 @@ def gangof4_plot(P, C, omega=None, **kwargs):
840954

841955
# Compute reasonable defaults for axes
842956
def _default_frequency_range(syslist, Hz=None, number_of_samples=None,
843-
feature_periphery_decades=None):
957+
feature_periphery_decades=None):
844958
"""Compute a reasonable default frequency range for frequency
845959
domain plots.
846960

0 commit comments

Comments
 (0)