"""
This example uses the streamline module to display field lines of an
electric dipole oriented along Z axis.

The electric field from a dipole is calculated from 
eqns (3.103), D. Griffiths "Introduction to Electrodynamics", 3rd ed.
In cartesian coordinates,

   Ex/E0 = 3/r^3 z/r x/r
   Ey/E0 = 3/r^3 z/r y/r
   Ez/E0 = 1/r^3 (3 z^2/r^2 - 1)

where E0 = q/(4\pi \epsilon_0 d^2) is the characteristic electric field
d is the characteristic length ("distance" between the charges), all distances
are measured in units of 'd'.

To get a prettier result, we use a fairly large grid to sample the
field. As a consequence, we need to clear temporary arrays as soon as
possible.

This scrip is a slightly modified version of 'magnetic_field_lines.py' by

 Author: Gael Varoquaux <gael.varoquaux@normalesup.org> 
 Copyright (c) 2007, Enthought, Inc.
 License: BSD Style.

which is distributed as an example with mayavi2
"""

import numpy as np

#### Calculate the field #####################################################

x, y, z = [e for e in np.ogrid[-2:2:31j, -2:2:31j, -2:2:31j] ]

rr = np.sqrt(x**2 + y**2 + z**2 + 4.)

rrinv3=1/rr**3

x_proj = x/rr
y_proj = y/rr
z_proj = z/rr

# Free memory early
del x, y, z, rr

Ex = 3 * rrinv3 * x_proj * z_proj
Ey = 3 * rrinv3 * y_proj * z_proj
Ez = rrinv3 * (3 * z_proj**2 - 1)

# Free memory early
del x_proj, y_proj, z_proj, rrinv3

#### Visualize the field #####################################################

from enthought.mayavi import mlab

fig = mlab.figure(1, size=(600, 600))

field = mlab.pipeline.vector_field(Ex, Ey, Ez)

# The above call makes a copy of the arrays, so we delete
# this copy to free memory.
del Ex, Ey, Ez

#vec = mlab.pipeline.vectors(field, scale_factor=3., mask_points=27)

magnitude = mlab.pipeline.extract_vector_norm(field)

electric_field_lines = mlab.pipeline.streamline(magnitude,seed_visible=False,seedtype='sphere',seed_scale=.9,
        seed_resolution=10,transparent=True,integration_direction='both')

#cut = mlab.pipeline.vector_cut_plane(magnitude, mask_points=2, scale_factor=3)

mlab.show()