예제 #1
0
class ContourLabeler:
    def __init__(self, ax):
        self.ax = ax

    def clabel(self, *args, **kwargs):
        """
        CLABEL(*args, **kwargs)

        Function signatures

        CLABEL(C) - plots contour labels,
                    C is the output of contour or a list of contours

        CLABEL(C,V) - creates labels only for those contours, given in
                      a list V

        CLABEL(C, **kwargs) - keyword args are explained below:



        * fontsize = None: as described in http://matplotlib.sf.net/fonts.html

        * colors = None:

           - a tuple of matplotlib color args (string, float, rgb, etc),
             different labels will be plotted in different colors in the order
             specified

           - one string color, e.g. colors = 'r' or colors = 'red', all labels
             will be plotted in this color

           - if colors == None, the color of each label matches the color
             of the corresponding contour

        * inline = 0: controls whether the underlying contour is removed
                     (inline = 1) or not

        * fmt = '%1.3f': a format string for the label

        """
        # todo, factor this out to a separate class and don't use hidden coll attrs

        if not self.ax.ishold(): self.ax.cla()

        fontsize = kwargs.get('fontsize', None)
        inline = kwargs.get('inline', 0)
        fmt = kwargs.get('fmt', '%1.3f')
        colors = kwargs.get('colors', None)

        if len(args) == 1:
            contours = args[0]
            levels = [con._label for con in contours]
        elif len(args) == 2:
            contours = args[0]
            levels = args[1]
        else:
            raise TypeError("Illegal arguments to clabel, see help(clabel)")

        self.fp = FontProperties()
        if fontsize == None:
            font_size = int(self.fp.get_size_in_points())
        else:
            if type(fontsize) not in [int, float, str]:
                raise TypeError("Font size must be an integer number.")
            else:
                if type(fontsize) == str:
                    font_size = int(self.fp.get_size_in_points())

                else:
                    self.fp.set_size(fontsize)
                    font_size = fontsize
        fslist = [font_size] * len(levels)

        if colors == None:
            colors = [c._colors[0] for c in contours]
        else:
            colors = colors * len(contours)

        if inline not in [0, 1]:
            raise TypeError("inline must be 0 or 1")

        self.cl = []
        self.cl_xy = []

        # we have a list of contours and each contour has a list of
        # segments.  We want changes in the contour color to be
        # reflected in changes in the label color.  This is a good use
        # for traits observers, but in the interim, until traits are
        # utilized, we'll create a dict mapping i,j to text instances.
        # i is the contour level index, j is the sement index
        self.labeld = {}
        if inline == 1:
            self.inline_labels(levels, contours, colors, fslist, fmt)
        else:
            self.labels(levels, contours, colors, fslist, fmt)

        for label in self.cl:
            self.ax.add_artist(label)

        ret = silent_list('Text', self.cl)
        ret.mappable = getattr(contours, 'mappable', None)
        # support colormapping for label
        if ret.mappable is not None:
            ret.mappable.labeld = self.labeld
        return ret

    def print_label(self, linecontour, labelwidth):
        "if contours are too short, don't plot a label"
        lcsize = len(linecontour)
        if lcsize > 10 * labelwidth:
            return 1

        xmax = amax(array(linecontour)[:, 0])
        xmin = amin(array(linecontour)[:, 0])
        ymax = amax(array(linecontour)[:, 1])
        ymin = amin(array(linecontour)[:, 1])

        lw = labelwidth
        if (xmax - xmin) > 1.2 * lw or (ymax - ymin) > 1.2 * lw:
            return 1
        else:
            return 0

    def too_close(self, x, y, lw):
        "if there's a label already nearby, find a better place"
        if self.cl_xy != []:
            dist = [
                sqrt((x - loc[0])**2 + (y - loc[1])**2) for loc in self.cl_xy
            ]
            for d in dist:
                if d < 1.2 * lw:
                    return 1
                else:
                    return 0
        else:
            return 0

    def get_label_coords(self, distances, XX, YY, ysize, lw):
        """ labels are ploted at a location with the smallest
        dispersion of the contour from a straight line
        unless there's another label nearby, in which case
        the second best place on the contour is picked up
        if there's no good place a label isplotted at the
        beginning of the contour
        """

        hysize = int(ysize / 2)
        adist = argsort(distances)

        for ind in adist:
            x, y = XX[ind][hysize], YY[ind][hysize]
            if self.too_close(x, y, lw):
                continue
            else:
                self.cl_xy.append((x, y))
                return x, y, ind

        ind = adist[0]
        x, y = XX[ind][hysize], YY[ind][hysize]
        self.cl_xy.append((x, y))
        return x, y, ind

    def get_label_width(self, lev, fmt, fsize):
        "get the width of the label in points"
        if is_string_like(lev):
            lw = (len(lev)) * fsize
        else:
            lw = (len(fmt % lev)) * fsize

        return lw

    def set_label_props(self, label, text, color):
        "set the label properties - color, fontsize, text"
        label.set_text(text)
        label.set_color(color)
        label.set_fontproperties(self.fp)
        label.set_clip_box(self.ax.bbox)

    def get_text(self, lev, fmt):
        "get the text of the label"
        if is_string_like(lev):
            return lev
        else:
            return fmt % lev

    def break_linecontour(self, linecontour, rot, labelwidth, ind):
        "break a contour in two contours at the location of the label"
        lcsize = len(linecontour)
        hlw = int(labelwidth / 2)

        #length of label in screen coords
        ylabel = abs(hlw * sin(rot * pi / 180))
        xlabel = abs(hlw * cos(rot * pi / 180))

        trans = self.ax.transData

        slc = trans.seq_xy_tups(linecontour)
        x, y = slc[ind]
        xx = array(slc)[:, 0].copy()
        yy = array(slc)[:, 1].copy()

        #indices which are under the label
        inds = nonzero(((xx < x + xlabel) & (xx > x - xlabel))
                       & ((yy < y + ylabel) & (yy > y - ylabel)))

        if len(inds) > 0:
            #if the label happens to be over the beginning of the
            #contour, the entire contour is removed, i.e.
            #indices to be removed are
            #inds= [0,1,2,3,305,306,307]
            #should rewrite this in a better way
            linds = nonzero(inds[1:] - inds[:-1] != 1)
            if inds[0] == 0 and len(linds) != 0:
                ii = inds[linds[0]]
                lc1 = linecontour[ii + 1:inds[ii + 1]]
                lc2 = []

            else:
                lc1 = linecontour[:inds[0]]
                lc2 = linecontour[inds[-1] + 1:]

        else:
            lc1 = linecontour[:ind]
            lc2 = linecontour[ind + 1:]

        if rot < 0:
            new_x1, new_y1 = x - xlabel, y + ylabel
            new_x2, new_y2 = x + xlabel, y - ylabel
        else:
            new_x1, new_y1 = x - xlabel, y - ylabel
            new_x2, new_y2 = x + xlabel, y + ylabel

        new_x1d, new_y1d = trans.inverse_xy_tup((new_x1, new_y1))
        new_x2d, new_y2d = trans.inverse_xy_tup((new_x2, new_y2))

        if rot > 0:
            if len(lc1) > 0 and (lc1[-1][0] <= new_x1d) and (lc1[-1][1] <=
                                                             new_y1d):
                lc1.append((new_x1d, new_y1d))

            if len(lc2) > 0 and (lc2[0][0] >= new_x2d) and (lc2[0][1] >=
                                                            new_y2d):
                lc2.insert(0, (new_x2d, new_y2d))
        else:
            if len(lc1) > 0 and ((lc1[-1][0] <= new_x1d) and
                                 (lc1[-1][1] >= new_y1d)):
                lc1.append((new_x1d, new_y1d))

            if len(lc2) > 0 and ((lc2[0][0] >= new_x2d) and
                                 (lc2[0][1] <= new_y2d)):
                lc2.insert(0, (new_x2d, new_y2d))

        return [lc1, lc2]

    def locate_label(self, linecontour, labelwidth):
        """find a good place to plot a label (relatively flat
        part of the contour) and the angle of rotation for the
        text object
        """

        nsize = len(linecontour)
        if labelwidth > 1:
            xsize = int(ceil(nsize / labelwidth))
        else:
            xsize = 1
        if xsize == 1:
            ysize = nsize
        else:
            ysize = labelwidth

        XX = resize(array(linecontour)[:, 0], (xsize, ysize))
        YY = resize(array(linecontour)[:, 1], (xsize, ysize))

        yfirst = YY[:, 0]
        ylast = YY[:, -1]
        xfirst = XX[:, 0]
        xlast = XX[:, -1]
        s = (reshape(yfirst,
                     (xsize, 1)) - YY) * (reshape(xlast, (xsize, 1)) - reshape(
                         xfirst,
                         (xsize, 1))) - (reshape(xfirst, (xsize, 1)) - XX) * (
                             reshape(ylast,
                                     (xsize, 1)) - reshape(yfirst, (xsize, 1)))
        L = sqrt((xlast - xfirst)**2 + (ylast - yfirst)**2)
        dist = add.reduce(([(abs(s)[i] / L[i]) for i in range(xsize)]), -1)
        x, y, ind = self.get_label_coords(dist, XX, YY, ysize, labelwidth)
        angle = arctan2(ylast - yfirst, xlast - xfirst)
        rotation = angle[ind] * 180 / pi
        if rotation > 90:
            rotation = rotation - 180
        if rotation < -90:
            rotation = 180 + rotation

        dind = list(linecontour).index((x, y))

        return x, y, rotation, dind

    def inline_labels(self, levels, contours, colors, fslist, fmt):
        trans = self.ax.transData
        contourNum = 0
        for lev, con, color, fsize in zip(levels, contours, colors, fslist):
            toremove = []
            toadd = []
            lw = self.get_label_width(lev, fmt, fsize)
            for segNum, linecontour in enumerate(con._segments):
                key = contourNum, segNum
                # for closed contours add one more point to
                # avoid division by zero
                if linecontour[0] == linecontour[-1]:
                    linecontour.append(linecontour[1])
                # transfer all data points to screen coordinates
                slc = trans.seq_xy_tups(linecontour)
                if self.print_label(slc, lw):
                    x, y, rotation, ind = self.locate_label(slc, lw)
                    # transfer the location of the label back to
                    # data coordinates
                    dx, dy = trans.inverse_xy_tup((x, y))
                    t = Text(dx,
                             dy,
                             rotation=rotation,
                             horizontalalignment='center',
                             verticalalignment='center')
                    self.labeld[key] = t
                    text = self.get_text(lev, fmt)
                    self.set_label_props(t, text, color)
                    self.cl.append(t)
                    new = self.break_linecontour(linecontour, rotation, lw,
                                                 ind)

                    for c in new:
                        toadd.append(c)
                    toremove.append(linecontour)
            for c in toremove:
                con._segments.remove(c)
            for c in toadd:
                con._segments.append(c)

            contourNum += 1

    def labels(self, levels, contours, colors, fslist, fmt):
        trans = self.ax.transData
        for lev, con, color, fsize in zip(levels, contours, colors, fslist):
            lw = self.get_label_width(lev, fmt, fsize)
            for linecontour in con._segments:
                # for closed contours add one more point
                if linecontour[0] == linecontour[-1]:
                    linecontour.append(linecontour[1])
                # transfer all data points to screen coordinates
                slc = trans.seq_xy_tups(linecontour)
                if self.print_label(slc, lw):
                    x, y, rotation, ind = self.locate_label(slc, lw)
                    # transfer the location of the label back into
                    # data coordinates
                    dx, dy = trans.inverse_xy_tup((x, y))
                    t = Text(dx,
                             dy,
                             rotation=rotation,
                             horizontalalignment='center',
                             verticalalignment='center')
                    text = self.get_text(lev, fmt)
                    self.set_label_props(t, text, color)
                    self.cl.append(t)
                else:
                    pass
