def cylinder(engine, radius, height, resolution): test = BuiltinSurface() test.source = 'cylinder' test.data_source.center = np.array([0., 0., 0.]) test.data_source.radius = radius test.data_source.height = height test.data_source.capping = False test.data_source.resolution = resolution test_surface = Surface() engine.add_source(test) engine.add_filter(test_surface)
def setUp(self): """Initial setting up of test fixture, automatically called by TestCase before any other test method is invoked""" e = NullEngine() # Uncomment to see visualization for debugging etc. #e = Engine() e.start() s = e.new_scene() poly_data = BuiltinSurface() e.add_source(poly_data) outline = Outline() e.add_module(outline) surface = Surface() e.add_module(surface) poly_data.data_source.shaft_radius = 0.05 poly_data.data_source.shaft_resolution = 7 poly_data.data_source.tip_radius = 0.1 self.e = e self.scene = e.current_scene return
def drawContinentsSphere(r_earth): ############################################################################### # Display continents outline, using the VTK Builtin surface 'Earth' from mayavi.sources.builtin_surface import BuiltinSurface continents_src = BuiltinSurface(source='earth', name='Continents') # The on_ratio of the Earth source controls the level of detail of the # continents outline. #r_earth = 6371.0 # km continents_src.data_source.on_ratio = 2 continents_src.data_source.radius = r_earth continents = mlab.pipeline.surface(continents_src, color=(0, 0, 0))
def main(): continents_src = BuiltinSurface(source="earth", name="Continents") continents = mlab.pipeline.surface(continents_src, color=(0, 0, 0)) mlab.figure(1, bgcolor=(0.48, 0.48, 0.48), fgcolor=(0, 0, 0), size=(400, 400)) mlab.view(63.4, 73.8, 4, [-0.05, 0, 0]) mlab.show() #TODO: implement pass
def continents(color=(0, 0, 0), linewidth=1, resolution=2, opacity=1, radius=MEAN_EARTH_RADIUS): """ Plot the outline of the continents. Parameters: * color : tuple RGB color of the lines. Default = black * linewidth : float The width of the continent lines * resolution : float The data_source.on_ratio parameter that controls the resolution of the continents * opacity : float The opacity of the lines. Must be between 0 and 1 * radius : float The radius of the sphere where the continents will be plotted. Defaults to the mean Earth radius Returns: * continents : Mayavi surface The Mayavi surface element of the continents """ _lazy_import_mlab() _lazy_import_BuiltinSurface() continents_src = BuiltinSurface(source='earth', name='Continents') continents_src.data_source.on_ratio = resolution continents_src.data_source.radius = MEAN_EARTH_RADIUS surf = mlab.pipeline.surface(continents_src, color=color) surf.actor.property.line_width = linewidth surf.actor.property.opacity = opacity return surf
def affichage_terre(): from mayavi import mlab from tvtk.tools import visual a = mlab.figure(1, bgcolor=(0.48, 0.48, 0.48), fgcolor=(0, 0, 0), size=(1000, 800)) mlab.clf() visual.set_viewer(a) from mayavi.sources.builtin_surface import BuiltinSurface continents_src = BuiltinSurface(source='earth', name='Countries') continents_src.data_source.on_ratio = 2 continents = mlab.pipeline.surface(continents_src, color=(0, 0, 0)) sphere = mlab.points3d(0, 0, 0, scale_mode='none', scale_factor=2, color=(0.67, 0.77, 0.93), resolution=50, opacity=0.7, name='Earth') sphere.actor.property.specular = 0.45 sphere.actor.property.specular_power = 5 sphere.actor.property.backface_culling = True def Arrow_From_A_to_B(x1, y1, z1, x2, y2, z2): ar1 = visual.arrow(x=x1, y=y1, z=z1) ar1.length_cone = 0.4 arrow_length = np.sqrt((x2 - x1)**2 + (y2 - y1)**2 + (z2 - z1)**2) ar1.actor.scale = [arrow_length, arrow_length, arrow_length] ar1.pos = ar1.pos / arrow_length ar1.axis = [x2 - x1, y2 - y1, z2 - z1] #ix = mlab.points3d(2,0,0,mode='arrow') PLUS ELEGANT ??? #http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html ! POINTS3D ! return ar1 axX = Arrow_From_A_to_B(0, 0, 0, 1, 0, 0) axY = Arrow_From_A_to_B(0, 0, 0, 0, 1, 0) axZ = Arrow_From_A_to_B(0, 0, 0, 0, 0, 1) '''coords = kepler2card(np.array([1,0,0]),satellites[0][14],-satellites[0][15]) sat1 = Arrow_From_A_to_B(0, 0, 0, coords[0], coords[1], coords[2]) coords2 = kepler2card(np.array([1,0,0]),satellites[0][14],0) sat2 = Arrow_From_A_to_B(0, 0, 0, coords[0], coords[1], coords[2]) ''' theta = np.linspace(0, 2 * np.pi, 100) for i in range(len(satellites)): U = kepler2card(np.array([1, 0, 0]), -satellites[i][15], satellites[i][14]) N = kepler2card(np.array([0, 0, 1]), -satellites[i][15], satellites[i][14]) #U = kepler2card(np.array([1,0,0]),-82.1335,54.1115) #N = kepler2card(np.array([0,0,1]),-82.1335,54.1115) UN = np.cross(U, N) #P = R*np.cos(theta)*U+R*np.sin(theta)*U.cross(N) x = R * np.cos(theta) * U[0] + R * np.sin(theta) * UN[0] y = R * np.cos(theta) * U[1] + R * np.sin(theta) * UN[1] z = R * np.cos(theta) * U[2] + R * np.sin(theta) * UN[2] mlab.plot3d(x, y, z, color=(1, 1, 1), opacity=0.2, tube_radius=None) for i in range(len(satellites)): U = kepler2card(np.array([1, 0, 0]), -cosatellites[i][15], satellites[i][14]) N = kepler2card(np.array([0, 0, 1]), -satellites[i][15], satellites[i][14]) #U = kepler2card(np.array([1,0,0]),-82.1335,54.1115) #N = kepler2card(np.array([0,0,1]),-82.1335,54.1115) UN = np.cross(U, N) theta = np.deg2rad(satellites[i][18]) * 0 n = satellites[i][18] d = ((mu**(1 / 3)) / ((2 * n * np.pi / 86400)**(2 / 3))) x = (d * np.cos(theta) * U[0] + d * np.sin(theta) * UN[0]) * R y = (d * np.cos(theta) * U[1] + d * np.sin(theta) * UN[1]) * R z = (d * np.cos(theta) * U[2] + d * np.sin(theta) * UN[2]) * R if (satellites[i][0].find("16") != -1 or satellites[i][0].find("26") != -1): points = mlab.points3d(x, y, z, scale_mode='none', scale_factor=0.1, color=(1, 0, 0)) x = (d * np.cos(theta) * U[0] + d * np.sin(theta) * UN[0]) y = (d * np.cos(theta) * U[1] + d * np.sin(theta) * UN[1]) z = (d * np.cos(theta) * U[2] + d * np.sin(theta) * UN[2]) if (satellites[i][0].find("16") != -1 or satellites[i][0].find("26") != -1): print(theta) points = mlab.points3d(x, y, z, scale_mode='none', scale_factor=0.1, color=(1, 0, 0)) return n, d
mlab.pipeline.surface(points, color=(1, 1, 1), representation='wireframe', line_width=4, name='Connections') ############################################################################### # Display city names for city, index in cities.iteritems(): label = mlab.text(x[index], y[index], city, z=z[index], width=0.016 * len(city), name=city) label.property.shadow = True ############################################################################### # Display continents outline, using the VTK Builtin surface 'Earth' from mayavi.sources.builtin_surface import BuiltinSurface continents_src = BuiltinSurface(source='earth', name='Continents') # The on_ratio of the Earth source controls the level of detail of the # continents outline. continents_src.data_source.on_ratio = 2 continents = mlab.pipeline.surface(continents_src, color=(0, 0, 0)) ############################################################################### # Display a semi-transparent sphere, for the surface of the Earth # We use a sphere Glyph, throught the points3d mlab function, rather than # building the mesh ourselves, because it gives a better transparent # rendering. sphere = mlab.points3d(0, 0, 0, scale_mode='none', scale_factor=2, color=(0.67, 0.77, 0.93), resolution=50,
# 6.绘制城市名字 for city, index in cities.items(): label = mlab.text( x[index], y[index], city, z=z[index], # x,y,city是城市名称,z坐标,width是文本宽度,name表示文本对象 width=0.016 * len(city), name=city) label.property.shadow = True # 7.绘制地球上大洲的边界 from mayavi.sources.builtin_surface import BuiltinSurface # 使用mlab的管线绘制表面函数对边界进行绘制 continents_src = BuiltinSurface(source="earth", name="Continents") # 8.设置模型LOD的层级,实现近细远粗 continents_src.data_source.on_ratio = 2 # 2级lod continents = mlab.pipeline.surface(continents_src, color=(0, 0, 0)) # 9.赤道线numpy数组的构造过程 theta = np.linspace(0, 2 * np.pi, 100) x = np.cos(theta) y = np.sin(theta) z = np.zeros_like(theta) # 10.绘制赤道线 mlab.plot3d(x, y, z, color=(1, 1, 1), opacity=0.2, tube_radius=None) # 11.上面效果不是太好,添加镜面反射等参数 # 调整镜面反射参数 sphere.actor.property.specular = 0.45
def sv_user_traj_3d_interactive(gps_ds, sv_pos, user_pos, ele=10., azim=20.): """ This method will provide an interactive 3d model plotted using mayavi to show all trajectories. :param gps_ds: :param sv_pos: :param user_pos: :return: """ mlab.figure(1, bgcolor=(0.48, 0.48, 0.48), fgcolor=(0, 0, 0), size=(400, 400)) mlab.clf() ########################################################################## # Display continents outline, using the VTK Builtin surface 'Earth' continents_src = BuiltinSurface(source='earth', name='Continents') # The on_ratio of the Earth source controls the level of detail of the # continents outline. continents_src.data_source.on_ratio = 1 continents = mlab.pipeline.surface(continents_src, color=(0, 0, 0)) for svn in range(0, len(sv_pos)): mlab.plot3d(sv_pos[svn, 0, :] / radius, sv_pos[svn, 1, :] / radius, sv_pos[svn, 2, :] / radius, color=(1, 1, 0.5), opacity=0.5, tube_radius=None) xml = len(sv_pos[svn, 0, :]) / 2 yml = len(sv_pos[svn, 1, :]) / 2 zml = len(sv_pos[svn, 2, :]) / 2 xm, ym, zm = sv_pos[svn, 0, int(xml)] / radius, sv_pos[ svn, 1, int(yml)] / radius, sv_pos[svn, 2, int(zml)] / radius label = mlab.text(xm, ym, gps_ds.Rx_sv_list[svn], z=zm, width=0.0155 * len(gps_ds.Rx_sv_list[svn])) label.property.shadow = True mlab.plot3d(user_pos[:, 0] / radius, user_pos[:, 1] / radius, user_pos[:, 2] / radius, color=(1, 1, 1), opacity=0.5, tube_radius=None) xml = len(user_pos[:, 0]) / 2 yml = len(user_pos[:, 1]) / 2 zml = len(user_pos[:, 2]) / 2 xm, ym, zm = user_pos[int(xml), 0] / radius, user_pos[ int(yml), 1] / radius, user_pos[int(zml), 2] / radius label = mlab.text(xm, ym, "User", z=zm, width=0.077) label.property.shadow = True ############################################################################### # Display a semi-transparent sphere, for the surface of the Earth # We use a sphere Glyph, throught the points3d mlab function, rather than # building the mesh ourselves, because it gives a better transparent # rendering. ocean_blue = (0.4, 0.5, 1.0) sphere = mlab.points3d( 0, 0, 0, scale_mode='none', scale_factor=2, # color=(0.67, 0.77, 0.93), color=ocean_blue, resolution=50, opacity=.85, name='Earth') # # These parameters, as well as the color, where tweaked through the GUI, # with the record mode to produce lines of code usable in a script. sphere.actor.property.specular = 0.45 sphere.actor.property.specular_power = 5 # Backface culling is necessary for more a beautiful transparent # rendering. sphere.actor.property.backface_culling = True # Plot the equator and the tropiques theta = np.linspace(0, 2 * np.pi, 100) for angle in (-np.pi / 6, 0, np.pi / 6): x = np.cos(theta) * np.cos(angle) y = np.sin(theta) * np.cos(angle) z = np.ones_like(theta) * np.sin(angle) mlab.plot3d(x, y, z, color=(1, 1, 1), opacity=0.2, tube_radius=None) mlab.view(azim, ele) mlab.show()
def mayaviBuiltinPlane(len_x=1.0, len_y=1.0, loc=(0.0,0.0,0.0), normal=(0.0,0.0,1.0), planeCol=(1.0,1.0,1.0), planeOpa=0.8, planeLighting=False, drawNormal=True, normVecScale=1.0, engine=None, scene=None): """Function to draw Mayavi builtin plane surface The function assumes that a figure has already been drawn and a scene is present. Parameters ---------- len_x : float side length of the plane along x axis (normally horizontal/width) len_y : float side length of the plane along y axis (normally vertical/height) loc : 3-tuple of floats (x,y,z) coordinates of the centroid of the plane in the world-coordinates normal : 3-tuple of floats (nx, ny, nz) normal vector of the plane surface planeCol : 3-tuple of floats (values between 0 and 1) Color of the plane planeOpa : float (between 0 and 1) opacity property of the plane planeLighting : boolean whether or not to use lighting on the plane (default is False) drawNormal : boolean whether or not to draw the plane normal normVecScale : float scale factor of the normal vector if drawn (default=1) engine : mayavi engine optional scene : mayavi scene optional Note: If the engine and the scene is not passed explicitly, it will grab the current engine and the first scene of that engine. Returns ------- if `drawNormal` is `True` : 2-tuple (plane_surface, normal_vector) if `drawNormal` is `False` : plane_surface """ if not engine: engine = mlab.get_engine() if not scene: scene = engine.scenes[0] # ensure the normal vector is of length 1 nvl = np.sqrt(normal[0]**2 + normal[1]**2 + normal[2]**2) if nvl != 1.00: normal = normal[0]/nvl, normal[1]/nvl, normal[2]/nvl plane_surf = BuiltinSurface() engine.add_source(plane_surf) plane_surf.source = 'plane' #plane_surf.data_source.center = np.array(loc) origin = loc[0] - len_x/2.0, loc[1] - len_y/2.0, 0.0 point1 = origin[0] + len_x, origin[1], 0.0 point2 = origin[0], origin[1] + len_y, 0.0 plane_surf.data_source.origin = np.array(origin) plane_surf.data_source.point1 = np.array(point1) plane_surf.data_source.point2 = np.array(point2) plane_surf.data_source.center = np.array(loc) # this needs to be set here plane_surf.data_source.normal = np.array(normal) surface = Surface() engine.add_filter(surface, plane_surf) # Add color and opacity and lighting properties surface.actor.property.color = planeCol surface.actor.property.opacity = planeOpa surface.actor.property.lighting = planeLighting if drawNormal: # Draw plane normal norm_drawn = mlab.quiver3d(loc[0], loc[1], loc[2], [normal[0]], [normal[1]], [normal[2]], # required to do this way for matching shape ... else Mayavi will throw error. This may also be related to the issue -- https://github.com/enthought/mayavi/issues/85 mode='arrow', resolution=16, color=(1,0,0), scale_factor=normVecScale, opacity=1.0) return plane_surf, norm_drawn else: return plane_surf
def _plot_rays_mayavi(inventory, catalog, phase_list=['P'], colorscheme='default', animate=False, savemovie=False, figsize=(800, 800), taup_model='iasp91', coastlines='internal', icol=0, event_labels=True, station_labels=True, fname_out=None, view_dict=None): """ internal mayavi plotting routine. Check out the plot_rays routine for more information on the parameters """ try: from mayavi import mlab except Exception as err: print(err) msg = ("obspy failed to import mayavi. " "You need to install the mayavi module " "(e.g. conda install mayavi, pip install mayavi). " "If it is installed and still doesn't work, " "try setting the environmental variable QT_API to " "pyqt (e.g. export QT_API=pyqt) before running the " "code. Another option is to avoid mayavi and " "directly use kind='vtk' for vtk file output of the " "radiation pattern that can be used by external " "software like paraview") raise ImportError(msg) if isinstance(taup_model, str): from obspy.taup import TauPyModel model = TauPyModel(model=taup_model) else: model = taup_model if fname_out is not None: offscreen = True else: offscreen = False nphases = len(phase_list) greatcircles = get_ray_paths(inventory=inventory, catalog=catalog, phase_list=phase_list, coordinate_system='XYZ', taup_model=model) if len(greatcircles) == 0: raise ValueError('no paths found for the input stations and events') # define colorschemes if colorscheme == 'dark' or colorscheme == 'default': # we use the color set that is used in taup, but adjust the lightness # to get nice shiny rays that are well visible in the dark 3d plots from obspy.taup.tau import COLORS ncolors = len(COLORS) color_hues = [rgb_to_hls(*hex2color(col))[0] for col in COLORS] # swap green and red to start with red: color_hues[2], color_hues[1] = color_hues[1], color_hues[2] # first color is for the continents etc: continent_hue = color_hues[0] event_hue = color_hues[-1] # the remaining colors are for the rays: ray_hues = color_hues[1:-1] # now convert all of the hues to rgb colors: ray_light = 0.45 ray_sat = 1.0 raycolors = [ hls_to_rgb(ray_hues[(iphase + icol) % (ncolors - 2)], ray_light, ray_sat) for iphase in range(nphases) ] stationcolor = hls_to_rgb(continent_hue, 0.7, 0.7) continentcolor = hls_to_rgb(continent_hue, 0.3, 0.2) eventcolor = hls_to_rgb(event_hue, 0.7, 0.7) cmbcolor = continentcolor bgcolor = (0, 0, 0) # sizes: sttextsize = (0.015, 0.015, 0.015) stmarkersize = 0.01 evtextsize = (0.015, 0.015, 0.015) evmarkersize = 0.05 elif colorscheme == 'bright': # colors (a distinct colors for each phase): lightness = 0.2 saturation = 1.0 ncolors = nphases + 2 # two extra colors for continents and events hues = np.linspace(0., 1. - 1. / ncolors, ncolors) raycolors = [ hls_to_rgb(hue, lightness, saturation) for hue in hues[2:] ] stationcolor = hls_to_rgb(hues[0], 0.2, 0.5) continentcolor = hls_to_rgb(hues[0], 0.6, 0.2) eventcolor = hls_to_rgb(hues[1], 0.2, 0.5) cmbcolor = continentcolor bgcolor = (1, 1, 1) # sizes: # everything has to be larger in bright background plot because it # is more difficult to read sttextsize = (0.02, 0.02, 0.02) stmarkersize = 0.02 evtextsize = (0.02, 0.02, 0.02) evmarkersize = 0.06 else: raise ValueError('colorscheme {:s} not recognized'.format(colorscheme)) # assemble each phase and all stations/events to plot them in a single call stations_loc = [] stations_lab = [] events_loc = [] events_lab = [] phases = [[] for iphase in range(nphases)] for gcircle, name, stlabel, time, magnitude, evid, _ in greatcircles: iphase = phase_list.index(name) phases[iphase].append(gcircle) date = UTCDateTime(time) evlabel = '{:s} | M{:.1f}'.format(str(date.date), magnitude) if stlabel not in stations_lab: x, y, z = gcircle[0, -1], gcircle[1, -1], gcircle[2, -1] stations_loc.append((x, y, z)) stations_lab.append(stlabel) if evlabel not in events_lab: x, y, z = gcircle[0, 0], gcircle[1, 0], gcircle[2, 0] events_loc.append((x, y, z)) events_lab.append(evlabel) # now begin mayavi plotting if offscreen: mlab.options.offscreen = True fig = mlab.figure(size=figsize, bgcolor=bgcolor) # define connectivity of each ray and add it to the scene for iphase, phase in enumerate(phases): index = 0 connects = [] for ray in phase: ndim, npoints = ray.shape connects.append( np.vstack([ np.arange(index, index + npoints - 1.5), np.arange(index + 1, index + npoints - .5) ]).T) index += npoints # collapse all points of the phase into a long array if len(phase) > 1: points = np.hstack(phase) elif len(phase) == 1: points = phase[0] else: continue connects = np.vstack(connects) # Create the points src = mlab.pipeline.scalar_scatter(*points) # Connect them src.mlab_source.dataset.lines = connects # The stripper filter cleans up connected lines lines = mlab.pipeline.stripper(src) color = raycolors[iphase] mlab.pipeline.surface(lines, line_width=0.5, color=color) # plot all stations fig.scene.disable_render = True # Super duper trick stxs, stys, stzs = zip(*stations_loc) mlab.points3d(stxs, stys, stzs, scale_factor=stmarkersize, color=stationcolor) if station_labels: for loc, stlabel in zip(stations_loc, stations_lab): mlab.text3d(loc[0], loc[1], loc[2], stlabel, scale=sttextsize, color=stationcolor) # plot all events evxs, evys, evzs = zip(*events_loc) evsource = mlab.pipeline.vector_scatter(evxs, evys, evzs, -np.array(evxs), -np.array(evys), -np.array(evzs)) evmarkers = mlab.pipeline.glyph(evsource, scale_factor=evmarkersize, scale_mode='none', color=eventcolor, mode='cone', resolution=8) evmarkers.glyph.glyph_source.glyph_position = 'head' if event_labels: for loc, evlabel in zip(events_loc, events_lab): mlab.text3d(loc[0], loc[1], loc[2], evlabel, scale=evtextsize, color=eventcolor) fig.scene.disable_render = False # Super duper trick # read and plot coastlines if coastlines == 'internal': from mayavi.sources.builtin_surface import BuiltinSurface data_source = BuiltinSurface(source='earth', name="Continents") data_source.data_source.on_ratio = 1 else: data_source = mlab.pipeline.open(coastlines) coastmesh = mlab.pipeline.surface(data_source, opacity=1.0, line_width=0.5, color=continentcolor) coastmesh.actor.actor.scale = np.array([1.02, 1.02, 1.02]) # plot block sphere that hides the backside of the continents rad = 0.99 phi, theta = np.mgrid[0:np.pi:51j, 0:2 * np.pi:51j] x = rad * np.sin(phi) * np.cos(theta) y = rad * np.sin(phi) * np.sin(theta) z = rad * np.cos(phi) blocksphere = mlab.mesh(x, y, z, color=bgcolor) blocksphere.actor.property.frontface_culling = True # front not rendered # make CMB sphere r_earth = model.model.radius_of_planet r_cmb = r_earth - model.model.cmb_depth rad = r_cmb / r_earth phi, theta = np.mgrid[0:np.pi:201j, 0:2 * np.pi:201j] x = rad * np.sin(phi) * np.cos(theta) y = rad * np.sin(phi) * np.sin(theta) z = rad * np.cos(phi) cmb = mlab.mesh(x, y, z, color=cmbcolor, opacity=0.3, line_width=0.5) cmb.actor.property.interpolation = 'gouraud' # cmb.actor.property.interpolation = 'flat' # make ICB sphere r_iocb = r_earth - model.model.iocb_depth rad = r_iocb / r_earth phi, theta = np.mgrid[0:np.pi:101j, 0:2 * np.pi:101j] x = rad * np.sin(phi) * np.cos(theta) y = rad * np.sin(phi) * np.sin(theta) z = rad * np.cos(phi) icb = mlab.mesh(x, y, z, color=cmbcolor, opacity=0.3, line_width=0.5) icb.actor.property.interpolation = 'gouraud' if view_dict is None: view_dict = { 'azimuth': 0., 'elevation': 90., 'distance': 4., 'focalpoint': (0., 0., 0.) } mlab.view(**view_dict) # to make a movie from the image files, you can use the command: # avconv -qscale 5 -r 20 -b 9600 -i %05d.png -vf scale=800:752 movie.mp4 if animate and not offscreen: @mlab.show @mlab.animate(delay=20) def anim(): iframe = 0 while 1: if savemovie and iframe < 360: mlab.savefig('{:05d}.png'.format(iframe)) # camera moves from East to West opposite of Earth's rotation fig.scene.camera.azimuth(-1.) fig.scene.render() iframe += 1 yield anim() # Starts the animation. else: if offscreen: mlab.savefig(fname_out) else: mlab.show()
def plot_3d(self, **kwargs): ''' plots a 3d earth model using mayavi kwargs: grid_x: Draw the x plane (default False) grid_y: Draw the y plane (default False) grid_z: Draw the z plane (default False) earth: Draw the earth outline and coastlines (default True) plot_quakes: Draw earthquakes in earthquake_list (default False) earthquake_list: a list of earthquake coordinates specified by (lat,lon,depth) ''' import mayavi from mayavi import mlab from tvtk.api import tvtk from mayavi.scripts import mayavi2 from spherical_section import generate from mayavi.sources.vtk_data_source import VTKDataSource from mayavi.sources.builtin_surface import BuiltinSurface from mayavi.modules.api import Outline, GridPlane, ScalarCutPlane draw_gridx = kwargs.get('grid_x', False) draw_gridy = kwargs.get('grid_y', False) draw_gridz = kwargs.get('grid_z', False) draw_earth = kwargs.get('earth', True) draw_quakes = kwargs.get('draw_quakes', False) earthquake_list = kwargs.get('earthquake_list', 'none') #build the spherical section dims = (len(self.rad) - 1, len(self.lon) - 1, len(self.colat) - 1) pts = generate(phi=np.radians(self.lon), theta=np.radians(self.colat), rad=self.rad) sgrid = tvtk.StructuredGrid(dimensions=dims) sgrid.points = pts s = np.zeros(len(pts)) #map data onto the grid count = 0 for i in range(0, len(self.colat) - 1): for j in range(0, len(self.lon) - 1): for k in range(0, len(self.rad) - 1): s[count] = self.data[k, i, j] sgrid.point_data.scalars = s sgrid.point_data.scalars.name = 'scalars' count += 1 #use vtk dataset src = VTKDataSource(data=sgrid) #set figure defaults mlab.figure(bgcolor=(0, 0, 0)) #outline mlab.pipeline.structured_grid_outline(src, opacity=0.3) #show grid planes if draw_gridx == True: gx = mlab.pipeline.grid_plane(src, color=(1, 1, 1), opacity=0.25) gx.grid_plane.axis = 'x' if draw_gridy == True: gy = mlab.pipeline.grid_plane(src, color=(1, 1, 1), opacity=0.25) gy.grid_plane.axis = 'y' if draw_gridz == True: gz = mlab.pipeline.grid_plane(src, color=(1, 1, 1), opacity=0.25) gz.grid_plane.axis = 'z' #cutplane mlab.pipeline.scalar_cut_plane(src, plane_orientation='y_axes', colormap='jet', view_controls=False) #draw earth and coastlines if draw_earth == True: coastline_src = BuiltinSurface(source='earth', name='Continents') coastline_src.data_source.radius = 6371.0 coastline_src.data_source.on_ratio = 1 mlab.pipeline.surface(coastline_src) #plot earthquakes if draw_quakes == True: print "Sorry, this doesn't work yet" lat_pts = earthquake_list[:, 0] lon_pts = earthquake_list[:, 1] rad_pts = 6371.0 - earthquake_list[:, 2] theta_pts = np.radians(lat_pts) phi_pts = np.radians(lon_pts) #convert point to cartesian x_pts = rad_pts * np.cos(phi_pts) * np.sin(theta_pts) y_pts = rad_pts * np.sin(phi_pts) * np.sin(theta_pts) z_pts = rad_pts * np.cos(theta_pts) eq_pts = mlab.points3d(x_pts, y_pts, z_pts, scale_mode='none', scale_factor=100.0, color=(1, 1, 0)) mlab.show()