1
2
3
4
5
6
7
8
9
10 """funcutils.py -- Subroutines that tabulate a function's values.
11
12 Convenience functions that evaluate a python function on a grid of
13 points and tabulate the output to be used with Gnuplot.
14
15 """
16
17 try:
18 import numpy
19 except ImportError:
20 import Numeric as numpy
21
22 import Gnuplot, utils
23
24
26 """Evaluate and tabulate a function on a 1- or 2-D grid of points.
27
28 f should be a function taking one or two floating-point
29 parameters.
30
31 If f takes one parameter, then xvals should be a 1-D array and
32 yvals should be None. The return value is a numpy array
33 '[f(x[0]), f(x[1]), ..., f(x[-1])]'.
34
35 If f takes two parameters, then 'xvals' and 'yvals' should each be
36 1-D arrays listing the values of x and y at which 'f' should be
37 tabulated. The return value is a matrix M where 'M[i,j] =
38 f(xvals[i],yvals[j])', which can for example be used in the
39 'GridData' constructor.
40
41 If 'ufunc=0', then 'f' is evaluated at each point using a Python
42 loop. This can be slow if the number of points is large. If
43 speed is an issue, you should write 'f' in terms of numpy ufuncs
44 and use the 'ufunc=1' feature described next.
45
46 If called with 'ufunc=1', then 'f' should be a function that is
47 composed entirely of ufuncs (i.e., a function that can operate
48 element-by-element on whole matrices). It will be passed the
49 xvals and yvals as rectangular matrices.
50
51 """
52
53 if yvals is None:
54
55 xvals = numpy.asarray(xvals, dtype)
56
57 if ufunc:
58 return f(xvals)
59 else:
60 if dtype is None:
61 try:
62 dtype = xvals.dtype.char
63 except AttributeError:
64 dtype = xvals.typecode()
65
66 m = numpy.zeros((len(xvals),), dtype)
67 for xi in range(len(xvals)):
68 x = xvals[xi]
69 m[xi] = f(x)
70 return m
71 else:
72
73 xvals = numpy.asarray(xvals, dtype)
74 yvals = numpy.asarray(yvals, dtype)
75
76 if ufunc:
77 return f(xvals[:,numpy.newaxis], yvals[numpy.newaxis,:])
78 else:
79 if dtype is None:
80
81
82
83 try:
84 dtype = (numpy.zeros((1,), xvals.dtype.char) +
85 numpy.zeros((1,), yvals.dtype.char)).dtype.char
86 except AttributeError:
87 dtype = (numpy.zeros((1,), xvals.typecode()) +
88 numpy.zeros((1,), xvals.typecode())).typecode()
89
90 m = numpy.zeros((len(xvals), len(yvals)), dtype)
91 for xi in range(len(xvals)):
92 x = xvals[xi]
93 for yi in range(len(yvals)):
94 y = yvals[yi]
95 m[xi,yi] = f(x,y)
96 return m
97
98
99
100 grid_function = tabulate_function
101
102
104 """Evaluate a function of 1 variable and store the results in a Data.
105
106 Computes a function f of one variable on a set of specified points
107 using 'tabulate_function', then store the results into a 'Data' so
108 that it can be plotted. After calculation, the data are written
109 to a file; no copy is kept in memory. Note that this is quite
110 different than 'Func' (which tells gnuplot to evaluate the
111 function).
112
113 Arguments:
114
115 'xvals' -- a 1-d array with dimension 'numx'
116
117 'f' -- the function to plot--a callable object for which
118 f(x) returns a number.
119
120 'ufunc=<bool>' -- evaluate 'f' as a ufunc?
121
122 Other keyword arguments are passed through to the Data
123 constructor.
124
125 'f' should be a callable object taking one argument. 'f(x)' will
126 be computed at all values in xvals.
127
128 If called with 'ufunc=1', then 'f' should be a function that is
129 composed entirely of ufuncs, and it will be passed the 'xvals' and
130 'yvals' as rectangular matrices.
131
132 Thus if you have a function 'f', a vector 'xvals', and a Gnuplot
133 instance called 'g', you can plot the function by typing
134 'g.splot(compute_Data(xvals, f))'.
135
136 """
137
138 xvals = utils.float_array(xvals)
139
140
141 data = tabulate_function(f, xvals, ufunc=ufunc)
142
143 return Gnuplot.Data(xvals, data, **keyw)
144
145
147 """Evaluate a function of 2 variables and store the results in a GridData.
148
149 Computes a function 'f' of two variables on a rectangular grid
150 using 'tabulate_function', then store the results into a
151 'GridData' so that it can be plotted. After calculation the data
152 are written to a file; no copy is kept in memory. Note that this
153 is quite different than 'Func' (which tells gnuplot to evaluate
154 the function).
155
156 Arguments:
157
158 'xvals' -- a 1-d array with dimension 'numx'
159
160 'yvals' -- a 1-d array with dimension 'numy'
161
162 'f' -- the function to plot--a callable object for which
163 'f(x,y)' returns a number.
164
165 'ufunc=<bool>' -- evaluate 'f' as a ufunc?
166
167 Other keyword arguments are passed to the 'GridData' constructor.
168
169 'f' should be a callable object taking two arguments.
170 'f(x,y)' will be computed at all grid points obtained by
171 combining elements from 'xvals' and 'yvals'.
172
173 If called with 'ufunc=1', then 'f' should be a function that is
174 composed entirely of ufuncs, and it will be passed the 'xvals' and
175 'yvals' as rectangular matrices.
176
177 Thus if you have a function 'f' and two vectors 'xvals' and
178 'yvals' and a Gnuplot instance called 'g', you can plot the
179 function by typing 'g.splot(compute_GridData(f, xvals, yvals))'.
180
181 """
182
183 xvals = utils.float_array(xvals)
184 yvals = utils.float_array(yvals)
185
186
187 data = tabulate_function(f, xvals, yvals, ufunc=ufunc)
188
189 return Gnuplot.GridData(data, xvals, yvals, **keyw)
190
191
192
195