ScatterPlotItem.py 35.4 KB
Newer Older
1
from pyqtgraph.Qt import QtGui, QtCore, USE_PYSIDE
2
3
from pyqtgraph.Point import Point
import pyqtgraph.functions as fn
4
5
from .GraphicsItem import GraphicsItem
from .GraphicsObject import GraphicsObject
6
7
import numpy as np
import scipy.stats
Luke Campagnola's avatar
Luke Campagnola committed
8
9
import weakref
import pyqtgraph.debug as debug
10
from pyqtgraph.pgcollections import OrderedDict
11
import pyqtgraph as pg
Luke Campagnola's avatar
Luke Campagnola committed
12
#import pyqtgraph as pg 
13
14

__all__ = ['ScatterPlotItem', 'SpotItem']
15
16
17


## Build all symbol paths
Luke Campagnola's avatar
Luke Campagnola committed
18
Symbols = OrderedDict([(name, QtGui.QPainterPath()) for name in ['o', 's', 't', 'd', '+']])
19
20
21
22
23
24
25
26
27
28
29
Symbols['o'].addEllipse(QtCore.QRectF(-0.5, -0.5, 1, 1))
Symbols['s'].addRect(QtCore.QRectF(-0.5, -0.5, 1, 1))
coords = {
    't': [(-0.5, -0.5), (0, 0.5), (0.5, -0.5)],
    'd': [(0., -0.5), (-0.4, 0.), (0, 0.5), (0.4, 0)],
    '+': [
        (-0.5, -0.05), (-0.5, 0.05), (-0.05, 0.05), (-0.05, 0.5),
        (0.05, 0.5), (0.05, 0.05), (0.5, 0.05), (0.5, -0.05), 
        (0.05, -0.05), (0.05, -0.5), (-0.05, -0.5), (-0.05, -0.05)
    ],
}
30
for k, c in coords.items():
31
32
33
34
35
    Symbols[k].moveTo(*c[0])
    for x,y in c[1:]:
        Symbols[k].lineTo(x, y)
    Symbols[k].closeSubpath()

36
37
38
39
40
41
42
43
44
45
    
def drawSymbol(painter, symbol, size, pen, brush):
    painter.scale(size, size)
    painter.setPen(pen)
    painter.setBrush(brush)
    if isinstance(symbol, basestring):
        symbol = Symbols[symbol]
    if np.isscalar(symbol):
        symbol = Symbols.values()[symbol % len(Symbols)]
    painter.drawPath(symbol)
46

47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
    
def renderSymbol(symbol, size, pen, brush, device=None):
    """
    Render a symbol specification to QImage.
    Symbol may be either a QPainterPath or one of the keys in the Symbols dict.
    If *device* is None, a new QPixmap will be returned. Otherwise,
    the symbol will be rendered into the device specified (See QPainter documentation 
    for more information).
    """
    ## see if this pixmap is already cached
    #global SymbolPixmapCache
    #key = (symbol, size, fn.colorTuple(pen.color()), pen.width(), pen.style(), fn.colorTuple(brush.color()))
    #if key in SymbolPixmapCache:
        #return SymbolPixmapCache[key]
        
Luke Campagnola's avatar
Luke Campagnola committed
62
    ## Render a spot with the given parameters to a pixmap
63
    penPxWidth = max(np.ceil(pen.width()), 1)
64
    image = QtGui.QImage(int(size+penPxWidth), int(size+penPxWidth), QtGui.QImage.Format_ARGB32)
Luke Campagnola's avatar
Luke Campagnola committed
65
66
67
    image.fill(0)
    p = QtGui.QPainter(image)
    p.setRenderHint(p.Antialiasing)
68
    p.translate(image.width()*0.5, image.height()*0.5)
69
    drawSymbol(p, symbol, size, pen, brush)
Luke Campagnola's avatar
Luke Campagnola committed
70
    p.end()
71
72
73
74
    return image
    #pixmap = QtGui.QPixmap(image)
    #SymbolPixmapCache[key] = pixmap
    #return pixmap
Luke Campagnola's avatar
Luke Campagnola committed
75

76
77
78
79
80
def makeSymbolPixmap(size, pen, brush, symbol):
    ## deprecated
    img = renderSymbol(symbol, size, pen, brush)
    return QtGui.QPixmap(img)
    
