python-chaco 4.1.0-1 / usr / share / pyshared / chaco / grid.py

""" Defines the PlotGrid class, and associated Traits UI View and validator
function.
"""

from __future__ import with_statement

from numpy import around, array, asarray, column_stack, float64, inf, zeros, zeros_like

# Enthought library imports
from enable.api import black_color_trait, LineStyle
from traits.api import Any, Bool, Callable, Enum, Float, Instance, \
        CInt, Trait, Property, TraitError, Tuple, on_trait_change
from traitsui.api import HGroup, Item, VGroup, View, TextEditor

# Local, relative imports
from abstract_overlay import AbstractOverlay
from abstract_mapper import AbstractMapper
from log_mapper import LogMapper
from ticks import AbstractTickGenerator, DefaultTickGenerator


def float_or_auto(val):
    """
    Validator function that returns *val* if *val* is either a number or
    the word 'auto'.  This is used as a validator for the text editor
    in the traits UI for the tick_interval trait.
    """
    try:
        return float(val)
    except:
        if isinstance(val, basestring) and val == "auto":
            return val
    raise TraitError, "Tick interval must be a number or 'auto'."

# View for setting grid properties.
GridView = View(VGroup(
                HGroup(Item("grid_interval", label="Interval", editor=TextEditor(evaluate=float_or_auto)),
                       Item("visible", label="Visible")),
                Item("line_color", label="Color", style="custom"),
                Item("line_style", label="Dash style"),
                Item("line_width", label="Thickness")
                ),
                buttons = ["OK", "Cancel"]
            )


def Alias(name):
    return Property(lambda obj: getattr(obj, name),
                    lambda obj, val: setattr(obj, name, val))


