mpfitexpr.py

""" 
	Copyright (C) 2009 Sergey Koposov
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.
    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import mpfit 
import re
import numpy
import scipy

def mpfitexpr(func, x, y, err , start_params, check=True, full_output=False,
						imports=None, **kw):
	"""Fit the used defined expression to the data
	Input:
	- func: string with the function definition 
	- x: x vector
	- y: y vector
	- err: vector with the errors of y
	- start_params: the starting parameters for the fit
	Output:
	- The tuple (params, yfit) with best-fit params and the values of func evaluated at x
	Keywords:
	- check: boolean parameter. If true(default) the function will be checked for sanity
	- full_output: boolean parameter. If True(default is False) then instead of best-fit parameters the mpfit object is returned
	- imports: list of strings, of optional modules to be imported, required to evaluate the function
	Example:
	params,yfit=mpfitexpr('p[0]+p[2]*(x-p[1])',x,y,err,[0,10,1])
	
	If you need to use numpy and scipy functions in your function, then
		you must to use the full names of these functions, e.g.:
		numpy.sin, numpy.cos etc.
	
	This function is motivated by mpfitexpr() from wonderful MPFIT IDL package
		written by Craig Markwardt	
	
	"""

	hash={}
	hash['numpy']=numpy
	hash['scipy']=scipy
	
	if imports is not None:
		for i in imports:
			#exec '%s=__import__("%s")'%(a,b) in globals(),locals()
			hash[i]= __import__(i)
	def myfunc(p,fjac=None,x=None, y=None, err=None):
		return [0, eval('(y-(%s))/err'%func,hash,locals())]

	myre = "(?:[^a-zA-Z_]|^)p\[(\d+)\]"
	r = re.compile(myre)
	maxp = -1
	for m in re.finditer(r,func):
		curp = int(m.group(1))
		maxp = curp if curp > maxp else maxp	
	if check:
		if maxp == -1: 
			raise Exception("wrong function format")
		if maxp + 1 != len(start_params):
			raise Exception("the length of the start_params != the length of the parameter verctor of the function")
	fa={'x' : x, 'y' : y,'err' : err}
	res = mpfit.mpfit(myfunc,start_params,functkw=fa,**kw)
	yfit = eval(func, hash, {'x':x, 'p': res.params})
	if full_output:
		return (res, yfit)
	else:
		return (res.params, yfit)