# \$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 or smooth approximation of a two-dimensional curve. Input parameters: - s: smoothing of the spline curve. The default is 0, which means plain interpolation, no no extra smoothing. If s > 0, then some smoothing is performed, and the curve represents an approximation of the input x,y curve. - w: the weight factor for each data point. """ # Important notes on spline CAVEATS: # - the x values better be sorted in ascending order, or else # the routine would return nonsense (i.e. NaN's). # - no two same values of x can be specified. def __init__(self, x, y, w=None, s=0): self.init(x,y) self.s = s self.w = w 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 # (No, now this is handled by splrep() properly.) #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.splrep(self.x, self.y, w=self.w, s=self.s) def spline(self, xnew): try: params = self.spline_params except: self.init_spline_params() return scipy.interpolate.splev(x=xnew, tck=self.spline_params, der=0) 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)