Home Original page

Merge pull request #953 from murrayrm/pzmap_plots-22Jul2023 · python-control/python-control@eb0f3f9

1-

import numpy as np

2-

from numpy import cos, sin, sqrt, linspace, pi, exp

1+

# grid.py - code to add gridlines to root locus and pole-zero diagrams

2+

#

3+

# This code generates grids for pole-zero diagrams (including root locus

4+

# diagrams). Rather than just draw a grid in place, it uses the AxisArtist

5+

# package to generate a custom grid that will scale with the figure.

6+

#

7+38

import matplotlib.pyplot as plt

4-

from mpl_toolkits.axisartist import SubplotHost

5-

from mpl_toolkits.axisartist.grid_helper_curvelinear \

6-

import GridHelperCurveLinear

79

import mpl_toolkits.axisartist.angle_helper as angle_helper

10+

import numpy as np

811

from matplotlib.projections import PolarAxes

912

from matplotlib.transforms import Affine2D

13+

from mpl_toolkits.axisartist import SubplotHost

14+

from mpl_toolkits.axisartist.grid_helper_curvelinear import \

15+

GridHelperCurveLinear

16+

from numpy import cos, exp, linspace, pi, sin, sqrt

17+18+

from .iosys import isdtime

101911201221

class FormatterDMS(object):

@@ -65,14 +74,15 @@ def __call__(self, transform_xy, x1, y1, x2, y2):

6574

return lon_min, lon_max, lat_min, lat_max

6675677668-

def sgrid():

77+

def sgrid(scaling=None):

6978

# From matplotlib demos:

7079

# https://matplotlib.org/gallery/axisartist/demo_curvelinear_grid.html

7180

# https://matplotlib.org/gallery/axisartist/demo_floating_axis.html

72817382

# PolarAxes.PolarTransform takes radian. However, we want our coordinate

74-

# system in degree

83+

# system in degrees

7584

tr = Affine2D().scale(np.pi/180., 1.) + PolarAxes.PolarTransform()

85+7686

# polar projection, which involves cycle, and also has limits in

7787

# its coordinates, needs a special method to find the extremes

7888

# (min, max of the coordinate within the view).

@@ -89,6 +99,7 @@ def sgrid():

8999

tr, extreme_finder=extreme_finder, grid_locator1=grid_locator1,

90100

tick_formatter1=tick_formatter1)

91101102+

# Set up an axes with a specialized grid helper

92103

fig = plt.gcf()

93104

ax = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)

94105

@@ -97,15 +108,20 @@ def sgrid():

97108

ax.axis[:].major_ticklabels.set_visible(visible)

98109

ax.axis[:].major_ticks.set_visible(False)

99110

ax.axis[:].invert_ticklabel_direction()

111+

ax.axis[:].major_ticklabels.set_color('gray')

100112113+

# Set up internal tickmarks and labels along the real/imag axes

101114

ax.axis["wnxneg"] = axis = ax.new_floating_axis(0, 180)

102115

axis.set_ticklabel_direction("-")

103116

axis.label.set_visible(False)

117+104118

ax.axis["wnxpos"] = axis = ax.new_floating_axis(0, 0)

105119

axis.label.set_visible(False)

120+106121

ax.axis["wnypos"] = axis = ax.new_floating_axis(0, 90)

107122

axis.label.set_visible(False)

108-

axis.set_axis_direction("left")

123+

axis.set_axis_direction("right")

124+109125

ax.axis["wnyneg"] = axis = ax.new_floating_axis(0, 270)

110126

axis.label.set_visible(False)

111127

axis.set_axis_direction("left")

@@ -119,43 +135,41 @@ def sgrid():

119135

ax.axis["bottom"].get_helper().nth_coord_ticks = 0

120136121137

fig.add_subplot(ax)

122-123-

# RECTANGULAR X Y AXES WITH SCALE

124-

# par2 = ax.twiny()

125-

# par2.axis["top"].toggle(all=False)

126-

# par2.axis["right"].toggle(all=False)

127-

# new_fixed_axis = par2.get_grid_helper().new_fixed_axis

128-

# par2.axis["left"] = new_fixed_axis(loc="left",

129-

# axes=par2,

130-

# offset=(0, 0))

131-

# par2.axis["bottom"] = new_fixed_axis(loc="bottom",

132-

# axes=par2,

133-

# offset=(0, 0))

134-

# FINISH RECTANGULAR

135-136138

ax.grid(True, zorder=0, linestyle='dotted')

137139138-

_final_setup(ax)

140+

_final_setup(ax, scaling=scaling)

139141

return ax, fig

140142141143142-

def _final_setup(ax):

144+

# Utility function used by all grid code

145+

def _final_setup(ax, scaling=None):

143146

ax.set_xlabel('Real')

144147

ax.set_ylabel('Imaginary')

145-

ax.axhline(y=0, color='black', lw=1)

146-

ax.axvline(x=0, color='black', lw=1)

147-

plt.axis('equal')

148+

ax.axhline(y=0, color='black', lw=0.25)

149+

ax.axvline(x=0, color='black', lw=0.25)

148150151+

# Set up the scaling for the axes

152+

scaling = 'equal' if scaling is None else scaling

153+

plt.axis(scaling)

149154150-

def nogrid():

151-

f = plt.gcf()

152-

ax = plt.axes()

153155154-

_final_setup(ax)

155-

return ax, f

156+

# If not grid is given, at least separate stable/unstable regions

157+

def nogrid(dt=None, ax=None, scaling=None):

158+

fig = plt.gcf()

159+

if ax is None:

160+

ax = fig.gca()

161+162+

# Draw the unit circle for discrete time systems

163+

if isdtime(dt=dt, strict=True):

164+

s = np.linspace(0, 2*pi, 100)

165+

ax.plot(np.cos(s), np.sin(s), 'k--', lw=0.5, dashes=(5, 5))

156166167+

_final_setup(ax, scaling=scaling)

168+

return ax, fig

157169158-

def zgrid(zetas=None, wns=None, ax=None):

170+

# Grid for discrete time system (drawn, not rendered by AxisArtist)

171+

# TODO (at some point): think about using customized grid generator?

172+

def zgrid(zetas=None, wns=None, ax=None, scaling=None):

159173

"""Draws discrete damping and frequency grid"""

160174161175

fig = plt.gcf()

@@ -206,5 +220,9 @@ def zgrid(zetas=None, wns=None, ax=None):

206220

ax.annotate(r"$\frac{"+num+r"\pi}{T}$", xy=(an_x, an_y),

207221

xytext=(an_x, an_y), size=9)

208222209-

_final_setup(ax)

223+

# Set default axes to allow some room around the unit circle

224+

ax.set_xlim([-1.1, 1.1])

225+

ax.set_ylim([-1.1, 1.1])

226+227+

_final_setup(ax, scaling=scaling)

210228

return ax, fig