81
class SymbolAtlas(object):
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
    """
    Used to efficiently construct a single QPixmap containing all rendered symbols
    for a ScatterPlotItem. This is required for fragment rendering.
    
    Use example:
        atlas = SymbolAtlas()
        sc1 = atlas.getSymbolCoords('o', 5, QPen(..), QBrush(..))
        sc2 = atlas.getSymbolCoords('t', 10, QPen(..), QBrush(..))
        pm = atlas.getAtlas()
        
    """
    class SymbolCoords(list):  ## needed because lists are not allowed in weak references.
        pass
    
    def __init__(self):
        # symbol key : [x, y, w, h] atlas coordinates
        # note that the coordinate list will always be the same list object as 
        # long as the symbol is in the atlas, but the coordinates may
        # change if the atlas is rebuilt.
        # weak value; if all external refs to this list disappear, 
        # the symbol will be forgotten.
        self.symbolMap = weakref.WeakValueDictionary()
        
        self.atlasData = None # numpy array of atlas image
        self.atlas = None     # atlas as QPixmap
        self.atlasValid = False
        
    def getSymbolCoords(self, opts):
        """
        Given a list of spot records, return an object representing the coordinates of that symbol within the atlas
        """
        coords = np.empty(len(opts), dtype=object)
        for i, rec in enumerate(opts):
            symbol, size, pen, brush = rec['symbol'], rec['size'], rec['pen'], rec['brush']
            pen = fn.mkPen(pen) if not isinstance(pen, QtGui.QPen) else pen
            brush = fn.mkBrush(brush) if not isinstance(pen, QtGui.QBrush) else brush
            key = (symbol, size, fn.colorTuple(pen.color()), pen.width(), pen.style(), fn.colorTuple(brush.color()))
            if key not in self.symbolMap:
                newCoords = SymbolAtlas.SymbolCoords()
                self.symbolMap[key] = newCoords
                self.atlasValid = False
                #try:
                    #self.addToAtlas(key)  ## squeeze this into the atlas if there is room
                #except:
                    #self.buildAtlas()  ## otherwise, we need to rebuild
            
            coords[i] = self.symbolMap[key]
        return coords
        
    def buildAtlas(self):
        # get rendered array for all symbols, keep track of avg/max width
        rendered = {}
        avgWidth = 0.0
        maxWidth = 0
        images = []
        for key, coords in self.symbolMap.items():
            if len(coords) == 0:
                pen = fn.mkPen(color=key[2], width=key[3], style=key[4])
                brush = fn.mkBrush(color=key[5])
                img = renderSymbol(key[0], key[1], pen, brush)
                images.append(img)  ## we only need this to prevent the images being garbage collected immediately
                arr = fn.imageToArray(img, copy=False, transpose=False)
            else:
                (x,y,w,h) = self.symbolMap[key]
                arr = self.atlasData[x:x+w, y:y+w]
            rendered[key] = arr
            w = arr.shape[0]
            avgWidth += w
            maxWidth = max(maxWidth, w)
            
        nSymbols = len(rendered)
        if nSymbols > 0:
            avgWidth /= nSymbols
            width = max(maxWidth, avgWidth * (nSymbols**0.5))
        else:
            avgWidth = 0
            width = 0
        
        # sort symbols by height
        symbols = sorted(rendered.keys(), key=lambda x: rendered[x].shape[1], reverse=True)
        
        self.atlasRows = []
Luke Campagnola's avatar
Luke Campagnola committed
164

165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
        x = width
        y = 0
        rowheight = 0
        for key in symbols:
            arr = rendered[key]
            w,h = arr.shape[:2]
            if x+w > width:
                y += rowheight
                x = 0
                rowheight = h
                self.atlasRows.append([y, rowheight, 0])
            self.symbolMap[key][:] = x, y, w, h
            x += w
            self.atlasRows[-1][2] = x
        height = y + rowheight
180

181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
        self.atlasData = np.zeros((width, height, 4), dtype=np.ubyte)
        for key in symbols:
            x, y, w, h = self.symbolMap[key]
            self.atlasData[x:x+w, y:y+h] = rendered[key]
        self.atlas = None
        self.atlasValid = True
    
    def getAtlas(self):
        if not self.atlasValid:
            self.buildAtlas()
        if self.atlas is None:
            if len(self.atlasData) == 0:
                return QtGui.QPixmap(0,0)
            img = fn.makeQImage(self.atlasData, copy=False, transpose=False)
            self.atlas = QtGui.QPixmap(img)
        return self.atlas
        
    
    
    
201
class ScatterPlotItem(GraphicsObject):
202
203
204
205
    """
    Displays a set of x/y points. Instances of this class are created
    automatically as part of PlotDataItem; these rarely need to be instantiated
    directly.
206
    
207
208
209
210
211
212
213
214
215
216
217
218
    The size, shape, pen, and fill brush may be set for each point individually 
    or for all points. 
    
    
    ========================  ===============================================
    **Signals:**
    sigPlotChanged(self)      Emitted when the data being plotted has changed
    sigClicked(self, points)  Emitted when the curve is clicked. Sends a list
                              of all the points under the mouse pointer.
    ========================  ===============================================
    
    """
219
220
221
    #sigPointClicked = QtCore.Signal(object, object)
    sigClicked = QtCore.Signal(object, object)  ## self, points
    sigPlotChanged = QtCore.Signal(object)
222
223
224
225
    def __init__(self, *args, **kargs):
        """
        Accepts the same arguments as setData()
        """
Luke Campagnola's avatar
Luke Campagnola committed
226
        prof = debug.Profiler('ScatterPlotItem.__init__', disabled=True)
227
        GraphicsObject.__init__(self)
228
229
230
231
232
233
        
        self.picture = None   # QPicture used for rendering when pxmode==False
        self.fragments = None # fragment specification for pxmode; updated every time the view changes.
        self.fragmentAtlas = SymbolAtlas()
        
        self.data = np.empty(0, dtype=[('x', float), ('y', float), ('size', float), ('symbol', object), ('pen', object), ('brush', object), ('data', object), ('fragCoords', object), ('item', object)])