예제 #2
0
class ContourLabeler:
    '''Mixin to provide labelling capability to ContourSet'''
    def clabel(self, *args, **kwargs):
        """
        clabel(CS, **kwargs) - add labels to line contours in CS,
               where CS is a ContourSet object returned by contour.

        clabel(CS, V, **kwargs) - only label contours listed in V

        keyword arguments:

        * fontsize = None: as described in http://matplotlib.sf.net/fonts.html

        * colors = None:

           - a tuple of matplotlib color args (string, float, rgb, etc),
             different labels will be plotted in different colors in the order
             specified

           - one string color, e.g. colors = 'r' or colors = 'red', all labels
             will be plotted in this color

           - if colors == None, the color of each label matches the color
             of the corresponding contour

        * inline = True: controls whether the underlying contour is removed
                     (inline = True) or not (False)

        * fmt = '%1.3f': a format string for the label

        """
        fontsize = kwargs.get('fontsize', None)
        inline = kwargs.get('inline', 1)
        self.fmt = kwargs.get('fmt', '%1.3f')
        colors = kwargs.get('colors', None)

        if len(args) == 0:
            levels = self.levels
            indices = range(len(self.levels))
        elif len(args) == 1:
            levlabs = list(args[0])
            indices, levels = [], []
            for i, lev in enumerate(self.levels):
                if lev in levlabs:
                    indices.append(i)
                    levels.append(lev)
            if len(levels) < len(levlabs):
                msg = "Specified levels " + str(levlabs)
                msg += "\n don't match available levels "
                msg += str(self.levels)
                raise ValueError(msg)
        else:
            raise TypeError("Illegal arguments to clabel, see help(clabel)")
        self.label_levels = levels
        self.label_indices = indices

        self.fp = FontProperties()
        if fontsize == None:
            font_size = int(self.fp.get_size_in_points())
        else:
            if type(fontsize) not in [int, float, str]:
                raise TypeError("Font size must be an integer number.")
                # Can't it be floating point, as indicated in line above?
            else:
                if type(fontsize) == str:
                    font_size = int(self.fp.get_size_in_points())
                else:
                    self.fp.set_size(fontsize)
                    font_size = fontsize
        self.fslist = [font_size] * len(levels)

        if colors == None:
            self.label_mappable = self
            self.label_cvalues = take(self.cvalues, self.label_indices)
        else:
            cmap = ListedColormap(colors, N=len(self.label_levels))
            self.label_cvalues = range(len(self.label_levels))
            self.label_mappable = ScalarMappable(cmap=cmap, norm=no_norm())

        #self.cl = []   # Initialized in ContourSet.__init__
        #self.cl_cvalues = [] # same
        self.cl_xy = []

        self.labels(inline)

        for label in self.cl:
            self.ax.add_artist(label)

        self.label_list = silent_list('Text', self.cl)
        return self.label_list

    def print_label(self, linecontour, labelwidth):
        "if contours are too short, don't plot a label"
        lcsize = len(linecontour)
        if lcsize > 10 * labelwidth:
            return 1

        xmax = amax(array(linecontour)[:, 0])
        xmin = amin(array(linecontour)[:, 0])
        ymax = amax(array(linecontour)[:, 1])
        ymin = amin(array(linecontour)[:, 1])

        lw = labelwidth
        if (xmax - xmin) > 1.2 * lw or (ymax - ymin) > 1.2 * lw:
            return 1
        else:
            return 0

    def too_close(self, x, y, lw):
        "if there's a label already nearby, find a better place"
        if self.cl_xy != []:
            dist = [
                sqrt((x - loc[0])**2 + (y - loc[1])**2) for loc in self.cl_xy
            ]
            for d in dist:
                if d < 1.2 * lw:
                    return 1
                else:
                    return 0
        else:
            return 0

    def get_label_coords(self, distances, XX, YY, ysize, lw):
        """ labels are ploted at a location with the smallest
        dispersion of the contour from a straight line
        unless there's another label nearby, in which case
        the second best place on the contour is picked up
        if there's no good place a label isplotted at the
        beginning of the contour
        """

        hysize = int(ysize / 2)
        adist = argsort(distances)

        for ind in adist:
            x, y = XX[ind][hysize], YY[ind][hysize]
            if self.too_close(x, y, lw):
                continue
            else:
                self.cl_xy.append((x, y))
                return x, y, ind

        ind = adist[0]
        x, y = XX[ind][hysize], YY[ind][hysize]
        self.cl_xy.append((x, y))
        return x, y, ind

    def get_label_width(self, lev, fmt, fsize):
        "get the width of the label in points"
        if is_string_like(lev):
            lw = (len(lev)) * fsize
        else:
            lw = (len(fmt % lev)) * fsize

        return lw

    def set_label_props(self, label, text, color):
        "set the label properties - color, fontsize, text"
        label.set_text(text)
        label.set_color(color)
        label.set_fontproperties(self.fp)
        label.set_clip_box(self.ax.bbox)

    def get_text(self, lev, fmt):
        "get the text of the label"
        if is_string_like(lev):
            return lev
        else:
            return fmt % lev

    def break_linecontour(self, linecontour, rot, labelwidth, ind):
        "break a contour in two contours at the location of the label"
        lcsize = len(linecontour)
        hlw = int(labelwidth / 2)

        #length of label in screen coords
        ylabel = abs(hlw * sin(rot * pi / 180))
        xlabel = abs(hlw * cos(rot * pi / 180))

        trans = self.ax.transData

        slc = trans.seq_xy_tups(linecontour)
        x, y = slc[ind]
        xx = asarray(slc)[:, 0].copy()
        yy = asarray(slc)[:, 1].copy()

        #indices which are under the label
        inds = nonzero(((xx < x + xlabel) & (xx > x - xlabel))
                       & ((yy < y + ylabel) & (yy > y - ylabel)))

        if len(inds) > 0:
            #if the label happens to be over the beginning of the
            #contour, the entire contour is removed, i.e.
            #indices to be removed are
            #inds= [0,1,2,3,305,306,307]
            #should rewrite this in a better way
            linds = nonzero(inds[1:] - inds[:-1] != 1)
            if inds[0] == 0 and len(linds) != 0:
                ii = inds[linds[0]]
                lc1 = linecontour[ii + 1:inds[ii + 1]]
                lc2 = []

            else:
                lc1 = linecontour[:inds[0]]
                lc2 = linecontour[inds[-1] + 1:]

        else:
            lc1 = linecontour[:ind]
            lc2 = linecontour[ind + 1:]

        if rot < 0:
            new_x1, new_y1 = x - xlabel, y + ylabel
            new_x2, new_y2 = x + xlabel, y - ylabel
        else:
            new_x1, new_y1 = x - xlabel, y - ylabel
            new_x2, new_y2 = x + xlabel, y + ylabel

        new_x1d, new_y1d = trans.inverse_xy_tup((new_x1, new_y1))
        new_x2d, new_y2d = trans.inverse_xy_tup((new_x2, new_y2))
        new_xy1 = array(((new_x1d, new_y1d), ))
        new_xy2 = array(((new_x2d, new_y2d), ))

        if rot > 0:
            if (len(lc1) > 0 and (lc1[-1][0] <= new_x1d)
                    and (lc1[-1][1] <= new_y1d)):
                lc1 = concatenate((lc1, new_xy1))
                #lc1.append((new_x1d, new_y1d))

            if (len(lc2) > 0 and (lc2[0][0] >= new_x2d)
                    and (lc2[0][1] >= new_y2d)):
                lc2 = concatenate((new_xy2, lc2))
                #lc2.insert(0, (new_x2d, new_y2d))
        else:
            if (len(lc1) > 0
                    and ((lc1[-1][0] <= new_x1d) and (lc1[-1][1] >= new_y1d))):
                lc1 = concatenate((lc1, new_xy1))
                #lc1.append((new_x1d, new_y1d))

            if (len(lc2) > 0
                    and ((lc2[0][0] >= new_x2d) and (lc2[0][1] <= new_y2d))):
                lc2 = concatenate((new_xy2, lc2))
                #lc2.insert(0, (new_x2d, new_y2d))

        return [lc1, lc2]

    def locate_label(self, linecontour, labelwidth):
        """find a good place to plot a label (relatively flat
        part of the contour) and the angle of rotation for the
        text object
        """

        nsize = len(linecontour)
        if labelwidth > 1:
            xsize = int(ceil(nsize / labelwidth))
        else:
            xsize = 1
        if xsize == 1:
            ysize = nsize
        else:
            ysize = labelwidth

        XX = resize(asarray(linecontour)[:, 0], (xsize, ysize))
        YY = resize(asarray(linecontour)[:, 1], (xsize, ysize))

        yfirst = YY[:, 0]
        ylast = YY[:, -1]
        xfirst = XX[:, 0]
        xlast = XX[:, -1]
        s = ((reshape(yfirst, (xsize, 1)) - YY) *
             (reshape(xlast, (xsize, 1)) - reshape(xfirst, (xsize, 1))) -
             (reshape(xfirst, (xsize, 1)) - XX) *
             (reshape(ylast, (xsize, 1)) - reshape(yfirst, (xsize, 1))))
        L = sqrt((xlast - xfirst)**2 + (ylast - yfirst)**2)
        dist = add.reduce(([(abs(s)[i] / L[i]) for i in range(xsize)]), -1)
        x, y, ind = self.get_label_coords(dist, XX, YY, ysize, labelwidth)
        #print 'ind, x, y', ind, x, y
        angle = arctan2(ylast - yfirst, xlast - xfirst)
        rotation = angle[ind] * 180 / pi
        if rotation > 90:
            rotation = rotation - 180
        if rotation < -90:
            rotation = 180 + rotation

        # There must be a more efficient way...
        lc = [tuple(l) for l in linecontour]
        dind = lc.index((x, y))
        #print 'dind', dind
        #dind = list(linecontour).index((x,y))

        return x, y, rotation, dind

    def labels(self, inline):
        levels = self.label_levels
        fslist = self.fslist
        trans = self.ax.transData
        colors = self.label_mappable.to_rgba(self.label_cvalues)
        fmt = self.fmt
        for icon, lev, color, cvalue, fsize in zip(self.label_indices,
                                                   self.label_levels, colors,
                                                   self.label_cvalues, fslist):
            con = self.collections[icon]
            lw = self.get_label_width(lev, fmt, fsize)
            additions = []
            for segNum, linecontour in enumerate(con._segments):
                # for closed contours add one more point to
                # avoid division by zero
                if all(linecontour[0] == linecontour[-1]):
                    linecontour = concatenate(
                        (linecontour, linecontour[1][newaxis, :]))
                    #linecontour.append(linecontour[1])
                # transfer all data points to screen coordinates
                slc = trans.seq_xy_tups(linecontour)
                if self.print_label(slc, lw):
                    x, y, rotation, ind = self.locate_label(slc, lw)
                    # transfer the location of the label back to
                    # data coordinates
                    dx, dy = trans.inverse_xy_tup((x, y))
                    t = Text(dx,
                             dy,
                             rotation=rotation,
                             horizontalalignment='center',
                             verticalalignment='center')
                    text = self.get_text(lev, fmt)
                    self.set_label_props(t, text, color)
                    self.cl.append(t)
                    self.cl_cvalues.append(cvalue)
                    if inline:
                        new = self.break_linecontour(linecontour, rotation, lw,
                                                     ind)
                        con._segments[segNum] = new[0]
                        additions.append(new[1])
            con._segments.extend(additions)
