tweaksHEADmaster

2 files changed, 80 insertions, 39 deletions

diff --git a/importdata.py b/importdata.py
index a72f71e..befb171 100644
--- a/importdata.py
+++ b/importdata.py
@@ -4,14 +4,17 @@
 
 import shutil, os, scipy, pylab;
 
-def dothing(arg,thisdir,files):
+def dothing(arg,thisdir,files,prefix="d"):
     global all_data
     for fname in files:
         if os.path.isfile(thisdir + '/' + fname):
             if (fname.find('.') == -1):
                 try:
                     dat = pylab.load(thisdir + '/' + fname)
-                    all_data[fname.split('.')[0]] = dat
+                    if(fname.split('.')[0][0].isalpha()):
+                        all_data[fname.split('.')[0]] = dat
+                    else:
+                        all_data[prefix + fname.split('.')[0]] = dat
                     print "Loaded (no type): %s/%s" % (thisdir,fname);
                 except:
                     print "Tried, but didn't load (no type):  %s/%s" % (thisdir,fname);
@@ -19,7 +22,10 @@ def dothing(arg,thisdir,files):
                 try:
                     dat = pylab.load(thisdir + '/' + fname)
                     #print all_data
-                    all_data[fname.split('.')[0]] = dat
+                    if(fname.split('.')[0][0].isalpha()):
+                        all_data[fname.split('.')[0]] = dat
+                    else:
+                        all_data[prefix + fname.split('.')[0]] = dat
                     print "Loaded .TKA: %s/%s" % (thisdir,fname);
                     #print "all_data has %d" % len(all_data)
                 except:
@@ -28,8 +34,14 @@ def dothing(arg,thisdir,files):
                 try:
                     dat = pylab.load(thisdir + '/' + fname,skiprows=1)
                     #print all_data
-                    all_data[fname.split('.')[0]] = dat
-                    print "Loaded .txt: %s/%s" % (thisdir,fname);
+                    if(fname.split('.')[0][0].isalpha()):
+                        all_data[fname.split('.')[0]] = dat
+                        print "Loaded .txt: %s/%s as %s" % (thisdir,fname,
+                            fname.split('.')[0]);
+                    else:
+                        all_data[prefix + fname.split('.')[0]] = dat
+                        print "Loaded .txt: %s/%s as %s" % (thisdir,fname,
+                            prefix + fname.split('.')[0]);
                     #print "all_data has %d" % len(all_data)
                 except:
                     print "Tried, but didn't load .txt: %s/%s" % (thisdir,fname);
@@ -42,7 +54,10 @@ def dothing(arg,thisdir,files):
 #                    print len(dat)
 #                    print dat
                     dat = scipy.array(dat);
-                    all_data[fname.split('.')[0]] = dat
+                    if(fname.split('.')[0][0].isalpha()):
+                        all_data[fname.split('.')[0]] = dat
+                    else:
+                        all_data[prefix + fname.split('.')[0]] = dat
                     print "Loaded .Spe: %s/%s" % (thisdir,fname);
 #                except:
                     print "Tried, but didn't load .Spe: %s/%s" % (thisdir,fname);
@@ -51,7 +66,7 @@ def dothing(arg,thisdir,files):
         else:
             print "Wasn't a file: %s/%s" % (thisdir,fname);
 
-def import_data(outfile='data.pydat',root='.',saveit=True,recurse=False):
+def import_data(outfile='data.pydat',root='.',saveit=True,recurse=False,prefix="d"):
     """Imports all that yummah files..."""
     global all_data
     if saveit:
diff --git a/stats.py b/stats.py
index 9acf3f4..eea94a2 100644
--- a/stats.py
+++ b/stats.py
@@ -50,7 +50,7 @@ def rchisqr_poisson(tp,ffunc,tx,ty,floor=1.):
     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):
+def fittodata(tx,ty,p0,ffunc,err=None,floor=1.,ret_all=False,param_names=None):
     """does a fit and returns a dict of shit. 
     ret_all=True gives more shit."""
     if err == None:
@@ -101,50 +101,70 @@ def fittodata(tx,ty,p0,ffunc,err=None,floor=1.,ret_all=False):
                 "err":err,
                 "x":tx,
                 "y":ty,
-                "success":success};
+                "success":success,
+                "param_names":param_names};
     else:
         return {"params":p1,
                 "param_error":p1_err,
-                "rchisqr":rcs}
+                "rchisqr":rcs,
+                "param_names":param_names};
 
 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)"""
+    second is stddev)"""
     #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);
