%auto
from sage.plot.primitive import GraphicPrimitive
from sage.misc.decorators import options
from sage.misc.misc import xsrange
# Below is the base class that is used to make 'field plots'.
# Its implementation is motivated by 'PlotField'.
# Currently it is used to make the functions 'plot_vector_field'
# and 'plot_slope_field'.
# TODO: use this to make these functions:
# 'plot_gradient_field' and 'plot_hamiltonian_field'
class PlotField_local(GraphicPrimitive):
"""
Primitive class that initializes the
PlotField graphics type
"""
def __init__(self, xpos_array, ypos_array, xvec_array, yvec_array, options):
"""
Create the graphics primitive PlotField. This sets options
and the array to be plotted as attributes.
EXAMPLES::
sage: x,y = var('x,y')
sage: R=plot_slope_field(x+y,(x,0,1),(y,0,1),plot_points=2)
sage: r=R[0]
sage: r.options()['headaxislength']
0
sage: r.xpos_array
[0.0, 0.0, 1.0, 1.0]
sage: r.yvec_array
masked_array(data = [0.0 0.707106781187 0.707106781187 0.894427191],
mask = [False False False False],
fill_value = 1e+20)
TESTS:
We test dumping and loading a plot::
sage: x,y = var('x,y')
sage: P = plot_vector_field((sin(x), cos(y)), (x,-3,3), (y,-3,3))
sage: Q = loads(dumps(P))
"""
self.xpos_array = xpos_array
self.ypos_array = ypos_array
self.xvec_array = xvec_array
self.yvec_array = yvec_array
GraphicPrimitive.__init__(self, options)
def get_minmax_data(self):
"""
Returns a dictionary with the bounding box data.
EXAMPLES::
sage: x,y = var('x,y')
sage: d = plot_vector_field((.01*x,x+y), (x,10,20), (y,10,20))[0].get_minmax_data()
sage: d['xmin']
10.0
sage: d['ymin']
10.0
"""
from sage.plot.plot import minmax_data
return minmax_data(self.xpos_array, self.ypos_array, dict=True)
def _allowed_options(self):
"""
Returns a dictionary with allowed options for PlotField.
EXAMPLES::
sage: x,y = var('x,y')
sage: P=plot_vector_field((sin(x), cos(y)), (x,-3,3), (y,-3,3))
sage: d=P[0]._allowed_options()
sage: d['pivot']
'Where the arrow should be placed in relation to the point (tail, middle, tip)'
"""
return {'plot_points':'How many points to use for plotting precision',
'pivot': 'Where the arrow should be placed in relation to the point (tail, middle, tip)',
'headwidth': 'Head width as multiple of shaft width, default is 3',
'headlength': 'head length as multiple of shaft width, default is 5',
'headaxislength': 'head length at shaft intersection, default is 4.5',
'zorder':'The layer level in which to draw',
'color':'The color of the arrows'}
def _repr_(self):
"""
String representation of PlotField graphics primitive.
EXAMPLES::
sage: x,y = var('x,y')
sage: P=plot_vector_field((sin(x), cos(y)), (x,-3,3), (y,-3,3))
sage: P[0]
PlotField defined by a 20 x 20 vector grid
"""
return "PlotField defined by a %s x %s vector grid"%(self.options()['plot_points'], self.options()['plot_points'])
def _render_on_subplot(self, subplot):
"""
TESTS::
sage: x,y = var('x,y')
sage: P=plot_vector_field((sin(x), cos(y)), (x,-3,3), (y,-3,3))
"""
options = self.options()
quiver_options = options.copy()
quiver_options.pop('plot_points')
subplot.quiver(self.xpos_array, self.ypos_array, self.xvec_array, self.yvec_array, angles='xy', **quiver_options)
@options(plot_points=20,frame=True)
def plot_vector_field_unscaled((f, g), xrange, yrange, **options):
r"""
``plot_vector_field`` takes two functions of two variables xvar and yvar
(for instance, if the variables are `x` and `y`, take `(f(x,y), g(x,y))`)
and plots vector arrows of the function over the specified ranges, with
xrange being of xvar between xmin and xmax, and yrange similarly (see below).
``plot_vector_field((f, g), (xvar, xmin, xmax), (yvar, ymin, ymax))``
EXAMPLES:
Plot some vector fields involving sin and cos::
sage: x,y = var('x y')
sage: plot_vector_field((sin(x), cos(y)), (x,-3,3), (y,-3,3))
::
sage: plot_vector_field(( y, (cos(x)-2)*sin(x)), (x,-pi,pi), (y,-pi,pi))
Plot a gradient field::
sage: u,v = var('u v')
sage: f = exp(-(u^2+v^2))
sage: plot_vector_field(f.gradient(), (u,-2,2), (v,-2,2), color='blue')
Plot two orthogonal vector fields::
sage: x,y = var('x,y')
sage: a=plot_vector_field((x,y), (x,-3,3),(y,-3,3),color='blue')
sage: b=plot_vector_field((y,-x),(x,-3,3),(y,-3,3),color='red')
sage: show(a+b)
We ignore function values that are infinite or NaN::
sage: x,y = var('x,y')
sage: plot_vector_field( (-x/sqrt(x^2+y^2), -y/sqrt(x^2+y^2)), (x, -10, 10), (y, -10, 10))
::
sage: x,y = var('x,y')
sage: plot_vector_field( (-x/sqrt(x+y), -y/sqrt(x+y)), (x, -10, 10), (y, -10, 10))
Extra options will get passed on to show(), as long as they are valid::
sage: plot_vector_field((x, y), (x, -2, 2), (y, -2, 2), xmax=10)
sage: plot_vector_field((x, y), (x, -2, 2), (y, -2, 2)).show(xmax=10) # These are equivalent
"""
from sage.plot.all import Graphics
from sage.plot.misc import setup_for_eval_on_grid
z, ranges = setup_for_eval_on_grid([f,g], [xrange, yrange], options['plot_points'])
f,g = z
xpos_array, ypos_array, xvec_array, yvec_array = [],[],[],[]
for x in xsrange(*ranges[0], include_endpoint=True):
for y in xsrange(*ranges[1], include_endpoint=True):
xpos_array.append(x)
ypos_array.append(y)
xvec_array.append(f(x,y))
yvec_array.append(g(x,y))
import numpy
xvec_array = numpy.ma.masked_invalid(numpy.array(xvec_array, dtype=float))
yvec_array = numpy.ma.masked_invalid(numpy.array(yvec_array, dtype=float))
g = Graphics()
g._set_extra_kwds(Graphics._extract_kwds_for_show(options))
g.add_primitive(PlotField_local(xpos_array, ypos_array, xvec_array, yvec_array, options))
return g