kmatplot/kmatplot/python/kplot.py

"""This module allows plotting via KMatplot.

As you need Numeric arrays to send any data you will have to import module
Numeric.

Instantiating class kplot starts and connects to a KMatplot session.
Plot via members plot, surface, polygon, pixmap and contour.
"""

__version__ = 0.3
__author__ = "Martin Wiechert <martin.wiechert@gmx.de>"
__date__ = "December 19, 2001"
__credits__ = """Guido van Rossum (Python)
Kamil Dobkowski (KMatplot)
Nathaniel Gray (gracePlot.py)
"""

import KMatplot as kmp
import _kplot

class Error (Exception):
    pass
            
class kplot:
    """Instantiating this class starts and connects to a KMatplot session.
    
    You can have multiple subplots.
    Attribute 'current' contains the number of the active subplot which all
    plotting is directed to. Add new subplots with the 'subplot' function.

    Function 'hold' toggles between overlay and overwrite mode. In the latter
    the current (sub)plot will be erased before new data are sent. Default is
    overlay mode.
    """

    def __init__ (self):
        self.proc = kmp.kmpProcess ()
        self.subplots = []
        self.subplot ()
        
    def subplot (self): # FIXME geometry management missing
        """Starts a new subplot and directs future drawing to it.

        Switch between subplots via attribute 'current'."""

        self.current = len (self.subplots)
        self.subplots.append (_kplot._ksubplot (self))

    def hold (self, switch = None):
        """Pass 1 for overlay mode, 0 for overwrite mode or nothing to toggle.

        Only the current subplot is affected.

        Obviously you cannot overlay 2d and 3d plots, so if you try this the
        hold status will be silently ignored."""
        
        if switch is None:
            self.subplots [self.current].hold = \
                          not self.subplots [self.current].hold
        else:
            self.subplots [self.current].hold = switch

    def plot (self, *args, **kwargs):
        """Plots a 2d curve, possibly with error bars.

        There are essentially three ways to pass data:
        
        1) By keyword arguments X, Y, dX, dY:
            One curve is plotted per row in Y. If dX or dY are specified
            they must have the same shape as Y.
            If X is specified it must either have the same shape as Y or
            be one-dimensional and be as long as the rows in Y.

        2) By keyword arguments XY / XYdX / XYdY / XYdXdY / YdX / YdY / YdXdY:
            One curve per 2 / 3 / 3 / 4 / 2 / 2 / 3 rows in the input array.
            Number of rows must be a multiple of 2 / 3 / 3 / 4 / 2 / 2 / 3.
            You may also use permutations like YdYX if your data happen to
            come in that order.

        3) By positional arguments:
            One argument is interpreted as Y.
            Two to four arguments are treated as X, Y [, dX [, dY]]
            Same shape requirements as in 1).
            If you give only one or two arguments you may in addition
            use dX / dY keyword arguments.
        """

        self.subplots [self.current].plot (args, kwargs)

    def arrows (self, X, Y, dX, dY):
        """Draws a set of plane arrows.

        Arrows start at X, Y and extend by dX, dY.
        
        arrows (X, Y, dX, dY)

        All arguments are required and must have same size, i.e. number of
        elements. Shapes are ignored.

        *** Not yet implemented. ***
        """

        pass # FIXME: implement it (and adjust doc string).
        # self.subplots [self.current].plot (
        # [X.flat, Y.flat, dX.flat, dY.flat], {}, kmp.PlotVectors ?)
        
    def surface (self, Z, X = None, Y = None):
        """Plots a surface.

        surface (Z [, X] [, Y])

        Z must be a matrix.
        X, Y must be vectors of sizes number of columns of Z,
                                      number of rows of Z respectively.
        """

        self.subplots [self.current].plane_map ([Z, X, Y], kmp.PlotSurface)

    def polygon (self, X, Y, Z):
        """Plots one or more polygons in space.

        polygon (X, Y, Z)
        
        All arguments are required and must be of same shape.
        Either vectors (one polygon) or matrices (one polygon per row).
        """

        if not (X.shape == Y.shape and X.shape ==Z.shape):
            raise Error, 'Arguments must have same shapes.'
        self.subplots [self.current].polygon (X, Y, Z)
        
    def contour (self, Z, X = None, Y = None):
        """Plots a colour-coded height map.

        contour (Z [, X] [, Y])
        
        Z must be a matrix.
        X, Y must be vectors of sizes number of columns of Z,
                                      number of rows of Z respectively.
        """

        self.subplots [self.current].plane_map ([Z, X, Y], kmp.PlotContour)

    def pixmap (self, red, green = None, blue = None, X = None, Y = None):
        """Plots a possibly scaled possibly coloured pixmap.

        pixmap (red [,green, blue] [, X] [, Y])

        If green and blue are not specified output is not red but greyscaled.

        red, green, blue must be a matrices of same shape.
        X, Y must be vectors of sizes number of columns of Z + 1,
                                      number of rows of Z + 1 respectively.
        """

        if green or blue:
            if not (green and blue):
                raise Error, 'Green and blue can only be specified together.'
            if not (red.shape == green.shape and red.shape == blue.shape):
                raise Error, 'Colours must have same shapes.'
        self.subplots [self.current].plane_map ([red, green, blue, None, X, Y],
                                                kmp.PlotImage)