234
235
236
        self.bounds = [None, None]  ## caches data bounds
        self._maxSpotWidth = 0      ## maximum size of the scale-variant portion of all spots
        self._maxSpotPxWidth = 0    ## maximum size of the scale-invariant portion of all spots
237
238
239
240
241
242
        self.opts = {
            'pxMode': True, 
            'useCache': True,  ## If useCache is False, symbols are re-drawn on every paint. 
            'antialias': pg.getConfigOption('antialias'),
        }   
        self.exportOpts = False
243
        
Luke Campagnola's avatar
Luke Campagnola committed
244
245
246
247
248
        self.setPen(200,200,200, update=False)
        self.setBrush(100,100,150, update=False)
        self.setSymbol('o', update=False)
        self.setSize(7, update=False)
        prof.mark('1')
249
        self.setData(*args, **kargs)
Luke Campagnola's avatar
Luke Campagnola committed
250
251
        prof.mark('setData')
        prof.finish()
252
        
253
254
        #self.setCacheMode(self.DeviceCoordinateCache)
        
255
    def setData(self, *args, **kargs):
256
        """
257
258
259
260
261
        **Ordered Arguments:**
        
        * If there is only one unnamed argument, it will be interpreted like the 'spots' argument.
        * If there are two unnamed arguments, they will be interpreted as sequences of x and y values.
        
262
        ====================== ===============================================================================================
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
        **Keyword Arguments:**
        *spots*                Optional list of dicts. Each dict specifies parameters for a single spot:
                               {'pos': (x,y), 'size', 'pen', 'brush', 'symbol'}. This is just an alternate method
                               of passing in data for the corresponding arguments.
        *x*,*y*                1D arrays of x,y values.
        *pos*                  2D structure of x,y pairs (such as Nx2 array or list of tuples)
        *pxMode*               If True, spots are always the same size regardless of scaling, and size is given in px.
                               Otherwise, size is in scene coordinates and the spots scale with the view.
                               Default is True
        *symbol*               can be one (or a list) of:
                               * 'o'  circle (default)
                               * 's'  square
                               * 't'  triangle
                               * 'd'  diamond
                               * '+'  plus
Luke Campagnola's avatar
Luke Campagnola committed
278
279
280
                               * any QPainterPath to specify custom symbol shapes. To properly obey the position and size,
                               custom symbols should be centered at (0,0) and width and height of 1.0. Note that it is also
                               possible to 'install' custom shapes by setting ScatterPlotItem.Symbols[key] = shape.
281
282
283
284
285
        *pen*                  The pen (or list of pens) to use for drawing spot outlines.
        *brush*                The brush (or list of brushes) to use for filling spots.
        *size*                 The size (or list of sizes) of spots. If *pxMode* is True, this value is in pixels. Otherwise,
                               it is in the item's local coordinate system.
        *data*                 a list of python objects used to uniquely identify each spot.
Luke Campagnola's avatar
Luke Campagnola committed
286
        *identical*            *Deprecated*. This functionality is handled automatically now.
287
288
289
        *antialias*            Whether to draw symbols with antialiasing. Note that if pxMode is True, symbols are 
                               always rendered with antialiasing (since the rendered symbols can be cached, this 
                               incurs very little performance cost)
290
        ====================== ===============================================================================================
291
        """
292
        oldData = self.data  ## this causes cached pixmaps to be preserved while new data is registered.
Luke Campagnola's avatar
Luke Campagnola committed
293
294
        self.clear()  ## clear out all old data
        self.addPoints(*args, **kargs)
295

Luke Campagnola's avatar
Luke Campagnola committed
296
297
298
299
300
    def addPoints(self, *args, **kargs):
        """
        Add new points to the scatter plot. 
        Arguments are the same as setData()
        """
301
        
302
303
304
305
306
307
308
309
        ## deal with non-keyword arguments
        if len(args) == 1:
            kargs['spots'] = args[0]
        elif len(args) == 2:
            kargs['x'] = args[0]
            kargs['y'] = args[1]
        elif len(args) > 2:
            raise Exception('Only accepts up to two non-keyword arguments.')
310
        
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
        ## convert 'pos' argument to 'x' and 'y'
        if 'pos' in kargs:
            pos = kargs['pos']
            if isinstance(pos, np.ndarray):
                kargs['x'] = pos[:,0]
                kargs['y'] = pos[:,1]
            else:
                x = []
                y = []
                for p in pos:
                    if isinstance(p, QtCore.QPointF):
                        x.append(p.x())
                        y.append(p.y())
                    else:
                        x.append(p[0])
                        y.append(p[1])
                kargs['x'] = x
                kargs['y'] = y
329
        
330
331
332
333
334
        ## determine how many spots we have
        if 'spots' in kargs:
            numPts = len(kargs['spots'])
        elif 'y' in kargs and kargs['y'] is not None:
            numPts = len(kargs['y'])
335
        else:
336
337
338
            kargs['x'] = []
            kargs['y'] = []
            numPts = 0
339
        
