1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187

#! /usr/bin/env python
# $Id: funcutils.py 297 20070330 11:25:28Z mhagger $
# Copyright (C) 19982003 Michael Haggerty <mhagger@alum.mit.edu>
#
# This file is licensed under the GNU Lesser General Public License
# (LGPL). See LICENSE.txt for details.
"""funcutils.py  Subroutines that tabulate a function's values.
Convenience functions that evaluate a python function on a grid of
points and tabulate the output to be used with Gnuplot.
"""
import numpy
import Gnuplot
from Gnuplot import utils
def tabulate_function(f, xvals, yvals=None, dtype=None, ufunc=0):
"""Evaluate and tabulate a function on a 1 or 2D grid of points.
f should be a function taking one or two floatingpoint
parameters.
If f takes one parameter, then xvals should be a 1D array and
yvals should be None. The return value is a numpy array
'[f(x[0]), f(x[1]), ..., f(x[1])]'.
If f takes two parameters, then 'xvals' and 'yvals' should each be
1D arrays listing the values of x and y at which 'f' should be
tabulated. The return value is a matrix M where 'M[i,j] =
f(xvals[i],yvals[j])', which can for example be used in the
'GridData' constructor.
If 'ufunc=0', then 'f' is evaluated at each point using a Python
loop. This can be slow if the number of points is large. If
speed is an issue, you should write 'f' in terms of numpy ufuncs
and use the 'ufunc=1' feature described next.
If called with 'ufunc=1', then 'f' should be a function that is
composed entirely of ufuncs (i.e., a function that can operate
elementbyelement on whole matrices). It will be passed the
xvals and yvals as rectangular matrices.
"""
if yvals is None:
# f is a function of only one variable:
xvals = numpy.asarray(xvals, dtype)
if ufunc:
return f(xvals)
else:
if dtype is None:
dtype = xvals.dtype.char
m = numpy.zeros((len(xvals),), dtype)
for xi in range(len(xvals)):
x = xvals[xi]
m[xi] = f(x)
return m
else:
# f is a function of two variables:
xvals = numpy.asarray(xvals, dtype)
yvals = numpy.asarray(yvals, dtype)
if ufunc:
return f(xvals[:,numpy.newaxis], yvals[numpy.newaxis,:])
else:
if dtype is None:
# choose a result dtype based on what '+' would return
# (yecch!):
dtype = (numpy.zeros((1,), xvals.dtype.char) +
numpy.zeros((1,), yvals.dtype.char)).dtype.char
m = numpy.zeros((len(xvals), len(yvals)), dtype)
for xi in range(len(xvals)):
x = xvals[xi]
for yi in range(len(yvals)):
y = yvals[yi]
m[xi,yi] = f(x,y)
return m
# For backwards compatibility:
grid_function = tabulate_function
def compute_Data(xvals, f, ufunc=0, **keyw):
"""Evaluate a function of 1 variable and store the results in a Data.
Computes a function f of one variable on a set of specified points
using 'tabulate_function', then store the results into a 'Data' so
that it can be plotted. After calculation, the data are written
to a file; no copy is kept in memory. Note that this is quite
different than 'Func' (which tells gnuplot to evaluate the
function).
Arguments:
'xvals'  a 1d array with dimension 'numx'
'f'  the function to plota callable object for which
f(x) returns a number.
'ufunc=<bool>'  evaluate 'f' as a ufunc?
Other keyword arguments are passed through to the Data
constructor.
'f' should be a callable object taking one argument. 'f(x)' will
be computed at all values in xvals.
If called with 'ufunc=1', then 'f' should be a function that is
composed entirely of ufuncs, and it will be passed the 'xvals' and
'yvals' as rectangular matrices.
Thus if you have a function 'f', a vector 'xvals', and a Gnuplot
instance called 'g', you can plot the function by typing
'g.splot(compute_Data(xvals, f))'.
"""
xvals = utils.float_array(xvals)
# evaluate function:
data = tabulate_function(f, xvals, ufunc=ufunc)
return Gnuplot.Data(xvals, data, **keyw)
def compute_GridData(xvals, yvals, f, ufunc=0, **keyw):
"""Evaluate a function of 2 variables and store the results in a GridData.
Computes a function 'f' of two variables on a rectangular grid
using 'tabulate_function', then store the results into a
'GridData' so that it can be plotted. After calculation the data
are written to a file; no copy is kept in memory. Note that this
is quite different than 'Func' (which tells gnuplot to evaluate
the function).
Arguments:
'xvals'  a 1d array with dimension 'numx'
'yvals'  a 1d array with dimension 'numy'
'f'  the function to plota callable object for which
'f(x,y)' returns a number.
'ufunc=<bool>'  evaluate 'f' as a ufunc?
Other keyword arguments are passed to the 'GridData' constructor.
'f' should be a callable object taking two arguments.
'f(x,y)' will be computed at all grid points obtained by
combining elements from 'xvals' and 'yvals'.
If called with 'ufunc=1', then 'f' should be a function that is
composed entirely of ufuncs, and it will be passed the 'xvals' and
'yvals' as rectangular matrices.
Thus if you have a function 'f' and two vectors 'xvals' and
'yvals' and a Gnuplot instance called 'g', you can plot the
function by typing 'g.splot(compute_GridData(f, xvals, yvals))'.
"""
xvals = utils.float_array(xvals)
yvals = utils.float_array(yvals)
# evaluate function:
data = tabulate_function(f, xvals, yvals, ufunc=ufunc)
return Gnuplot.GridData(data, xvals, yvals, **keyw)
# For backwards compatibility:
def GridFunc(f, xvals, yvals, **keyw):
return compute_GridData(xvals, yvals, f, **keyw)
