class Bar: """ Bar (histogram) """ def __init__(self, ax, transform, Z, facecolors="white", edgecolors="black", linewidth=0, clip=False): """ """ self.Z = Z if isinstance(facecolors, np.ndarray): shape = facecolors.shape facecolors = facecolors.reshape(-1, shape[-1]) facecolors = mpl.colors.to_rgba_array(facecolors) self.facecolors = facecolors.reshape(shape[0], shape[1], 4) else: shape = Z.shape self.facecolors = np.zeros((shape[0], shape[1], 4)) self.facecolors[...] = mpl.colors.to_rgba(facecolors) if isinstance(edgecolors, np.ndarray): shape = edgecolors.shape edgecolors = edgecolors.reshape(-1, shape[-1]) edgecolors = mpl.colors.to_rgba_array(edgecolors) self.edgecolors = edgecolors.reshape(shape[0], shape[1], 4) else: shape = Z.shape self.edgecolors = np.zeros((shape[0], shape[1], 4)) self.edgecolors[...] = mpl.colors.to_rgba(edgecolors) self.linewidth = linewidth self.xlim = -0.5, +0.50 self.ylim = -0.5, +0.50 self.zlim = -0.5, +0.50 self.clip = clip # Because all the bars have the same orientation, we can use a hack to # shade each face at once instead of computing individual face lighting. self.shade = np.array([[1.00, 1.00, 0.75, 1.00, 0.50, 1.00]]) self.collection = PolyCollection([], clip_on=self.clip, snap=False) self.update(transform) ax.add_collection(self.collection, autolim=False) def update(self, transform): """ """ Z = self.Z xmin, xmax = self.xlim ymin, ymax = self.ylim zmin, zmax = self.zlim dx, dy = 0.5 * 1 / Z.shape[0], 0.5 * 1 / Z.shape[1] # Each bar is described by 8 vertices and 6 faces V = np.zeros((Z.shape[0], Z.shape[1], 8, 3)) F = np.zeros((Z.shape[0], Z.shape[1], 6, 4), dtype=int) # Face and edge colors for the six faces FC = np.zeros((Z.shape[0], Z.shape[1], 6, 4)) FC[:, :] = self.facecolors.reshape(Z.shape[0], Z.shape[1], 1, 4) FC *= self.shade.T FC[:, :, :, 3] = 1 EC = np.zeros((Z.shape[0], Z.shape[1], 6, 4)) EC[:, :] = self.edgecolors.reshape(Z.shape[0], Z.shape[1], 1, 4) # Build vertices X, Y = np.meshgrid(np.linspace(xmin, xmax, Z.shape[0]), np.linspace(ymin, ymax, Z.shape[1])) V[..., 0] = X.reshape(Z.shape[0], Z.shape[1], 1) V[..., 1] = Y.reshape(Z.shape[0], Z.shape[1], 1) V[:, :, 0] += [+dx, +dy, zmin] V[:, :, 1] += [+dx, -dy, zmin] V[:, :, 2] += [-dx, -dy, zmin] V[:, :, 3] += [-dx, +dy, zmin] V[:, :, 4] += [+dx, +dy, zmin] V[:, :, 5] += [+dx, -dy, zmin] V[:, :, 6] += [-dx, -dy, zmin] V[:, :, 7] += [-dx, +dy, zmin] V[:, :, 4:, 2] += Z.reshape(Z.shape[0], Z.shape[1], 1) # Build faces I = 8 * np.arange(Z.shape[0] * Z.shape[1]) F[:, :] = I.reshape(Z.shape[0], Z.shape[1], 1, 1) F[:, :] += [ [0, 1, 2, 3], # -Z [0, 1, 5, 4], # +X [2, 3, 7, 6], # -X [1, 2, 6, 5], # -Y [0, 3, 7, 4], # +Y [4, 5, 6, 7] ] # +Z # Actual transformation V = V.reshape(-1, 3) V = glm.transform(V[F], transform) #[...,:2] # Depth computation # We combine the global "depth" of the bar (depth of the bottom face) # and the local depth of each face. This trick avoids problems when # sorting all the different faces. Z1 = (V[:, :, 0, :, 2].mean(axis=2)).reshape(Z.shape[0], Z.shape[1], 1) Z2 = (V[..., 2].mean(axis=3) + 10 * Z1).ravel() # Sorting I = np.argsort(-Z2) V = (V[..., :2].reshape(Z.shape[0] * Z.shape[1] * 6, 4, 2)) self.collection.set_verts(V[I]) self.collection.set_facecolors(FC.reshape(-1, 4)[I]) self.collection.set_edgecolors(EC.reshape(-1, 4)[I]) self.collection.set_linewidths(self.linewidth) if self.linewidth == 0.0: self.collection.set_antialiased(False) else: self.collection.set_antialiased(True)
class Mesh(): """ Mesh described by vertices and faces """ def __init__(self, ax, transform, vertices, faces, cmap=None, facecolors="white", edgecolors="black", linewidths=0.5, mode="front"): """ """ self.collection = PolyCollection([], clip_on=False, snap=False) self.vertices = vertices self.faces = faces self.cmap = cmap self.facecolors = mpl.colors.to_rgba_array(facecolors) self.edgecolors = mpl.colors.to_rgba_array(edgecolors) self.linewidths = linewidths self.mode = mode self.update(transform) ax.add_collection(self.collection, autolim=False) def update(self, transform): """ Update mesh according to transform (4x4 array) """ T = glm.transform(self.vertices, transform)[self.faces] Z = -T[:, :, 2].mean(axis=1) if self.cmap is not None: # Facecolors using depth buffer norm = mpl.colors.Normalize(vmin=Z.min(), vmax=Z.max()) facecolors = self.cmap(norm(Z)) else: facecolors = self.facecolors edgecolors = self.edgecolors linewidths = self.linewidths # Back face culling if self.mode == "front": front, back = glm.frontback(T) T, Z = T[front], Z[front] if len(facecolors) == len(self.faces): facecolors = facecolors[front] if len(edgecolors) == len(self.faces): edgecolors = edgecolors[front] # Front face culling elif self.mode == "back": front, back = glm.frontback(T) T, Z = T[back], Z[back] if len(facecolors) == len(faces): facecolors = facecolors[back] if len(edgecolor) == len(faces): edgecolors = edgecolors[back] # Separate 2d triangles from zbuffer triangles = T[:, :, :2] antialiased = linewidths > 0 # Sort triangles according to z buffer I = np.argsort(Z) triangles = triangles[I, :] if len(facecolors) == len(I): facecolors = facecolors[I, :] if len(edgecolors) == len(I): edgecolors = edgecolors[I, :] self.collection.set_verts(triangles) self.collection.set_linewidths(linewidths) self.collection.set_facecolors(facecolors) self.collection.set_edgecolors(edgecolors) self.collection.set_antialiased(antialiased)
[target_dist,target_dist],linewidth=4,linestyle="--",color="k",alpha=0.5) else: if lat_distance: if geomag_distance: target_mag=aacgm.Convert(target_lat,target_lon,0,year=2010,flag=0) p1color_ax.plot([p.date2num(startday),p.date2num(endday)], [target_mag[0],target_mag[0]],linewidth=4,linestyle="--",color="k",alpha=0.5) else: p1color_ax.plot([p.date2num(startday),p.date2num(endday)], [target_lat,target_lat],linewidth=4,linestyle="--",color="k",alpha=0.5) p.axes(p2color_ax) pixs=N.array(p2pixels) colors=p2values coll = PolyCollection(pixs,edgecolors='none',linewidths=0.0,zorder=10) coll.set_antialiased(False) coll.set_facecolor(colors) coll.set_alpha(1) p2color_ax.add_collection(coll) p2color_ax.autoscale_view() p2color_ax.xaxis.set_major_formatter(dateFmt) p2color_ax.set_xlim(p.date2num(startday), p.date2num(endday)) p2color_ax.set_ylim(plotdict["min_rang"],plotdict["max_rang"]) p2color_ax.xaxis.set_major_locator(tick2_minute_interval) p2color_locator=MaxNLocator(nbins=6,prune='both',integer=True) p2color_ax.yaxis.set_major_locator(p2color_locator) p2color_ax.set_xticklabels([]) # p2color_ax.set_xticks([]) p2color_ax.yaxis.grid(True, linestyle='-', which='major', color='lightgrey',alpha=0.5) p2color_ax.xaxis.grid(True, linestyle='-', which='major', color='lightgrey',alpha=0.5) print "pcolor done"