Luke Campagnola's avatar
Luke Campagnola committed
340
341
342
343
344
345
        ## Extend record array
        oldData = self.data
        self.data = np.empty(len(oldData)+numPts, dtype=self.data.dtype)
        ## note that np.empty initializes object fields to None and string fields to ''
        
        self.data[:len(oldData)] = oldData
346
347
        #for i in range(len(oldData)):
            #oldData[i]['item']._data = self.data[i]  ## Make sure items have proper reference to new array
Luke Campagnola's avatar
Luke Campagnola committed
348
349
350
            
        newData = self.data[len(oldData):]
        newData['size'] = -1  ## indicates to use default size
351
        
352
353
        if 'spots' in kargs:
            spots = kargs['spots']
354
            for i in range(len(spots)):
355
356
                spot = spots[i]
                for k in spot:
Luke Campagnola's avatar
Luke Campagnola committed
357
358
359
360
361
                    #if k == 'pen':
                        #newData[k] = fn.mkPen(spot[k])
                    #elif k == 'brush':
                        #newData[k] = fn.mkBrush(spot[k])
                    if k == 'pos':
362
363
364
365
366
                        pos = spot[k]
                        if isinstance(pos, QtCore.QPointF):
                            x,y = pos.x(), pos.y()
                        else:
                            x,y = pos[0], pos[1]
Luke Campagnola's avatar
Luke Campagnola committed
367
368
369
370
371
372
                        newData[i]['x'] = x
                        newData[i]['y'] = y
                    elif k in ['x', 'y', 'size', 'symbol', 'pen', 'brush', 'data']:
                        newData[i][k] = spot[k]
                    #elif k == 'data':
                        #self.pointData[i] = spot[k]
373
374
375
                    else:
                        raise Exception("Unknown spot parameter: %s" % k)
        elif 'y' in kargs:
Luke Campagnola's avatar
Luke Campagnola committed
376
377
            newData['x'] = kargs['x']
            newData['y'] = kargs['y']
378
        
Luke Campagnola's avatar
Luke Campagnola committed
379
        if 'pxMode' in kargs:
380
            self.setPxMode(kargs['pxMode'])
381
382
        if 'antialias' in kargs:
            self.opts['antialias'] = kargs['antialias']
Luke Campagnola's avatar
Luke Campagnola committed
383
            
384
        ## Set any extra parameters provided in keyword arguments
Luke Campagnola's avatar
Luke Campagnola committed
385
        for k in ['pen', 'brush', 'symbol', 'size']:
386
387
            if k in kargs:
                setMethod = getattr(self, 'set' + k[0].upper() + k[1:])
Luke Campagnola's avatar
Luke Campagnola committed
388
                setMethod(kargs[k], update=False, dataSet=newData)
389
        
Luke Campagnola's avatar
Luke Campagnola committed
390
391
        if 'data' in kargs:
            self.setPointData(kargs['data'], dataSet=newData)
392
            
Luke Campagnola's avatar
Luke Campagnola committed
393
        self.prepareGeometryChange()
394
        self.bounds = [None, None]
395
396
        self.invalidate()
        self.updateSpots(newData)
Luke Campagnola's avatar
Luke Campagnola committed
397
        self.sigPlotChanged.emit(self)
398
        
399
400
401
402
403
404
    def invalidate(self):
        ## clear any cached drawing state
        self.picture = None
        self.fragments = None
        self.update()
        
Luke Campagnola's avatar
Luke Campagnola committed
405
406
407
    def getData(self):
        return self.data['x'], self.data['y']
    
408
409
        
    def setPoints(self, *args, **kargs):
410
        ##Deprecated; use setData
411
412
        return self.setData(*args, **kargs)
        
413
414
415
416
417
418
    def implements(self, interface=None):
        ints = ['plotData']
        if interface is None:
            return ints
        return interface in ints
    
419
    def setPen(self, *args, **kargs):
Luke Campagnola's avatar
Luke Campagnola committed
420
421
422
423
424
425
426
        """Set the pen(s) used to draw the outline around each spot. 
        If a list or array is provided, then the pen for each spot will be set separately.
        Otherwise, the arguments are passed to pg.mkPen and used as the default pen for 
        all spots which do not have a pen explicitly set."""
        update = kargs.pop('update', True)
        dataSet = kargs.pop('dataSet', self.data)
        
427
428
        if len(args) == 1 and (isinstance(args[0], np.ndarray) or isinstance(args[0], list)):
            pens = args[0]
Luke Campagnola's avatar
Luke Campagnola committed
429
430
431
            if len(pens) != len(dataSet):
                raise Exception("Number of pens does not match number of points (%d != %d)" % (len(pens), len(dataSet)))
            dataSet['pen'] = pens
432
433
        else:
            self.opts['pen'] = fn.mkPen(*args, **kargs)
Luke Campagnola's avatar
Luke Campagnola committed
434
        
435
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
436
437
        if update:
            self.updateSpots(dataSet)
438
439
        
    def setBrush(self, *args, **kargs):
Luke Campagnola's avatar
Luke Campagnola committed
440
441
442
443
444
445
446
        """Set the brush(es) used to fill the interior of each spot. 
        If a list or array is provided, then the brush for each spot will be set separately.
        Otherwise, the arguments are passed to pg.mkBrush and used as the default brush for 
        all spots which do not have a brush explicitly set."""
        update = kargs.pop('update', True)
        dataSet = kargs.pop('dataSet', self.data)
            