class PlotGrid(AbstractOverlay):
    """ An overlay that represents a grid.

    A grid is a set of parallel lines, horizontal or vertical. You can use
    multiple grids with different settings for the horizontal and vertical
    lines in a plot.
    """

    #------------------------------------------------------------------------
    # Data-related traits
    #------------------------------------------------------------------------

    # The mapper (and associated range) that drive this PlotGrid.
    mapper = Instance(AbstractMapper)

    # The dataspace interval between grid lines.
    grid_interval = Trait('auto', 'auto', Float)

    # The dataspace value at which to start this grid.  If None, then
    # uses the mapper.range.low.
    data_min = Trait(None, None, Float)

    # The dataspace value at which to end this grid.  If None, then uses
    # the mapper.range.high.
    data_max = Trait(None, None, Float)

    # A callable that implements the AbstractTickGenerator Interface.
    tick_generator = Instance(AbstractTickGenerator)

    #------------------------------------------------------------------------
    # Layout traits
    #------------------------------------------------------------------------

    # The orientation of the grid lines.  "horizontal" means that the grid
    # lines are parallel to the X axis and the ticker and grid interval
    # refer to the Y axis.
    orientation = Enum('horizontal', 'vertical')

    # Draw the ticks starting at the end of the mapper range? If False, the
    # ticks are drawn starting at 0. This setting can be useful to keep the
    # grid from from "flashing" as the user resizes the plot area.
    flip_axis = Bool(False)

    # Optional specification of the grid bounds in the dimension transverse
    # to the ticking/gridding dimension, i.e. along the direction specified
    # by self.orientation.  If this is specified but transverse_mapper is
    # not specified, then there is no effect.
    #
    #   None : use self.bounds or self.component.bounds (if overlay)
    #   Tuple : (low, high) extents, used for every grid line
    #   Callable : Function that takes an array of dataspace grid ticks
    #              and returns either an array of shape (N,2) of (starts,ends)
    #              for each grid point or a single tuple (low, high)
    transverse_bounds = Trait(None, Tuple, Callable)

    # Mapper in the direction corresponding to self.orientation, i.e. transverse
    # to the direction of self.mapper.  This is used to compute the screen
    # position of transverse_bounds.  If this is not specified, then
    # transverse_bounds has no effect, and vice versa.
    transverse_mapper = Instance(AbstractMapper)

    # Dimensions that the grid is resizable in (overrides PlotComponent).
    resizable = "hv"

    #------------------------------------------------------------------------
    # Appearance traits
    #------------------------------------------------------------------------

    # The color of the grid lines.
    line_color = black_color_trait

    # The style (i.e., dash pattern) of the grid lines.
    line_style = LineStyle('solid')

    # The thickness, in pixels, of the grid lines.
    line_width = CInt(1)
    line_weight = Alias("line_width")

    # Default Traits UI View for modifying grid attributes.
    traits_view = GridView

    #------------------------------------------------------------------------
    # Private traits; mostly cached information
    #------------------------------------------------------------------------

    _cache_valid = Bool(False)
    _tick_list = Any
    _tick_positions = Any

    # An array (N,2) of start,end positions in the transverse direction
    # i.e. the direction corresponding to self.orientation
    _tick_extents = Any

    #_length = Float(0.0)


    #------------------------------------------------------------------------
    # Public methods
    #------------------------------------------------------------------------

    def __init__(self, **traits):
        # TODO: change this back to a factory in the instance trait some day
        self.tick_generator = DefaultTickGenerator()
        super(PlotGrid, self).__init__(**traits)
        self.bgcolor = "none" #make sure we're transparent
        return

    @on_trait_change("bounds,bounds_items,position,position_items")
    def invalidate(self):
        """ Invalidate cached information about the grid.
        """
        self._reset_cache()
        return


    #------------------------------------------------------------------------
    # PlotComponent and AbstractOverlay interface
    #------------------------------------------------------------------------

    def do_layout(self, *args, **kw):
        """ Tells this component to do layout at a given size.

        Overrides PlotComponent.
        """
        if self.use_draw_order and self.component is not None:
            self._layout_as_overlay(*args, **kw)
        else:
            super(PlotGrid, self).do_layout(*args, **kw)
        return

    #------------------------------------------------------------------------
    # Private methods
    #------------------------------------------------------------------------

    def _do_layout(self):
        """ Performs a layout.

        Overrides PlotComponent.
        """
        return

    def _layout_as_overlay(self, size=None, force=False):
        """ Lays out the axis as an overlay on another component.
        """
        if self.component is not None:
            self.position = self.component.position
            self.bounds = self.component.bounds
        return

    def _reset_cache(self):
        """ Clears the cached tick positions.
        """
        self._tick_positions = array([], dtype=float)
        self._tick_extents = array([], dtype=float)
        self._cache_valid = False
        return

    def _compute_ticks(self, component=None):
        """ Calculates the positions for the grid lines.
        """
        if (self.mapper is None):
            self._reset_cache()
            self._cache_valid = True
            return

        if self.data_min is None:
            datalow = self.mapper.range.low
        else:
            datalow = self.data_min
        if self.data_max is None:
            datahigh = self.mapper.range.high
        else:
            datahigh = self.data_max

        # Map the low and high data points
        screenhigh = self.mapper.map_screen(datalow)
        screenlow = self.mapper.map_screen(datahigh)

        if (datalow == datahigh) or (screenlow == screenhigh) or \
           (datalow in [inf, -inf]) or (datahigh in [inf, -inf]):
            self._reset_cache()
            self._cache_valid = True
            return

        if component is None:
            component = self.component

        if component is not None:
            bounds = component.bounds
            position = component.position
        else:
            bounds = self.bounds
            position = self.position

        if isinstance(self.mapper, LogMapper):
            scale = 'log'
        else:
            scale = 'linear'

        ticks = self.tick_generator.get_ticks(datalow, datahigh, datalow, datahigh,
                                              self.grid_interval, use_endpoints = False,
                                              scale=scale)
        tick_positions = self.mapper.map_screen(array(ticks, float64))

        if self.orientation == 'horizontal':
            self._tick_positions = around(column_stack((zeros_like(tick_positions) + position[0],
                                           tick_positions)))
        elif self.orientation == 'vertical':
            self._tick_positions = around(column_stack((tick_positions,
                                           zeros_like(tick_positions) + position[1])))
        else:
            raise self.NotImplementedError

        # Compute the transverse direction extents
        self._tick_extents = zeros((len(ticks), 2), dtype=float)
        if self.transverse_bounds is None or self.transverse_mapper is None:
            # No mapping needed, just use the extents
            if self.orientation == 'horizontal':
                extents = (position[0], position[0] + bounds[0])
            elif self.orientation == 'vertical':
                extents = (position[1], position[1] + bounds[1])
            self._tick_extents[:] = extents
        elif callable(self.transverse_bounds):
            data_extents = self.transverse_bounds(ticks)
            tmapper = self.transverse_mapper
            if isinstance(data_extents, tuple):
                self._tick_extents[:] = tmapper.map_screen(asarray(data_extents))
            else:
                extents = array([tmapper.map_screen(data_extents[:,0]),
                                 tmapper.map_screen(data_extents[:,1])]).T
                self._tick_extents = extents
        else:
            # Already a tuple
            self._tick_extents[:] = self.transverse_mapper.map_screen(asarray(self.transverse_bounds))

        self._cache_valid = True


    def _draw_overlay(self, gc, view_bounds=None, mode='normal'):
        """ Draws the overlay layer of a component.

        Overrides PlotComponent.
        """
        self._draw_component(gc, view_bounds, mode)
        return

    def overlay(self, other_component, gc, view_bounds=None, mode="normal"):
        """ Draws this component overlaid on another component.

        Overrides AbstractOverlay.
        """
        if not self.visible:
            return
        self._compute_ticks(other_component)
        self._draw_component(gc, view_bounds, mode)
        self._cache_valid = False
        return

    def _draw_component(self, gc, view_bounds=None, mode="normal"):
        """ Draws the component.

        This method is preserved for backwards compatibility. Overrides
        PlotComponent.
        """
        # What we're really trying to do with a grid is plot contour lines in
        # the space of the plot.  In a rectangular plot, these will always be
        # straight lines.
        if not self.visible:
            return

        if not self._cache_valid:
            self._compute_ticks()

        if len(self._tick_positions) == 0:
            return

        with gc:
            gc.set_line_width(self.line_weight)
            gc.set_line_dash(self.line_style_)
            gc.set_stroke_color(self.line_color_)
            gc.set_antialias(False)

            if self.component is not None:
                gc.clip_to_rect(*(self.component.position + self.component.bounds))
            else:
                gc.clip_to_rect(*(self.position + self.bounds))

            gc.begin_path()
            if self.orientation == "horizontal":
                starts = self._tick_positions.copy()
                starts[:,0] = self._tick_extents[:,0]
                ends = self._tick_positions.copy()
                ends[:,0] = self._tick_extents[:,1]
            else:
                starts = self._tick_positions.copy()
                starts[:,1] = self._tick_extents[:,0]
                ends = self._tick_positions.copy()
                ends[:,1] = self._tick_extents[:,1]
            if self.flip_axis:
                starts, ends = ends, starts
            gc.line_set(starts, ends)
            gc.stroke_path()
        return

    def _mapper_changed(self, old, new):
        if old is not None:
            old.on_trait_change(self.mapper_updated, "updated", remove=True)
        if new is not None:
            new.on_trait_change(self.mapper_updated, "updated")
        self.invalidate()
        return

    def mapper_updated(self):
        """
        Event handler that is bound to this mapper's **updated** event.
        """
        self.invalidate()
        return

    def _position_changed_for_component(self):
        self.invalidate()

    def _position_items_changed_for_component(self):
        self.invalidate()

    def _bounds_changed_for_component(self):
        self.invalidate()

    def _bounds_items_changed_for_component(self):
        self.invalidate()

    #------------------------------------------------------------------------
    # Event handlers for visual attributes.  These mostly just call request_redraw()
    #------------------------------------------------------------------------

    @on_trait_change("visible,line_color,line_style,line_weight")
    def visual_attr_changed(self):
        """ Called when an attribute that affects the appearance of the grid
        is changed.
        """
        if self.component:
            self.component.invalidate_draw()
            self.component.request_redraw()
        else:
            self.invalidate_draw()
            self.request_redraw()

    def _grid_interval_changed(self):
        self.invalidate()
        self.visual_attr_changed()

    def _orientation_changed(self):
        self.invalidate()
        self.visual_attr_changed()
        return


    ### Persistence ###########################################################
    #_pickles = ("orientation", "line_color", "line_style", "line_weight",
    #            "grid_interval", "mapper")

    def __getstate__(self):
        state = super(PlotGrid,self).__getstate__()
        for key in ['_cache_valid', '_tick_list', '_tick_positions', '_tick_extents']:
            if state.has_key(key):
                del state[key]

        return state

    def _post_load(self):
        super(PlotGrid, self)._post_load()
        self._mapper_changed(None, self.mapper)
        self._reset_cache()
        self._cache_valid = False
        return


# EOF