Hello,
I needed to use "stem" for complex impulse vectors and found that it
was too much work to make it behave. I therefore wrote "cstem" which
is intended for complex impulse vectors. It takes x (optional) and z
vectors, and a small number of special keyword parameters -- remaining
keywords are handed off to "plot" during plotting.
I've tried making it as similar to "stem" as I found useful, but it is
not a drop-in replacement. It is written as a function, not an axis
method, but if found useful to others, changing it over should not be
difficult.
The code has ReST docstring including example of use:
>>> z = np.random.uniform(size=10) + 1j*np.random.uniform(size=10)
>>> cstem(z)
or, using explicit x axis and keeping handles
>>> x = np.r_[1:3:10j]
>>> r_mark, i_mark, r_line, i_line, baseline = cstem(x, z)
License: anything BSD-ish that is broadly used in matplotlib already.
Enjoy! Comments/suggestions welcome.
--
Tom Grydeland
#
# Licensed under BSD-style terms, exact text TBD based on inputs from matplotlib list.
#
# (c) 2010 Tom Grydeland and Northern Research Institute Tromsø AS
import numpy as np
import pylab as pl
def _cstem_style(defs, kws, re, line):
ts = defs.copy()
ts.update(kws)
if 'color' in ts and 'markercolor' not in ts:
ts['markercolor'] = ts['color']
style = {}
for k in ts:
if k is 'basefmt': continue
if k is 'marker' and line: continue
if k is 'linestyle' and line: continue
if k is 'markercolor':
if line: continue
if re:
style['markerfacecolor'] = ts[k]
style['markeredgecolor'] = 'black'
else:
style['markerfacecolor'] = 'white'
style['markeredgecolor'] = ts[k]
continue
if k is 'reallinestyle':
if line and re: style['linestyle'] = ts[k]
continue
if k is 'imaglinestyle':
if line and not re: style['linestyle'] = ts[k]
continue
style[k] = ts[k]
return style
def cstem(x, z=None, **kwargs):
"""stemplot for complex impulse vectors.
Stem plot for a complex vector, using the same color, but distinct line
styles and marker shading, for the real and imaginary parts.
Parameters
--
x : vector of x values, optional (indices into z used if not given)
z : complex vector of values
Returns
---
r_mark, i_mark, r_line, i_line, baseline : handles for real/imag markers,
stemlines and the base (x axis) line, respectively
Other keyword parameters
color : used for stem line and marker (unless markercolor is specified
separately). If not given, taken from axes' color cycle.
marker : ('o') shape of marker (see pylab.plot() for available shapes).
Markers with area (i.e. one of 'ov^<>sphHDd*') are recommended.
markercolor: (same as 'color' if not given) used for 'markerfacecolor'
when plotting the real part (edges black), and for
'markeredgecolor' when plotting the imaginary part (face white).
reallinestyle, imaglinestyle : ('solid'/'dashed', respectively)
style of stem lines for real/imaginary part.
basefmt : line format used to plot a base line (x axis).
All other keyword parameters are passed to pylab.plot(), which will
barf on unknown keywords
See Also
plot : underlying plot function
stem : stem plot for real impulse vectors
Examples
>>> z = np.random.uniform(size=10) + 1j*np.random.uniform(size=10)
>>> cstem(z)
or, using explicit x axis and keeping handles
>>> x = np.r_[1:3:10j]
>>> r_mark, i_mark, r_line, i_line, baseline = cstem(x, z)
"""
if z is None:
z = x
x = np.arange(len(z))
defaults = {
'marker': 'o',
'linestyle' : 'None',
'reallinestyle' : 'solid',
'imaglinestyle' : 'dashed',
}
basefmt = 'k-'
if 'basefmt' in kwargs:
basefmt = kwargs.pop('basefmt')
ax = pl.gca()
remember_hold = ax._hold
if not ax._hold: ax.cla()
ax.hold(True)
if 'color' not in kwargs:
kwargs['color'] = ax._get_lines._get_next_cycle_color()
rsline = []; isline = []
immark, = ax.plot(x, z.imag, **_cstem_style(defaults, kwargs, False, False))
rmmark, = ax.plot(x, z.real, **_cstem_style(defaults, kwargs, True, False))
r_style = _cstem_style(defaults, kwargs, True, True)
i_style = _cstem_style(defaults, kwargs, False, True)
for xi, zi in zip(x, z):
il, = ax.plot([xi, xi], [0, zi.imag], **i_style)
rl, = ax.plot([xi, xi], [0, zi.real], **r_style)
rsline.append(rl)
isline.append(il)
baseline, = ax.plot([np.amin(x), np.amax(x)], [0,0], basefmt)
ax.hold(remember_hold)
return rmmark, immark, rsline, is
