ScatterPlotItem.py 34.3 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
Luke Campagnola's avatar
Luke Campagnola committed
11
#import pyqtgraph as pg 
12
13

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


## Build all symbol paths
Luke Campagnola's avatar
Luke Campagnola committed
17
Symbols = OrderedDict([(name, QtGui.QPainterPath()) for name in ['o', 's', 't', 'd', '+']])
18
19
20
21
22
23
24
25
26
27
28
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)
    ],
}
29
for k, c in coords.items():
30
31
32
33
34
    Symbols[k].moveTo(*c[0])
    for x,y in c[1:]:
        Symbols[k].lineTo(x, y)
    Symbols[k].closeSubpath()

35
36
37
38
39
40
41
42
43
44
    
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)
45

46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
    
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
61
    ## Render a spot with the given parameters to a pixmap
62
    penPxWidth = max(np.ceil(pen.width()), 1)
63
    image = QtGui.QImage(int(size+penPxWidth), int(size+penPxWidth), QtGui.QImage.Format_ARGB32)
Luke Campagnola's avatar
Luke Campagnola committed
64
65
66
    image.fill(0)
    p = QtGui.QPainter(image)
    p.setRenderHint(p.Antialiasing)
67
    p.translate(image.width()*0.5, image.height()*0.5)
68
    drawSymbol(p, symbol, size, pen, brush)
Luke Campagnola's avatar
Luke Campagnola committed
69
    p.end()
70
71
72
73
    return image
    #pixmap = QtGui.QPixmap(image)
    #SymbolPixmapCache[key] = pixmap
    #return pixmap
Luke Campagnola's avatar
Luke Campagnola committed
74

75
76
77
78
79
def makeSymbolPixmap(size, pen, brush, symbol):
    ## deprecated
    img = renderSymbol(symbol, size, pen, brush)
    return QtGui.QPixmap(img)
    
80
class SymbolAtlas(object):
81
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
    """
    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
163

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
        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
179

180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        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
        
    
    
    
200
class ScatterPlotItem(GraphicsObject):
201
202
203
204
    """
    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.
205
    
206
207
208
209
210
211
212
213
214
215
216
217
    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.
    ========================  ===============================================
    
    """
218
219
220
    #sigPointClicked = QtCore.Signal(object, object)
    sigClicked = QtCore.Signal(object, object)  ## self, points
    sigPlotChanged = QtCore.Signal(object)
221
222
223
224
    def __init__(self, *args, **kargs):
        """
        Accepts the same arguments as setData()
        """
Luke Campagnola's avatar
Luke Campagnola committed
225
        prof = debug.Profiler('ScatterPlotItem.__init__', disabled=True)
226
        GraphicsObject.__init__(self)
227
228
229
230
231
232
        
        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)])
233
234
235
        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
236
        self.opts = {'pxMode': True, 'useCache': True, 'exportMode': False}   ## If useCache is False, symbols are re-drawn on every paint.
237
        
Luke Campagnola's avatar
Luke Campagnola committed
238
239
240
241
242
        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')
243
        self.setData(*args, **kargs)
Luke Campagnola's avatar
Luke Campagnola committed
244
245
        prof.mark('setData')
        prof.finish()
246
        
247
248
        #self.setCacheMode(self.DeviceCoordinateCache)
        
249
    def setData(self, *args, **kargs):
250
        """
251
252
253
254
255
        **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.
        
256
        ====================== ===============================================================================================
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
        **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
272
273
274
                               * 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.
275
276
277
278
279
        *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
280
        *identical*            *Deprecated*. This functionality is handled automatically now.
281
        ====================== ===============================================================================================