예제 #3
0
class ContourLabeler:
    '''Mixin to provide labelling capability to ContourSet'''

    def clabel(self, *args, **kwargs):
        """
        clabel(CS, **kwargs) - add labels to line contours in CS,
               where CS is a ContourSet object returned by contour.

        clabel(CS, V, **kwargs) - only label contours listed in V

        keyword arguments:

        * fontsize = None: as described in http://matplotlib.sf.net/fonts.html

        * colors = None:

           - a tuple of matplotlib color args (string, float, rgb, etc),
             different labels will be plotted in different colors in the order
             specified

           - one string color, e.g. colors = 'r' or colors = 'red', all labels
             will be plotted in this color

           - if colors == None, the color of each label matches the color
             of the corresponding contour

        * inline = True: controls whether the underlying contour is removed
                     (inline = True) or not (False)

        * fmt = '%1.3f': a format string for the label

        """
        fontsize = kwargs.get('fontsize', None)
        inline = kwargs.get('inline', 1)
        self.fmt = kwargs.get('fmt', '%1.3f')
        colors = kwargs.get('colors', None)



        if len(args) == 0:
            levels = self.levels
            indices = range(len(self.levels))
        elif len(args) == 1:
            levlabs = list(args[0])
            indices, levels = [], []
            for i, lev in enumerate(self.levels):
                if lev in levlabs:
                    indices.append(i)
                    levels.append(lev)
            if len(levels) < len(levlabs):
                msg = "Specified levels " + str(levlabs)
                msg += "\n don't match available levels "
                msg += str(self.levels)
                raise ValueError(msg)
        else:
            raise TypeError("Illegal arguments to clabel, see help(clabel)")
        self.label_levels = levels
        self.label_indices = indices

        self.fp = FontProperties()
        if fontsize == None:
            font_size = int(self.fp.get_size_in_points())
        else:
            if type(fontsize) not in [int, float, str]:
                raise TypeError("Font size must be an integer number.")
                # Can't it be floating point, as indicated in line above?
            else:
                if type(fontsize) == str:
                    font_size = int(self.fp.get_size_in_points())
                else:
                    self.fp.set_size(fontsize)
                    font_size = fontsize
        self.fslist = [font_size] * len(levels)

        if colors == None:
            self.label_mappable = self
            self.label_cvalues = take(self.cvalues, self.label_indices)
        else:
            cmap = ListedColormap(colors, N=len(self.label_levels))
            self.label_cvalues = range(len(self.label_levels))
            self.label_mappable = ScalarMappable(cmap = cmap,
                                                 norm = no_norm())

        #self.cl = []   # Initialized in ContourSet.__init__
        #self.cl_cvalues = [] # same
        self.cl_xy = []

        self.labels(inline)

        for label in self.cl:
            self.ax.add_artist(label)

        self.label_list =  silent_list('Text', self.cl)
        return self.label_list


    def print_label(self, linecontour,labelwidth):
        "if contours are too short, don't plot a label"
        lcsize = len(linecontour)
        if lcsize > 10 * labelwidth:
            return 1

        xmax = amax(array(linecontour)[:,0])
        xmin = amin(array(linecontour)[:,0])
        ymax = amax(array(linecontour)[:,1])
        ymin = amin(array(linecontour)[:,1])

        lw = labelwidth
        if (xmax - xmin) > 1.2* lw or (ymax - ymin) > 1.2 * lw:
            return 1
        else:
            return 0

    def too_close(self, x,y, lw):
        "if there's a label already nearby, find a better place"
        if self.cl_xy != []:
            dist = [sqrt((x-loc[0]) ** 2 + (y-loc[1]) ** 2)
                    for loc in self.cl_xy]
            for d in dist:
                if d < 1.2*lw:
                    return 1
                else: return 0
        else: return 0

    def get_label_coords(self, distances, XX, YY, ysize, lw):
        """ labels are ploted at a location with the smallest
        dispersion of the contour from a straight line
        unless there's another label nearby, in which case
        the second best place on the contour is picked up
        if there's no good place a label isplotted at the
        beginning of the contour
        """

        hysize = int(ysize/2)
        adist = argsort(distances)

        for ind in adist:
            x, y = XX[ind][hysize], YY[ind][hysize]
            if self.too_close(x,y, lw):
                continue
            else:
                self.cl_xy.append((x,y))
                return x,y, ind

        ind = adist[0]
        x, y = XX[ind][hysize], YY[ind][hysize]
        self.cl_xy.append((x,y))
        return x,y, ind

    def get_label_width(self, lev, fmt, fsize):
        "get the width of the label in points"
        if is_string_like(lev):
            lw = (len(lev)) * fsize
        else:
            lw = (len(fmt%lev)) * fsize

        return lw


    def set_label_props(self, label, text, color):
        "set the label properties - color, fontsize, text"
        label.set_text(text)
        label.set_color(color)
        label.set_fontproperties(self.fp)
        label.set_clip_box(self.ax.bbox)

    def get_text(self, lev, fmt):
        "get the text of the label"
        if is_string_like(lev):
            return lev
        else:
            return fmt%lev


    def break_linecontour(self, linecontour, rot, labelwidth, ind):
        "break a contour in two contours at the location of the label"
        lcsize = len(linecontour)
        hlw = int(labelwidth/2)

        #length of label in screen coords
        ylabel = abs(hlw * sin(rot*pi/180))
        xlabel = abs(hlw * cos(rot*pi/180))

        trans = self.ax.transData

        slc = trans.seq_xy_tups(linecontour)
        x,y = slc[ind]
        xx= array(slc)[:,0].copy()
        yy=array(slc)[:,1].copy()

        #indices which are under the label
        inds=nonzero(((xx < x+xlabel) & (xx > x-xlabel)) &
                     ((yy < y+ylabel) & (yy > y-ylabel)))

        if len(inds) >0:
            #if the label happens to be over the beginning of the
            #contour, the entire contour is removed, i.e.
            #indices to be removed are
            #inds= [0,1,2,3,305,306,307]
            #should rewrite this in a better way
            linds = nonzero(inds[1:]- inds[:-1] != 1)
            if inds[0] == 0 and len(linds) != 0:
                ii = inds[linds[0]]
                lc1 =linecontour[ii+1:inds[ii+1]]
                lc2 = []

            else:
                lc1=linecontour[:inds[0]]
                lc2= linecontour[inds[-1]+1:]

        else:
            lc1=linecontour[:ind]
            lc2 = linecontour[ind+1:]


        if rot <0:
            new_x1, new_y1 = x-xlabel, y+ylabel
            new_x2, new_y2 = x+xlabel, y-ylabel
        else:
            new_x1, new_y1 = x-xlabel, y-ylabel
            new_x2, new_y2 = x+xlabel, y+ylabel

        new_x1d, new_y1d = trans.inverse_xy_tup((new_x1, new_y1))
        new_x2d, new_y2d = trans.inverse_xy_tup((new_x2, new_y2))

        if rot > 0:
            if len(lc1) > 0 and (lc1[-1][0] <= new_x1d) and (lc1[-1][1] <= new_y1d):
                lc1.append((new_x1d, new_y1d))

            if len(lc2) > 0 and (lc2[0][0] >= new_x2d) and (lc2[0][1] >= new_y2d):
                lc2.insert(0, (new_x2d, new_y2d))
        else:
            if len(lc1) > 0 and ((lc1[-1][0] <= new_x1d) and (lc1[-1][1] >= new_y1d)):
                lc1.append((new_x1d, new_y1d))

            if len(lc2) > 0 and ((lc2[0][0] >= new_x2d) and (lc2[0][1] <= new_y2d)):
                lc2.insert(0, (new_x2d, new_y2d))

        return [lc1,lc2]


    def locate_label(self, linecontour, labelwidth):
        """find a good place to plot a label (relatively flat
        part of the contour) and the angle of rotation for the
        text object
        """

        nsize= len(linecontour)
        if labelwidth > 1:
            xsize = int(ceil(nsize/labelwidth))
        else:
            xsize = 1
        if xsize == 1:
            ysize = nsize
        else:
            ysize = labelwidth

        XX = resize(array(linecontour)[:,0],(xsize, ysize))
        YY = resize(array(linecontour)[:,1],(xsize,ysize))

        yfirst = YY[:,0]
        ylast = YY[:,-1]
        xfirst = XX[:,0]
        xlast = XX[:,-1]
        s = ( (reshape(yfirst, (xsize,1))-YY) *
              (reshape(xlast,(xsize,1)) - reshape(xfirst,(xsize,1)))
              - (reshape(xfirst,(xsize,1))-XX)
              * (reshape(ylast,(xsize,1)) - reshape(yfirst,(xsize,1))) )
        L=sqrt((xlast-xfirst)**2+(ylast-yfirst)**2)
        dist = add.reduce(([(abs(s)[i]/L[i]) for i in range(xsize)]),-1)
        x,y,ind = self.get_label_coords(dist, XX, YY, ysize, labelwidth)
        angle = arctan2(ylast - yfirst, xlast - xfirst)
        rotation = angle[ind]*180/pi
        if rotation > 90:
            rotation = rotation -180
        if rotation < -90:
            rotation = 180 + rotation

        dind = list(linecontour).index((x,y))

        return x,y, rotation, dind

    def labels(self, inline):
        levels = self.label_levels
        fslist = self.fslist
        trans = self.ax.transData
        colors = self.label_mappable.to_rgba(self.label_cvalues)
        fmt = self.fmt
        for icon, lev, color, cvalue, fsize in zip(self.label_indices,
                                          self.label_levels,
                                          colors,
                                          self.label_cvalues, fslist):
            con = self.collections[icon]
            toremove = []
            toadd = []
            lw = self.get_label_width(lev, fmt, fsize)
            for segNum, linecontour in enumerate(con._segments):
                # for closed contours add one more point to
                # avoid division by zero
                if linecontour[0] == linecontour[-1]:
                    linecontour.append(linecontour[1])
                # transfer all data points to screen coordinates
                slc = trans.seq_xy_tups(linecontour)
                if self.print_label(slc,lw):
                    x,y, rotation, ind  = self.locate_label(slc, lw)
                    # transfer the location of the label back to
                    # data coordinates
                    dx,dy = trans.inverse_xy_tup((x,y))
                    t = Text(dx, dy, rotation = rotation,
                             horizontalalignment='center',
                             verticalalignment='center')
                    text = self.get_text(lev,fmt)
                    self.set_label_props(t, text, color)
                    self.cl.append(t)
                    self.cl_cvalues.append(cvalue)
                    if inline:
                        new = self.break_linecontour(linecontour, rotation,
                                                       lw, ind)
                        toadd.extend(new)
                        #for c in new: toadd.append(c)
                        toremove.append(linecontour)
            for c in toremove:
                con._segments.remove(c)
            for c in toadd:
                con._segments.append(c)
