def plot_poly_collection(ax, verts, zs=None, color=None, cmap=None, vmin=None, vmax=None, **kwargs): from matplotlib.collections import PolyCollection # if None in zs: # zs = None # color=None overwrites specified facecolor/edgecolor with default color if color is not None: kwargs["color"] = color import matplotlib as mpl norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) poly = PolyCollection(verts, **kwargs) if zs is not None: poly.set_array(np.asarray(zs)) poly.set_cmap(cmap) poly.set_clim(vmin, vmax) ax.add_collection3d(poly, zs=zs, zdir="y") # ax.autoscale_view() return poly
def main_2d(vertices, faces, facecolors, edges, edgecolors, verts, vertexcolors, args, canvassize=4): fig = plt.figure() fig.set_size_inches(canvassize, canvassize) ax = fig.add_subplot(111) ax.set_aspect("equal") plt.axis('off') these_faces = [vertices[i] for i in faces] pd = PolyCollection(these_faces, edgecolors=edgecolors) try: pd.set_facecolor(facecolors) except ValueError: pd.set_array(np.array(facecolors)) pd.set_cmap(plt.get_cmap(args.cmap)) ax.add_collection(pd) these_edges = vertices[edges] ld = LineCollection(these_edges, edgecolor=edgecolors) ax.add_collection(ld) these_vertices = vertices[verts] ax.scatter(these_vertices[..., 0], these_vertices[..., 1], c=vertexcolors)
units=var_atts[varname]['units'] data=utils.get_packed_data(dt, tsvert, dtc) data=numpy.ma.masked_equal(data,fill_val,copy=False) # apply the scale_factor if scale_factor is not None: data=data.astype(scale_factor.dtype.char)*scale_factor # apply the add_offset if add_offset is not None: data+=add_offset qvert_list=mesh.getEntAdj(quads,iBase.Type.vertex) poly_list=[] for qv in qvert_list: cds=mesh.getVtxCoords(qv) poly_list.append(cds[:,[0,1]].tolist()) pcoll=PolyCollection(poly_list, edgecolor='none') pcoll.set_array(data) pcoll.set_cmap(cm.jet) a=fig.add_subplot(ncol,nrow,i) a.add_collection(pcoll, autolim=True) a.autoscale_view() colorbar(pcoll,orientation='vertical') title("%s (%s)" % (varname,units)) i+=1 show(0)
def tripcolor(ax, *args, **kwargs): """ Create a pseudocolor plot of an unstructured triangular grid to the :class:`~matplotlib.axes.Axes`. The triangulation can be specified in one of two ways; either:: tripcolor(triangulation, ...) where triangulation is a :class:`~matplotlib.tri.Triangulation` object, or :: tripcolor(x, y, ...) tripcolor(x, y, triangles, ...) tripcolor(x, y, triangles=triangles, ...) tripcolor(x, y, mask=mask, ...) tripcolor(x, y, triangles, mask=mask, ...) in which case a Triangulation object will be created. See :class:`~matplotlib.tri.Triangulation` for a explanation of these possibilities. The next argument must be *C*, the array of color values, one per point in the triangulation. The colors used for each triangle are from the mean C of the triangle's three points. The remaining kwargs are the same as for :meth:`~matplotlib.axes.Axes.pcolor`. **Example:** .. plot:: mpl_examples/pylab_examples/tripcolor_demo.py """ if not ax._hold: ax.cla() alpha = kwargs.pop('alpha', 1.0) norm = kwargs.pop('norm', None) cmap = kwargs.pop('cmap', None) vmin = kwargs.pop('vmin', None) vmax = kwargs.pop('vmax', None) shading = kwargs.pop('shading', 'flat') tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs) x = tri.x y = tri.y triangles = tri.get_masked_triangles() # Vertices of triangles. verts = np.concatenate((x[triangles][...,np.newaxis], y[triangles][...,np.newaxis]), axis=2) C = np.asarray(args[0]) if C.shape != x.shape: raise ValueError('C array must have same length as triangulation x and' ' y arrays') # Color values, one per triangle, mean of the 3 vertex color values. C = C[triangles].mean(axis=1) if shading == 'faceted': edgecolors = (0,0,0,1), linewidths = (0.25,) else: edgecolors = 'face' linewidths = (1.0,) kwargs.setdefault('edgecolors', edgecolors) kwargs.setdefault('antialiaseds', (0,)) kwargs.setdefault('linewidths', linewidths) collection = PolyCollection(verts, **kwargs) collection.set_alpha(alpha) collection.set_array(C) if norm is not None: assert(isinstance(norm, Normalize)) collection.set_cmap(cmap) collection.set_norm(norm) if vmin is not None or vmax is not None: collection.set_clim(vmin, vmax) else: collection.autoscale_None() ax.grid(False) minx = tri.x.min() maxx = tri.x.max() miny = tri.y.min() maxy = tri.y.max() corners = (minx, miny), (maxx, maxy) ax.update_datalim( corners) ax.autoscale_view() ax.add_collection(collection) return collection
def tripcolor(ax, *args, **kwargs): """ Create a pseudocolor plot of an unstructured triangular grid to the :class:`~matplotlib.axes.Axes`. The triangulation can be specified in one of two ways; either:: tripcolor(triangulation, ...) where triangulation is a :class:`~matplotlib.tri.Triangulation` object, or :: tripcolor(x, y, ...) tripcolor(x, y, triangles, ...) tripcolor(x, y, triangles=triangles, ...) tripcolor(x, y, mask, ...) tripcolor(x, y, mask=mask, ...) tripcolor(x, y, triangles, mask, ...) tripcolor(x, y, triangles, mask=mask, ...) in which case a Triangulation object will be created. See :class:`~matplotlib.tri.Triangulation` for a explanation of these possibilities. The next argument must be *C*, the array of color values, one per point in the triangulation. The colors used for each triangle are from the mean C of the triangle's three points. The remaining kwargs are the same as for :meth:`~matplotlib.axes.Axes.pcolor`. **Example:** .. plot:: mpl_examples/pylab_examples/tripcolor_demo.py """ if not ax._hold: ax.cla() alpha = kwargs.pop('alpha', 1.0) norm = kwargs.pop('norm', None) cmap = kwargs.pop('cmap', None) vmin = kwargs.pop('vmin', None) vmax = kwargs.pop('vmax', None) shading = kwargs.pop('shading', 'flat') tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs) x = tri.x y = tri.y triangles = tri.get_masked_triangles() # Vertices of triangles. verts = np.concatenate( (x[triangles][..., np.newaxis], y[triangles][..., np.newaxis]), axis=2) C = np.asarray(args[0]) if C.shape != x.shape: raise ValueError('C array must have same length as triangulation x and' ' y arrays') # Color values, one per triangle, mean of the 3 vertex color values. C = C[triangles].mean(axis=1) if shading == 'faceted': edgecolors = (0, 0, 0, 1), linewidths = (0.25, ) else: edgecolors = 'face' linewidths = (1.0, ) kwargs.setdefault('edgecolors', edgecolors) kwargs.setdefault('antialiaseds', (0, )) kwargs.setdefault('linewidths', linewidths) collection = PolyCollection(verts, **kwargs) collection.set_alpha(alpha) collection.set_array(C) if norm is not None: assert (isinstance(norm, Normalize)) collection.set_cmap(cmap) collection.set_norm(norm) if vmin is not None or vmax is not None: collection.set_clim(vmin, vmax) else: collection.autoscale_None() ax.grid(False) minx = tri.x.min() maxx = tri.x.max() miny = tri.y.min() maxy = tri.y.max() corners = (minx, miny), (maxx, maxy) ax.update_datalim(corners) ax.autoscale_view() ax.add_collection(collection) return collection
def draw_nx_tapered_edges(G, pos, edgelist=None, width=0.5, edge_color='k', style='solid', alpha=1.0, edge_cmap=None, edge_vmin=None, edge_vmax=None, ax=None, label=None, highlight=None, tapered=False, **kwds): """Draw the edges of the graph G. This draws only the edges of the graph G. Parameters ---------- G : graph A networkx graph pos : dictionary A dictionary with nodes as keys and positions as values. Positions should be sequences of length 2. edgelist : collection of edge tuples Draw only specified edges(default=G.edges()) width : float, or array of floats Line width of edges (default=1.0) edge_color : color string, or array of floats Edge color. Can be a single color format string (default='r'), or a sequence of colors with the same length as edgelist. If numeric values are specified they will be mapped to colors using the edge_cmap and edge_vmin,edge_vmax parameters. style : string Edge line style (default='solid') (solid|dashed|dotted,dashdot) alpha : float The edge transparency (default=1.0) edge_ cmap : Matplotlib colormap Colormap for mapping intensities of edges (default=None) edge_vmin,edge_vmax : floats Minimum and maximum for edge colormap scaling (default=None) ax : Matplotlib Axes object, optional Draw the graph in the specified Matplotlib axes. label : [None| string] Label for legend Returns ------- matplotlib.collection.LineCollection `LineCollection` of the edges Examples -------- >>> G=nx.dodecahedral_graph() >>> edges=nx.draw_networkx_edges(G,pos=nx.spring_layout(G)) Also see the NetworkX drawing examples at http://networkx.github.io/documentation/latest/gallery.html See Also -------- draw() draw_networkx() draw_networkx_nodes() draw_networkx_labels() draw_networkx_edge_labels() """ if ax is None: ax = plt.gca() if edgelist is None: edgelist = list(G.edges()) if not edgelist or len(edgelist) == 0: # no edges! return None if highlight is not None and (isinstance(edge_color, basestring) or not cb.iterable(edge_color)): idMap = {} nodes = G.nodes() for i in range(len(nodes)): idMap[nodes[i]] = i ecol = [edge_color] * len(edgelist) eHighlight = [ highlight[idMap[edge[0]]] or highlight[idMap[edge[1]]] for edge in edgelist ] for i in range(len(eHighlight)): if eHighlight[i]: ecol[i] = '0.0' edge_color = ecol # set edge positions if not cb.iterable(width): lw = np.full(len(edgelist), width) else: lw = width edge_pos = [] wdScale = 0.01 for i in range(len(edgelist)): e = edgelist[i] w = wdScale * lw[i] / 2 p0 = pos[e[0]] p1 = pos[e[1]] dx = p1[0] - p0[0] dy = p1[1] - p0[1] l = math.sqrt(dx * dx + dy * dy) edge_pos.append( ((p0[0] + w * dy / l, p0[1] - w * dx / l), (p0[0] - w * dy / l, p0[1] + w * dx / l), (p1[0], p1[1]))) edge_vertices = np.asarray(edge_pos) if not isinstance(edge_color, basestring) \ and cb.iterable(edge_color) \ and len(edge_color) == len(edge_vertices): if np.alltrue([isinstance(c, basestring) for c in edge_color]): # (should check ALL elements) # list of color letters such as ['k','r','k',...] edge_colors = tuple( [colorConverter.to_rgba(c, alpha) for c in edge_color]) elif np.alltrue([not isinstance(c, basestring) for c in edge_color]): # If color specs are given as (rgb) or (rgba) tuples, we're OK if np.alltrue( [cb.iterable(c) and len(c) in (3, 4) for c in edge_color]): edge_colors = tuple(edge_color) else: # numbers (which are going to be mapped with a colormap) edge_colors = None else: raise ValueError( 'edge_color must consist of either color names or numbers') else: if isinstance(edge_color, basestring) or len(edge_color) == 1: edge_colors = (colorConverter.to_rgba(edge_color, alpha), ) else: raise ValueError( 'edge_color must be a single color or list of exactly m colors where m is the number or edges' ) if tapered: edge_collection = PolyCollection( edge_vertices, facecolors=edge_colors, linewidths=0, antialiaseds=(1, ), transOffset=ax.transData, ) else: edge_collection = LineCollection( edge_pos, colors=edge_colors, linewidths=lw, antialiaseds=(1, ), linestyle=style, transOffset=ax.transData, ) edge_collection.set_zorder(1) # edges go behind nodes edge_collection.set_label(label) ax.add_collection(edge_collection) # Note: there was a bug in mpl regarding the handling of alpha values for # each line in a LineCollection. It was fixed in matplotlib in r7184 and # r7189 (June 6 2009). We should then not set the alpha value globally, # since the user can instead provide per-edge alphas now. Only set it # globally if provided as a scalar. if cb.is_numlike(alpha): edge_collection.set_alpha(alpha) if edge_colors is None: if edge_cmap is not None: assert (isinstance(edge_cmap, Colormap)) edge_collection.set_array(np.asarray(edge_color)) edge_collection.set_cmap(edge_cmap) if edge_vmin is not None or edge_vmax is not None: edge_collection.set_clim(edge_vmin, edge_vmax) else: edge_collection.autoscale() # update view minx = np.amin(np.ravel(edge_vertices[:, :, 0])) maxx = np.amax(np.ravel(edge_vertices[:, :, 0])) miny = np.amin(np.ravel(edge_vertices[:, :, 1])) maxy = np.amax(np.ravel(edge_vertices[:, :, 1])) w = maxx - minx h = maxy - miny padx, pady = 0.05 * w, 0.05 * h corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady) ax.update_datalim(corners) ax.autoscale_view() return edge_collection
y = ncv['node_y'][:] x = ncv['node_x'][:] s = ncv['hzg_ecosmo_sed2'][:] nv = ncv['nv'][:, :3] - 1 verts = [] for nvi in nv: verts.append([(x[i], y[i]) for i in nvi]) verts = asarray(verts) f = figure(figsize=(15, 5)) p = PolyCollection(verts, closed=True, edgecolor='none') p.set_array(s[0]) p.set_cmap(cm.RdYlBu_r) pn = PolyCollection(verts, closed=True, edgecolor='none') pn.set_array(arange(len(s[0]))) pn.set_cmap(cm.RdYlBu_r) ax = axes([0.1, 0.75, 0.88, 0.2]) ax.add_collection(p, autolim=True) autoscale() colorbar(p) ax = axes([0.1, 0.25, 0.88, 0.2]) ax.add_collection(pn, autolim=True) autoscale() colorbar(pn)
def tile(self, x, y, w, h, color=None, anchor='center', edgecolors='face', linewidth=0.8, **kwargs): """Plot rectanguler tiles based onto these `Axes`. ``x`` and ``y`` give the anchor point for each tile, with ``w`` and ``h`` giving the extent in the X and Y axis respectively. Parameters ---------- x, y, w, h : `array_like`, shape (n, ) Input data color : `array_like`, shape (n, ) Array of amplitudes for tile color anchor : `str`, optional Anchor point for tiles relative to ``(x, y)`` coordinates, one of - ``'center'`` - center tile on ``(x, y)`` - ``'ll'`` - ``(x, y)`` defines lower-left corner of tile - ``'lr'`` - ``(x, y)`` defines lower-right corner of tile - ``'ul'`` - ``(x, y)`` defines upper-left corner of tile - ``'ur'`` - ``(x, y)`` defines upper-right corner of tile **kwargs Other keywords are passed to :meth:`~matplotlib.collections.PolyCollection` Returns ------- collection : `~matplotlib.collections.PolyCollection` the collection of tiles drawn Examples -------- >>> import numpy >>> from matplotlib import pyplot >>> import gwpy.plot # to get gwpy's Axes >>> x = numpy.arange(10) >>> y = numpy.arange(x.size) >>> w = numpy.ones_like(x) * .8 >>> h = numpy.ones_like(x) * .8 >>> fig = pyplot.figure() >>> ax = fig.gca() >>> ax.tile(x, y, w, h, anchor='ll') >>> pyplot.show() """ # get color and sort if color is not None and kwargs.get('c_sort', True): sortidx = color.argsort() x = x[sortidx] y = y[sortidx] w = w[sortidx] h = h[sortidx] color = color[sortidx] # define how to make a polygon for each tile if anchor == 'll': def _poly(x, y, w, h): return ((x, y), (x, y+h), (x+w, y+h), (x+w, y)) elif anchor == 'lr': def _poly(x, y, w, h): return ((x-w, y), (x-w, y+h), (x, y+h), (x, y)) elif anchor == 'ul': def _poly(x, y, w, h): return ((x, y-h), (x, y), (x+w, y), (x+w, y-h)) elif anchor == 'ur': def _poly(x, y, w, h): return ((x-w, y-h), (x-w, y), (x, y), (x, y-h)) elif anchor == 'center': def _poly(x, y, w, h): return ((x-w/2., y-h/2.), (x-w/2., y+h/2.), (x+w/2., y+h/2.), (x+w/2., y-h/2.)) else: raise ValueError("Unrecognised tile anchor {!r}".format(anchor)) # build collection cmap = kwargs.pop('cmap', rcParams['image.cmap']) coll = PolyCollection((_poly(*tile) for tile in zip(x, y, w, h)), edgecolors=edgecolors, linewidth=linewidth, **kwargs) if color is not None: coll.set_array(color) coll.set_cmap(cmap) out = self.add_collection(coll) self.autoscale_view() return out
return np.array([vert1, vert2, vert3, vert4]) DPI = 72 fig = plt.figure(figsize=(700 / DPI, 500 / DPI), dpi=DPI) ax = plt.subplot() ax.set_xlim(LEFT, RIGHT) ax.set_ylim(BOTTOM, TOP) # Get the boxes for the valid comparisons verts_valid = [get_polygon(bound) for i,bound in enumerate(C) \ if not C_masked.mask[i]] C_valid = [C_masked.data[i] for i in range(len(C)) if not C_masked.mask[i]] c = PolyCollection(verts_valid) c.set_array(C_masked) c.set_cmap(colormap) # ax.add_collection(c) # Get the boxes for absent actual data verts_invalid = [get_polygon(bound) for i,bound in enumerate(C) \ if C_masked.mask[i]] c = PolyCollection(verts_invalid, hatch=r"./", facecolor='white') ax.add_collection(c) # Draw the box partitions qtree.draw(ax) for area in area_data: plot_gdf(ax, area_data[area]['df_lines'], area_data[area]['df_buses'], 'orangered') plot_gdf(ax, area_data[area]['df_synth'], area_data[area]['df_cords'],
xdict = dict(zip(ni,x)) ydict = dict(zip(ni,y)) s = ncv['hzg_ecosmo_sed2'][:] nv = ncv['nv'][:,:3]-1 verts=[] for nvi in nv: verts.append([(xdict[i+1],ydict[i+1]) for i in nvi]) verts=asarray(verts) f=figure(figsize=(10,10)) p = PolyCollection(verts,closed=True,edgecolor='none') p.set_array(s[0]) p.set_cmap(cm.RdYlBu_r) pn = PolyCollection(verts,closed=True,edgecolor='none') pn.set_array(arange(len(s[0]))) pn.set_cmap(cm.RdYlBu_r) ax=axes([0.1,0.75,0.88,0.2]) ax.add_collection(p,autolim=True) autoscale() colorbar(p) ax=axes([0.1,0.25,0.88,0.2]) ax.add_collection(pn,autolim=True) autoscale() colorbar(pn)
def _gen2d3d(*args, **pltkwargs): # UPDATE """ Abstract layout for 2d plot. For convienence, a few special labels, colorbar and background keywords have been implemented. If these are not adequate, it one can add custom colorbars, linelabels background images etc... easily just by using respecitve calls to plot (plt.colorbar(), plt.imshow(), plt.clabel() ); my implementations are only for convienences and some predefined, cool styles. countours: Number of desired contours from output. label: Predefined label types. For now, only integer values 1,2. Use plt.clabel to add a custom label. background: Integers 1,2 will add gray or autumn colormap under contour plot. Use plt.imgshow() to generate custom background, or pass a PIL-opened image (note, proper image scaling not yet implemented). c_mesh, r_mesh: These control how man column and row iso lines will be projected onto the 3d plot. For example, if c_mesh=10, then 10 isolines will be plotted as columns, despite the actual length of the columns. Alternatively, one can pass in c_stride directly, which is a column step size rather than an absolute number, and c_mesh will be disregarded. fill: bool (False) Fill between contour lines. **pltkwargs: Will be passed directly to plt.contour(). Returns ------- tuple: (Axes, SurfaceFunction) Returns axes object and the surface function (e.g. contours for contour plot. Surface for surface plot. """ # Use a label mapper to allow for datetimes in any plot x/y axis _x_dti = _ = _y_dti = False # Passed Spectra if len(args) == 1: ts = args[0] try: index = np.array([dates.date2num(x) for x in ts.index]) _x_dti = True except AttributeError: index = ts.index.values #VALUES NECESSARY FOR POLY CMAP try: cols = np.array([dates.date2num(x) for x in ts.columns]) _y_dti = True except AttributeError: cols = ts.columns.values #VALUES NECESSARY FOR POLY CMAP yy, xx = np.meshgrid(cols, index) # Passed xx, yy, ts/zz elif len(args) == 3: xx, yy, ts = args cols, index = ts.columns.values, ts.index.values else: raise PlotError( "Please pass a single spectra, or xx, yy, zz. Got %s args" % len(args)) # Boilerplate from basic_plots._genplot(); could refactor as decorator xlabel = pltkwargs.pop('xlabel', '') ylabel = pltkwargs.pop('ylabel', '') zlabel = pltkwargs.pop('zlabel', '') title = pltkwargs.pop('title', '') labelsize = pltkwargs.pop('labelsize', pvconfig.LABELSIZE) #Can also be ints titlesize = pltkwargs.pop('titlesize', pvconfig.TITLESIZE) # Choose plot kind kind = pltkwargs.pop('kind', 'contour') grid = pltkwargs.pop('grid', True) # pltkwargs.setdefault('legend', False) #(any purpose in 2d?) # LEGEND FOR 2D PLOT: http://stackoverflow.com/questions/10490302/how-do-you-create-a-legend-for-a-contour-plot-in-matplotlib pltkwargs.setdefault('linewidth', 1) cbar = pltkwargs.pop('cbar', False) outline = pltkwargs.pop('outline', None) if outline: if kind != 'surf' and kind != 'waterfall': raise PlotError( '"outline" is only valid for "surf" and "waterfall"' ' plots. Please use color/cmap for all other color' ' designations.') fig = pltkwargs.pop('fig', None) ax = pltkwargs.pop('ax', None) fill = pltkwargs.pop('fill', pvconfig.FILL_CONTOUR) xlim = pltkwargs.pop('xlim', None) ylim = pltkwargs.pop('ylim', None) zlim = pltkwargs.pop('zlim', None) #Private attributes _modifyax = pltkwargs.pop('_modifyax', True) contours = pltkwargs.pop('contours', pvconfig.NUM_CONTOURS) label = pltkwargs.pop('label', None) projection = None if kind in PLOTPARSER.plots_3d: projection = '3d' elev = pltkwargs.pop('elev', 35) azim = pltkwargs.pop('azim', -135) view = pltkwargs.pop('view', None) if view: if view == 1: elev, azim = 35, -135 elif view == 2: elev, azim = 35, -45 elif view == 3: elev, azim = 20, -10 # Side view elif view == 4: elev, azim = 20, -170 elif view == 5: elev, azim = 0, -90 elif view == 6: elev, azim = 65, -90 else: raise PlotError('View must be between 1 and 6; otherwise set' ' "elev" and "azim" keywords.') # Orientation of zlabel (doesn't work...) _zlabel_rotation = 0.0 if azim < 0: _zlabel_rotation = 90.0 if 'mesh' in pltkwargs: pltkwargs['c_mesh'] = pltkwargs['r_mesh'] = pltkwargs.pop('mesh') # Defaults will be ignored if mesh or ciso in kwargs ciso_default = pvconfig.C_MESH if len(ts.columns) < ciso_default: ciso_default = len(ts.columns) riso_default = pvconfig.R_MESH if len(ts.index) < riso_default: riso_default = len(ts.index) c_mesh = pltkwargs.pop('c_mesh', ciso_default) r_mesh = pltkwargs.pop('r_mesh', riso_default) if c_mesh > ts.shape[1] or c_mesh < 0: raise PlotError( '"c_mesh/column mesh" must be between 0 and %s, got "%s"' % (ts.shape[1], c_mesh)) if r_mesh > ts.shape[0] or r_mesh < 0: raise PlotError( '"r_mesh/row mesh" must be between 0 and %s, got "%s"' % (ts.shape[0], r_mesh)) if c_mesh == 0: cstride = 0 else: cstride = _ir(ts.shape[1] / float(c_mesh)) if r_mesh == 0: rstride = 0 else: rstride = _ir(ts.shape[0] / float(r_mesh)) pltkwargs.setdefault('cstride', cstride) pltkwargs.setdefault('rstride', rstride) elif kind == 'contour': pass else: raise PlotError('_gen2d3d invalid kind: "%s". ' 'Choose from %s' % (kind, PLOTPARSER.plots_2d_3d)) # Is this the best logic for 2d/3d fig? if not ax: f = plt.figure() # ax = f.gca(projection=projection) ax = f.add_subplot(111, projection=projection) if not fig: fig = f # PLT.CONTOUR() doesn't take 'color'; rather, takes 'colors' for now if 'color' in pltkwargs: if kind == 'contour' or kind == 'contour3d': pltkwargs['colors'] = pltkwargs.pop('color') # Convienence method to pass in string colors if 'colormap' in pltkwargs: pltkwargs['cmap'] = pltkwargs.pop('colormap') if 'cmap' in pltkwargs: if isinstance(pltkwargs['cmap'], basestring): pltkwargs['cmap'] = pu.cmget(pltkwargs['cmap']) # Contour Plots # ------------- # Broken background image ### More here http://matplotlib.org/examples/pylab_examples/image_demo3.html ### # Refactored with xx, yy instead of df.columns/index UNTESTED #if background: #xmin, xmax, ymin, ymax = xx.min(), xx.max(), yy.min(), yy.max() ## Could try rescaling contour rather than image: ##http://stackoverflow.com/questions/10850882/pyqt-matplotlib-plot-contour-data-on-top-of-picture-scaling-issue #if background==1: #im = ax.imshow(ts, interpolation='bilinear', origin='lower', #cmap=cm.gray, extent=(xmin, xmax, ymin, ymax)) #### This will take a custom image opened in PIL or it will take plt.imshow() returned from somewhere else #else: #try: #im = ax.imshow(background) #### Perhaps image was not correctly opened #except Exception: #raise badvalue_error(background, 'integer 1,2 or a PIL-opened image') # Note this overwrites the 'contours' variable from an int to array if kind == 'contour' or kind == 'contour3d': # Cornercase datetimeindex and offest from add projection hack try: pltkwargs['offset'] = dates.date2num(pltkwargs['offset']) except Exception: pass if fill: #Values of DTI doesn't work mappable = ax.contourf(xx, yy, ts.values, contours, **pltkwargs) #linewidths is a pltkwargs arg else: mappable = ax.contour(xx, yy, ts.values, contours, **pltkwargs) ### Pick a few label styles to choose from. if label: if label == 1: ax.clabel(inline=1, fontsize=10) elif label == 2: ax.clabel(levels[1::2], inline=1, fontsize=10) #label every second line else: raise PlotError(label, 'integer of value 1 or 2') elif kind == 'surf': mappable = ax.plot_surface(xx, yy, ts, **pltkwargs) if outline: try: pltkwargs['cmap'] = pu.cmget(outline) except Exception: #Don't change; attribute error fails when outline=None pltkwargs['color'] = outline pltkwargs.pop('cmap') custom_wireframe(ax, xx, yy, ts, **pltkwargs) # Wires are thrown out, since mappable is the surface, and only mappable returned elif kind == 'wire': pltkwargs.setdefault('color', 'black') mappable = custom_wireframe(ax, xx, yy, ts, **pltkwargs) elif kind == 'waterfall': # Parse outline color (if colormap, error!) try: pu.cmget(outline) except Exception: pltkwargs['edgecolors'] = outline else: raise PlotError( 'Waterfall "outline" must be a solid color, not colormap.') pltkwargs.setdefault('closed', False) alpha = pltkwargs.setdefault('alpha', None) # Need to handle cmap/colors a bit differently for PolyCollection API if 'color' in pltkwargs: pltkwargs['facecolors'] = pltkwargs.pop('color') cmap = pltkwargs.setdefault('cmap', None) if alpha is None: #as opposed to 0 alpha = 0.6 * (13.0 / ts.shape[1]) if alpha > 0.6: alpha = 0.6 #Delete stride keywords (waterfall doesn't have strides: not a surface!) for key in ['cstride', 'rstride']: try: del pltkwargs[key] except KeyError: pass # Verts are index dotted with data verts = [] for col in ts.columns: values = ts[col] values[0], values[-1] = values.min().min(), values.min().min() verts.append(list(zip(ts.index, values))) mappable = PolyCollection(verts, **pltkwargs) if cmap: mappable.set_array( cols) #If set array in __init__, autogens a cmap! mappable.set_cmap(pltkwargs['cmap']) mappable.set_alpha(alpha) #zdir is the direction used to plot; dont' fully understand ax.add_collection3d(mappable, zs=cols, zdir='x') # custom limits/labels polygon plot (reverse x,y) if not ylim: ylim = (max(index), min(index)) #REVERSE AXIS FOR VIEWING PURPOSES if not xlim: xlim = (min(cols), max(cols)) #x if not zlim: zlim = (min(ts.min()), max(ts.max()) ) #How to get absolute min/max of ts values # Reverse labels/DTI call for correct orientaion HACK HACK HACK xlabel, ylabel = ylabel, xlabel _x_dti, _y_dti = _y_dti, _x_dti azim = -1 * azim # General Features # ---------------- # Some applications (like add_projection) shouldn't alther axes features if not _modifyax: return (ax, mappable) if cbar: # Do I want colorbar outside of fig? Wouldn't be better on axes? try: cb = fig.colorbar(mappable, ax=ax) # Label colorbar on contour since no 3d-zlabel if kind == 'contour': cb.set_label(zlabel) except Exception: raise PlotError("Colorbar failed; did you pass a colormap?") if grid: if grid == True: ax.grid() else: ax.grid(color=grid) #Let's any supported color in # Format datetime axis # ------------------- if _x_dti: ax.xaxis.set_major_formatter(mplticker.FuncFormatter(format_date)) # Uncomment for custom 3d timestamp orientation # if projection: # for t1 in ax.yaxis.get_ticklabels(): # t1.set_ha('right') # t1.set_rotation(30) # ax.yaxis._axinfo['label']['space_factor'] = _TIMESTAMPPADDING if _y_dti: ax.yaxis.set_major_formatter(mplticker.FuncFormatter(format_date)) # Uncomment for custom 3d timestamp orientation # if projection: # for t1 in ax.yaxis.get_ticklabels(): # t1.set_ha('right') # t1.set_rotation(30) # ax.yaxis._axinfo['label']['space_factor'] = _TIMESTAMPPADDING if xlim: ax.set_xlim3d(xlim) if ylim: ax.set_ylim3d(ylim) if zlim: ax.set_zlim3d(zlim) # Set elevation/azimuth for 3d plots if projection: ax.view_init(elev, azim) ax.set_zlabel(zlabel, fontsize=labelsize, rotation=_zlabel_rotation) ax.set_xlabel(xlabel, fontsize=labelsize) ax.set_ylabel(ylabel, fontsize=labelsize) ax.set_title(title, fontsize=titlesize) # Return Ax, contours/surface/polygons etc... return (ax, mappable)
def plot_probe(probe, ax=None, contacts_colors=None, with_channel_index=False, with_contact_id=False, with_device_index=False, text_on_contact=None, first_index='auto', contacts_values=None, cmap='viridis', title=True, contacts_kargs={}, probe_shape_kwargs={}, xlims=None, ylims=None, zlims=None, show_channel_on_click=False): """Plot a Probe object. Generates a 2D or 3D axis, depending on Probe.ndim Parameters ---------- probe : Probe The probe object ax : matplotlib.axis, optional The axis to plot the probe on. If None, an axis is created, by default None contacts_colors : matplotlib color, optional The color of the contacts, by default None with_channel_index : bool, optional If True, channel indices are displayed on top of the channels, by default False with_contact_id : bool, optional If True, channel ids are displayed on top of the channels, by default False with_device_index : bool, optional If True, device channel indices are displayed on top of the channels, by default False text_on_contact: None or list or numpy.array Addintional text to plot on each contact first_index : str, optional The first index of the contacts, by default 'auto' (taken from channel ids) contacts_values : np.array, optional Values to color the contacts with, by default None cmap : str, optional [description], by default 'viridis' title : bool, optional If True, the axis title is set to the probe name, by default True contacts_kargs : dict, optional Dict with kwargs for contacts (e.g. alpha, edgecolor, lw), by default {} probe_shape_kwargs : dict, optional Dict with kwargs for probe shape (e.g. alpha, edgecolor, lw), by default {} xlims : tuple, optional Limits for x dimension, by default None ylims : tuple, optional Limits for y dimension, by default None zlims : tuple, optional Limits for z dimension, by default None show_channel_on_click : bool, optional If True, the channel information is shown upon click, by default False Returns ------- poly : PolyCollection The polygon collection for contacts poly_contour : PolyCollection The polygon collection for the probe shape """ import matplotlib.pyplot as plt if probe.ndim == 2: from matplotlib.collections import PolyCollection elif probe.ndim == 3: from mpl_toolkits.mplot3d.art3d import Poly3DCollection if ax is None: if probe.ndim == 2: fig, ax = plt.subplots() ax.set_aspect('equal') else: fig = plt.figure() ax = fig.add_subplot(1, 1, 1, projection='3d') else: fig = ax.get_figure() if first_index == 'auto': if 'first_index' in probe.annotations: first_index = probe.annotations['first_index'] elif probe.annotations.get('manufacturer', None) == 'neuronexus': # neuronexus is one based indexing first_index = 1 else: first_index = 0 assert first_index in (0, 1) _probe_shape_kwargs = dict( facecolor='green', edgecolor='k', lw=0.5, alpha=0.3) _probe_shape_kwargs.update(probe_shape_kwargs) _contacts_kargs = dict(alpha=0.7, edgecolor=[0.3, 0.3, 0.3], lw=0.5) _contacts_kargs.update(contacts_kargs) n = probe.get_contact_count() if contacts_colors is None and contacts_values is None: contacts_colors = ['orange'] * n elif contacts_colors is not None: contacts_colors = contacts_colors elif contacts_values is not None: contacts_colors = None vertices = probe.get_contact_vertices() if probe.ndim == 2: poly = PolyCollection( vertices, color=contacts_colors, **_contacts_kargs) ax.add_collection(poly) elif probe.ndim == 3: poly = Poly3DCollection( vertices, color=contacts_colors, **_contacts_kargs) ax.add_collection3d(poly) if contacts_values is not None: poly.set_array(contacts_values) poly.set_cmap(cmap) if show_channel_on_click: assert probe.ndim == 2, 'show_channel_on_click works only for ndim=2' def on_press(event): return _on_press(probe, event) fig.canvas.mpl_connect('button_press_event', on_press) fig.canvas.mpl_connect('button_release_event', on_release) # probe shape planar_contour = probe.probe_planar_contour if planar_contour is not None: if probe.ndim == 2: poly_contour = PolyCollection( [planar_contour], **_probe_shape_kwargs) ax.add_collection(poly_contour) elif probe.ndim == 3: poly_contour = Poly3DCollection( [planar_contour], **_probe_shape_kwargs) ax.add_collection3d(poly_contour) else: poly_contour = None if text_on_contact is not None: text_on_contact = np.asarray(text_on_contact) assert text_on_contact.size == probe.get_contact_count() if with_channel_index or with_contact_id or with_device_index or text_on_contact is not None: if probe.ndim == 3: raise NotImplementedError('Channel index is 2d only') for i in range(n): txt = [] if with_channel_index: txt.append(f'{i + first_index}') if with_contact_id and probe.contact_ids is not None: contact_id = probe.contact_ids[i] txt.append(f'id{contact_id}') if with_device_index and probe.device_channel_indices is not None: chan_ind = probe.device_channel_indices[i] txt.append(f'dev{chan_ind}') if text_on_contact is not None: txt.append(f'{text_on_contact[i]}') txt = '\n'.join(txt) x, y = probe.contact_positions[i] ax.text(x, y, txt, ha='center', va='center', clip_on=True) if xlims is None or ylims is None or (zlims is None and probe.ndim == 3): xlims, ylims, zlims = get_auto_lims(probe) ax.set_xlim(*xlims) ax.set_ylim(*ylims) if probe.si_units == "um": unit_str = "($\mu m$)" else: unit_str = f"({probe.si_units})" ax.set_xlabel(f'x {unit_str}', fontsize=15) ax.set_ylabel(f'y {unit_str}', fontsize=15) if probe.ndim == 3: ax.set_zlim(zlims) ax.set_zlabel('z') if probe.ndim == 2: ax.set_aspect('equal') if title: ax.set_title(probe.get_title()) return poly, poly_contour
def Plot(self, filename=None, edgeColor='none'): """ Plot the optimized shape Parameters ---------- None Returns ------- None """ fig = plt.figure("Result", figsize=(12,12), clear=True) if self.fem.nDof == 2: collection = PolyCollection(self.fem.nodes[self.fem.elements], edgecolors=edgeColor) collection.set_array(self.densities) collection.set_cmap('gray_r') collection.set_clim(vmin=0, vmax=1) ax = fig.gca() ax.add_collection(collection) ax.set_xlim(self.fem.nodes[:,0].min(), self.fem.nodes[:,0].max()) ax.set_ylim(self.fem.nodes[:,1].min(), self.fem.nodes[:,1].max()) ratio = ((ax.get_ylim()[1] - ax.get_ylim()[0]) / (ax.get_xlim()[1] - ax.get_xlim()[0])) if ratio < 1: fig.set_figheight(ratio * fig.get_figwidth()) else: fig.set_figwidth(fig.get_figheight() / ratio) ax.axis('off') elif self.fem.nDof == 3: if not hasattr(self, 'facePairs'): face = np.array([0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 5, 4, 3, 2, 6, 7, 0, 4, 7, 3, 1, 5, 6, 2]).reshape(6,4) faces = np.concatenate([el[face] for el in self.fem.elements]) order = np.arange(faces.shape[0]) for i in range(faces.shape[1]): ind = np.argsort(faces[:,i], kind='stable') faces = faces[ind,:] order = order[ind] elements = np.tile(np.arange(self.fem.nElem).reshape(-1, 1), (1,6)).ravel()[order] self.facePairs = [] self.faceNodes = [] i = 0 while i < faces.shape[0] - 1: if (faces[i,:] == faces[i+1,:]).all(): self.facePairs.append([elements[i], elements[i+1]]) self.faceNodes.append(faces[i]) i += 2 else: self.facePairs.append([elements[i], self.fem.nElem]) self.faceNodes.append(faces[i]) i += 1 if i < faces.shape[0]: self.facePairs.append([elements[i], self.fem.nElem]) self.faceNodes.append(faces[i]) self.facePairs = np.array(self.facePairs) self.faceNodes = np.array(self.faceNodes) densities = np.append(self.densities, [0]) faces = np.logical_xor(densities[self.facePairs[:,0]] > self.dens_thresh, densities[self.facePairs[:,1]] > self.dens_thresh) print("Plotting %i faces" % (self.faceNodes[faces].size // 4)) collection = Poly3DCollection(self.fem.nodes[self.faceNodes[faces].reshape(-1,4)], edgecolors=edgeColor) densities = densities[self.facePairs].max(axis=1)[faces] collection.set_array(densities) # elements = self.fem.elements[self.densities > self.dens_thresh,:] # collection = Poly3DCollection(self.fem.nodes[elements[:,face].reshape(-1,4)], # edgecolors="k") # densities = np.tile(self.densities[self.densities > self.dens_thresh], (6,1)).T.ravel() # collection.set_array(densities) collection.set_cmap('gray_r') collection.set_clim(vmin=0, vmax=1) ax = fig.gca(projection='3d') ax.add_collection3d(collection) ax.set_xlim(np.min(self.fem.nodes), np.max(self.fem.nodes)) ax.set_ylim(np.min(self.fem.nodes), np.max(self.fem.nodes)) ax.set_zlim(np.min(self.fem.nodes), np.max(self.fem.nodes)) plt.draw() plt.pause(0.01) if filename: plt.tight_layout() plt.savefig(filename, bbox_inches='tight', pad_inches=0)
def _gen2d3d(*args, **pltkwargs): # UPDATE """ Abstract layout for 2d plot. For convienence, a few special labels, colorbar and background keywords have been implemented. If these are not adequate, it one can add custom colorbars, linelabels background images etc... easily just by using respecitve calls to plot (plt.colorbar(), plt.imshow(), plt.clabel() ); my implementations are only for convienences and some predefined, cool styles. countours: Number of desired contours from output. label: Predefined label types. For now, only integer values 1,2. Use plt.clabel to add a custom label. background: Integers 1,2 will add gray or autumn colormap under contour plot. Use plt.imgshow() to generate custom background, or pass a PIL-opened image (note, proper image scaling not yet implemented). c_iso, r_iso: These control how man column and row iso lines will be projected onto the 3d plot. For example, if c_iso=10, then 10 isolines will be plotted as columns, despite the actual length of the columns. Alternatively, one can pass in c_stride directly, which is a column step size rather than an absolute number, and c_iso will be disregarded. fill: bool (False) Fill between contour lines. **pltkwargs: Will be passed directly to plt.contour(). Returns ------- tuple: (Axes, SurfaceFunction) Returns axes object and the surface function (e.g. contours for contour plot. Surface for surface plot. """ # Use a label mapper to allow for datetimes in any plot x/y axis _x_dti = _ = _y_dti = False # Passed Spectra if len(args) == 1: ts = args[0] try: index = np.array([dates.date2num(x) for x in ts.index]) _x_dti = True except AttributeError: index = ts.index.values #VALUES NECESSARY FOR POLY CMAP try: cols = np.array([dates.date2num(x) for x in ts.columns]) _y_dti = True except AttributeError: cols = ts.columns.values #VALUES NECESSARY FOR POLY CMAP yy, xx = np.meshgrid(cols, index) # Passed xx, yy, ts/zz elif len(args) == 3: xx, yy, ts = args cols, index = ts.columns.values, ts.index.values else: raise PlotError("Please pass a single spectra, or xx, yy, zz. Got %s args" % len(args)) # Boilerplate from basic_plots._genplot(); could refactor as decorator xlabel = pltkwargs.pop('xlabel', '') ylabel = pltkwargs.pop('ylabel', '') zlabel = pltkwargs.pop('zlabel', '') title = pltkwargs.pop('title', '') # Choose plot kind kind = pltkwargs.pop('kind', 'contour') grid = pltkwargs.pop('grid', '') # pltkwargs.setdefault('legend', False) #(any purpose in 2d?) # LEGEND FOR 2D PLOT: http://stackoverflow.com/questions/10490302/how-do-you-create-a-legend-for-a-contour-plot-in-matplotlib pltkwargs.setdefault('linewidth', 1) cbar = pltkwargs.pop('cbar', False) fig = pltkwargs.pop('fig', None) ax = pltkwargs.pop('ax', None) fill = pltkwargs.pop('fill', False) xlim = pltkwargs.pop('xlim', None) ylim = pltkwargs.pop('ylim', None) zlim = pltkwargs.pop('zlim', None) #Private attributes _modifyax = pltkwargs.pop('_modifyax', True) contours = pltkwargs.pop('contours', 6) label = pltkwargs.pop('label', None) projection = None if kind in PLOTPARSER.plots_3d: projection = '3d' elev = pltkwargs.pop('elev', 35) azim = pltkwargs.pop('azim', -135) view = pltkwargs.pop('view', None) if view: if view == 1: elev, azim = 35, -135 elif view == 2: elev, azim = 35, -45 elif view == 3: elev, azim = 20, -10 # Side view elif view == 4: elev, azim = 20, -170 elif view == 5: elev, azim = 0,-90 elif view == 6: elev, azim = 65, -90 else: raise PlotError('View must be between 1 and 6; otherwise set' ' "elev" and "azim" keywords.') # Orientation of zlabel (doesn't work...) _zlabel_rotation = 0.0 if azim < 0: _zlabel_rotation = 90.0 c_iso = pltkwargs.pop('c_iso', 10) r_iso = pltkwargs.pop('r_iso', 10) if c_iso > ts.shape[1] or c_iso < 0: raise PlotError('"c_iso" must be between 0 and %s, got "%s"' % (ts.shape[1], c_iso)) if r_iso > ts.shape[0] or r_iso < 0: raise PlotError('"r_iso" must be between 0 and %s, got "%s"' % (ts.shape[0], r_iso)) if c_iso == 0: cstride = 0 else: cstride = _ir(ts.shape[1]/float(c_iso) ) if r_iso == 0: rstride = 0 else: rstride = _ir(ts.shape[0]/float(r_iso) ) pltkwargs.setdefault('cstride', cstride) pltkwargs.setdefault('rstride', rstride) elif kind == 'contour': pass else: raise PlotError('_gen2d3d invalid kind: "%s". ' 'Choose from %s' % (kind, PLOTPARSER.plots_2d_3d)) # Is this the best logic for 2d/3d fig? if not ax: f = plt.figure() # ax = f.gca(projection=projection) ax = f.add_subplot(111, projection=projection) if not fig: fig = f labelsize = pltkwargs.pop('labelsize', 'medium') #Can also be ints titlesize = pltkwargs.pop('titlesize', 'large') ticksize = pltkwargs.pop('ticksize', '') #Put in default and remove bool gate # PLT.CONTOUR() doesn't take 'color'; rather, takes 'colors' for now if 'color' in pltkwargs: if kind == 'contour': pltkwargs['colors'] = pltkwargs.pop('color') # Convienence method to pass in string colors if 'colormap' in pltkwargs: pltkwargs['cmap'] = pltkwargs.pop('colormap') if 'cmap' in pltkwargs: if isinstance(pltkwargs['cmap'], basestring): pltkwargs['cmap'] = pu.cmget(pltkwargs['cmap']) # Contour Plots # ------------- # Broken background image ### More here http://matplotlib.org/examples/pylab_examples/image_demo3.html ### # Refactored with xx, yy instead of df.columns/index UNTESTED #if background: #xmin, xmax, ymin, ymax = xx.min(), xx.max(), yy.min(), yy.max() ## Could try rescaling contour rather than image: ##http://stackoverflow.com/questions/10850882/pyqt-matplotlib-plot-contour-data-on-top-of-picture-scaling-issue #if background==1: #im = ax.imshow(ts, interpolation='bilinear', origin='lower', #cmap=cm.gray, extent=(xmin, xmax, ymin, ymax)) #### This will take a custom image opened in PIL or it will take plt.imshow() returned from somewhere else #else: #try: #im = ax.imshow(background) #### Perhaps image was not correctly opened #except Exception: #raise badvalue_error(background, 'integer 1,2 or a PIL-opened image') # Note this overwrites the 'contours' variable from an int to array if kind == 'contour' or kind == 'contour3d': if fill: mappable = ax.contourf(xx, yy, ts, contours, **pltkwargs) #linewidths is a pltkwargs arg else: mappable = ax.contour(xx, yy, ts, contours, **pltkwargs) ### Pick a few label styles to choose from. if label: if label==1: ax.clabel(inline=1, fontsize=10) elif label==2: ax.clabel(levels[1::2], inline=1, fontsize=10) #label every second line else: raise PlotError(label, 'integer of value 1 or 2') elif kind == 'surf': mappable = ax.plot_surface(xx, yy, ts, **pltkwargs) # pltkwargs.pop('edgecolors') wires = overload_plot_wireframe(ax, xx, yy, ts, **pltkwargs) # print np.shape(wires._segments3d) wires = wire_cmap(wires, ax, cmap='jet') elif kind == 'wire': pltkwargs.setdefault('color', 'black') mappable = overload_plot_wireframe(ax, xx, yy, ts, **pltkwargs) elif kind == 'waterfall': edgecolors = pltkwargs.setdefault('edgecolors', None) pltkwargs.setdefault('closed', False) alpha = pltkwargs.setdefault('alpha', None) # Need to handle cmap/colors a bit differently for PolyCollection API if 'color' in pltkwargs: pltkwargs['facecolors']=pltkwargs.pop('color') cmap = pltkwargs.setdefault('cmap', None) if alpha is None: #as opposed to 0 alpha = 0.6 * (13.0/ts.shape[1]) if alpha > 0.6: alpha = 0.6 #Delete stride keywords for key in ['cstride', 'rstride']: try: del pltkwargs[key] except KeyError: pass # Verts are index dotted with data verts = [] for col in ts.columns: values = ts[col] values[0], values[-1] = values.min().min(), values.min().min() verts.append(list(zip(ts.index, values))) mappable = PolyCollection(verts, **pltkwargs) if cmap: mappable.set_array(cols) #If set array in __init__, autogens a cmap! mappable.set_cmap(pltkwargs['cmap']) mappable.set_alpha(alpha) #zdir is the direction used to plot; dont' fully understand ax.add_collection3d(mappable, zs=cols, zdir='x' ) # custom limits/labels polygon plot (reverse x,y) if not ylim: ylim = (max(index), min(index)) #REVERSE AXIS FOR VIEWING PURPOSES if not xlim: xlim = (min(cols), max(cols)) #x if not zlim: zlim = (min(ts.min()), max(ts.max())) #How to get absolute min/max of ts values # Reverse labels/DTI call for correct orientaion HACK HACK HACK xlabel, ylabel = ylabel, xlabel _x_dti, _y_dti = _y_dti, _x_dti azim = -1 * azim # General Features # ---------------- # Some applications (like add_projection) shouldn't alther axes features if not _modifyax: return (ax, mappable) if cbar: # Do I want colorbar outside of fig? Wouldn't be better on axes? try: fig.colorbar(mappable, ax=ax) except Exception: raise PlotError("Colorbar failed; did you pass a colormap?") if grid: ax.grid() # Format datetime axis # ------------------- if _x_dti: ax.xaxis.set_major_formatter(mplticker.FuncFormatter(format_date)) # Uncomment for custom 3d timestamp orientation # if projection: # for t1 in ax.yaxis.get_ticklabels(): # t1.set_ha('right') # t1.set_rotation(30) # ax.yaxis._axinfo['label']['space_factor'] = _TIMESTAMPPADDING if _y_dti: ax.yaxis.set_major_formatter(mplticker.FuncFormatter(format_date)) # Uncomment for custom 3d timestamp orientation # if projection: # for t1 in ax.yaxis.get_ticklabels(): # t1.set_ha('right') # t1.set_rotation(30) # ax.yaxis._axinfo['label']['space_factor'] = _TIMESTAMPPADDING if xlim: ax.set_xlim3d(xlim) if ylim: ax.set_ylim3d(ylim) if zlim: ax.set_zlim3d(zlim) # Set elevation/azimuth for 3d plots if projection: ax.view_init(elev, azim) ax.set_zlabel(zlabel, fontsize=labelsize, rotation= _zlabel_rotation) ax.set_xlabel(xlabel, fontsize=labelsize) ax.set_ylabel(ylabel, fontsize=labelsize) ax.set_title(title, fontsize=titlesize) # Return Ax, contours/surface/polygons etc... return (ax, mappable)
def plot_probe(probe, ax=None, electrode_colors=None, with_channel_index=False, first_index='auto', electrode_values=None, cmap='viridis', title=True, electrodes_kargs={}, probe_shape_kwargs={}): """ plot one probe. switch 2d 3d depending the Probe.ndim """ import matplotlib.pyplot as plt if probe.ndim == 2: from matplotlib.collections import PolyCollection elif probe.ndim == 3: from mpl_toolkits.mplot3d.art3d import Poly3DCollection if ax is None: if probe.ndim == 2: fig, ax = plt.subplots() ax.set_aspect('equal') else: fig = plt.figure() ax = fig.add_subplot(1, 1, 1, projection='3d') if first_index == 'auto': if 'first_index' in probe.annotations: first_index = probe.annotations['first_index'] elif probe.annotations.get('manufacturer', None) == 'neuronexus': # neuronexus is one based indexing first_index = 1 else: first_index = 0 assert first_index in (0, 1) _probe_shape_kwargs = dict(facecolor='green', edgecolor='k', lw=0.5, alpha=0.3) _probe_shape_kwargs.update(probe_shape_kwargs) _electrodes_kargs = dict(alpha=0.7, edgecolor=[0.3, 0.3, 0.3], lw=0.5) _electrodes_kargs.update(electrodes_kargs) n = probe.get_electrode_count() if electrode_colors is None and electrode_values is None: electrode_colors = ['orange'] * n elif electrode_colors is not None: electrode_colors = electrode_colors elif electrode_values is not None: electrode_colors = None # electrodes positions = probe.electrode_positions min_, max_ = np.min(positions), np.max(positions) vertices = probe.get_electrodes_vertices() if probe.ndim == 2: poly = PolyCollection(vertices, color=electrode_colors, **_electrodes_kargs) ax.add_collection(poly) elif probe.ndim == 3: poly = poly3d = Poly3DCollection(vertices, color=electrode_colors, **_electrodes_kargs) ax.add_collection3d(poly3d) if electrode_values is not None: poly.set_array(electrode_values) poly.set_cmap(cmap) # probe shape vertices = probe.probe_planar_contour if vertices is not None: if probe.ndim == 2: poly_contour = PolyCollection([vertices], **_probe_shape_kwargs) ax.add_collection(poly_contour) elif probe.ndim == 3: poly_contour = Poly3DCollection([vertices], **_probe_shape_kwargs) ax.add_collection3d(poly_contour) min_, max_ = np.min(vertices), np.max(vertices) else: poly_contour = None if with_channel_index: if probe.ndim == 3: raise NotImplementedError('Channel index is 2d only') for i in range(n): x, y = probe.electrode_positions[i] if probe.device_channel_indices is None: txt = f'{i + first_index}' else: chan_ind = probe.device_channel_indices[i] txt = f'prb{i + first_index}\ndev{chan_ind}' ax.text(x, y, txt, ha='center', va='center') min_ -= 40 max_ += 40 ax.set_xlim(min_, max_) ax.set_ylim(min_, max_) ax.set_xlabel('x') ax.set_ylabel('y') if probe.ndim == 3: ax.set_zlim(min_, max_) ax.set_zlabel('z') if probe.ndim == 2: ax.set_aspect('equal') if title: ax.set_title(probe.get_title()) return poly, poly_contour