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