#!/usr/bin/env python
#coding=utf-8

try:
    from mayavi import mlab
except ImportError:
    from enthought.mayavi import mlab

import numpy as np

XX, YY = np.mgrid[-1:1:10j, -1:1:10j]

def saddle(x,y):
    return x**2-y**2

ZZ = saddle(XX, YY)

mlab.surf(XX, YY, ZZ)

mlab.figure()
def lorenz(x,y,z):
    dxdt = 10*(y-x)
    dydt = 28*x-y-x*z
    dzdt = x*y - 8*z/3
    return dxdt, dydt, dzdt

XX, YY, ZZ = np.mgrid[-50:50:20j,-50:50:20j, -10:60:20j]
u,v,w = lorenz(XX, YY, ZZ)
mlab.quiver3d(XX, YY, ZZ, u, v, w)

from scipy.integrate import odeint
ics = (10, 50, 50)
dt = 0.001
t = np.arange(0, 50, dt)
lorenz_ode = lambda r, t: lorenz(*r)
trajectory = odeint(lorenz_ode, ics, t)

traj_x, traj_y, traj_z = trajectory.T
mlab.plot3d(traj_x, traj_y, traj_z, t, tube_radius=None)

mlab.figure()
mlab.quiver3d(XX, YY, ZZ, u, v, w)
mlab.flow(XX, YY, ZZ, u, v, w)
mlab.show()