+    return fittodata(tx,ty,p0,gauss_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["Mean","StdDev"]);
+
+def fittodata_gaussian_amp(tx,ty,p0,err=None,floor=1.,ret_all=False):
+    """runs fittodata for a gaussian with amplitude parameter (first parameter 
+    is mean, second is stddev, third is amplitude)"""
+    bigsum = scipy.integrate.simps(ty,x=tx)
+    gauss_func = lambda a,b: a[2] * 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,
+            param_names=["Mean","StdDev","Amplitude"]);
+
+def fittodata_gaussian_plus(tx,ty,p0,err=None,floor=1.,ret_all=False):
+    """runs fittodata for a gaussian with extra parameters (first parameter 
+    is mean, second is stddev, third is amplitude, fourth is constant offset, 
+    fifth is linear trend)"""
+    bigsum = scipy.integrate.simps(ty,x=tx)
+    gauss_func = lambda a,b: a[2] * bigsum/(a[1] * scipy.sqrt(2*pi)) * scipy.exp(-.5 * ((b - a[0])/a[1])**2) + a[3] + a[4]*b
+    return fittodata(tx,ty,p0,gauss_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["Mean","StdDev","Amplitude","VertOffset","Slope"]);
 
 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);
+    return fittodata(tx,ty,p0,lin_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["VertOffset (m)","Slope (b)"]);
+
+def fittodata_quadratic(tx,ty,p0,err=None,floor=1.,ret_all=False):
+    """runs fittodata for a quadratic curve (first param is a,
+    second is b, third is c: ax^2+bx+c)"""
+    quad_func= lambda a,b: a[0] * b*b  + a[1]*b + a[2]
+    return fittodata(tx,ty,p0,quad_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["QuadaticTerm (a)","VertOffset (b)","Slope (c)"]);
 
 def fittodata_poisson(tx,ty,p0,err=None,floor=.5,ret_all=False):
-    """runs fittodata for a poisson (first param is the mean"""
+    """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);
+    return fittodata(tx,ty,p0,poissfunc,err=err,floor=floor,ret_all=ret_all,
+            param_names=["Mean"]);
 
 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);
+    return fittodata(tx,ty,p0,exp_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["x0","tau","VertOffset"]);
 
 def fittodata_lorentz(tx,ty,p0,err=None,floor=1.,ret_all=False):
     """runs fittodata for a lorentzian (first param is gamma,
@@ -153,7 +173,19 @@ def fittodata_lorentz(tx,ty,p0,err=None,floor=1.,ret_all=False):
     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);
+    return fittodata(tx,ty,p0,lorentz_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["Gamma","Mean"]);
+
+def fittodata_lorentz_plus(tx,ty,p0,err=None,floor=1.,ret_all=False):
+    """runs fittodata for a lorentzian with extra parameters (first param 
+    is gamma, second param is mean, third is constant offset, fourth is
+    linear trend, fifth is amplitute)"""
+    bigsum = scipy.integrate.simps(ty,x=tx)
+    if type(p0) == type(0.1):
+        p0 = [p0,]
+    lorentz_func = lambda a,b,: a[4]*bigsum * a[0]/(2.*pi*((a[1]-b)**2+(a[0]/2.)**2. )) + a[2] + a[3]*b
+    return fittodata(tx,ty,p0,lorentz_func,err=err,floor=floor,ret_all=ret_all,
+            param_names=["Gamma","Mean","VertOffset","Slope","Amplitude"]);
 
 def plotfit(indict,newfig=True):
     if not indict.has_key('x'):
@@ -164,6 +196,7 @@ def plotfit(indict,newfig=True):
     y = indict['y']
     err = indict['err']
     ffunc = indict['ffunc']
+    param_names = indict['param_names']
     if newfig: pylab.figure();
     pylab.errorbar(x,y,fmt='b.',yerr=err)
     d=pylab.plot(x,y,'k.')
@@ -179,19 +212,12 @@ def plotfit(indict,newfig=True):
     p1_err = indict['param_error']
     
     ax = pylab.gca()
-    extra_txt = 'reduced chi squared: %g'\
-        % (indict['rchisqr'])
+    extra_txt = 'reduced chi^2: %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])
-    
+        if len(param_names)>i:
+            extra_txt += '\n%s: %.3g +/- %.2g' % (param_names[i],p1[i],p1_err[i]);
+        else:
+            extra_txt += '\np[%i]: %.3g +/- %.2g' % (i,p1[i],p1_err[i]);
     t = pylab.text(.77, .75, extra_txt,
         horizontalalignment='center',
         verticalalignment='center',