plot.py 文件源码

def plot_trace(xy, ids = None, depth = 0, colormap = 'rainbow', line_color = 'k', line_width = 1, point_size = 5, title = None):
  """Plot trajectories with positions color coded according to discrete ids"""

  #if ids is not None:
  uids = np.unique(ids);

  cmap = cm.get_cmap(colormap);
  n = len(uids);
  colors = cmap(range(n), bytes = True);

  #lines
  if line_width is not None:
    #plt.plot(xy[:,0], xy[:,1], color = lines);    
    plot = pg.plot(xy[:,0], xy[:,1], pen = pg.mkPen(color = line_color, width = line_width))    
  else:
    plot = pg.plot(title = title);

  if ids is None:
    sp = pg.ScatterPlotItem(pos = xy, size=point_size, pen=pg.mkPen(colors[0])); #, pxMode=True);
  else:
    sp = pg.ScatterPlotItem(size=point_size); #, pxMode=True);
    spots = [];
    for j,i in enumerate(uids):
      idx = ids == i;
      spots.append({'pos': xy[idx,:].T, 'data': 1, 'brush':pg.mkBrush(colors[j])}); #, 'size': point_size});
    sp.addPoints(spots)

  plot.addItem(sp);

  return plot;


#  legs = [];
#  for k,i in enumerate(uids):
#    ii = np.where(ids == i)[0];
#    if depth > 0:
#      ii = [ii-d for d in range(depth)];
#      ii = np.unique(np.concatenate(ii));
#    
#    plt.plot(data[ii, 0], data[ii, 1], '.', color = color[k]);
#
#    legs.append(mpatches.Patch(color=color[k], label= str(i)));
#  
#  plt.legend(handles=legs);