File contents

'''Fit 1D and 2D polynomials of order k to a set of data points.  Actually, they're
McLaren series in 1 and 2 variables.

Written by Chris Burns
Obfuscated by Dan Kelson.
'''
from Numeric import *
from LinearAlgebra import singular_value_decomposition

def divz(x, y=1, repl=0.0,out=Float32, tol=0):
   if len(shape(y)) or len(shape(x)):
      if len(shape(y)):  bad = less_equal(absolute(y), tol)
      else:  band = ones(x.shape)*(abs(y)<=tol)
      not_bad = 1-bad
      numer = (x*not_bad)
      denom = (y+bad)
      a = (repl*bad).astype(out)
      b = (numer/denom).astype(out)
      return(a + b).astype(out)
   else:
      if abs(y) <= tol:  return(repl)
      else:  return(x/y)

def fitsvd(A, b):
   decomp = singular_value_decomposition(A)
   sol1 = transpose(decomp[2])
   sol2 = divz(1.0,decomp[1], tol=1e-10)
   sol3 = dot(transpose(decomp[0]),b)
   if sometrue(sol3):
      solr = (sol2*sol3)
      soll = dot(sol1,solr)
      err = sqrt(sum(power(sol1*sol2, 2)))
   else:
      soll = zeros(sol3.shape)
      err = zeros(sol3.shape)
   return soll,err


def fac(N):
   if N > 1:
      return(N*fac(N-1))
   else:
      return(1)

def fitpoly(x, y, w, k=1, x0=0):
   '''Fit a 1D McLaren series of degree k to some data (x,y) with weights w, 
   centered on the point x=x0.  Returns (coeff,err) which are the 
   coefficients of the series:  f(x0), f'(x0), f''(x0), ... and the 
   associated errors.'''

   A = [1.0/fac(i)*power(x-x0, i)*w for i in range(k+1)]
   A = transpose(array(A))
   b = y*w

   soll,err = fitsvd(A,b)
   return(soll, err)

def fit2Dpoly(x, y, z, w, k=1, x0=0, y0=0):
   '''Fit a 2D McLaren series of degree k to some data (x,y,z) with weights w,
   centered on the point (x=x0, y=y0).  Returns the tuple (coeff, err).  Coeff
   are the coefficients of the series in the following order.  Let f[i,j](x0,y0)
   be the ith partial derivative of f w.r.t. x and jth partial derivative of f
   w.r.t y evaluated at (x0,y0).  The Coeff is
      f[0,0](x0,y0), f[1,0](x0,y0), ..., f[k,0](x0,y0),
      f[0,1](x0,y0), f[1,1](x0,y0), ..., f[k,1](x0,y0),
      ...
      f[0,k](x0,y0), f[1,k](x0,y0), ..., f[k,k](x0,y0)'''

   A = [power(x-x0,i)*power(y-y0,j)/fac(i)/fac(j) for j in range(k+1) for i in range(k+1) if i+j <= k]
   soll,err = fitsvd(transpose(A)*w[::,NewAxis],z*w)
   return (soll, err)