282
        """
283
        oldData = self.data  ## this causes cached pixmaps to be preserved while new data is registered.
Luke Campagnola's avatar
Luke Campagnola committed
284
285
        self.clear()  ## clear out all old data
        self.addPoints(*args, **kargs)
286

Luke Campagnola's avatar
Luke Campagnola committed
287
288
289
290
291
    def addPoints(self, *args, **kargs):
        """
        Add new points to the scatter plot. 
        Arguments are the same as setData()
        """
292
        
293
294
295
296
297
298
299
300
        ## 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.')
301
        
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        ## 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
320
        
321
322
323
324
325
        ## 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'])
326
        else:
327
328
329
            kargs['x'] = []
            kargs['y'] = []
            numPts = 0
330
        
Luke Campagnola's avatar
Luke Campagnola committed
331
332
333
334
335
336
        ## 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
337
338
        #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
339
340
341
            
        newData = self.data[len(oldData):]
        newData['size'] = -1  ## indicates to use default size
342
        
343
344
        if 'spots' in kargs:
            spots = kargs['spots']
345
            for i in range(len(spots)):
346
347
                spot = spots[i]
                for k in spot:
Luke Campagnola's avatar
Luke Campagnola committed
348
349
350
351
352
                    #if k == 'pen':
                        #newData[k] = fn.mkPen(spot[k])
                    #elif k == 'brush':
                        #newData[k] = fn.mkBrush(spot[k])
                    if k == 'pos':
353
354
355
356
357
                        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
358
359
360
361
362
363
                        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]
364
365
366
                    else:
                        raise Exception("Unknown spot parameter: %s" % k)
        elif 'y' in kargs:
Luke Campagnola's avatar
Luke Campagnola committed
367
368
            newData['x'] = kargs['x']
            newData['y'] = kargs['y']
369
        
Luke Campagnola's avatar
Luke Campagnola committed
370
        if 'pxMode' in kargs:
371
            self.setPxMode(kargs['pxMode'])
Luke Campagnola's avatar
Luke Campagnola committed
372
            
373
        ## Set any extra parameters provided in keyword arguments
Luke Campagnola's avatar
Luke Campagnola committed
374
        for k in ['pen', 'brush', 'symbol', 'size']:
375
376
            if k in kargs:
                setMethod = getattr(self, 'set' + k[0].upper() + k[1:])
Luke Campagnola's avatar
Luke Campagnola committed
377
                setMethod(kargs[k], update=False, dataSet=newData)
378
        
Luke Campagnola's avatar
Luke Campagnola committed
379
380
        if 'data' in kargs:
            self.setPointData(kargs['data'], dataSet=newData)
381
        
Luke Campagnola's avatar
Luke Campagnola committed
382
        self.prepareGeometryChange()
383
        self.bounds = [None, None]
384
385
        self.invalidate()
        self.updateSpots(newData)
Luke Campagnola's avatar
Luke Campagnola committed
386
        self.sigPlotChanged.emit(self)
387
        
388
389
390
391
392
393
    def invalidate(self):
        ## clear any cached drawing state
        self.picture = None
        self.fragments = None
        self.update()
        
Luke Campagnola's avatar
Luke Campagnola committed
394
395
396
    def getData(self):
        return self.data['x'], self.data['y']
    
397
398
        
    def setPoints(self, *args, **kargs):
399
        ##Deprecated; use setData
400
401
        return self.setData(*args, **kargs)
        
402
403
404
405
406
407
    def implements(self, interface=None):
        ints = ['plotData']
        if interface is None:
            return ints
        return interface in ints
    
408
    def setPen(self, *args, **kargs):
Luke Campagnola's avatar
Luke Campagnola committed
409
410
411
412
413
414
415
        """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)
        
416
417
        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
418
419
420
            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
421
422
        else:
            self.opts['pen'] = fn.mkPen(*args, **kargs)
Luke Campagnola's avatar
Luke Campagnola committed
423
        
424
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
425
426
        if update:
            self.updateSpots(dataSet)
427
428
        
    def setBrush(self, *args, **kargs):
