def drawArrows(self, figure, bbox, number, partner): """ draw an 'arrow' from the cell 'number' to 'partner' """ #load the center of mass position n = np.array(number) p = np.array(partner) n.shape = (n.size, ) p.shape = (p.size, ) if p.size > 0 and n.size > 0: com1 = self.file1["features"][str(n[0])]["com"][:] com2 = self.file2["features"][str(p[0])]["com"][:] #write the cell label as text mlab.text3d(com1[2]-bbox[2]+1,com1[1]-bbox[1]+1,com1[0]-bbox[0]+1, str(number), color=(1,1,1), figure=figure) #plot a point where the current cell is mlab.points3d([com1[2]-bbox[2]+1],[com1[1]-bbox[1]+1],[com1[0]-bbox[0]+1],color=(0,0,1), figure=figure) #plot a line to the descendant's center mlab.plot3d([com1[2]-bbox[2]+1,com2[2]-bbox[2]+1], [com1[1]-bbox[1]+1,com2[1]-bbox[1]+1], [com1[0]-bbox[0]+1,com2[0]-bbox[0]+1], tube_radius=0.2, color=(1,0,0), figure=figure) #plot a second line, if there is a split if p.size == 2: com3 = self.file2["features"][str(p[1])]["com"][:] mlab.plot3d([com1[2]-bbox[2]+1,com3[2]-bbox[2]+1], [com1[1]-bbox[1]+1,com3[1]-bbox[1]+1], [com1[0]-bbox[0]+1,com3[0]-bbox[0]+1], tube_radius=0.2, color=(1,0,0), figure=figure)
def topo(j_data, file_name=None, size=(600,600)): u"""Creates the topological view of the molecule. ** Parameters ** j_data : dict Data on the molecule, as deserialized from the scanlog format. file_name : str, optional Base name of the file in which to save the image. size : tuple(int, int), optional The size of the image to save. ** Returns ** figure : mayavi.core.scene.Scene The MayaVi scene containing the visualization. """ figure, normal = _init_scene(j_data) geom = np.array(j_data["results"]["geometry"]["elements_3D_coords_converged"]).reshape((-1,3))/A_to_a0 ## Show labels and numbers ( = indices + 1 ) for i, atom in enumerate(j_data["molecule"]["atoms_Z"]): P, label = geom[i], tab[atom][1] mlab.text3d(P[0] - normal[0], P[1] - normal[1], P[2] - normal[2], label + str(i + 1), color=(0,0,0), scale=0.5, figure=figure) if file_name is not None: mlab.savefig("{}-TOPO.png".format(file_name), figure=figure, size=size) return figure
def plot_3d(lats,lons,depth,vs,runlat,runlon,vtkname='rayleigh_c_u.vtk', annotate_depth=False,coastline=False,annotate_lat=True,annotate_lon=True): """ plot 3d results using mayavi """ if coastline: plot_canada_map() sgrid = get_tvtk_grid(lats,lons,depth,vs) d = mlab.pipeline.add_dataset(sgrid) #sf = mlab.pipeline.surface(d) #gx = mlab.pipeline.grid_plane(d) #gx.grid_plane.axis = 'x' gy = mlab.pipeline.grid_plane(d) gy.grid_plane.axis = 'x' gy.module_manager.scalar_lut_manager.show_scalar_bar = True gy.module_manager.scalar_lut_manager.lut_mode = 'jet' gy.module_manager.scalar_lut_manager.data_range = np.array([ 2. , 4.8]) gy.module_manager.scalar_lut_manager.scalar_bar_representation.maximum_size = np.array([100000, 100000]) gy.module_manager.scalar_lut_manager.scalar_bar_representation.minimum_size = np.array([1, 1]) gy.module_manager.scalar_lut_manager.scalar_bar_representation.position2 = np.array([ 0.08796009, 0.56264591]) gy.module_manager.scalar_lut_manager.scalar_bar_representation.position = np.array([ 0.03396896, 0.39182879]) gy.actor.mapper.progress = 1.0 gy.actor.mapper.scalar_range = np.array([ 0., 1.]) gy.actor.mapper.scalar_visibility = True gy.actor.property.representation = 'surface' gy.grid_plane.position = 6 #gz = mlab.pipeline.grid_plane(d) #gz.grid_plane.axis = 'z' if annotate_lat: for lat in runlat: x,y,z = convert_pt(lat,-58.,10.) txt = mlab.text3d(x,y,z,'%d'%(lat),color=(0,0,0),line_width=10.0) txt.scale = [20,20,20] if annotate_lon: for lon in runlon[1::]: x,y,z = convert_pt(49.,lon,10.) txt = mlab.text3d(x,y,z,'%d'%(lon),color=(0,0,0),line_width=10.0) txt.scale = [20,20,20] if annotate_depth: for dp in [-10,-40,-80,-120]: x,y,z = convert_pt(49,-68.,dp) txt = mlab.text3d(x,y,z,'%d km'%(dp),color=(0,0,0),line_width=10.0) txt.scale = [20,20,20] ### Include 3D screenshot in matplotlib #arr = mlab.screenshot() #import pylab as pl #pl.imshow(arr) #pl.show() mlab.text(0.76,0.86,'49N',width=0.1) mlab.show()
def initialize_rope(label_img, xyz,bgr, plotting=False): # XXX that sucks rope_mask = (label_img==1) | ndi.morphology.binary_dilation(label_img==2,np.ones((5,5))) rope_mask = ndi.binary_opening(rope_mask, np.ones((3,3))) rope_mask = ndi.binary_dilation(rope_mask, np.ones((15,15))) skel = mip.skeletonize(rope_mask) if plotting: cv2.imshow('bgr',bgr.copy()) cv2.imshow('labels',label_img.astype('uint8')*50) cv2.imshow("skel", skel.astype('uint8')*50) cv2.imshow("rope_mask", rope_mask.astype('uint8')*50) cv2.waitKey(5) #start_end = (uv_start, uv_end) = get_cc_centers(label_img==2,2) G = skel2graph(skel) path2d = longest_shortest_path(G) path3d = to3d(path2d,xyz) xyzs_unif = curves.unif_resample(path3d,N_CTRL_PTS,tol=.0025) #us,vs = xyz2uv(xyzs_unif).T #labels = label_img[us,vs] labels = np.ones(xyzs_unif.shape[0]-1,'int') labels[0] = labels[-1] = 2 if plotting: xs,ys,zs = xyzs_unif.T us_skel, vs_skel = np.nonzero(skel) xs_skel, ys_skel, zs_skel = to3d(np.c_[us_skel, vs_skel], xyz).T import matplotlib.pyplot as plt, enthought.mayavi.mlab as mlab mlab.plot3d(xs,ys,zs,np.arange(len(xs)),tube_radius=.0025, colormap='spectral') mlab.points3d(xs_skel, ys_skel, zs_skel, scale_factor=.02, color=(1,1,1),opacity=.1) for (i,path) in enumerate(path3d): mlab.text3d(path[0,0],path[0,1],path[0,2],str(2*i),scale=.01,color=(0,0,0)) mlab.text3d(path[-1,0],path[-1,1],path[-1,2],str(2*i+1),scale=.01,color=(0,0,0)) mlab.show() return xyzs_unif, labels
def test_text3d(self): """ Test the text3d module. """ data = np.random.random((3, 3, 3)) src = mlab.pipeline.scalar_field(data) t = mlab.text3d(0, 0, 0, 'foo', opacity=0.5, scale=2, orient_to_camera=False, color=(0, 0, 0), orientation=(90, 0, 0))
def text3d( x, y, z, s, bcolor=None, bwidth=0.5, bn=16, **kwargs ): """ Mayavi text3d command augmented with poor man's bold. """ from enthought.mayavi import mlab h = [] if bcolor is not None: args = kwargs.copy() args['color'] = bcolor for i in range( bn ): phi = 2.0 * np.pi * i / bn x_ = x + bwidth * np.cos( phi ) y_ = y + bwidth * np.sin( phi ) h += [ mlab.text3d( x_, y_, z, s, **args ) ] h[-1].actor.property.lighting = False h += [ mlab.text3d( x_, y_, z, s, **kwargs ) ] h[-1].actor.property.lighting = False return h
def drawArrows(self, figure, bbox, number, partner): """ draw an 'arrow' from the cell 'number' to 'partner' """ #load the center of mass position n = np.array(number) p = np.array(partner) n.shape = (n.size, ) p.shape = (p.size, ) if p.size > 0 and n.size > 0: com1 = self.file1["features"][str(n[0])]["com"][:] com2 = self.file2["features"][str(p[0])]["com"][:] #write the cell label as text mlab.text3d(com1[2] - bbox[2] + 1, com1[1] - bbox[1] + 1, com1[0] - bbox[0] + 1, str(number), color=(1, 1, 1), figure=figure) #plot a point where the current cell is mlab.points3d([com1[2] - bbox[2] + 1], [com1[1] - bbox[1] + 1], [com1[0] - bbox[0] + 1], color=(0, 0, 1), figure=figure) #plot a line to the descendant's center mlab.plot3d([com1[2] - bbox[2] + 1, com2[2] - bbox[2] + 1], [com1[1] - bbox[1] + 1, com2[1] - bbox[1] + 1], [com1[0] - bbox[0] + 1, com2[0] - bbox[0] + 1], tube_radius=0.2, color=(1, 0, 0), figure=figure) #plot a second line, if there is a split if p.size == 2: com3 = self.file2["features"][str(p[1])]["com"][:] mlab.plot3d([com1[2] - bbox[2] + 1, com3[2] - bbox[2] + 1], [com1[1] - bbox[1] + 1, com3[1] - bbox[1] + 1], [com1[0] - bbox[0] + 1, com3[0] - bbox[0] + 1], tube_radius=0.2, color=(1, 0, 0), figure=figure)
def test_den(): # P = np.random.random((10,10)) # P=np.load(r'c:\maxdenP.np.npy') # myfilestr=r'c:\structfactors_density.dat' # x,y,z=np.loadtxt(myfilestr).T # P=z.reshape((101,101)) # print P.shape fig = mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1)) x, y, z = np.array([[1, 0, 1], [0, -1, 1], [-1, 0, -1], [0, 1, -1], [1, 1, 0], [-1, -1, 0]], "float64").T # view along z-axis pts_as = mlab.points3d(x, y, z, color=(1, 0, 0), colormap="gist_rainbow", figure=fig, scale_factor=0.3) # x,y,z=gen_fe() print "x", x, x.shape print "y", y, y.shape print "z", z, z.shape x, y, z = np.array([[0, 1, 1], [0, -1, -1]], "float64").T pts_FE = mlab.points3d(x, y, z, color=(0, 1, 0), colormap="gist_rainbow", figure=fig, scale_factor=0.3) mlab.text3d(0, 1, 1, "011", scale=0.2, color=(1, 1, 0)) mlab.text3d(1, 1, 0, "110", scale=0.2) mlab.text3d(-1, -1, 0, "-1-10", scale=0.2) mlab.text3d(1, 0, 1, "101", scale=0.2) mlab.text3d(0, -1, 1, "0-11", scale=0.2) mlab.text3d(-1, 0, -1, "-10-1", scale=0.2) mlab.text3d(0, 1, -1, "01-1", scale=0.2) # pts_fe=mlab.points3d(x,y,z-.125,color=(0,1,0),colormap='gist_rainbow',figure=fig,scale_factor=.02) # x,y,z=gen_sr() # print 'x',x,x.shape # pts_sr=mlab.points3d(x,y,z-.125,color=(0,0,1),colormap='gist_rainbow',figure=fig) # outline=mlab.outline(figure=fig,extent=[0,1,0,1,-1,0]) outline = mlab.outline(figure=fig) mlab.orientation_axes(figure=fig, xlabel="a", ylabel="b", zlabel="c") # print 'shape',P.shape # P = P[:,:, np.newaxis] # print 'after',P.shape # src = mlab.pipeline.scalar_field(P) # src = mlab.pipeline.array2d_source(P) # surf = mlab.pipeline.surface(src,figure=fig,extent=[0,1,0,1,-1,0],name='surf2',opacity=0.4) # surf.transform.GetTransform().RotateX(90) print "done"
def sensor_connectivity_3d(raw, picks, con, idx, n_con=20, min_dist=0.05, scale_factor=0.005, tube_radius=0.001): """ Function to plot sensor connectivity showing strongest connections(n_con) excluding sensors that are less than min_dist apart. https://github.com/mne-tools/mne-python/blob/master/examples/connectivity/plot_sensor_connectivity.py Parameters ---------- raw : Raw object Instance of mne.io.Raw picks : list Picks to be included. con : ndarray (n_channels, n_channels) Connectivity matrix. idx : list List of indices of sensors of interest. n_con : int Number of connections of interest. min_dist : float Minimum distance between sensors allowed. Note: Please modify scale factor and tube radius to appropriate sizes if the plot looks scrambled. """ # Now, visualize the connectivity in 3D try: from enthought.mayavi import mlab except: from mayavi import mlab mlab.figure(size=(600, 600), bgcolor=(0.5, 0.5, 0.5)) # Plot the sensor location sens_loc = [raw.info['chs'][picks[i]]['loc'][:3] for i in idx] sens_loc = np.array(sens_loc) pts = mlab.points3d(sens_loc[:, 0], sens_loc[:, 1], sens_loc[:, 2], color=(1, 1, 1), opacity=1, scale_factor=scale_factor) # Get the strongest connections threshold = np.sort(con, axis=None)[-n_con] ii, jj = np.where(con >= threshold) # Remove close connections con_nodes = list() con_val = list() for i, j in zip(ii, jj): if sci.linalg.norm(sens_loc[i] - sens_loc[j]) > min_dist: con_nodes.append((i, j)) con_val.append(con[i, j]) con_val = np.array(con_val) # Show the connections as tubes between sensors vmax = np.max(con_val) vmin = np.min(con_val) for val, nodes in zip(con_val, con_nodes): x1, y1, z1 = sens_loc[nodes[0]] x2, y2, z2 = sens_loc[nodes[1]] points = mlab.plot3d([x1, x2], [y1, y2], [z1, z2], [val, val], vmin=vmin, vmax=vmax, tube_radius=tube_radius, colormap='RdBu') points.module_manager.scalar_lut_manager.reverse_lut = True mlab.scalarbar(title='Phase Lag Index (PLI)', nb_labels=4) # Add the sensor names for the connections shown nodes_shown = list(set([n[0] for n in con_nodes] + [n[1] for n in con_nodes])) for node in nodes_shown: x, y, z = sens_loc[node] mlab.text3d(x, y, z, raw.ch_names[picks[node]], scale=0.005, color=(0, 0, 0)) view = (-88.7, 40.8, 0.76, np.array([-3.9e-4, -8.5e-3, -1e-2])) mlab.view(*view)
def test_den(): #P = np.random.random((10,10)) #P=np.load(r'c:\maxdenP.np.npy') #myfilestr=r'c:\structfactors_density.dat' #x,y,z=np.loadtxt(myfilestr).T #P=z.reshape((101,101)) #print P.shape fig = mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1)) x, y, z = np.array([[1, 0, 1], [0, -1, 1], [-1, 0, -1], [0, 1, -1], [1, 1, 0], [-1, -1, 0]], 'float64').T #view along z-axis pts_as = mlab.points3d(x, y, z, color=(1, 0, 0), colormap='gist_rainbow', figure=fig, scale_factor=.3) #x,y,z=gen_fe() print 'x', x, x.shape print 'y', y, y.shape print 'z', z, z.shape x, y, z = np.array([ [0, 1, 1], [0, -1, -1], ], 'float64').T pts_FE = mlab.points3d(x, y, z, color=(0, 1, 0), colormap='gist_rainbow', figure=fig, scale_factor=.3) mlab.text3d(0, 1, 1, '011', scale=.2, color=(1, 1, 0)) mlab.text3d(1, 1, 0, '110', scale=.2) mlab.text3d(-1, -1, 0, '-1-10', scale=.2) mlab.text3d(1, 0, 1, '101', scale=.2) mlab.text3d(0, -1, 1, '0-11', scale=.2) mlab.text3d(-1, 0, -1, '-10-1', scale=.2) mlab.text3d(0, 1, -1, '01-1', scale=.2) #pts_fe=mlab.points3d(x,y,z-.125,color=(0,1,0),colormap='gist_rainbow',figure=fig,scale_factor=.02) #x,y,z=gen_sr() #print 'x',x,x.shape #pts_sr=mlab.points3d(x,y,z-.125,color=(0,0,1),colormap='gist_rainbow',figure=fig) #outline=mlab.outline(figure=fig,extent=[0,1,0,1,-1,0]) outline = mlab.outline(figure=fig) mlab.orientation_axes(figure=fig, xlabel='a', ylabel='b', zlabel='c') #print 'shape',P.shape #P = P[:,:, np.newaxis] #print 'after',P.shape #src = mlab.pipeline.scalar_field(P) #src = mlab.pipeline.array2d_source(P) #surf = mlab.pipeline.surface(src,figure=fig,extent=[0,1,0,1,-1,0],name='surf2',opacity=0.4) #surf.transform.GetTransform().RotateX(90) print 'done'
# Show the connections as tubes between sensors vmax = np.max(con_val) vmin = np.min(con_val) for val, nodes in zip(con_val, con_nodes): x1, y1, z1 = sens_loc[nodes[0]] x2, y2, z2 = sens_loc[nodes[1]] points = mlab.plot3d([x1, x2], [y1, y2], [z1, z2], [val, val], vmin=vmin, vmax=vmax, tube_radius=0.001, colormap='RdBu') points.module_manager.scalar_lut_manager.reverse_lut = True mlab.scalarbar(title='Phase Lag Index (PLI)', nb_labels=4) # Add the sensor names for the connections shown nodes_shown = list(set([n[0] for n in con_nodes] + [n[1] for n in con_nodes])) for node in nodes_shown: x, y, z = sens_loc[node] mlab.text3d(x, y, z, raw.ch_names[picks[node]], scale=0.005, color=(0, 0, 0)) view = (-88.7, 40.8, 0.76, np.array([-3.9e-4, -8.5e-3, -1e-2])) mlab.view(*view)
def plot_3d_spectrum_mayavi(fs, fignum=None, vmin=None, vmax=None, pop_ids=None): """ Logarithmic heatmap of single 3d FS. This method relies on MayaVi2's mlab interface. See http://code.enthought.com/projects/mayavi/docs/development/html/mayavi/mlab.html . To edit plot properties, click leftmost icon in the toolbar. If you get an ImportError upon calling this function, it is likely that you don't have mayavi installed. fs: FS to plot vmin: Values in fs below vmin are masked in plot. vmax: Values in fs above vmax saturate the color spectrum. fignum: Figure number to plot into. If None, a new figure will be created. Note that these are MayaVi figures, which are separate from matplotlib figures. pop_ids: If not None, override pop_ids stored in Spectrum. """ from enthought.mayavi import mlab fig = mlab.figure(fignum, bgcolor=(1,1,1)) mlab.clf(fig) if vmin is None: vmin = fs.min() if vmax is None: vmax = fs.max() # Which entries should I plot? toplot = numpy.logical_not(fs.mask) toplot = numpy.logical_and(toplot, fs.data >= vmin) # For the color mapping normalized = (numpy.log(fs)-numpy.log(vmin))\ /(numpy.log(vmax)-numpy.log(vmin)) normalized = numpy.minimum(normalized, 1) xs,ys,zs = numpy.indices(fs.shape) flat_xs = xs.flatten() flat_ys = ys.flatten() flat_zs = zs.flatten() flat_toplot = toplot.flatten() mlab.barchart(flat_xs[flat_toplot], flat_ys[flat_toplot], flat_zs[flat_toplot], normalized.flatten()[flat_toplot], colormap='hsv', scale_mode='none', lateral_scale=1, figure=fig) if pop_ids is None: if fs.pop_ids is not None: pop_ids = fs.pop_ids else: pop_ids = ['pop0','pop1','pop2'] a = mlab.axes(xlabel=pop_ids[0],ylabel=pop_ids[1],zlabel=pop_ids[2], figure=fig, color=(0,0,0)) a.axes.label_format = "" a.title_text_property.color = (0,0,0) mlab.text3d(fs.sample_sizes[0],fs.sample_sizes[1],fs.sample_sizes[2]+1, '(%i,%i,%i)'%tuple(fs.sample_sizes), scale=0.75, figure=fig, color=(0,0,0)) mlab.view(azimuth=-40, elevation=65, distance='auto', focalpoint='auto') mlab.show()
con_nodes.append((i, j)) con_val.append(con[i, j]) con_val = np.array(con_val) # Show the connections as tubes between sensors vmax = np.max(con_val) vmin = np.min(con_val) for val, nodes in zip(con_val, con_nodes): x1, y1, z1 = sens_loc[nodes[0]] x2, y2, z2 = sens_loc[nodes[1]] points = mlab.plot3d([x1, x2], [y1, y2], [z1, z2], [val, val], vmin=vmin, vmax=vmax, tube_radius=0.001, colormap='RdBu') points.module_manager.scalar_lut_manager.reverse_lut = True mlab.scalarbar(title='Phase Lag Index (PLI)', nb_labels=4) # Add the sensor names for the connections shown nodes_shown = list(set([n[0] for n in con_nodes] + [n[1] for n in con_nodes])) for node in nodes_shown: x, y, z = sens_loc[node] mlab.text3d(x, y, z, raw.ch_names[picks[node]], scale=0.005, color=(0, 0, 0)) view = (-88.7, 40.8, 0.76, np.array([-3.9e-4, -8.5e-3, -1e-2])) mlab.view(*view)
def initialize_rope(label_img, xyz,bgr, plotting=False): rope_mask = (label_img==1) | ndi.morphology.binary_dilation(label_img==2,np.ones((5,5))) rope_mask = ndi.binary_opening(rope_mask, np.ones((3,3))) if plotting: cv2.imshow('bgr',bgr.copy()) cv2.imshow('labels',label_img.astype('uint8')*50) cv2.waitKey(5) start_end = (uv_start, uv_end) = get_cc_centers(label_img==2,2) skel = mip.skeletonize(rope_mask) skel = skeletons.removeSpurs1(skel,10) if plotting: cv2.imshow('skel',skel*100) paths2d = skeletons.get_paths_2d(skel) endpts2d = [] for path in paths2d: endpts2d.append(path[0]) endpts2d.append(path[-1]) i_start = nearest_neighbor(uv_start, endpts2d) i_end = nearest_neighbor(uv_end, endpts2d) print "start,end",i_start,i_end paths3d = [to3d(path,xyz) for path in paths2d] if plotting: starts_ends = [] for path in paths2d: starts_ends.append(path[0]) starts_ends.append(path[-1]) plot_colorskel(paths2d,bgr,start_end,starts_ends) C = skeletons.make_cost_matrix(paths3d) PG = skeletons.make_path_graph(C, [len(path) for path in paths3d]) (_,nodes) = skeletons.longest_path_between(PG,i_start, set([]), i_end) if nodes is None: raise Exception("longest path failed") print nodes total_path = [] for node in nodes[::2]: if node%2 == 0: total_path.extend(paths3d[node//2]) else: total_path.extend(paths3d[node//2][::-1]) total_path = np.array(total_path) xyzs_unif = curves.unif_resample(total_path,N_SEGS+1,tol=.0025) #us,vs = xyz2uv(xyzs_unif).T #labels = label_img[us,vs] labels = np.ones(xyzs_unif.shape[0]-1,'int') labels[0] = labels[-1] = 2 if plotting: xs,ys,zs = xyzs_unif.T us_skel, vs_skel = np.nonzero(skel) xs_skel, ys_skel, zs_skel = to3d(np.c_[us_skel, vs_skel], xyz).T import matplotlib.pyplot as plt, enthought.mayavi.mlab as mlab mlab.plot3d(xs,ys,zs,np.arange(len(xs)),tube_radius=.0025, colormap='spectral') mlab.points3d(xs_skel, ys_skel, zs_skel, scale_factor=.02, color=(1,1,1),opacity=.1) for (i,path) in enumerate(paths3d): mlab.text3d(path[0,0],path[0,1],path[0,2],str(2*i),scale=.01,color=(0,0,0)) mlab.text3d(path[-1,0],path[-1,1],path[-1,2],str(2*i+1),scale=.01,color=(0,0,0)) mlab.show() return xyzs_unif, labels
## For surface rendering field=mlab.contour3d(sregion,colormap='gist_ncar') # Generate a scalar field #field.contour.maximum_contour = 13000 field.contour.minimum_contour = 2000 field.actor.property.opacity = 0.3 mlab.outline() mlab.xlabel('RA(J2000)') mlab.ylabel('DEC(J2000)') mlab.zlabel('Velocity') #mlab.axes() #mlab.colorbar() ## Add 3-d text in the cube mlab.text3d(30,80,23,'Outflows in L1551 IRS5',scale=2.) #mlab.text(45,45,"outflows!",width=0.2,z=20) ## Draw arrows to show outflows etc. vector1=mlab.quiver3d(55,55,30,1,1,0.2,mode='arrow',scale_factor=10.0,color=(0.0,0.0,1.0)) vector2=mlab.quiver3d(55,40,30,1,-1.2,0.2,mode='arrow',scale_factor=10.0,color=(0.0,0.0,1.0)) vector3=mlab.quiver3d(45,50,30,-0.8,1,0.2,mode='2dthick_arrow',scale_factor=10.0,color=(0.0,0.0,1.0)) vector4=mlab.quiver3d(49,48,30,0.2,0,1,mode='2darrow',scale_factor=8.0,color=(0.0,0.0,1.0)) vector5=mlab.quiver3d(47,58,20,-1,1.8,-0.2,mode='arrow',scale_factor=10.0,color=(1.0,0.0,0.0)) vector6=mlab.quiver3d(45,45,20,-1.2,-1,-0.2,mode='arrow',scale_factor=10.0,color=(1.0,0.0,0.0)) vector7=mlab.quiver3d(55,45,20,1,-1.4,-0.2,mode='2dthick_arrow',scale_factor=10.0,color=(1.0,0.0,0.0)) vector8=mlab.quiver3d(48,48,20,-0.2,0,-1,mode='2darrow',scale_factor=8.0,color=(1.0,0.0,0.0)) # Insert the continuum image as a background img=pyfits.open('L1551.cont.fits') # Read the fits data
def atom_chooser(file_name = "C:/LanthanideSMMs/GdCu SMMs/7.cif", atoms = [], fract_coords = []): """display a list of atoms using a list of fractional coordinate arrays unit cell parameters will be extracted from the CIF-file inputs: file_name: a string with front slashes for file tree delimiters that directs to the location of the CIF file atoms: string list of atoms index in same manner as fract_coords fract_coords: a list of numpy arrays containing the fractional coordinates output: chosen_atom_indices """ import CifFile #existing library for reading in CIF format import os #operating system import re #regular expressions import numpy as np #numeric python from enthought.mayavi import mlab #convert file_name to acceptable format file_name = "file://" + file_name file_name = os.path.normpath(file_name) #read contents of CIF file cf = CifFile.ReadCif(file_name) sample_name = cf.dictionary.keys()[0] #unit cell in meters (converted from Angstroms) a = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_length_a']))) * 1e-10 b = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_length_b']))) * 1e-10 c = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_length_c']))) * 1e-10 #angles in degrees alpha = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_angle_alpha']))) beta = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_angle_beta']))) gamma = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_angle_gamma']))) #angles in radians alpha_ = alpha*np.pi/180 beta_ = beta*np.pi/180 gamma_ = gamma*np.pi/180 #need to generate real space vectors... a1 = np.array([a, 0, 0]) a2 = np.array([b*np.cos(gamma_), b*np.sin(gamma_)*np.sqrt(1-(np.cos(gamma_)*np.cos(beta_)-np.cos(alpha_))**2/(np.sin(gamma_)**2*np.sin(beta_)**2)), b*np.sin(gamma_)*(np.cos(gamma_)*np.cos(beta_)-np.cos(alpha_))/(np.sin(gamma_)*np.sin(beta_))]) a3 = np.array([c*np.cos(beta_), 0, c*np.sin(beta_)]) a1_A = a1*1e10 a2_A = a2*1e10 a3_A = a3*1e10 #start the mayavi pipeline # mlab.figure(fgcolor = (0,0,0), bgcolor = (1,1,1)) #display the bounding box of the unit cell mlab.plot3d([0,a1_A[0]], [0,a1_A[1]], [0,a1_A[2]], line_width = 4) mlab.plot3d([0,a2_A[0]], [0,a2_A[1]], [0,a2_A[2]], line_width = 4) mlab.plot3d([a1_A[0], a1_A[0]+a2_A[0]], [a1_A[1], a1_A[1]+a2_A[1]], [a1_A[2], a1_A[2]+a2_A[2]], line_width = 4) mlab.plot3d([a2_A[0], a1_A[0]+a2_A[0]], [a2_A[1], a1_A[1]+a2_A[1]], [a2_A[2], a1_A[2]+a2_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+0,a1_A[0]+a3_A[0]], [a3_A[1]+0,a1_A[1]+a3_A[1]], [a3_A[2]+0,a1_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+0,a2_A[0]+a3_A[0]], [a3_A[1]+0,a2_A[1]+a3_A[1]], [a3_A[2]+0,a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+a1_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a1_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a1_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+a2_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a2_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a2_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([0,a3_A[0]], [0,a3_A[1]], [0,a3_A[2]], line_width = 4) mlab.plot3d([a1_A[0],a1_A[0]+a3_A[0]], [a1_A[1],a1_A[1]+a3_A[1]], [a1_A[2],a1_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a2_A[0],a2_A[0]+a3_A[0]], [a2_A[1],a2_A[1]+a3_A[1]], [a2_A[2],a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a2_A[0]+a1_A[0],a2_A[0]+a1_A[0]+a3_A[0]], [a2_A[1]+a1_A[1],a2_A[1]+a1_A[1]+a3_A[1]], [a2_A[2]+a1_A[2],a2_A[2]+a1_A[2]+a3_A[2]], line_width = 4) r_coords = np.zeros(np.shape(fract_coords)) #take the fractional coordinates and generate real coordinates using the basis vectors for i in range(np.size(atoms)): r_coords[i] = a1_A * fract_coords[i][0] + a2_A * fract_coords[i][1] + a3_A * fract_coords[i][2] for i in range(np.size(atoms)): mlab.points3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.15, scale_factor = 1, color = (0,0,1), mode = 'cube') mlab.text3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.35, scale = 0.5, text = atoms[i], color = (0,0,1)) # placeholder, delete when ready chosen_atom_indices = [] mlab.show() return chosen_atom_indices
def _nodepicked_callback(self, picker_obj, evt): """ The callback function to determine what to do with a picked node """ from enthought.tvtk.api import tvtk picker_obj = tvtk.to_tvtk(picker_obj) picked = picker_obj.actors # if picked actor are the glyphs if self.nodes.actor.actor._vtk_obj in [o._vtk_obj for o in picked]: pos = picker_obj.picked_positions[0] dist_small = sys.maxint # the biggest integer for i, points in enumerate(self.nodes.mlab_source.points): # if picked positions ON a node dist = np.linalg.norm((np.array(pos)- points)) # save the smallest distance node if dist <= dist_small: j = i dist_small = dist # get the id of the picked node picknodeid = 'n' + str(j+1) if self.nodepicker.show_info: # generate a nice view for the information nodedata = self.network.graph.node[picknodeid] nodedata['dn_internal_id'] = picknodeid deg = self.network.graph.degree(picknodeid) from types import IntType, StringType if type(deg) == IntType or type(deg) == StringType: nodedata['degree'] = deg # does this node has a self-loop? if self.network.graph.has_edge(picknodeid, picknodeid): nodedata['has_selfloop'] = True else: nodedata['has_selfloop'] = True mynode = Node(nodedata=nodedata) mynode.configure_traits() elif self.nodepicker.toggle_selection: # toggles the selection of a node if self.datasourcemanager._srcobj.selected_nodes[j] == 1: # unpick self._select_nodes(selection_node_array = np.array([j]), \ first_is_selected = True, change_vector = False) else: # pick self._select_nodes(selection_node_array = np.array([j]), activate = True, \ first_is_selected = True, change_vector = False) elif self.nodepicker.select_node: # toggle the picking state if self.datasourcemanager._srcobj.selected_nodes[j] == 1: # unpick self._select_nodes(selection_node_array = np.array([j]), \ first_is_selected = True) else: # pick self._select_nodes(selection_node_array = np.array([j]), activate = True, \ first_is_selected = True) elif self.nodepicker.select_node2: # get the neighbors of the picked node in an array nbrs = self.datasourcemanager._srcobj.relabled_graph.neighbors(j) # put the two lists together cur_nodes = np.append(j, np.array(nbrs)) if self.datasourcemanager._srcobj.selected_nodes[j] == 1: # unpick self._select_nodes(selection_node_array = cur_nodes, \ first_is_selected = True) else: # pick self._select_nodes(selection_node_array = cur_nodes, activate = True ,\ first_is_selected = True) elif self.nodepicker.toggle_label: if self.node_labels.has_key(j): # we want to remove the node label from the pipeline first # and then delete the entry in the dictionary a = self.node_labels[j] a.remove() del self.node_labels[j] else: srcobj = self.datasourcemanager._srcobj # create the text3d instance and add it to the dictionary from enthought.mayavi import mlab a = mlab.text3d(srcobj.positions[j,0], srcobj.positions[j,1], \ srcobj.positions[j,2], \ ' ' + srcobj.labels[j], # white spaces are hackish name = 'Node ' + srcobj.labels[j]) self.node_labels[j] = a else: logger.info("No valid pick operation.")
else: C.append('g') fig = mlab.figure(1, bgcolor=(1, 1, 1), size=(500, 500)) fig_res = 100 S = [1 for n in xrange(len(X))] mlab.points3d(X, Y, Z, resolution=fig_res) for x,y,z,l in zip(X,Y,Z,labels): if int(l)<8: d = .17 mlab.text3d(x-d,y+d*1.25,z,l,line_width=3.5, scale=.15,color=(0,0,0), opacity=.9) elif int(l) in [9,8]: mlab.text3d(x-d/2,y+d*1.3,z,l,line_width=3.5, scale=.15,color=(0,0,0), opacity=.9) else: mlab.text3d(x-d/2,y+d*1.5,z,l,line_width=3.5, scale=.15,color=(0,0,0), opacity=.9) def connection_plot(r1,r2): segments = 10 dr = (r2 - r1)/segments ''' for n in xrange(segments): if not n%3: x1,y1,z1 = [k[0] for k in zip(r1+n*dr)]
pts,s = latlondep2geo(depth,lat,lon,vs) sgrid = tvtk.StructuredGrid(dimensions=(9,13,81),points=pts) sgrid.points = pts sgrid.point_data.scalars = ravel(s.copy()) sgrid.point_data.scalars.name = 'scalars' d = mlab.pipeline.add_dataset(sgrid) sf = mlab.pipeline.surface(d) gx = mlab.pipeline.grid_plane(d) gx.grid_plane.axis = 'x' gy = mlab.pipeline.grid_plane(d) gy.grid_plane.axis = 'y' gz = mlab.pipeline.grid_plane(d) gz.grid_plane.axis = 'z' for lat in runlat: x,y,z = convert_pt(lat,-68.,10.) txt = mlab.text3d(x,y,z,'%d'%(lat),color=(0,0,0),line_width=10.0) txt.scale = [30,30,30] for lon in runlon[1::]: x,y,z = convert_pt(43.,lon,10.) txt = mlab.text3d(x,y,z,'%d'%(lon),color=(0,0,0),line_width=10.0) txt.scale = [30,30,30] for dp in [-10,-40]: x,y,z = convert_pt(43.,-68.5,dp) txt = mlab.text3d(x,y,z,'%d km'%dp,color=(0,0,0),line_width=10.0) txt.scale = [30,30,30] #plot_canada_map() mlab.show() #arr = mlab.screenshot() #import pylab as pl #pl.imshow(arr)
def mshplot3D(mesh, scale=[1, 1, 1], elm_ids=[], labels=[0, 0, 0]): ''' Plot 3D meshes (and optionally label vertices, elements, and edges) @param mesh A 3D mesh object @param elm_ids Numpy array of elements that will be plotted. If left empty, all elements are plotted. @param labels (\c bool) List of three elements indicateting what should be labeled. By default, nothing is labeled labels[0]=True labels vertices labels[1]=True labels elements labels[2]=True labels edges @param scale (\c int) List indicating the scaling of the vertices -- this is used to scale the z-direction for a thin mesh. By default, there is no scaling, i.e. scale = [1, 1, 1] @see msh.mesh.Mesh3D @author Matt Ueckermann ''' if type(elm_ids) is not int: if len(elm_ids) == 0: elm_ids = np.arange(len(mesh.elm)) #Figure out how many vertices are in each element type n_vert_in_type = [len(int_el_pqr(element=element, dim=3)[0]) \ for element in mesh.u_elm_type] #Now create the TVTK data-structures for the 'cells' and 'offsets' #cells give [#verts, vert1id, vert2id...,vertnid, #verts ...] #offsets gives the id's in the cells list where #verts are listed. So above # it is [0, n+1] cells = np.array([], dtype=int) offset = np.array([], dtype=int) for elm, elm_type in zip(mesh.elm[elm_ids, :], mesh.elm_type[elm_ids]): n_vt = n_vert_in_type[elm_type] offset = np.append(offset, len(cells)) cells = np.append(cells, [n_vt] + elm[:n_vt].tolist()) #Also have to create a list of element-types -- to do that we have to #convert from my numbering system to the TVTK numbering system if mesh.dim == 3: type_convert = np.array([tvtk.Tetra().cell_type,\ tvtk.Hexahedron().cell_type, tvtk.Wedge().cell_type]) elif mesh.dim == 2: type_convert = np.array([tvtk.Triangle().cell_type,\ tvtk.Quad().cell_type]) cell_types = mesh.u_elm_type[mesh.elm_type[elm_ids]] cell_types = type_convert[cell_types] #To help visualize the mesh, we color it be the distance from the origin if mesh.dim == 3: x, y, z = mesh.vert[:].T elif mesh.dim == 2: x, y = mesh.vert[:, :2].T z = np.zeros_like(x) dist = np.sqrt(x**2 + y**2 + z**2) #and we scale the x, y, z, coordinates according the the assigned 'scale' points = np.column_stack((x * scale[0], y * scale[1], z * scale[2])) #Now we create the data-structures cell_array = tvtk.CellArray() cell_array.set_cells(len(offset), cells) ug = tvtk.UnstructuredGrid(points=points) ug.set_cells(cell_types, offset, cell_array) ug.point_data.scalars = dist.ravel() ug.point_data.scalars.name = 'Distance from Origin' #Next we set the new data-structures as a mayavi source src = sources.vtk_data_source.VTKDataSource(data=ug) #Extract the edges from the grid edges = mlab.pipeline.extract_edges(src) #Use a shorter name for the elements elm = mesh.elm[elm_ids, :] if any(labels) and len(elm) > 20: string = "WARNING:\nAre you sure you want to label more than 20" + \ "elements in 3D? -- Labels are very memory" + \ "(apparently -- who would have guessed?) \nY(es) to proceed:" answer = raw_input(string) if answer.lower() not in ['yes', 'y', 'yes']: labels = [0, 0, 0] #Next add labels if desired: maxdist = np.max(dist) / 2. if labels[0]: #Vertex labels for i in xrange(len(offset)): for vnum in elm[i, :]: if vnum >= 0: mlab.text3d(points[vnum, 0], points[vnum, 1], \ points[vnum, 2], '%d' % vnum, \ color=(0, 0, 0), scale=0.02*maxdist) if labels[1]: #element labels for i in xrange(len(offset)): if elm_ids[i] < 0: elm_label = mesh.numelm + elm_ids[i] else: elm_label = elm_ids[i] x = np.mean(points[elm[i, :cells[offset[i]]], 0]) y = np.mean(points[elm[i, :cells[offset[i]]], 1]) z = np.mean(points[elm[i, :cells[offset[i]]], 2]) mlab.text3d(x, y, z, '%d' % elm_label, color=(0.3, 0.3, 0.9), \ scale=0.03*maxdist) if labels[2]: #edge labels if len(elm_ids) < len(mesh.elm): elm2ed = connect_elm2ed(mesh.elm2elm[:], mesh.ed2ed) ed_ids = np.unique(elm2ed[elm_ids]) ed_ids = ed_ids[ed_ids >= 0] ed2ed = mesh.ed2ed[ed_ids, :] else: ed_ids = np.arange(len(mesh.ed2ed)) ed2ed = mesh.ed2ed[:] for i in xrange(len(ed2ed)): n = 8 if ed2ed[i, -1] < 0: n = 7 x = np.mean(points[ed2ed[i, 4:n], 0]) y = np.mean(points[ed2ed[i, 4:n], 1]) z = np.mean(points[ed2ed[i, 4:n], 2]) mlab.text3d(x, y, z, '%d' % ed_ids[i], color=(0.3, 0.9, 0.3), \ scale=0.03*maxdist) #And plot them! mlab.pipeline.surface(edges, opacity=0.4, line_width=2) mlab.axes() mlab.title('3D mesh', size=0.5, height=0.95) #mlab.show() return cells, offset, cell_types
#display the bounding box of the unit cell mlab.plot3d([0,a1_A[0]], [0,a1_A[1]], [0,a1_A[2]], line_width = 4) mlab.plot3d([0,a2_A[0]], [0,a2_A[1]], [0,a2_A[2]], line_width = 4) mlab.plot3d([a1_A[0], a1_A[0]+a2_A[0]], [a1_A[1], a1_A[1]+a2_A[1]], [a1_A[2], a1_A[2]+a2_A[2]], line_width = 4) mlab.plot3d([a2_A[0], a1_A[0]+a2_A[0]], [a2_A[1], a1_A[1]+a2_A[1]], [a2_A[2], a1_A[2]+a2_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+0,a1_A[0]+a3_A[0]], [a3_A[1]+0,a1_A[1]+a3_A[1]], [a3_A[2]+0,a1_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+0,a2_A[0]+a3_A[0]], [a3_A[1]+0,a2_A[1]+a3_A[1]], [a3_A[2]+0,a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+a1_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a1_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a1_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a3_A[0]+a2_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a2_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a2_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([0,a3_A[0]], [0,a3_A[1]], [0,a3_A[2]], line_width = 4) mlab.plot3d([a1_A[0],a1_A[0]+a3_A[0]], [a1_A[1],a1_A[1]+a3_A[1]], [a1_A[2],a1_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a2_A[0],a2_A[0]+a3_A[0]], [a2_A[1],a2_A[1]+a3_A[1]], [a2_A[2],a2_A[2]+a3_A[2]], line_width = 4) mlab.plot3d([a2_A[0]+a1_A[0],a2_A[0]+a1_A[0]+a3_A[0]], [a2_A[1]+a1_A[1],a2_A[1]+a1_A[1]+a3_A[1]], [a2_A[2]+a1_A[2],a2_A[2]+a1_A[2]+a3_A[2]], line_width = 4) r_coords = np.zeros(np.shape(fract_coords)) #take the fractional coordinates and generate real coordinates using the basis vectors for i in range(np.size(atoms)): r_coords[i] = a1_A * fract_coords[i][0] + a2_A * fract_coords[i][1] + a3_A * fract_coords[i][2] for i in range(np.size(atoms)): mlab.points3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.15, scale_factor = 1, color = color_dict[atoms[i]], mode = 'cube') mlab.text3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.35, scale = 0.5, text = atoms[i]+', '+str(i), color = color_dict[atoms[i]]) # placeholder, delete when ready chosen_atom_indices = [] mlab.show()