Fit a charge-charge interaction tail

fit1overr.py

#!/usr/bin/env python


''' fit a simple equation '''


import os, sys
from math import *


fnin = None # input file
exponent = -1
xmin = -1e30
xmax = 1e30



def usage():
  ''' print usage and die '''
  print '''fit.py [OPTIONS] input.dat
  Linear regression
  Options:
    -e:           fit a + b x^e
    -x, -xmin=:   minimal x for fitting
    -X, -xmax=:   maximal x for fitting
    -v:           be verbose
    -vv:          be more verbose
    -h, --help:   print this message

  Examples:
    fit.py -e-1 my.dat
  '''
  exit(1)



def doargs():
  ''' handle command line options '''
  import getopt, glob
  try:
    opts, args = getopt.gnu_getopt(sys.argv[1:], "e:x:X",
        ["xmin=", "xmax=", "verbose=", "help"])
  except getopt.GetoptError, err:
    print str(err)
    usage()

  global fnin, exponent, xmin, xmax

  for o, a in opts:
    if o in ("-e",):
      exponent = float(a)
    elif o in ("-x", "--xmin"):
      xmin = float(a)
    elif o in ("-X", "--xmax"):
      xmax = float(a)
    elif o in ("-v",):
      verbose += 1
    elif o in ("-h", "--help",):
      usage()

  ls = args
  if len(ls) > 0:
    fnin = ls[0]
  else:
    fnin = glob.glob("*.dat")[0]



def fitlow(ls):
  n = len(ls)
  sz = sy = szz = syy = szy = 0
  for i in range(n):
    x, y = ls[i]
    z = pow(x, exponent)
    sy += y
    sz += z
    szz += z*z
    szy += z*y
    syy += y*y
  ''' now solve
    n a + sz b = sy --> sz a + sz*sz/n b = sy*sz/n
    sz a + szz b = szy
  '''
  var = szz - sz*sz/n
  cov = szy - sy*sz/n
  b = cov/var
  a = (sy - b*sz)/n
  vary = syy - sy*sy/n
  res = (vary - cov*cov/var)/n
  print "best fit: %g%+g*x^(%s), residue %g (%d points)" % (
      a, b, exponent, sqrt(res), n)
  unitb = 0.00013893545782*1e7
  invepsr = b / unitb
  print invepsr



def dofit(fn):
  ls = []
  for ln in open(fn).readlines():
    xy = ln.strip()
    if not xy: continue
    xy = xy.split()
    xi = float(xy[0])
    yi = float(xy[1])
    if xi < xmin or xi > xmax: continue
    ls += [(xi, yi),]
  fitlow(ls)



if __name__ == "__main__":
  doargs()
  dofit(fnin)