diff --git a/math/spline_2d.py b/math/spline_2d.py new file mode 100644 index 0000000..431ed66 --- /dev/null +++ b/math/spline_2d.py @@ -0,0 +1,91 @@ +# $Id: spline_2d.py,v 1.1 2010-01-22 18:48:19 wirawan Exp $ +# +# wpylib.math.spline_2d module +# Created: 20091204 +# Wirawan Purwanto +# + +import numpy +import scipy.interpolate + +class spline_2d: + """Simple interpolation for two-dimensional curve. + Created to handle the quirks of current SciPy interpolation + routines.""" + # Important notes on spline CAVEATS: + # - for some reason we HAVE to make a copy of the 'x' array + # (to make it contiguous, probably?) + # - also, the x values better be sorted in ascending order, or else + # the routine would return nonsense (i.e. NaN's). + def __init__(self, x, y): + self.init(x,y) + + def init(self, x, y): + # First, the x must be sorted, so we make a private copy of + # the data: + self.data = numpy.array(zip(x, y), dtype=[('x', float), ('y', float)]) + # Quirk 1: The x axis data must be sorted ascending + self.data.sort(order=['x']) + self.x = self.data['x'] + self.y = self.data['y'] + # Quirk 2: the x data for spline function must be contiguous + # See below in init_spline_params() + #self.x_copy = self.x.copy() + + try: + del self.spline_params + except: + pass + + def init_spline_params(self): + """Initialize spline params with default params. + You can call something to initialize the spline params before + calling the first spline function if you want different, non-default + parameters.""" + self.spline_params \ + = scipy.interpolate.splmake(self.x.copy(), self.y) + + def spline(self, xnew): + try: + params = self.spline_params + except: + self.init_spline_params() + return scipy.interpolate.spleval(self.spline_params, xnew) + + + +class spline_2d_piecewise: + """Simple spline_2d interpolator with piecewise datasets. + Interpolation is possible only in the ranges defined by the piecewise + datasets. + No checking is done whether the pieces are overlapping, discontinuous, etc. + The first piece found enclosing the coordinate will be taken for + interpolation.""" + def __init__(self, *datasets): + self.init(*datasets) + + def init(self, *datasets): + #if len(dsets) % 2: + # raise ValueError, "The input datasets must be given in x, y pairs + self.pieces = [] + for dset in datasets: + x = dset[0] + y = dset[1] + xmin = numpy.min(x) + xmax = numpy.max(x) + piece = spline_2d(x, y) + piece.xmin = xmin + piece.xmax = xmax + self.pieces.append(piece) + + def in_range(self, piece, x): + return piece.xmin <= x and x <= piece.xmax + + def get_piece(self, x): + for p in self.pieces: + if self.in_range(p, x): + return p + raise ValueError, "Out-of-range x value = %g" % x + + def spline(self, x): + return self.get_piece(x).spline(x)