#!/net/python/bin/python import pylab as p from scipy import optimize,linspace import numpy as num from numpy import random import sys import types ################################################### # One Dimensional Fitting # def fourier_basis1D(x,i,xmin=0.0,xmax=100.0): xmax = xmax+4*(xmax-xmin) if i%2 == 0: #even: use a cosine return num.cos( i * num.pi*(x-xmin)/(xmax-xmin) ) else: #odd: use a sine return num.sin( 0.5*(i+1) * num.pi*(x-xmin)/(xmax-xmin) ) def AnalyticFit1D(x,y,dy,N=20,basis=fourier_basis1D): xmin = min(x) xmax = max(x) if N > len(x): N=len(x) p0 = [1.0 for i in range(N)] fitfunc = lambda p, x: sum( p[i]*basis(x,i,xmin,xmax) for i in range(len(p)) ) if len(dy) == len(x): errfunc = lambda p, x, y, dy: (fitfunc(p,x)-y)/dy p1,success = optimize.leastsq( errfunc, p0[:], args=(x,y,dy) ) #args are extra args for errfunc else: errfunc = lambda p, x, y: (fitfunc(p,x)-y) p1,success = optimize.leastsq( errfunc, p0[:], args=(x,y) ) #args are extra args for errfunc return lambda x: fitfunc(p1,x) ################################################### # Two Dimensional Fitting # def fourier_basis2D(x,y,i,xmin=0.0,xmax=100.0,ymin=0.0,ymax=100.0): """ returns the ith basis function of x and y basis functions are sin and cos in each spatial direction multiplied by each other """ order = int( num.sqrt(i) ) ext = i - order**2 if ext % 2 == 0: xi = ext/2 yi = order else: xi = order yi = int( ext/2 ) xmax = xmax+4*(xmax-xmin) if xi%2 == 0: #even: use a cosine xfunc = num.cos( xi * num.pi*(x-xmin)/(xmax-xmin) ) else: #odd: use a sine xfunc = num.sin( 0.5*(xi+1) * num.pi*(x-xmin)/(xmax-xmin) ) ymax = ymax+4*(ymax-ymin) if yi%2 == 0: #even: use a cosine yfunc = num.cos( yi * num.pi*(y-ymin)/(ymax-ymin) ) else: #odd: use a sine yfunc = num.sin( 0.5*(yi+1) * num.pi*(y-ymin)/(ymax-ymin) ) return xfunc*yfunc def tolist(x): """ returns x as a list if x is already iterable, returns a copy of x """ try: return [xi for xi in x] except: return [x] def samesize(x,y): x_c = tolist(x) y_c = tolist(y) for i in range(len(x_c)): if len(tolist(x_c[i])) != len(tolist(y_c[i])): return False return True def fitfunc2D(p,x,y,xmin,xmax,ymin,ymax,basis=fourier_basis2D): """ p are the coefficients corresponding to the basis functions Note: this will fail unless x and y are correctly formatted: x should be an array of numbers y should be an array of arrays, the top level array corresponds to given x values returns a single array with the values of p * basis(x,y) """ numtypes = (types.IntType,types.FloatType,types.LongType) if x in numtypes: x = [x] if y in numtypes: y = [[y]] assert len(x) == len(y) output = [] for i in range(len(x)): for j in range( len(y[i]) ): S = sum( [ p[k] *\ basis(x[i],y[i][j],k,xmin,xmax,ymin,ymax)\ for k in range(len(p))\ ] ) output.append(S) return output def errfunc2D(p,x,y,z,dz,xmin,xmax,ymin,ymax,basis=fourier_basis2D): """ p are the coefficients corresponding to the basis functions Note: this will fail unless x and y are correctly formatted: x should be an array of numbers y should be an array of arrays, the top level array corresponds to given x values z & dz should be same shape as y returns a single array with the values of (p * basis(x,y)-z)/dz """ output = [] #S = (fitfunc(p,x,y)-z)/dz for i in range(len(x)): for j in range( len(y[i]) ): S = sum( [ (p[k] * basis(x[i],y[i][j],k,xmin,xmax,ymin,ymax)\ -z[i][j])/dz[i][j]\ for k in range(len(p))] ) output.append(S) return output def AnalyticFit2D(x,y,z,dz,N=25,basis=fourier_basis2D): #---------------------------------- xmin = 1.0*min(x) xmax = 1.0*max(x) ymin = 1.0*min( [min(yi) for yi in y] ) ymax = 1.0*max( [max(yi) for yi in y] ) #---------------------------------- print "xmin:",xmin print "xmax:",xmax print "ymin:",ymin print "ymax:",ymax num_pts = sum( [len(yi) for yi in y] ) if N > num_pts: N=num_pts print "AnalyticFit2D: fitting for %i d.o.f" % N if N>50: print " (This may take a while...)" p0 = [1.0 for i in range(N)] #if no dz is provided, this assures that the fitting routine # will still work if not samesize(z,dz): print "making dz match z" dz = [[1.0 for zij in zi] for zi in z] fitfunc = lambda p,xx,yy:\ fitfunc2D(p,xx,yy,xmin,xmax,ymin,ymax,basis) errfunc = lambda p,xx,yy,zz,dzz:\ errfunc2D(p,xx,yy,zz,dzz,xmin,xmax,ymin,ymax,basis) p1,success = optimize.leastsq( errfunc, p0[:], args=(x,y,z,dz) ) #args are extra args for errfunc print "Chi2: ", Chi2_2D(p1,x,y,z,dz,xmin,xmax,ymin,ymax,basis) return p1 def Chi2_2D(p,x,y,z,dz,xmin,xmax,ymin,ymax,basis): output = errfunc2D(p,x,y,z,dz,xmin,xmax,ymin,ymax,basis) return sum( [x**2 for x in output] ) if __name__ == '__main__': #------Testing 1D stuff: default-------------------------- if len(sys.argv) == 1: x = num.sort([random.random()*100 for i in range(10)]) y = [(num.sin(10*num.sqrt(xi)*num.pi/100)\ + 0.1*(random.random()-0.5)) for xi in x] dy = [0.1*(random.random()-0.5) for xi in x] p.errorbar(x,y,dy,fmt='.k') func = AnalyticFit1D(x,y,dy) x_t = linspace(x.min(),x.max(),100) p.plot(x_t,func(x_t),"r-") # Plot of the data and the fit p.show() #------Testing 2D stuff: happens if any args are passed--- else: x = [1,2,3,4,5,6,7,8,9,10] y = [num.sort([random.random()*100\ for i in range(5+int(10*random.random()))])\ for xi in x] z = [ [y[i][j]*( (num.sin(10*num.sqrt(x[i])*num.pi/10)\ +0.1*(random.random()-0.5)) )\ for j in range(len(y[i]))]\ for i in range(len(x))] dz = [[0.1*(random.random()-0.5) for zij in zi] for zi in z] p1 = AnalyticFit2D(x,y,z,dz) print "###########################################\nfitted." print p1.tolist()