def cross_section(self,func,anchor, zrot,yrot,xrot, length,width, Nl=20,Nw=20, cmap=tcm.slip, lw=1.0, clim=None, func_args=None, func_kwargs=None, **kwargs): if func_args is None: func_args = () if func_kwargs is None: func_kwargs = {} x = np.linspace(0,length,Nl) y = np.linspace(0,-width,Nw) x,y = np.meshgrid(x,y) z = 0.0*x R = rotation3D(zrot,yrot,xrot) p = np.concatenate((x[None,:,:], y[None,:,:], z[None,:,:]), axis=0) p = np.einsum('ij,jkl->kli',R,p) p += anchor p = np.reshape(p,(Nl*Nw,3)) c = func(p,*func_args,**func_kwargs) c = np.reshape(c,(Nl,Nw)) p = np.reshape(p,(Nl,Nw,3)) if not self.cross_section_clim_set: self.sm = cm.ScalarMappable(cmap=cmap) self.sm.set_array(c) if (clim is not None): self.sm.set_clim(clim[0],clim[1]) self.cross_section_clim_set = True cnorm = self.sm.norm(c)#(c - clim[0])/(clim[1] - clim[0]) s = self.plot_surface(p[:,:,0], p[:,:,1], p[:,:,2], shade=False, facecolors=cmap(cnorm), rstride=1,cstride=1, **kwargs) idx1 = np.array([[0,-1],[0,-1]]) idx2 = np.array([[0,0],[-1,-1]]) w = self.plot_wireframe(p[idx1,idx2,0], p[idx1,idx2,1], p[idx1,idx2,2], color='k',lw=lw) return self.sm
def vector_cross_section(self,func,anchor, zrot,yrot,xrot, length,width, Nl=20,Nw=20, lw=1.0, minmag=0.05, arrow_length=1000.0, func_args=None, func_kwargs=None, **kwargs): if func_args is None: func_args = () if func_kwargs is None: func_kwargs = {} x = np.linspace(0,length,Nl) y = np.linspace(0,-width,Nw) x,y = np.meshgrid(x,y) z = 0.0*x R = rotation3D(zrot,yrot,xrot) p = np.concatenate((x[None,:,:], y[None,:,:], z[None,:,:]), axis=0) p = np.einsum('ij,jkl->kli',R,p) p += anchor p = np.reshape(p,(Nl*Nw,3)) c = func(p,*func_args,**func_kwargs) cmag = np.sqrt(c[:,0]**2 + c[:,1]**2 + c[:,2]**2) #c = np.ma.masked_array(c,mask=cmag<minmag) c *= arrow_length ends = p + c mscale = kwargs.pop('mutation_scale',1.0) for i in range(len(p)): if cmag[i] > minmag: artist = Arrow3D(*zip(p[i],ends[i]),mutation_scale=cmag[i]*mscale,**kwargs) self.add_artist(artist)