# bnewbold_stats.py # helper functions! # @author: bryan newbold # @email: bnewbold@mit.edu # @date: oct 2006 import scipy, pylab def floored_val(x,min=1.): if pylab.size(x) > 1: ret = scipy.array(pylab.copy.copy(x)); for i in range(0,pylab.size(x)): ret[i] = floored_val(x[i],min=min); return ret; if(abs(x) < min): return min; else: return abs(x); def floored_root(x,min=1.): return scipy.sqrt(floored_val(x,min=min)); def rchisqr(tp,ffunc,tx,ty,terr): """finds reduced chi squared for a function compared to data""" if not (pylab.size(tx) == pylab.size(ty) == pylab.size(terr)): print "ERROR: input matrix mismatch... (rchisqr)"; return; try: ffunc(tp,tx[-1]); except: print "problem with passed function and/or parameters!"; try: eta_err = (ty - ffunc(tp,tx))**2 / (terr**2) except: print "i think you have a zero error in there..."; eta_err = (ty - ffunc(tp,tx))**2 / (terr**2) return sum(eta_err)/(pylab.size(tx)-pylab.size(tp)); def rchisqr_poisson(tp,ffunc,tx,ty,floor=1.): """finds reduced chi squared for a function compared to data, assuming each data point has a poisson distribution and thus an error of sqrt(val). The error value is floored at 1.0 to prevent stupid val=0 cases...""" if not (pylab.size(tx) == pylab.size(ty)): print "ERROR: input matrix mismatch... (rchisqr_poisson)"; return; terr = floored_root(tx,min=floor); return rchisqr(tp,ffunc,tx,ty,terr); def fittodata(tx,ty,p0,ffunc,err=None,floor=1.,ret_all=False): """does a fit and returns a dict of shit. ret_all=True gives more shit.""" if err == None: err = floored_root(ty,min=floor); if not (pylab.size(tx) == pylab.size(ty) == pylab.size(err)): print "ERROR: input matrix mismatch... (fittodata)"; return; errfunc = lambda n,m,w: ffunc(n,m) - w; tx = scipy.array(tx); #tx = tx.tolist(); tx = scipy.array(tx); ty = scipy.array(ty); #ty = ty.tolist(); ty = scipy.array(ty); err = err.tolist(); err = scipy.array(err); p0 = scipy.array(p0).tolist(); p0 = scipy.array(p0); p1,cov_x,extra1,extra2,success = scipy.optimize.leastsq(errfunc, p0.copy(), args=(tx,ty), full_output=1) if size(p1) == 1: p1 = [p1,] print "DEBUG: done with fit" print "p1: " + str(p1) rcs = rchisqr(p1,ffunc,tx,ty,err); try: p1_err = p1.copy() except: p1_err = [0,] for i in range(scipy.size(p1)): try: scipy.size(cov_x); except: p1_err = array([0.0,]) p1 = array([p1,]) break; if scipy.size(cov_x) == 1: p1_err = array([scipy.sqrt(cov_x),]); p1 = array([p1,]); else: p1_err[i] = scipy.sqrt(cov_x[i,i]); if ret_all: return {"params":p1, "param_error":p1_err, "rchisqr":rcs, "ffunc":ffunc, "err":err, "x":tx, "y":ty, "success":success}; else: return {"params":p1, "param_error":p1_err, "rchisqr":rcs} def fittodata_gaussian(tx,ty,p0,err=None,floor=1.,ret_all=False): """runs fittodata for a gaussian (first parameter is mean, second is stddev, third is amplitude)""" #gauss_func = lambda a,b: a[2]/(a[1] * sqrt(2*pi)) * exp(-.5 * ((b - a[0])/a[1])**2) #thing = std(ty); bigsum = scipy.integrate.simps(ty,x=tx) gauss_func = lambda a,b: bigsum/(a[1] * scipy.sqrt(2*pi)) * scipy.exp(-.5 * ((b - a[0])/a[1])**2) return fittodata(tx,ty,p0,gauss_func,err=err,floor=floor,ret_all=ret_all); def fittodata_linear(tx,ty,p0,err=None,floor=1.,ret_all=False): """runs fittodata for a linear curve (first param is slope, second is vertical offset)""" lin_func = lambda a,b: a[0] * b + a[1] return fittodata(tx,ty,p0,lin_func,err=err,floor=floor,ret_all=ret_all); def fittodata_poisson(tx,ty,p0,err=None,floor=.5,ret_all=False): """runs fittodata for a poisson (first param is the mean""" bigsum = scipy.integrate.simps(ty,x=tx) if type(p0) == type(0.1): p0 = [p0,] poissfunc = lambda a,b: bigsum * scipy.power(a[0],b) / scipy.factorial(b) * scipy.exp(-1 * a[0]) #def poissfunc(a,b): # print "--------" # print a # print b # print bigsum # print scipy.power(a[0],b) # print scipy.exp(a[0]) # print bigsum * scipy.power(a[0],b) / scipy.factorial(b) * scipy.exp(-1 * a[0]) # return bigsum * scipy.power(a[0],b) / scipy.factorial(b) * scipy.exp(-1 * a[0]) return fittodata(tx,ty,p0,poissfunc,err=err,floor=floor,ret_all=ret_all); def fittodata_exp(tx,ty,p0,err=None,floor=1.,ret_all=False): """runs fittodata for an exponential (first param is x0, second param is tau, third param is vertical offset)""" #exp_func = lambda a,b: a[0] * scipy.exp(- b/a[1]) + a[2] exp_func = lambda a,b: a[0] * e ** (- b/a[1]) + a[2] return fittodata(tx,ty,p0,exp_func,err=err,floor=floor,ret_all=ret_all); def fittodata_lorentz(tx,ty,p0,err=None,floor=1.,ret_all=False): """runs fittodata for a lorentzian (first param is gamma, second param is mean)""" bigsum = scipy.integrate.simps(ty,x=tx) if type(p0) == type(0.1): p0 = [p0,] lorentz_func = lambda a,b: bigsum * a[0]/(2.*pi*((a[1]-b)**2+(a[0]/2.)**2. )) return fittodata(tx,ty,p0,lorentz_func,err=err,floor=floor,ret_all=ret_all); def plotfit(indict,newfig=True): if not indict.has_key('x'): print "forgot to use ret_all on the fittodata?"; return; p1 = indict['params']; x = indict['x'] y = indict['y'] err = indict['err'] ffunc = indict['ffunc'] if newfig: pylab.figure(); pylab.errorbar(x,y,fmt='b.',yerr=err) d=pylab.plot(x,y,'k.') fx = x if pylab.size(fx) < 100: fx = r_[min(x):max(x):(max(x)-min(x))/150.] f=pylab.plot(fx,ffunc(p1,fx),'r-') #fy = indict['ffunc'](indict['params'],indict['x']) pylab.title("Rough Fit of Data") pylab.xlabel("X") pylab.ylabel("Y") p1 = indict['params'] p1_err = indict['param_error'] ax = pylab.gca() extra_txt = 'reduced chi squared: %g'\ % (indict['rchisqr']) for i in range(scipy.size(p1)): extra_txt += '\np[%i]: %g +/- %g' \ % (i,p1[i],p1_err[i]); # if(pylab.size(p1) == 2): # extra_txt = 'reduced chi squared: %g\np[0]: %g\np[1]: %g' \ # % (indict['rchisqr'],p1[0],p1[1]) # else: # extra_txt= 'reduced chi squared:%g\np[0]:%g, +/- %g\n p[1]:%g\np[2]:%g' \ # % (indict['rchisqr'],p1[0],p1[1],p1[2]) t = pylab.text(.77, .75, extra_txt, horizontalalignment='center', verticalalignment='center', transform = ax.transAxes) pylab.legend((d,f),('Data', 'Fit'))