# plotting3d.py
'''
Introduction into mayavi2.
'''

# Start ipython with
# ipython --gui=qt


import numpy as np
from mayavi import mlab

#
# 0d and 1d plots.
#

# points.
fig = mlab.figure()
x = np.linspace(0, 2*np.pi, 10)
y = np.sin(x)
z = x**2
mlab.points3d(x, y, z, color='k', mode='sphere', scale_factor=0.5)

s = x*y*z
mlab.points3d(x, y, z, s, scale_mode='scalar')

# lines.
c = np.linspace(0, 10*np.pi, 1000)
x = np.cos(c)
y = np.sin(c)
z = c/5
mlab.plot3d(x, y, z, c)


#
# 2d plots.
#

# imshow.
x = np.linspace(0, 2*np.pi, 100)
y = np.linspace(0, 2*np.pi, 100)
xx, yy = np.meshgrid(x, y, indexing='ij')
zz = np.sin(xx) + np.cos(yy)
mlab.imshow(zz, colormap='GnBu')

# surf
mlab.surf(10*zz)

# mesh
phi = np.linspace(0, 2*np.pi, 10)
theta = np.linspace(-np.pi/2, np.pi/2, 10)
pp, tt = np.meshgrid(phi, theta)
x = np.cos(pp)*np.cos(tt)*(1+np.cos(3*tt)/5)
y = np.sin(pp)*np.cos(tt)*(1+np.cos(3*tt)/5)
z = np.sin(tt)*(1+np.cos(3*tt)/5)
mlab.mesh(x, y, z, representation='fancymesh')


#
# 3d plots.
#

# contours in 3d.
x = np.linspace(-5, 5, 100)
y = x
z = y
xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
ff = xx**2*0.5 + yy**2 + 2*zz**2
con = mlab.contour3d(ff, contours=4, transparent=True)

# vector field.
rr2 = xx**2+yy**2+zz**2
cc = 1/(1+rr2)**3
bb = np.zeros(cc.shape + (3, ))
bb[..., 0] = 2*(yy - xx*zz)*cc
bb[..., 1] = -2*(xx + yy*zz)*cc
bb[..., 2] = (-1 + xx**2 + yy**2 - zz**2)*cc
mlab.quiver3d(bb[..., 0], bb[..., 1], bb[..., 2], mask_points=10)

# flows.
mlab.flow(bb[..., 0], bb[..., 1], bb[..., 2], integration_direction='both', seed_resolution=10)
mlab.flow(xx, yy, zz, bb[..., 0], bb[..., 1], bb[..., 2], integration_direction='both', seed_resolution=10, extent=[xx.min(), xx.max(), yy.min(), yy.max(), zz.min(), zz.max()], seedtype='plane')

# volume rendering.
ff = np.sin(xx*yy*zz)/(xx*yy*zz)
mlab.pipeline.volume(mlab.pipeline.scalar_field(ff))

#
# Handling figures.
#

mlab.clf()
mlab.gcf()
mlab.savefig()
mlab.draw()


#
# Making the plot pretty.
#

mlab.axes(color='k', line_width=2)
mlab.xlabel('x')
mlab.outline(color=(0, 0, 0))

# Camera control.
f = mlab.gcf()
camera = f.scene.camera
camera.yaw(45)