216.sagews

def plot3d_vector_field(f,xrange,yrange,zrange,grid=[5,5,5],thick_arrows=False,**kwds):
    """
    Return a 3-dimensional representation of a vector field contained
    in the specified frame.

    INPUT:
        f -- a 3-tuple (or list) of functions or expressions
        xrange, yrange, zrange -- 3-tuples of the form (var,var_min,var_max)
            defining a box in which to display the plot
        grid -- (default: [5,5,5]) a list specifying how many arrows
            to draw in each dimension
        thick_arrows -- (default: False) setting this to True may slow
            things down a bit...
    """
    from sage.ext.fast_eval import fast_float
    from sage.plot.plot3d.shapes2 import frame3d

    framepts = [(xrange[1],yrange[1],zrange[1]),(xrange[2],yrange[2],zrange[2])]
    x = xrange[0]
    y = yrange[0]
    z = zrange[0]
    xstep = (framepts[1][0]-framepts[0][0])/grid[0]
    ystep = (framepts[1][1]-framepts[0][1])/grid[1]
    zstep = (framepts[1][2]-framepts[0][2])/grid[2]
    gridpoints = [(framepts[0][0] + xstep*(i+0.5),\
        framepts[0][1] + ystep*(j+0.5), framepts[0][2] + zstep*(k+0.5))\
        for i in range(grid[0]) for j in range(grid[1]) for k in range(grid[2])]

    f_fast = fast_float(tuple(f),str(x),str(y),str(z))
    f_vals = [vector((f_fast[0](*point),f_fast[1](*point),f_fast[2](*point))) for point in gridpoints]
    fmax = max([val.norm() for val in f_vals])

    theplot = line3d(framepts,opacity=0) # stop the bounding box from shrinking
    theplot += text3d(str(x),(.5*(framepts[0][0]+framepts[1][0]),framepts[0][1],framepts[0][2])) +\
        text3d(str(y),(framepts[1][0],.5*(framepts[0][1]+framepts[1][1]),framepts[0][2])) +\
        text3d(str(z),(framepts[0][0],framepts[0][1],.5*(framepts[0][2]+framepts[1][2])))

    if thick_arrows:
        radius = lambda offset: 0.1*abs(offset)
    else:
        radius = lambda offset: None
    for i in range(len(gridpoints)):
        offset = 0.5/fmax*vector((f_vals[i][0]*xstep,f_vals[i][1]*ystep,f_vals[i][2]*zstep))
        tail = vector(gridpoints[i]) - offset
        head = vector(gridpoints[i]) + offset
        theplot += line3d([tail,head],arrow_head=True,radius=radius(offset),**kwds)

    return theplot

# A simple example illustrating the syntax of the plot3d_vector_field function.
# A particle starting near the origin will start out flowing in the positive
# x, y, and z directions, but will hook back to the zy-plane.

var('x,y,z')
plot3d_vector_field((cos(y),sin(x),sin(z)),(x,0,pi),(y,0,pi),(z,0,pi))

# This example is taken from #18 in 16.1.

var('x,y,z')
F = (x,y,z)
plot3d_vector_field(F,(x,-1,1),(y,-1,1),(z,-1,1),color='red')

# The vector field (y,-x,z) swirls upwards from the xy-plane.
# The grid=[6,6,6] option makes the function plot 6 arrows in each direction.

var('x,y,z')
plot3d_vector_field((y,-x,z),(x,-1,1),(y,-1,1),(z,0,2),grid=[6,6,6],color='gray')

# A plot of the gradient vector field of a function f together with some level surfaces of f.
# Note how the gradient vectors are perpendicular to the level surfaces.

def spherical(rho,theta,phi):
    return (rho*sin(phi)*cos(theta),rho*sin(phi)*sin(theta),rho*cos(phi))

var('theta,phi')
f(x,y,z) = x^2 + y^2 + z^2
demo = plot3d_vector_field(f.gradient(),(x,0,2),(y,-2,2),(z,-2,2),color='brown')
num_spheres = 4
for k in range(num_spheres):
    demo += parametric_plot3d(spherical(2*(k+1)/num_spheres,theta,phi),(theta,-pi/2,pi/2),(phi,0,pi),\
        opacity=(1-k/num_spheres),rgbcolor=(.5,.5*((k+1)/num_spheres),1))
demo.show(aspect_ratio=[1,1,1])

# An example from the worksheet.

var('x,y,z,t')
H = (z,y,x)
r = (cos(6*t),sin(6*t),t)
vfield = plot3d_vector_field(H,(x,-1,1),(y,-1,1),(z,0,pi),color='orange')
curve = parametric_plot3d(r,(t,0,pi),color='red')
vfield + curve