diff -r 000000000000 -r 9243d75024cc day2/demo.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/day2/demo.py Fri Oct 02 16:42:15 2009 +0530 @@ -0,0 +1,118 @@ +import numpy as np +import pylab as P +from enthought.mayavi import mlab +from scipy.integrate import odeint + + +###################################################################### +def lorenz(x, y, z, s=10.,r=28., b=8./3.): + """The Lorenz system.""" + u = s*(y-x) + v = r*x -y - x*z + w = x*y - b*z + return u, v, w + +def lorenz_ode(r, t): + x, y, z = r + u, v, w = lorenz(x, y, z) + return np.array([u, v, w]) + +def show_lorenz(new_fig=True): + x, y, z = np.mgrid[-50:50:100j,-50:50:100j,-10:60:70j] + u, v, w = lorenz(x, y, z) + + if new_fig: + fig = mlab.figure(size=(600,600), bgcolor=(0, 0, 0)) + + # Plot the flow of trajectories with suitable parameters. + f = mlab.flow(x, y, z, u, v, w, line_width=3, colormap='Paired') + f.module_manager.scalar_lut_manager.reverse_lut = True + f.stream_tracer.integration_direction = 'both' + f.stream_tracer.maximum_propagation = 200 + +def show_lorenz_traj(start=(0.,0.,0.)): + t = np.linspace(0., 50., 2000) + ret = lorenz_traj(t, np.asarray(start)) + x, y, z = ret[:,0], ret[:,1], ret[:,2] + mlab.plot3d(x, y, z, t, tube_radius=None) + + +def lorenz_traj(t, start=(0.,0.,0.)): + return odeint(lorenz_ode, start, t) + + +###################################################################### + +def convection_cell(x, y, z): + """velocity field of a multi-axis convection cell [1], in + hydrodynamics, as defined by its components sampled on a grid. + """ + u = np.sin(np.pi*x) * np.cos(np.pi*z) + v = -2*np.sin(np.pi*y) * np.cos(2*np.pi*z) + w = np.cos(np.pi*x)*np.sin(np.pi*z) + np.cos(np.pi*y)*np.sin(2*np.pi*z) + return u, v, w + +def show_convection1(new_fig=True): + x, y, z = np.mgrid[0:1:20j, 0:1:20j, 0:1:20j] + u, v, w = convection_cell(x, y, z) + if new_fig: + mlab.figure(size=(600,600)) + + mlab.quiver3d(u, v, w) + mlab.outline() + +def show_convection2(new_fig=True): + x, y, z = np.mgrid[0:1:20j, 0:1:20j, 0:1:20j] + u, v, w = convection_cell(x, y, z) + if new_fig: + mlab.figure(size=(600,600)) + + src = mlab.pipeline.vector_field(u, v, w) + mlab.pipeline.vectors(src, mask_points=20, scale_factor=3.) + mlab.outline() + +def show_convection3(new_fig=True): + x, y, z = np.mgrid[0:1:20j, 0:1:20j, 0:1:20j] + u, v, w = convection_cell(x, y, z) + + if new_fig: + mlab.figure(fgcolor=(0., 0., 0.), bgcolor=(0.5, 0.5, 0.5), + size=(600,600)) + src = mlab.pipeline.vector_field(u, v, w) + magnitude = mlab.pipeline.extract_vector_norm(src) + mlab.outline() + + # We apply the following modules on the magnitude object, in order to + # be able to display the norm of the vectors, eg as the color. + iso = mlab.pipeline.iso_surface(magnitude, contours=[1.9, ], + opacity=0.3) + + vec = mlab.pipeline.vectors(magnitude, + mask_points=40, + line_width=1, + color=(1, 1, 1), + scale_factor=4.) + + flow = mlab.pipeline.streamline(magnitude, + seedtype='plane', + seed_visible=True, + seed_scale=0.5, + seed_resolution=1, + linetype='ribbon',) + + vcp = mlab.pipeline.vector_cut_plane(magnitude, mask_points=2, + scale_factor=4, + colormap='jet', + plane_orientation='x_axes') + + +def test(): + show_lorenz() + show_lorenz_traj((10, 50, 50)) + show_convection1() + show_convection2() + show_convection3() + +if __name__ == '__main__': + test() +