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
  <Tom.Grydeland@(gmail.com)>
#
# 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, isline, baseline

------------------------------------------------------------------------------
Beautiful is writing same markup. Internet Explorer 9 supports
standards for HTML5, CSS3, SVG 1.1,  ECMAScript5, and DOM L2 & L3.
Spend less time writing and  rewriting code and more time creating great
experiences on the web. Be a part of the beta today.
http://p.sf.net/sfu/beautyoftheweb
_______________________________________________
Matplotlib-devel mailing list
Matplotlib-devel@lists.sourceforge.net
https://lists.sourceforge.net/lists/listinfo/matplotlib-devel

Reply via email to