Luke Campagnola's avatar
Luke Campagnola committed
429
430
431
432
433
434
435
        """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)
            
436
437
        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
438
439
440
441
442
            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
443
444
        else:
            self.opts['brush'] = fn.mkBrush(*args, **kargs)
445
            #self._spotPixmap = None
Luke Campagnola's avatar
Luke Campagnola committed
446
        
447
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
448
449
        if update:
            self.updateSpots(dataSet)
450

Luke Campagnola's avatar
Luke Campagnola committed
451
452
453
454
455
456
457
458
    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
            
459
460
        if isinstance(symbol, np.ndarray) or isinstance(symbol, list):
            symbols = symbol
Luke Campagnola's avatar
Luke Campagnola committed
461
462
463
            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
464
465
        else:
            self.opts['symbol'] = symbol
Luke Campagnola's avatar
Luke Campagnola committed
466
            self._spotPixmap = None
467
        
468
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
469
470
471
472
473
474
475
476
477
478
479
        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
            
480
481
        if isinstance(size, np.ndarray) or isinstance(size, list):
            sizes = size
Luke Campagnola's avatar
Luke Campagnola committed
482
483
484
            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
485
486
        else:
            self.opts['size'] = size
Luke Campagnola's avatar
Luke Campagnola committed
487
488
            self._spotPixmap = None
            
489
        dataSet['fragCoords'] = None
Luke Campagnola's avatar
Luke Campagnola committed
490
491
        if update:
            self.updateSpots(dataSet)
492
        
Luke Campagnola's avatar
Luke Campagnola committed
493
494
495
496
    def setPointData(self, data, dataSet=None):
        if dataSet is None:
            dataSet = self.data
            
497
        if isinstance(data, np.ndarray) or isinstance(data, list):
Luke Campagnola's avatar
Luke Campagnola committed
498
499
            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
500
501
502
503
504
505
506
507
        
        ## 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
508
        
509
    def setPxMode(self, mode):
Luke Campagnola's avatar
Luke Campagnola committed
510
511
512
        if self.opts['pxMode'] == mode:
            return
            
513
        self.opts['pxMode'] = mode
514
        self.invalidate()
515
        
Luke Campagnola's avatar
Luke Campagnola committed
516
517
518
    def updateSpots(self, dataSet=None):
        if dataSet is None:
            dataSet = self.data
519
520
        self._maxSpotWidth = 0
        self._maxSpotPxWidth = 0
521
        invalidate = False
522
        self.measureSpotSizes(dataSet)
523
524
525
526
527
528
529
530
531
532
533
534
535
536
        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()
537

538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
    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
            
            
        
564
    def measureSpotSizes(self, dataSet):
565
        for rec in dataSet:
566
            ## keep track of the maximum spot size and pixel size
567
            symbol, size, pen, brush = self.getSpotOpts(rec)
568
569
            width = 0
            pxWidth = 0
570
            if self.opts['pxMode']:
571
                pxWidth = size + pen.width()
572
            else:
573
                width = size
574
575
576
577
                if pen.isCosmetic():
                    pxWidth += pen.width()
                else:
                    width += pen.width()
578
579
            self._maxSpotWidth = max(self._maxSpotWidth, width)
            self._maxSpotPxWidth = max(self._maxSpotPxWidth, pxWidth)
580
        self.bounds = [None, None]
581
582
    
    
583
    def clear(self):
Luke Campagnola's avatar
Luke Campagnola committed
584
        """Remove all spots from the scatter plot"""
585
        #self.clearItems()
Luke Campagnola's avatar
Luke Campagnola committed
586
587
        self.data = np.empty(0, dtype=self.data.dtype)
        self.bounds = [None, None]
588
        self.invalidate()
Luke Campagnola's avatar
Luke Campagnola committed
589

590
    def dataBounds(self, ax, frac=1.0, orthoRange=None):
591
592
        if frac >= 1.0 and self.bounds[ax] is not None:
            return self.bounds[ax]
593
        
Luke Campagnola's avatar
Luke Campagnola committed
594
        #self.prepareGeometryChange()
595
596
597
598
599
        if self.data is None or len(self.data) == 0:
            return (None, None)
        
        if ax == 0:
            d = self.data['x']
600
            d2 = self.data['y']
601
602
        elif ax == 1:
            d = self.data['y']
603
604
605
606
607
608
            d2 = self.data['x']
        
        if orthoRange is not None:
            mask = (d2 >= orthoRange[0]) * (d2 <= orthoRange[1])
            d = d[mask]
            d2 = d2[mask]
609
610
            
        if frac >= 1.0:
611
612
613
614
            ## 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
615
616
617
618
            minIndex = np.argmin(d)
            maxIndex = np.argmax(d)
            minVal = d[minIndex]
            maxVal = d[maxIndex]
619
620
            spotSize = 0.5 * (self._maxSpotWidth + px * self._maxSpotPxWidth)
            self.bounds[ax] = (minVal-spotSize, maxVal+spotSize)
621
622
623
624
625
626
            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)))
            
627

628
629
630
631
632
    #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
633
634

    def boundingRect(self):
635
636
637
638
639
640
641
642
        (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
643
        return QtCore.QRectF(xmn, ymn, xmx-xmn, ymx-ymn)
644

645
    def viewTransformChanged(self):
Luke Campagnola's avatar
Luke Campagnola committed
646
        self.prepareGeometryChange()
647
        GraphicsObject.viewTransformChanged(self)
648
        self.bounds = [None, None]
649
        self.fragments = None
650
        
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
    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))
            
667
668
669
670
    def setExportMode(self, enabled, opts):
        self.opts['exportMode'] = enabled
            
            
671
    def paint(self, p, *args):
672
673
        #p.setPen(fn.mkPen('r'))
        #p.drawRect(self.boundingRect())
674
        if self.opts['pxMode'] is True:
675
676
677
678
679
680
681
682
683
684
685
686
687
            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()
            
688
            if not USE_PYSIDE and self.opts['useCache'] and self.opts['exportMode'] is False:
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
                p.drawPixmapFragments(self.fragments, atlas)
            else:
                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()
                
            self.picture.play(p)
709

710
        
711
    def points(self):
712
713
714
        for rec in self.data:
            if rec['item'] is None:
                rec['item'] = SpotItem(rec, self)
Luke Campagnola's avatar
Luke Campagnola committed
715
716
        return self.data['item']
        
717
718
719
720
721
722
    def pointsAt(self, pos):
        x = pos.x()
        y = pos.y()
        pw = self.pixelWidth()
        ph = self.pixelHeight()
        pts = []
Luke Campagnola's avatar
Luke Campagnola committed
723
        for s in self.points():
724
            sp = s.pos()
Luke Campagnola's avatar
Luke Campagnola committed
725
            ss = s.size()
726
727
728
            sx = sp.x()
            sy = sp.y()
            s2x = s2y = ss * 0.5
729
            if self.opts['pxMode']:
730
731
732
733
734
735
736
737
                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)
738
739
        #pts.sort(lambda a,b: cmp(b.zValue(), a.zValue()))
        return pts[::-1]
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
            

    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()


756
class SpotItem(object):
Luke Campagnola's avatar
Luke Campagnola committed
757
758
759
760
761
    """
    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.
    """
762

Luke Campagnola's avatar
Luke Campagnola committed
763
    def __init__(self, data, plot):
764
        #GraphicsItem.__init__(self, register=False)
Luke Campagnola's avatar
Luke Campagnola committed
765
766
        self._data = data
        self._plot = plot
767
768
769
        #self.setParentItem(plot)
        #self.setPos(QtCore.QPointF(data['x'], data['y']))
        #self.updateItem()
770
    
Luke Campagnola's avatar
Luke Campagnola committed
771
772
773
774
775
776
777
778
779
780
781
782
    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']
    
783
784
785
786
787
788
    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
789
790
791
792
793
794
795
796
797
798
799
800
    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']
801
        if symbol is None:
Luke Campagnola's avatar
Luke Campagnola committed
802
            symbol = self._plot.opts['symbol']
803
        try:
Luke Campagnola's avatar
Luke Campagnola committed
804
            n = int(symbol)
805
            symbol = list(Symbols.keys())[n % len(Symbols)]
Luke Campagnola's avatar
Luke Campagnola committed
806
        except:
807
            pass
Luke Campagnola's avatar
Luke Campagnola committed
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
        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

854
855
856
857
    def updateItem(self):
        self._data['fragCoords'] = None
        self._plot.updateSpots([self._data])
        self._plot.invalidate()
Luke Campagnola's avatar
Luke Campagnola committed
858

859
860
861
862
863
#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
864
    
865
866
867
    #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
868
    
869
870
871
872
873
874
875
876
877
    #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)
878
879


880
881
882
883
#class PathSpotItem(SpotItem, QtGui.QGraphicsPathItem):
    #def __init__(self, data, plot):
        #QtGui.QGraphicsPathItem.__init__(self)
        #SpotItem.__init__(self, data, plot)
884

885
886
887
888
889
890
891
    #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)