summaryrefslogtreecommitdiffstats
path: root/stats.py
diff options
context:
space:
mode:
Diffstat (limited to 'stats.py')
-rw-r--r--stats.py196
1 files changed, 196 insertions, 0 deletions
diff --git a/stats.py b/stats.py
new file mode 100644
index 0000000..3315968
--- /dev/null
+++ b/stats.py
@@ -0,0 +1,196 @@
+# 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)"""
+ lorentz_func = lambda a,b: 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'))