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

var('x,y,z')
plot3d_vector_field((y^2,x^2,1),(x,-3,3),(y,-3,3),(z,-3,3))

var('x,y,z')
plot3d_vector_field((x+y,-x,0),(x,-3,3),(y,-3,3),(z,0,0))