day2/demo.py
changeset 0 9243d75024cc
equal deleted inserted replaced
-1:000000000000 0:9243d75024cc
       
     1 import numpy as np
       
     2 import pylab as P
       
     3 from enthought.mayavi import mlab
       
     4 from scipy.integrate import odeint
       
     5 
       
     6 
       
     7 ######################################################################
       
     8 def lorenz(x, y, z, s=10.,r=28., b=8./3.):
       
     9     """The Lorenz system."""
       
    10     u = s*(y-x)
       
    11     v = r*x -y - x*z
       
    12     w = x*y - b*z
       
    13     return u, v, w
       
    14 
       
    15 def lorenz_ode(r, t):
       
    16     x, y, z = r
       
    17     u, v, w = lorenz(x, y, z)
       
    18     return np.array([u, v, w])
       
    19 
       
    20 def show_lorenz(new_fig=True):
       
    21     x, y, z = np.mgrid[-50:50:100j,-50:50:100j,-10:60:70j]
       
    22     u, v, w = lorenz(x, y, z)
       
    23 
       
    24     if new_fig:
       
    25         fig = mlab.figure(size=(600,600), bgcolor=(0, 0, 0))
       
    26 
       
    27     # Plot the flow of trajectories with suitable parameters.
       
    28     f = mlab.flow(x, y, z, u, v, w, line_width=3, colormap='Paired')
       
    29     f.module_manager.scalar_lut_manager.reverse_lut = True
       
    30     f.stream_tracer.integration_direction = 'both'
       
    31     f.stream_tracer.maximum_propagation = 200
       
    32 
       
    33 def show_lorenz_traj(start=(0.,0.,0.)):
       
    34     t = np.linspace(0., 50., 2000)
       
    35     ret = lorenz_traj(t, np.asarray(start))
       
    36     x, y, z = ret[:,0], ret[:,1], ret[:,2]
       
    37     mlab.plot3d(x, y, z, t, tube_radius=None)
       
    38 
       
    39 
       
    40 def lorenz_traj(t, start=(0.,0.,0.)):
       
    41     return odeint(lorenz_ode, start, t)
       
    42     
       
    43 
       
    44 ######################################################################
       
    45 
       
    46 def convection_cell(x, y, z):
       
    47     """velocity field of a multi-axis convection cell [1], in
       
    48     hydrodynamics, as defined by its components sampled on a grid. 
       
    49     """
       
    50     u =  np.sin(np.pi*x) * np.cos(np.pi*z)
       
    51     v = -2*np.sin(np.pi*y) * np.cos(2*np.pi*z)
       
    52     w = np.cos(np.pi*x)*np.sin(np.pi*z) + np.cos(np.pi*y)*np.sin(2*np.pi*z)
       
    53     return u, v, w
       
    54 
       
    55 def show_convection1(new_fig=True):
       
    56     x, y, z = np.mgrid[0:1:20j, 0:1:20j, 0:1:20j]
       
    57     u, v, w = convection_cell(x, y, z)
       
    58     if new_fig:
       
    59         mlab.figure(size=(600,600))
       
    60 
       
    61     mlab.quiver3d(u, v, w)
       
    62     mlab.outline()
       
    63 
       
    64 def show_convection2(new_fig=True):
       
    65     x, y, z = np.mgrid[0:1:20j, 0:1:20j, 0:1:20j]
       
    66     u, v, w = convection_cell(x, y, z)
       
    67     if new_fig:
       
    68         mlab.figure(size=(600,600))
       
    69 
       
    70     src = mlab.pipeline.vector_field(u, v, w)
       
    71     mlab.pipeline.vectors(src, mask_points=20, scale_factor=3.)
       
    72     mlab.outline()
       
    73 
       
    74 def show_convection3(new_fig=True):
       
    75     x, y, z = np.mgrid[0:1:20j, 0:1:20j, 0:1:20j]
       
    76     u, v, w = convection_cell(x, y, z)
       
    77 
       
    78     if new_fig:
       
    79         mlab.figure(fgcolor=(0., 0., 0.), bgcolor=(0.5, 0.5, 0.5),
       
    80                     size=(600,600))
       
    81     src = mlab.pipeline.vector_field(u, v, w)
       
    82     magnitude = mlab.pipeline.extract_vector_norm(src)
       
    83     mlab.outline()
       
    84 
       
    85     # We apply the following modules on the magnitude object, in order to
       
    86     # be able to display the norm of the vectors, eg as the color.
       
    87     iso = mlab.pipeline.iso_surface(magnitude, contours=[1.9, ], 
       
    88                                     opacity=0.3)
       
    89 
       
    90     vec = mlab.pipeline.vectors(magnitude, 
       
    91                                 mask_points=40,
       
    92                                 line_width=1,
       
    93                                 color=(1, 1, 1),
       
    94                                 scale_factor=4.)
       
    95 
       
    96     flow = mlab.pipeline.streamline(magnitude, 
       
    97                                     seedtype='plane',
       
    98                                     seed_visible=True,
       
    99                                     seed_scale=0.5,
       
   100                                     seed_resolution=1,
       
   101                                     linetype='ribbon',)
       
   102 
       
   103     vcp = mlab.pipeline.vector_cut_plane(magnitude, mask_points=2,
       
   104                                          scale_factor=4,
       
   105                                          colormap='jet',
       
   106                                          plane_orientation='x_axes')
       
   107 
       
   108 
       
   109 def test():
       
   110     show_lorenz()
       
   111     show_lorenz_traj((10, 50, 50))
       
   112     show_convection1()
       
   113     show_convection2()
       
   114     show_convection3()
       
   115 
       
   116 if __name__ == '__main__':
       
   117     test()
       
   118