447
448
        if len(args) == 1 and (isinstance(args[0], np.ndarray) or isinstance(args[0], list)):
            brushes = args[0]
Luke Campagnola's avatar
Luke Campagnola committed
449
450
451
452
453
            if len(brushes) != len(dataSet):
                raise Exception("Number of brushes does not match number of points (%d != %d)" % (len(brushes), len(dataSet)))
            #for i in xrange(len(brushes)):
                #self.data[i]['brush'] = fn.mkBrush(brushes[i], **kargs)
            dataSet['brush'] = brushes
454
455
        else:
            self.opts['brush'] = fn.mkBrush(*args, **kargs)
456
            #self._spotPixmap = None
Luke Campagnola's avatar
Luke Campagnola committed
457
        
458
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
459
460
        if update:
            self.updateSpots(dataSet)
461

Luke Campagnola's avatar
Luke Campagnola committed
462
463
464
465
466
467
468
469
    def setSymbol(self, symbol, update=True, dataSet=None):
        """Set the symbol(s) used to draw each spot. 
        If a list or array is provided, then the symbol for each spot will be set separately.
        Otherwise, the argument will be used as the default symbol for 
        all spots which do not have a symbol explicitly set."""
        if dataSet is None:
            dataSet = self.data
            
470
471
        if isinstance(symbol, np.ndarray) or isinstance(symbol, list):
            symbols = symbol
Luke Campagnola's avatar
Luke Campagnola committed
472
473
474
            if len(symbols) != len(dataSet):
                raise Exception("Number of symbols does not match number of points (%d != %d)" % (len(symbols), len(dataSet)))
            dataSet['symbol'] = symbols
475
476
        else:
            self.opts['symbol'] = symbol
Luke Campagnola's avatar
Luke Campagnola committed
477
            self._spotPixmap = None
478
        
479
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
480
481
482
483
484
485
486
487
488
489
490
        if update:
            self.updateSpots(dataSet)
    
    def setSize(self, size, update=True, dataSet=None):
        """Set the size(s) used to draw each spot. 
        If a list or array is provided, then the size for each spot will be set separately.
        Otherwise, the argument will be used as the default size for 
        all spots which do not have a size explicitly set."""
        if dataSet is None:
            dataSet = self.data
            
491
492
        if isinstance(size, np.ndarray) or isinstance(size, list):
            sizes = size
Luke Campagnola's avatar
Luke Campagnola committed
493
494
495
            if len(sizes) != len(dataSet):
                raise Exception("Number of sizes does not match number of points (%d != %d)" % (len(sizes), len(dataSet)))
            dataSet['size'] = sizes
496
497
        else:
            self.opts['size'] = size
Luke Campagnola's avatar
Luke Campagnola committed
498
499
            self._spotPixmap = None
            
500
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
501
502
        if update:
            self.updateSpots(dataSet)
503
        
Luke Campagnola's avatar
Luke Campagnola committed
504
505
506
507
    def setPointData(self, data, dataSet=None):
        if dataSet is None:
            dataSet = self.data
            
508
        if isinstance(data, np.ndarray) or isinstance(data, list):
Luke Campagnola's avatar
Luke Campagnola committed
509
510
            if len(data) != len(dataSet):
                raise Exception("Length of meta data does not match number of points (%d != %d)" % (len(data), len(dataSet)))
Luke Campagnola's avatar
Luke Campagnola committed
511
512
513
514
515
516
517
518
        
        ## Bug: If data is a numpy record array, then items from that array must be copied to dataSet one at a time.
        ## (otherwise they are converted to tuples and thus lose their field names.
        if isinstance(data, np.ndarray) and len(data.dtype.fields) > 1:
            for i, rec in enumerate(data):
                dataSet['data'][i] = rec
        else:
            dataSet['data'] = data
519
        
520
    def setPxMode(self, mode):
Luke Campagnola's avatar
Luke Campagnola committed
521
522
523
        if self.opts['pxMode'] == mode:
            return
            
524
        self.opts['pxMode'] = mode
525
        self.invalidate()
526
        
Luke Campagnola's avatar
Luke Campagnola committed
527
528
529
    def updateSpots(self, dataSet=None):
        if dataSet is None:
            dataSet = self.data
530
531
        self._maxSpotWidth = 0
        self._maxSpotPxWidth = 0
532
        invalidate = False
533
        self.measureSpotSizes(dataSet)
534
535
536
537
538
539
540
541
542
543
544
545
546
547
        if self.opts['pxMode']:
            mask = np.equal(dataSet['fragCoords'], None)
            if np.any(mask):
                invalidate = True
                opts = self.getSpotOpts(dataSet[mask])
                coords = self.fragmentAtlas.getSymbolCoords(opts)
                dataSet['fragCoords'][mask] = coords
                
            #for rec in dataSet:
                #if rec['fragCoords'] is None:
                    #invalidate = True
                    #rec['fragCoords'] = self.fragmentAtlas.getSymbolCoords(*self.getSpotOpts(rec))
        if invalidate:
            self.invalidate()
