''' Created on 20. jan. 2011 @author: pab license BSD ''' from __future__ import division import warnings import numpy as np from wafo.plotbackend import plotbackend from matplotlib import mlab __all__ = ['cltext', 'tallibing', 'test_docstrings'] _TALLIBING_GID = 'TALLIBING' _CLTEXT_GID = 'CLTEXT' def _matchfun(x, gidtxt): if hasattr(x, 'get_gid'): return x.get_gid() == gidtxt return False def delete_text_object(gidtxt, figure=None, axis=None, verbose=False): ''' Delete all text objects matching the gidtxt if it exists Parameters ---------- gidtxt : string figure, axis : objects current figure and current axis, respectively. verbose : bool If true print warnings when trying to delete non-existent objects ''' if figure is None: figure = plotbackend.gcf() if axis is None: axis = figure.gca() def lmatchfun(x): return _matchfun(x, gidtxt) objs = axis.findobj(lmatchfun) for obj in objs: try: axis.texts.remove(obj) except: if verbose: warnings.warn( 'Tried to delete a non-existing %s from axis' % gidtxt) objs = figure.findobj(lmatchfun) for obj in objs: try: figure.texts.remove(obj) except: if verbose: warnings.warn( 'Tried to delete a non-existing %s from figure' % gidtxt) def cltext(levels, percent=False, n=4, xs=0.036, ys=0.94, zs=0, figure=None, axis=None): ''' Places contour level text in the current window Parameters ---------- levels : vector contour levels or the corresponding percent which the contour line encloses percent : bool False if levels are the actual contour levels (default) True if levels are the corresponding percent which the contour line encloses n : integer maximum N digits of precision (default 4) figure, axis : objects current figure and current axis, respectively. default figure = plotbackend.gcf(), axis = plotbackend.gca() Returns ------- h = handles to the text objects. Notes ----- CLTEXT creates text objects in the current figure and prints "Level curves at:" if percent is False and "Level curves enclosing:" otherwise and the contour levels or percent. The handles to the lines of text may also be found by h = findobj(gcf,'gid','CLTEXT','type','text'); h = findobj(gca,'gid','CLTEXT','type','text'); To make the text objects follow the data in the axes set the units for the text objects 'data' by set(h,'unit','data') Examples -------- >>> import wafo.graphutil as wg >>> import wafo.demos as wd >>> import pylab as plt >>> x,y,z = wd.peaks(); >>> h = plt.contour(x,y,z) >>> h = wg.cltext(h.levels) >>> plt.show() ''' # TODO : Make it work like legend does (but without the box): include # position options etc... if figure is None: figure = plotbackend.gcf() if axis is None: axis = figure.gca() clevels = np.atleast_1d(levels) axpos = axis.get_position() xint = axpos.intervalx yint = axpos.intervaly xss = xint[0] + xs * (xint[1] - xint[0]) yss = yint[0] + ys * (yint[1] - yint[0]) # delete cltext object if it exists delete_text_object(_CLTEXT_GID, axis=axis) charHeight = 1.0 / 33.0 delta_y = charHeight if percent: titletxt = 'Level curves enclosing:' else: titletxt = 'Level curves at:' format_ = '%0.' + ('%d' % n) + 'g\n' cltxt = ''.join([format_ % level for level in clevels.tolist()]) titleProp = dict(gid=_CLTEXT_GID, horizontalalignment='left', verticalalignment='center', fontweight='bold', axes=axis) ha1 = figure.text(xss, yss, titletxt, **titleProp) yss -= delta_y txtProp = dict(gid=_CLTEXT_GID, horizontalalignment='left', verticalalignment='top', axes=axis) ha2 = figure.text(xss, yss, cltxt, **txtProp) plotbackend.draw_if_interactive() return ha1, ha2 def tallibing(*args, **kwds): ''' TALLIBING Display numbers on field-plot CALL h=tallibing(x,y,n,size,color) Parameters ---------- x, y : array position matrices n : array corresponding matrix of the values to be written (non-integers are rounded) mid_points : bool (default True) data-point-positions are in the middle of bins instead of the corners size : int, (default=8) font size (optional) color : str, (default='white') color of text (optional) Returns ------- h : list handles to TEXT objects TALLIBING writes the numbers in a 2D array as text at the positions given by the x and y coordinate matrices. When plotting binned results, the number of datapoints in each bin can be written on the bins in the plot. Example ------- >>> import wafo.graphutil as wg >>> import wafo.demos as wd >>> [x,y,z] = wd.peaks(n=20) >>> h0 = wg.pcolor(x,y,z) >>> h1 = wg.tallibing(x,y,z) See also -------- text ''' axis = kwds.pop('axis', None) if axis is None: axis = plotbackend.gca() x, y, n = _parse_data(*args, **kwds) if mlab.isvector(x) or mlab.isvector(y): x, y = np.meshgrid(x, y) n = np.round(n) # delete tallibing object if it exists delete_text_object(_TALLIBING_GID, axis=axis) txtProp = dict(gid=_TALLIBING_GID, size=8, color='w', horizontalalignment='center', verticalalignment='center', fontweight='demi', axes=axis) txtProp.update(**kwds) h = [] for xi, yi, ni in zip(x.ravel(), y.ravel(), n.ravel()): if ni: h.append(axis.text(xi, yi, str(ni), **txtProp)) plotbackend.draw_if_interactive() return h def _parse_data(*args, **kwds): nargin = len(args) data = np.atleast_2d(args[-1]).copy() M, N = data.shape if nargin == 1: x = np.arange(N) y = np.arange(M) elif nargin == 3: x, y = np.atleast_1d(*args[:-1]) if min(x.shape) != 1: x = x[0] if min(y.shape) != 1: y = y[:, 0] else: raise ValueError( 'Requires 3 or 1 in arguments! (x,y,data) or (data)') if kwds.pop('mid_point', True): xx = _find_mid_points(x) yy = _find_mid_points(y) return xx, yy, data return x, y, data pcolor = getattr(plotbackend, 'pcolor', None) pcolormesh = getattr(plotbackend, 'pcolormesh', None) def _find_mid_points(x): ''' Return points half way between all values of X and outside the endpoints. The outer limits have same distance from X's endpoints as the limits just inside. ''' dx = np.diff(x) * 0.5 dx = np.hstack((dx, dx[-1])) return x + dx def test_docstrings(): import doctest print('Testing docstrings in %s' % __file__) doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE) if __name__ == '__main__': test_docstrings()