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()