548

549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
    def getSpotOpts(self, recs):
        if recs.ndim == 0:
            rec = recs
            symbol = rec['symbol']
            if symbol is None:
                symbol = self.opts['symbol']
            size = rec['size']
            if size < 0:
                size = self.opts['size']
            pen = rec['pen']
            if pen is None:
                pen = self.opts['pen']
            brush = rec['brush']
            if brush is None:
                brush = self.opts['brush']
            return (symbol, size, fn.mkPen(pen), fn.mkBrush(brush))
        else:
            recs = recs.copy()
            recs['symbol'][np.equal(recs['symbol'], None)] = self.opts['symbol']
            recs['size'][np.equal(recs['size'], -1)] = self.opts['size']
            recs['pen'][np.equal(recs['pen'], None)] = fn.mkPen(self.opts['pen'])
            recs['brush'][np.equal(recs['brush'], None)] = fn.mkBrush(self.opts['brush'])
            return recs
            
            
        
575
    def measureSpotSizes(self, dataSet):
576
        for rec in dataSet:
577
            ## keep track of the maximum spot size and pixel size
578
            symbol, size, pen, brush = self.getSpotOpts(rec)
579
580
            width = 0
            pxWidth = 0
581
            if self.opts['pxMode']:
582
                pxWidth = size + pen.width()
583
            else:
584
                width = size
585
586
587
588
                if pen.isCosmetic():
                    pxWidth += pen.width()
                else:
                    width += pen.width()
589
590
            self._maxSpotWidth = max(self._maxSpotWidth, width)
            self._maxSpotPxWidth = max(self._maxSpotPxWidth, pxWidth)
591
        self.bounds = [None, None]
592
593
    
    
594
    def clear(self):
Luke Campagnola's avatar
Luke Campagnola committed
595
        """Remove all spots from the scatter plot"""
596
        #self.clearItems()
Luke Campagnola's avatar
Luke Campagnola committed
597
598
        self.data = np.empty(0, dtype=self.data.dtype)
        self.bounds = [None, None]
599
        self.invalidate()
Luke Campagnola's avatar
Luke Campagnola committed
600

601
    def dataBounds(self, ax, frac=1.0, orthoRange=None):
602
603
        if frac >= 1.0 and self.bounds[ax] is not None:
            return self.bounds[ax]
604
        
Luke Campagnola's avatar
Luke Campagnola committed
605
        #self.prepareGeometryChange()
606
607
608
609
610
        if self.data is None or len(self.data) == 0:
            return (None, None)
        
        if ax == 0:
            d = self.data['x']
611
            d2 = self.data['y']
612
613
        elif ax == 1:
            d = self.data['y']
614
615
616
617
618
619
            d2 = self.data['x']
        
        if orthoRange is not None:
            mask = (d2 >= orthoRange[0]) * (d2 <= orthoRange[1])
            d = d[mask]
            d2 = d2[mask]
620
621
            
        if frac >= 1.0:
622
623
624
625
            ## increase size of bounds based on spot size and pen width
            px = self.pixelLength(Point(1, 0) if ax == 0 else Point(0, 1))  ## determine length of pixel along this axis
            if px is None:
                px = 0
626
627
628
629
            minIndex = np.argmin(d)
            maxIndex = np.argmax(d)
            minVal = d[minIndex]
            maxVal = d[maxIndex]
630
631
            spotSize = 0.5 * (self._maxSpotWidth + px * self._maxSpotPxWidth)
            self.bounds[ax] = (minVal-spotSize, maxVal+spotSize)
632
633
634
635
636
637
            return self.bounds[ax]
        elif frac <= 0.0:
            raise Exception("Value for parameter 'frac' must be > 0. (got %s)" % str(frac))
        else:
            return (scipy.stats.scoreatpercentile(d, 50 - (frac * 50)), scipy.stats.scoreatpercentile(d, 50 + (frac * 50)))
            
638

639
640
641
642
643
    #def defaultSpotPixmap(self):
        ### Return the default spot pixmap
        #if self._spotPixmap is None:
            #self._spotPixmap = makeSymbolPixmap(size=self.opts['size'], brush=self.opts['brush'], pen=self.opts['pen'], symbol=self.opts['symbol'])
        #return self._spotPixmap
644
645

    def boundingRect(self):
646
647
648
649
650
651
652
653
        (xmn, xmx) = self.dataBounds(ax=0)
        (ymn, ymx) = self.dataBounds(ax=1)
        if xmn is None or xmx is None:
            xmn = 0
            xmx = 0
        if ymn is None or ymx is None:
            ymn = 0
            ymx = 0
654
        return QtCore.QRectF(xmn, ymn, xmx-xmn, ymx-ymn)
655

656
    def viewTransformChanged(self):
Luke Campagnola's avatar
Luke Campagnola committed
657
        self.prepareGeometryChange()
658
        GraphicsObject.viewTransformChanged(self)
659
        self.bounds = [None, None]
660
        self.fragments = None