예제 #4
0
class ContourLabeler:
    def __init__(self, ax):
        self.ax = ax

    def clabel(self, *args, **kwargs):
        """
        CLABEL(*args, **kwargs)

        Function signatures

        CLABEL(C) - plots contour labels,
                    C is the output of contour or a list of contours

        CLABEL(C,V) - creates labels only for those contours, given in
                      a list V

        CLABEL(C, **kwargs) - keyword args are explained below:



        * fontsize = None: as described in http://matplotlib.sf.net/fonts.html

        * colors = None:

           - a tuple of matplotlib color args (string, float, rgb, etc),
             different labels will be plotted in different colors in the order
             specified

           - one string color, e.g. colors = 'r' or colors = 'red', all labels
             will be plotted in this color

           - if colors == None, the color of each label matches the color
             of the corresponding contour

        * inline = 0: controls whether the underlying contour is removed
                     (inline = 1) or not

        * fmt = '%1.3f': a format string for the label

        """
        # todo, factor this out to a separate class and don't use hidden coll attrs

        if not self.ax.ishold(): self.ax.cla()

        fontsize = kwargs.get('fontsize', None)
        inline = kwargs.get('inline', 0)
        fmt = kwargs.get('fmt', '%1.3f')
        colors = kwargs.get('colors', None)



        if len(args) == 1:
            contours = args[0]
            levels = [con._label for con in contours]
        elif len(args) == 2:
            contours = args[0]
            levels = args[1]
        else:
            raise TypeError("Illegal arguments to clabel, see help(clabel)")



        self.fp = FontProperties()
        if fontsize == None:
            font_size = int(self.fp.get_size_in_points())
        else:
            if type(fontsize) not in [int, float, str]:
                raise TypeError("Font size must be an integer number.")
            else:
                if type(fontsize) == str:
                    font_size = int(self.fp.get_size_in_points())

                else:
                    self.fp.set_size(fontsize)
                    font_size = fontsize
        fslist = [font_size] * len(levels)

        if colors == None:
            colors = [c._colors[0] for c in contours]
        else:
            colors = colors * len(contours)

        if inline not in [0,1]:
            raise TypeError("inline must be 0 or 1")


        self.cl = []
        self.cl_xy = []

        # we have a list of contours and each contour has a list of
        # segments.  We want changes in the contour color to be
        # reflected in changes in the label color.  This is a good use
        # for traits observers, but in the interim, until traits are
        # utilized, we'll create a dict mapping i,j to text instances.
        # i is the contour level index, j is the sement index
        self.labeld = {}
        if inline == 1:
            self.inline_labels(levels, contours, colors, fslist, fmt)
        else:
            self.labels(levels, contours, colors, fslist, fmt)

        for label in self.cl:
            self.ax.add_artist(label)

        ret =  silent_list('Text', self.cl)
        ret.mappable = getattr(contours, 'mappable', None)
        # support colormapping for label
        if ret.mappable is not None:
            ret.mappable.labeld = self.labeld
        return ret



    def print_label(self, linecontour,labelwidth):
        "if contours are too short, don't plot a label"
        lcsize = len(linecontour)
        if lcsize > 10 * labelwidth:
            return 1

        xmax = amax(array(linecontour)[:,0])
        xmin = amin(array(linecontour)[:,0])
        ymax = amax(array(linecontour)[:,1])
        ymin = amin(array(linecontour)[:,1])

        lw = labelwidth
        if (xmax - xmin) > 1.2* lw or (ymax - ymin) > 1.2 * lw:
            return 1
        else:
            return 0

    def too_close(self, x,y, lw):
        "if there's a label already nearby, find a better place"
        if self.cl_xy != []:
            dist = [sqrt((x-loc[0]) ** 2 + (y-loc[1]) ** 2) for loc in self.cl_xy]
            for d in dist:
                if d < 1.2*lw:
                    return 1
                else: return 0
        else: return 0

    def get_label_coords(self, distances, XX, YY, ysize, lw):
        """ labels are ploted at a location with the smallest
        dispersion of the contour from a straight line
        unless there's another label nearby, in which case
        the second best place on the contour is picked up
        if there's no good place a label isplotted at the
        beginning of the contour
        """

        hysize = int(ysize/2)
        adist = argsort(distances)

        for ind in adist:
            x, y = XX[ind][hysize], YY[ind][hysize]
            if self.too_close(x,y, lw):
                continue
            else:
                self.cl_xy.append((x,y))
                return x,y, ind

        ind = adist[0]
        x, y = XX[ind][hysize], YY[ind][hysize]
        self.cl_xy.append((x,y))
        return x,y, ind

    def get_label_width(self, lev, fmt, fsize):
        "get the width of the label in points"
        if is_string_like(lev):
            lw = (len(lev)) * fsize
        else:
            lw = (len(fmt%lev)) * fsize

        return lw


    def set_label_props(self, label,text, color):
        "set the label properties - color, fontsize, text"
        label.set_text(text)
        label.set_color(color)
        label.set_fontproperties(self.fp)
        label.set_clip_box(self.ax.bbox)

    def get_text(self, lev, fmt):
        "get the text of the label"
        if is_string_like(lev):
            return lev
        else:
            return fmt%lev


    def break_linecontour(self, linecontour, rot, labelwidth, ind):
        "break a contour in two contours at the location of the label"
        lcsize = len(linecontour)
        hlw = int(labelwidth/2)

        #length of label in screen coords
        ylabel = abs(hlw * sin(rot*pi/180))
        xlabel = abs(hlw * cos(rot*pi/180))

        trans = self.ax.transData

        slc = trans.seq_xy_tups(linecontour)
        x,y = slc[ind]
        xx= array(slc)[:,0].copy()
        yy=array(slc)[:,1].copy()

        #indices which are under the label
        inds=nonzero(((xx < x+xlabel) & (xx > x-xlabel)) & ((yy < y+ylabel) & (yy > y-ylabel)))

        if len(inds) >0:
            #if the label happens to be over the beginning of the
            #contour, the entire contour is removed, i.e.
            #indices to be removed are
            #inds= [0,1,2,3,305,306,307]
            #should rewrite this in a better way
            linds = nonzero(inds[1:]- inds[:-1] != 1)
            if inds[0] == 0 and len(linds) != 0:
                ii = inds[linds[0]]
                lc1 =linecontour[ii+1:inds[ii+1]]
                lc2 = []

            else:
                lc1=linecontour[:inds[0]]
                lc2= linecontour[inds[-1]+1:]

        else:
            lc1=linecontour[:ind]
            lc2 = linecontour[ind+1:]


        if rot <0:
            new_x1, new_y1 = x-xlabel, y+ylabel
            new_x2, new_y2 = x+xlabel, y-ylabel
        else:
            new_x1, new_y1 = x-xlabel, y-ylabel
            new_x2, new_y2 = x+xlabel, y+ylabel

        new_x1d, new_y1d = trans.inverse_xy_tup((new_x1, new_y1))
        new_x2d, new_y2d = trans.inverse_xy_tup((new_x2, new_y2))

        if rot > 0:
            if len(lc1) > 0 and (lc1[-1][0] <= new_x1d) and (lc1[-1][1] <= new_y1d):
                lc1.append((new_x1d, new_y1d))

            if len(lc2) > 0 and (lc2[0][0] >= new_x2d) and (lc2[0][1] >= new_y2d):
                lc2.insert(0, (new_x2d, new_y2d))
        else:
            if len(lc1) > 0 and ((lc1[-1][0] <= new_x1d) and (lc1[-1][1] >= new_y1d)):
                lc1.append((new_x1d, new_y1d))

            if len(lc2) > 0 and ((lc2[0][0] >= new_x2d) and (lc2[0][1] <= new_y2d)):
                lc2.insert(0, (new_x2d, new_y2d))

        return [lc1,lc2]


    def locate_label(self, linecontour, labelwidth):
        """find a good place to plot a label (relatively flat
        part of the contour) and the angle of rotation for the
        text object
        """

        nsize= len(linecontour)
        if labelwidth > 1:
            xsize = int(ceil(nsize/labelwidth))
        else:
            xsize = 1
        if xsize == 1:
            ysize = nsize
        else:
            ysize = labelwidth

        XX = resize(array(linecontour)[:,0],(xsize, ysize))
        YY = resize(array(linecontour)[:,1],(xsize,ysize))

        yfirst = YY[:,0]
        ylast = YY[:,-1]
        xfirst = XX[:,0]
        xlast = XX[:,-1]
        s = (reshape(yfirst, (xsize,1))-YY)*(reshape(xlast,(xsize,1))-reshape(xfirst,(xsize,1)))-(reshape(xfirst,(xsize,1))-XX)*(reshape(ylast,(xsize,1))-reshape(yfirst,(xsize,1)))
        L=sqrt((xlast-xfirst)**2+(ylast-yfirst)**2)
        dist = add.reduce(([(abs(s)[i]/L[i]) for i in range(xsize)]),-1)
        x,y,ind = self.get_label_coords(dist, XX, YY, ysize, labelwidth)
        angle = arctan2(ylast - yfirst, xlast - xfirst)
        rotation = angle[ind]*180/pi
        if rotation > 90:
            rotation = rotation -180
        if rotation < -90:
            rotation = 180 + rotation

        dind = list(linecontour).index((x,y))

        return x,y, rotation, dind

    def inline_labels(self, levels, contours, colors, fslist, fmt):
        trans = self.ax.transData
        contourNum = 0
        for lev, con, color, fsize in zip(levels, contours, colors, fslist):
            toremove = []
            toadd = []
            lw = self.get_label_width(lev, fmt, fsize)
            for segNum, linecontour in enumerate(con._segments):
                key = contourNum, segNum
                # for closed contours add one more point to
                # avoid division by zero
                if linecontour[0] == linecontour[-1]:
                    linecontour.append(linecontour[1])
                # transfer all data points to screen coordinates
                slc = trans.seq_xy_tups(linecontour)
                if self.print_label(slc,lw):
                    x,y, rotation, ind  = self.locate_label(slc, lw)
                    # transfer the location of the label back to
                    # data coordinates
                    dx,dy = trans.inverse_xy_tup((x,y))
                    t = Text(dx, dy, rotation = rotation, horizontalalignment='center', verticalalignment='center')
                    self.labeld[key] = t
                    text = self.get_text(lev,fmt)
                    self.set_label_props(t, text, color)
                    self.cl.append(t)
                    new  =  self.break_linecontour(linecontour, rotation, lw, ind)

                    for c in new: toadd.append(c)
                    toremove.append(linecontour)
            for c in toremove:
                con._segments.remove(c)
            for c in toadd: con._segments.append(c)

            contourNum += 1


    def labels(self, levels, contours, colors, fslist, fmt):
        trans = self.ax.transData
        for lev, con, color, fsize in zip(levels, contours, colors, fslist):
            lw = self.get_label_width(lev, fmt, fsize)
            for linecontour in con._segments:
                # for closed contours add one more point
                if linecontour[0] == linecontour[-1]:
                    linecontour.append(linecontour[1])
                # transfer all data points to screen coordinates
                slc = trans.seq_xy_tups(linecontour)
                if self.print_label(slc,lw):
                    x,y, rotation, ind  = self.locate_label(slc, lw)
                    # transfer the location of the label back into
                    # data coordinates
                    dx,dy = trans.inverse_xy_tup((x,y))
                    t = Text(dx, dy, rotation = rotation, horizontalalignment='center', verticalalignment='center')
                    text = self.get_text(lev, fmt)
                    self.set_label_props(t, text, color)
                    self.cl.append(t)
                else:
                    pass