import numpy as np from .Qt import QtGui, QtCore class ColorMap(object): """ A ColorMap defines a relationship between a scalar value and a range of colors. ColorMaps are commonly used for false-coloring monochromatic images, coloring scatter-plot points, and coloring surface plots by height. Each color map is defined by a set of colors, each corresponding to a particular scalar value. For example: | 0.0 -> black | 0.2 -> red | 0.6 -> yellow | 1.0 -> white The colors for intermediate values are determined by interpolating between the two nearest colors in either RGB or HSV color space. To provide user-defined color mappings, see :class:`GradientWidget `. """ ## color interpolation modes RGB = 1 HSV_POS = 2 HSV_NEG = 3 ## boundary modes CLIP = 1 REPEAT = 2 MIRROR = 3 ## return types BYTE = 1 FLOAT = 2 QCOLOR = 3 enumMap = { 'rgb': RGB, 'hsv+': HSV_POS, 'hsv-': HSV_NEG, 'clip': CLIP, 'repeat': REPEAT, 'mirror': MIRROR, 'byte': BYTE, 'float': FLOAT, 'qcolor': QCOLOR, } def __init__(self, pos, color, mode=None): """ =============== ============================================================== **Arguments:** pos Array of positions where each color is defined color Array of RGBA colors. Integer data types are interpreted as 0-255; float data types are interpreted as 0.0-1.0 mode Array of color modes (ColorMap.RGB, HSV_POS, or HSV_NEG) indicating the color space that should be used when interpolating between stops. Note that the last mode value is ignored. By default, the mode is entirely RGB. =============== ============================================================== """ self.pos = np.array(pos) self.color = np.array(color) if mode is None: mode = np.ones(len(pos)) self.mode = mode self.stopsCache = {} def map(self, data, mode='byte'): """ Return an array of colors corresponding to the values in *data*. Data must be either a scalar position or an array (any shape) of positions. The *mode* argument determines the type of data returned: =========== =============================================================== byte (default) Values are returned as 0-255 unsigned bytes. float Values are returned as 0.0-1.0 floats. qcolor Values are returned as an array of QColor objects. =========== =============================================================== """ if isinstance(mode, basestring): mode = self.enumMap[mode.lower()] if mode == self.QCOLOR: pos, color = self.getStops(self.BYTE) else: pos, color = self.getStops(mode) # don't need this--np.interp takes care of it. #data = np.clip(data, pos.min(), pos.max()) # Interpolate # TODO: is griddata faster? # interp = scipy.interpolate.griddata(pos, color, data) if np.isscalar(data): interp = np.empty((color.shape[1],), dtype=color.dtype) else: if not isinstance(data, np.ndarray): data = np.array(data) interp = np.empty(data.shape + (color.shape[1],), dtype=color.dtype) for i in range(color.shape[1]): interp[...,i] = np.interp(data, pos, color[:,i]) # Convert to QColor if requested if mode == self.QCOLOR: if np.isscalar(data): return QtGui.QColor(*interp) else: return [QtGui.QColor(*x) for x in interp] else: return interp def mapToQColor(self, data): """Convenience function; see :func:`map() `.""" return self.map(data, mode=self.QCOLOR) def mapToByte(self, data): """Convenience function; see :func:`map() `.""" return self.map(data, mode=self.BYTE) def mapToFloat(self, data): """Convenience function; see :func:`map() `.""" return self.map(data, mode=self.FLOAT) def getGradient(self, p1=None, p2=None): """Return a QLinearGradient object spanning from QPoints p1 to p2.""" if p1 == None: p1 = QtCore.QPointF(0,0) if p2 == None: p2 = QtCore.QPointF(self.pos.max()-self.pos.min(),0) g = QtGui.QLinearGradient(p1, p2) pos, color = self.getStops(mode=self.BYTE) color = [QtGui.QColor(*x) for x in color] g.setStops(zip(pos, color)) #if self.colorMode == 'rgb': #ticks = self.listTicks() #g.setStops([(x, QtGui.QColor(t.color)) for t,x in ticks]) #elif self.colorMode == 'hsv': ## HSV mode is approximated for display by interpolating 10 points between each stop #ticks = self.listTicks() #stops = [] #stops.append((ticks[0][1], ticks[0][0].color)) #for i in range(1,len(ticks)): #x1 = ticks[i-1][1] #x2 = ticks[i][1] #dx = (x2-x1) / 10. #for j in range(1,10): #x = x1 + dx*j #stops.append((x, self.getColor(x))) #stops.append((x2, self.getColor(x2))) #g.setStops(stops) return g def getColors(self, mode=None): """Return list of all color stops converted to the specified mode. If mode is None, then no conversion is done.""" if isinstance(mode, basestring): mode = self.enumMap[mode.lower()] color = self.color if mode in [self.BYTE, self.QCOLOR] and color.dtype.kind == 'f': color = (color * 255).astype(np.ubyte) elif mode == self.FLOAT and color.dtype.kind != 'f': color = color.astype(float) / 255. if mode == self.QCOLOR: color = [QtGui.QColor(*x) for x in color] return color def getStops(self, mode): ## Get fully-expanded set of RGBA stops in either float or byte mode. if mode not in self.stopsCache: color = self.color if mode == self.BYTE and color.dtype.kind == 'f': color = (color * 255).astype(np.ubyte) elif mode == self.FLOAT and color.dtype.kind != 'f': color = color.astype(float) / 255. ## to support HSV mode, we need to do a little more work.. #stops = [] #for i in range(len(self.pos)): #pos = self.pos[i] #color = color[i] #imode = self.mode[i] #if imode == self.RGB: #stops.append((x,color)) #else: #ns = self.stopsCache[mode] = (self.pos, color) return self.stopsCache[mode] def getLookupTable(self, start=0.0, stop=1.0, nPts=512, alpha=None, mode='byte'): """ Return an RGB(A) lookup table (ndarray). =============== ============================================================================= **Arguments:** start The starting value in the lookup table (default=0.0) stop The final value in the lookup table (default=1.0) nPts The number of points in the returned lookup table. alpha True, False, or None - Specifies whether or not alpha values are included in the table. If alpha is None, it will be automatically determined. mode Determines return type: 'byte' (0-255), 'float' (0.0-1.0), or 'qcolor'. See :func:`map() `. =============== ============================================================================= """ if isinstance(mode, basestring): mode = self.enumMap[mode.lower()] if alpha is None: alpha = self.usesAlpha() x = np.linspace(start, stop, nPts) table = self.map(x, mode) if not alpha: return table[:,:3] else: return table def usesAlpha(self): """Return True if any stops have an alpha < 255""" max = 1.0 if self.color.dtype.kind == 'f' else 255 return np.any(self.color[:,3] != max) def isMapTrivial(self): """ Return True if the gradient has exactly two stops in it: black at 0.0 and white at 1.0. """ if len(self.pos) != 2: return False if self.pos[0] != 0.0 or self.pos[1] != 1.0: return False if self.color.dtype.kind == 'f': return np.all(self.color == np.array([[0.,0.,0.,1.], [1.,1.,1.,1.]])) else: return np.all(self.color == np.array([[0,0,0,255], [255,255,255,255]]))