661
        
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
    def generateFragments(self):
        tr = self.deviceTransform()
        if tr is None:
            return
        pts = np.empty((2,len(self.data['x'])))
        pts[0] = self.data['x']
        pts[1] = self.data['y']
        pts = fn.transformCoordinates(tr, pts)
        self.fragments = []
        for i in xrange(len(self.data)):
            rec = self.data[i]
            pos = QtCore.QPointF(pts[0,i], pts[1,i])
            x,y,w,h = rec['fragCoords']
            rect = QtCore.QRectF(y, x, h, w)
            self.fragments.append(QtGui.QPainter.PixmapFragment.create(pos, rect))
            
678
679
680
681
682
683
684
    def setExportMode(self, export, opts):
        if export:
            self.exportOpts = opts
            if 'antialias' not in opts:
                self.exportOpts['antialias'] = True
        else:
            self.exportOpts = False
685
686
            
            
687
    def paint(self, p, *args):
688
689
        #p.setPen(fn.mkPen('r'))
        #p.drawRect(self.boundingRect())
690
        if self.opts['pxMode'] is True:
691
692
693
694
695
696
697
698
699
700
701
702
703
            atlas = self.fragmentAtlas.getAtlas()
            #arr = fn.imageToArray(atlas.toImage(), copy=True)
            #if hasattr(self, 'lastAtlas'):
                #if np.any(self.lastAtlas != arr):
                    #print "Atlas changed:", arr
            #self.lastAtlas = arr
            
            if self.fragments is None:
                self.updateSpots()
                self.generateFragments()
                    
            p.resetTransform()
            
704
            if not USE_PYSIDE and self.opts['useCache'] and self.exportOpts is False:
705
706
                p.drawPixmapFragments(self.fragments, atlas)
            else:
707
708
709
710
711
712
                if self.exportOpts is not False:
                    aa = self.exportOpts['antialias']
                else:
                    aa = self.opts['antialias']
                p.setRenderHint(p.Antialiasing, aa)
                
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
                for i in range(len(self.data)):
                    rec = self.data[i]
                    frag = self.fragments[i]
                    p.resetTransform()
                    p.translate(frag.x, frag.y)
                    drawSymbol(p, *self.getSpotOpts(rec))
        else:
            if self.picture is None:
                self.picture = QtGui.QPicture()
                p2 = QtGui.QPainter(self.picture)
                for rec in self.data:
                    
                    p2.resetTransform()
                    p2.translate(rec['x'], rec['y'])
                    drawSymbol(p2, *self.getSpotOpts(rec))
                p2.end()
                
730
731
732
733
734
            if self.exportOpts is not False:
                aa = self.exportOpts['antialias']
            else:
                aa = self.opts['antialias']
            p.setRenderHint(p.Antialiasing, aa)
735
            self.picture.play(p)
736

737
        
738
    def points(self):
739
740
741
        for rec in self.data:
            if rec['item'] is None:
                rec['item'] = SpotItem(rec, self)
Luke Campagnola's avatar
Luke Campagnola committed
742
743
        return self.data['item']
        
744
745
746
747
748
749
    def pointsAt(self, pos):
        x = pos.x()
        y = pos.y()
        pw = self.pixelWidth()
        ph = self.pixelHeight()
        pts = []
Luke Campagnola's avatar
Luke Campagnola committed
750
        for s in self.points():
751
            sp = s.pos()
Luke Campagnola's avatar
Luke Campagnola committed
752
            ss = s.size()
753
754
755
            sx = sp.x()
            sy = sp.y()
            s2x = s2y = ss * 0.5
756
            if self.opts['pxMode']:
757
758
759
760
761
762
763
764
                s2x *= pw
                s2y *= ph
            if x > sx-s2x and x < sx+s2x and y > sy-s2y and y < sy+s2y:
                pts.append(s)
                #print "HIT:", x, y, sx, sy, s2x, s2y
            #else:
                #print "No hit:", (x, y), (sx, sy)
                #print "       ", (sx-s2x, sy-s2y), (sx+s2x, sy+s2y)
765
766
        #pts.sort(lambda a,b: cmp(b.zValue(), a.zValue()))
        return pts[::-1]
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
            

    def mouseClickEvent(self, ev):
        if ev.button() == QtCore.Qt.LeftButton:
            pts = self.pointsAt(ev.pos())
            if len(pts) > 0:
                self.ptsClicked = pts
                self.sigClicked.emit(self, self.ptsClicked)
                ev.accept()
            else:
                #print "no spots"
                ev.ignore()
        else:
            ev.ignore()


783
class SpotItem(object):
Luke Campagnola's avatar
Luke Campagnola committed
784
785
786
787
788
    """
    Class referring to individual spots in a scatter plot.
    These can be retrieved by calling ScatterPlotItem.points() or 
    by connecting to the ScatterPlotItem's click signals.
    """
789

Luke Campagnola's avatar
Luke Campagnola committed
790
    def __init__(self, data, plot):
791
        #GraphicsItem.__init__(self, register=False)
Luke Campagnola's avatar
Luke Campagnola committed
792
793
        self._data = data
        self._plot = plot
794
795
796
        #self.setParentItem(plot)
        #self.setPos(QtCore.QPointF(data['x'], data['y']))
        #self.updateItem()
797
    
