"""
Display electric field lines of a uniformly polarized dielectric sphere.

Polarization P0 is oriented along Z axis.

In cartesian coordinates, the electric field outside of the sphere 

   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 = P0/(3 \epsilon_0) is the magnitude of the electric field inside 
the sphere

   Ex = 0, Ey = 0, Ez/E0 = 1.

All distances are measured in units of 'R', the radius of the sphere
(i.e. the radius of the sphere in dimensionless units is one).

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.
"""

import numpy as np

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

def dipole(x, y, z):
    """ Electric field of a dipole. Returns Ex,Ey,Ez """
    rr = np.sqrt(x**2 + y**2 + z**2)
    rrinv3 = 1/rr**3
    x_proj = x/rr
    y_proj = y/rr
    z_proj = z/rr
    return 3*rrinv3*x_proj*z_proj, 3*rrinv3*y_proj*z_proj, rrinv3*(3*z_proj**2-1)

def inside_polarized_sphere(x, y, z, E, val):
    """ Electric field inside a polarized dielectric sphere. """
    rr = np.sqrt(x**2 + y**2 + z**2)
    return np.where(rr > 1., E, val)

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

Ex, Ey, Ez = dipole(x,y,z)
Ex = inside_polarized_sphere(x,y,z,Ex,0)
Ey = inside_polarized_sphere(x,y,z,Ey,0)
Ez = inside_polarized_sphere(x,y,z,Ez,-1)

# Free memory early
del x, y, z

#### 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=8)

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=1, scale_factor=1)

mlab.show()