Luke Campagnola's avatar
Luke Campagnola committed
798
799
800
801
802
803
804
805
806
807
808
809
    def data(self):
        """Return the user data associated with this spot."""
        return self._data['data']
    
    def size(self):
        """Return the size of this spot. 
        If the spot has no explicit size set, then return the ScatterPlotItem's default size instead."""
        if self._data['size'] == -1:
            return self._plot.opts['size']
        else:
            return self._data['size']
    
810
811
812
813
814
815
    def pos(self):
        return Point(self._data['x'], self._data['y'])
        
    def viewPos(self):
        return self._plot.mapToView(self.pos())
    
Luke Campagnola's avatar
Luke Campagnola committed
816
817
818
819
820
821
822
823
824
825
826
827
    def setSize(self, size):
        """Set the size of this spot. 
        If the size is set to -1, then the ScatterPlotItem's default size 
        will be used instead."""
        self._data['size'] = size
        self.updateItem()
    
    def symbol(self):
        """Return the symbol of this spot. 
        If the spot has no explicit symbol set, then return the ScatterPlotItem's default symbol instead.
        """
        symbol = self._data['symbol']
828
        if symbol is None:
Luke Campagnola's avatar
Luke Campagnola committed
829
            symbol = self._plot.opts['symbol']
830
        try:
Luke Campagnola's avatar
Luke Campagnola committed
831
            n = int(symbol)
832
            symbol = list(Symbols.keys())[n % len(Symbols)]
Luke Campagnola's avatar
Luke Campagnola committed
833
        except:
834
            pass
Luke Campagnola's avatar
Luke Campagnola committed
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
        return symbol
    
    def setSymbol(self, symbol):
        """Set the symbol for this spot.
        If the symbol is set to '', then the ScatterPlotItem's default symbol will be used instead."""
        self._data['symbol'] = symbol
        self.updateItem()

    def pen(self):
        pen = self._data['pen']
        if pen is None:
            pen = self._plot.opts['pen']
        return fn.mkPen(pen)
    
    def setPen(self, *args, **kargs):
        """Set the outline pen for this spot"""
        pen = fn.mkPen(*args, **kargs)
        self._data['pen'] = pen
        self.updateItem()
    
    def resetPen(self):
        """Remove the pen set for this spot; the scatter plot's default pen will be used instead."""
        self._data['pen'] = None  ## Note this is NOT the same as calling setPen(None)
        self.updateItem()
    
    def brush(self):
        brush = self._data['brush']
        if brush is None:
            brush = self._plot.opts['brush']
        return fn.mkBrush(brush)
    
    def setBrush(self, *args, **kargs):
        """Set the fill brush for this spot"""
        brush = fn.mkBrush(*args, **kargs)
        self._data['brush'] = brush
        self.updateItem()
    
    def resetBrush(self):
        """Remove the brush set for this spot; the scatter plot's default brush will be used instead."""
        self._data['brush'] = None  ## Note this is NOT the same as calling setBrush(None)
        self.updateItem()
    
    def setData(self, data):
        """Set the user-data associated with this spot"""
        self._data['data'] = data

881
882
883
884
    def updateItem(self):
        self._data['fragCoords'] = None
        self._plot.updateSpots([self._data])
        self._plot.invalidate()
Luke Campagnola's avatar
Luke Campagnola committed
885

886
887
888
889
890
#class PixmapSpotItem(SpotItem, QtGui.QGraphicsPixmapItem):
    #def __init__(self, data, plot):
        #QtGui.QGraphicsPixmapItem.__init__(self)
        #self.setFlags(self.flags() | self.ItemIgnoresTransformations)
        #SpotItem.__init__(self, data, plot)
Luke Campagnola's avatar
Luke Campagnola committed
891
    
892
893
894
    #def setPixmap(self, pixmap):
        #QtGui.QGraphicsPixmapItem.setPixmap(self, pixmap)
        #self.setOffset(-pixmap.width()/2.+0.5, -pixmap.height()/2.)
Luke Campagnola's avatar
Luke Campagnola committed
895
    
896
897
898
899
900
901
902
903
904
    #def updateItem(self):
        #symbolOpts = (self._data['pen'], self._data['brush'], self._data['size'], self._data['symbol'])
        
        ### If all symbol options are default, use default pixmap
        #if symbolOpts == (None, None, -1, ''):
            #pixmap = self._plot.defaultSpotPixmap()
        #else:
            #pixmap = makeSymbolPixmap(size=self.size(), pen=self.pen(), brush=self.brush(), symbol=self.symbol())
        #self.setPixmap(pixmap)
905
906


907
908
909
910
#class PathSpotItem(SpotItem, QtGui.QGraphicsPathItem):
    #def __init__(self, data, plot):
        #QtGui.QGraphicsPathItem.__init__(self)
        #SpotItem.__init__(self, data, plot)
911

912
913
914
915
916
917
918
    #def updateItem(self):
        #QtGui.QGraphicsPathItem.setPath(self, Symbols[self.symbol()])
        #QtGui.QGraphicsPathItem.setPen(self, self.pen())
        #QtGui.QGraphicsPathItem.setBrush(self, self.brush())
        #size = self.size()
        #self.resetTransform()
        #self.scale(size, size)