; From Dick French 
; rgfcurvefit4.pro
; Revisions:
;	17 Jan 1996 - add optional output of number of iterations
;	18 Jan 1996 - add optional logfile
;	1 July 1996 - change printout to give actual change of param,
;			not just absolute value
;	9 July 1997 - update to have the same format as IDL5.0 version of
;			curvefit, which has a /NODERIVATIVE keyword that
;			needs to be handled properly. Eliminate COMMON.
;	10 July 1997 - handle case where weights are zero or one - set
;			number of pts = number of pts with unit wt.
;		     - change convention so that sigma is for ALL input
;			parameters, whether fitted or not.
;	13 Aug 1997  - add a keyword FINITE_DIFFS that, if set, lets you 
;			specify finite differences to use for each parameter
;			instead of the machine precision calculation
;			when doing numerical derivatives. Has no effect
;			if NODERIVATIVE is not set.
;	28 Mar 1998  - add returned parameter 'converged'(0,1 for no,yes)
;		     - fix BUG where newest parameters were not being
;			used in the lambda loop!
;		     - take c = stuff out of inner loop.
;		     - compute chisqr at end to make sure is consistent
;			with returned parameters
;		     - bail out of inner lambda loop if more than 30 reps
;		     - add LAMBDACOUNTMAX iteration
;	10 Apr 1998  - define niter_used=0 if no fit performed

function rgfcurvefit4,x,y,weights,a_all,sigma,chisqr,correl_matrix,niter_used,$
	converged,$
	FUNCTION_NAME = FUNCTION_NAME, NODERIVATIVE=NODERIVATIVE, $
	FINITE_DIFFS = FINITE_DIFFS,$
	NITER = NITER, CONVCRIT = CONVCRIT, FITLIST = FITLIST, $
	CONVLIST = CONVLIST,PRINTEACH = PRINTEACH,LOGFILE = LOGFILE, $
	LAMBDACOUNTMAX = LAMBDACOUNTMAX
;+
; NAME:
;	RGFCURVEFIT4
;
; PURPOSE:
;	Non-linear least squares fit to a function of an arbitrary 
;	number of parameters.  The function may be any non-linear 
;	function. If available, partial derivatives can be calculated by
;	the user function, esle this routine will estimate partial
;	derivatives with a forward difference approximation.
;	approximated.
;
; CATEGORY:
;	E2 - Curve and Surface Fitting.
;
; CALLING SEQUENCE:
;	Result =  RGFCURVEFIT3(x,y,weights,a_all,sigma,chisqr,correl_matrix,niter_used,converged,$
;	FUNCTION_NAME = function_name, $
;	Niter = Niter, ConvCrit = ConvCrit, FitList = FitList, $
;	ConvList = ConvList,PrintEach = PrintEach,LogFile = LogFile,$
;	LambdaCountMax = LambdaCountMax)
;
; INPUTS:
;	X:  A row vector of independent variables. This routine does not
;		manipulate or use values in X, it simply passes X to
;		the user routine. This makes it possible to encode information
;		that the user routine needs without resorting to a labelled common.
;
;	Y:  A row vector containing the dependent variable.
;
;	weights:  A row vector of weights, the same length as x and y.
;		For no weighting,
;			Weights(i) = 1.0.
;		For instrumental (Gaussian)  weighting,
;			Weights(i) = 1.0/sigma(i)^2
;		For statistical (Poisson) weighting,
;			Weights(i) = 1.0/y(i)
;
;	a_all:  A vector, with as many elements as the number of terms, that 
;	    contains the initial estimate for each parameter.  If A is double-
;	    precision, calculations are performed in double precision, 
;	    otherwise they are performed in single precision. Final parameters
;	    are returned in A_all
;
; KEYWORDS:
;	FUNCTION_NAME:  The name of the function (actually, a procedure) to 
;	fit.  If omitted, "FUNCT" is used. The procedure must be written as
;	described under RESTRICTIONS, below.
;	NITER:  maximum number of iterations (20 if omitted)
;	CONVCRIT: fractional convergence criterion - the max fractional
;		change in any parameter allowed (if omitted, then 0.1%
;		change in chisquare is the default convergence criterion)
;		if CONVLIST is present, then CONVCRIT is ignored
;	FITLIST: integer array of size=# of parameters. An element set
;		to 0 indicates that that parameter is to be held
;		fixed; 1 indicates that it is to be fitted. This array,
;		if present, is passed to the fitting function
;	PRINTEACH: if 1, print each iteration, if 2, just the parameters
;		and correlation matrix
;	CONVLIST: list of acceptable absolute value of changes of fitted
;		parameters between successive iterations. This is to
;		enable convergence for parameters that slop around.
;	LOGFILE: if present, intermediate results are stored in this file
;	NODERIVATIVE: If this keyword is set then the user procedure will not
;		be requested to provide partial derivatives. The partial
;		derivatives will be estimated in RGFCURVEFIT3 using forward
;		differences. If analytical derivatives are available they 
;		should always be use.
;	FINITE_DIFFS: If NODERIVATIVE is set, and this keyword is present, then
;		the elements of the array will be used to specify the finite
;		differences used to compute numerical derivatives.
;	LAMBDACOUNTMAX: max number of times lambda loop is attempted.
;		Default is 30


;
; OUTPUTS:
;	Returns a vector of calculated function values.
;	a_all:  A vector of final parameters.
;
; OUTPUT PARAMETERS:
;	sigma:  A vector of standard deviations for the parameters in A.
;	chisqr:  Chisq at end of fit
;	correl_matrix: correlation matrix
;	niter_used: # of iterations actually performed
;	
; COMMON BLOCKS:
;	None.
;
; SIDE EFFECTS:
;	Y is made double precision.
;
; RESTRICTIONS:
;	The function to be fit must be defined and called FUNCT,
;	unless the FUNCTION_NAME keyword is supplied.  This function,
;	(actually written as a procedure) must accept values of
;	X (the independent variable), and A_all (the fitted function's
;	parameter values), and return F (the function's value at
;	X), and PDER (a 2D array of partial derivatives).
;	For an example, see FUNCT in the IDL User's Libaray.
;	A call to FUNCT is entered as:
;	FUNCT, X, A_all, F, PDER,FITLIST=fitlist
; where:
;	X = Vector of NPOINT independent variables, input.
;	A_all = Vector of function parameters, input.
;	F = Vector of NPOINT values of function, y(i) = funct(x(i)), output.
;	PDER = Array, (NPOINT, NTERMS), of partial derivatives of funct.
;		PDER(I,J) = Derivative of function at ith point with
;		respect to jth parameter.  Optional output parameter.
;		PDER should not be calculated if the parameter is not
;		supplied in call. NTERMS is the number of FITTED parameters
;		which may not always be the same as the size of A_all
;	FITLIST keyword corresponds to an integer array with dimension
;		equal to the number of parameters, and a value of 1
;		if the parameter is fitted and a value of 0 if the
;		parameter is not fitted.
;
; PROCEDURE:
;	Copied from "CURFIT", least squares fit to a non-linear
;	function, pages 237-239, Bevington, Data Reduction and Error
;	Analysis for the Physical Sciences.
;
;	"This method is the Gradient-expansion algorithm which
;	combines the best features of the gradient search with
;	the method of linearizing the fitting function."
;
;	Iterations are performed until the chi square changes by
;	only 0.1% or until 20 iterations have been performed.
;
;	The initial guess of the parameter values should be
;	as close to the actual values as possible or the solution
;	may not converge.
;
; MODIFICATION HISTORY:
;	Written, DMS, RSI, September, 1982.
;	Does not iterate if the first guess is good.  DMS, Oct, 1990.
;	Added CALL_PROCEDURE to make the function's name a parameter.
;		(Nov 1990)
;	12/14/92 - modified to reflect the changes in the 1991
;		   edition of Bevington (eq. II-27) (jiy-suggested by CreaSo)
;	11/21/93 - rfrench@cfa.harvard.edu - double precision, add number
; 		of iterations as optional keyword, add optional index
;		list of parameters to be fitted, fix code for
;		1-parameter fits
;	11/22/93 - fix error calculation of parameters from curvefit version
;		add printing of correlation matrix
;	11/23/93 - if all weights are 1 on input, then scale the errors
;		by the sigma of the fit
;		 - allow an array of independent variables, instead of
;		just a vector.
;	12/8/93  - make sure that correlation matrix works if only one parm
;	1/15/96 - renamed from newrgfcurvefit.pro since PDN changed that
;		program to have a different storage method for sigmas
;		and partial derivatives. 
;	1/16/96 - rgfcurvefit2 - allow convlist 
;	6/25/96 - handle case with no parameters fitted
;	7/2/97  - rgfcurvefit3 - clean up and double-check errors.
;-
	T_start=systime(1)
        on_error,0		;stay here if ERROR
				;Name of function to fit
	if n_elements(function_name) le 0 then function_name = "FUNCT"
	if n_elements(niter) le 0 then niter=20 else niter=niter(0)
	if n_elements(fitlist) le 0 then fitlist=intarr(n_elements(a_all))+1
	if n_elements(convcrit) le 0 then convcrit=0.001d0 else $
		convcrit=convcrit(0)
	if n_elements(convcrit) le 0 then default_conv=1 else default_conv=0
	if n_elements(convlist) gt 0 then begin
				default_conv=-1
				list=where(fitlist,nfit)
				if nfit gt 0 then $
					clist=convlist(where(fitlist)) $
				else default_conv=1
	endif
	if n_elements(lambdacountmax) le 0 then lambdacountmax=30 $
		else lambdacountmax=lambdacountmax(0)
	converged=1	; set convergence flag to 1 (yes)
	y=double(y)

	if n_elements(printeach) le 0 then printeach=0 else $
		printeach=printeach(0)
	if n_elements(logfile) le 0 then logfile='' else logfile=logfile(0)

	nterms = total(fitlist)	;# of FITTED parameters
	if nterms LE 0 then begin ; evaluate function and return
		call_procedure, Function_name,X,a_all,yfit,/NODERIVATIVE
		niter_used=0
		return,yfit ; no fit to perform
	endif
	log=0
	if logfile ne '' then begin
		log=1
		get_lun,lun_log
		openw,lun_log,logfile
	endif

; If we will be estimating partial derivatives then compute machine precision
       if keyword_set(NODERIVATIVE) then begin
	  if keyword_set(FINITE_DIFFS) then eps=FINITE_DIFFS else begin
           res = machar(/DOUBLE)
           eps = sqrt(res.eps)
	   eps = replicate(eps,n_elements(A_all))
	  endelse
       endif

	a=a_all(where(fitlist)) ; grab only the fitted parms
	a = double(a)		; make double precision
; handle case of some zero-weighted points, which we should not include
; in the total number of data points
	npts=n_elements(y)
	lwzero=where(weights eq 0,nlwzero)
	lwone=where(weights eq 1.,nlwone)
; set npts to the number of unit-weighted points if weights are (0|1)
	if (nlwzero+nlwone) eq npts then npts=nlwone
	nfree = npts-nterms ;Degs of freedom
	if nfree lt 0 then message,'rgfcurvefit4 - not enough data points.'
	flambda = 0.001d0		;Initial lambda
	diag = lindgen(nterms)*(nterms+1) ;subscripts of diagonal elements
;
;      Define the partial derivative array
       pder = dblarr(n_elements(y), nterms)

	for iter = 1L,niter do begin	;Iteration loop
		T_iter=systime(1)
;
; evaluate alpha and beta matrices.
;
          if keyword_set(NODERIVATIVE) then begin
; Evaluate function and estimate partial derivatives
             call_procedure, Function_name, x, a_all, yfit
	     lfit=where(fitlist)
             for term=0L, nterms-1 do begin
		iterm=lfit(term)
                p = a_all       ; Copy current parameters
; Increment size for forward difference derivative
                inc = eps[iterm] * abs(p[iterm])
                if (inc eq 0.) then inc = eps[iterm]
                p[iterm] = p[iterm] + inc
                call_procedure, function_name, x, p, yfit1
                pder[0,term] = (yfit1-yfit)/inc

;	plot,x,pder[*,term],title='partial derivative for term '+ $
;		strtrim(string(term),2),ystyle=2
             endfor
          endif else begin
; The user's procedure will return partial derivatives
		call_procedure, function_name,x,a_all,yfit,pder,fitlist=fitlist
          endelse
	if nterms eq 1 then pder=reform(pder,n_elements(y),1)
	beta = (y-yfit)*weights # pder
	alpha = transpose(pder) # (weights # (dblarr(nterms)+1)*pder)
;
	chisq1 = total(weights*(y-yfit)^2)/nfree ;present chi squared.
	if(printeach eq 1) then print,'chisq1=',chisq1
				; If a good fit, no need to iterate
; make this test tighter - rfrench
	if chisq1 lt total(abs(y))/1e14/NFREE then begin
		print,'NO FIT PERFORMED - ADEQUATE FIT AT START'
		if log then begin
		  printf,lun_log,'NO FIT PERFORMED - ADEQUATE FIT AT START'
		  close,lun_log
		  free_lun,lun_log
		endif
		sigma = fltarr(nterms)	;Return all 0's
		return, yfit
		endif
;
;	invert modified curvature matrix to find new parameters.
;
; the following two terms moved out of loop in rgfcurvefit4
	c = sqrt(alpha[diag])
	c = c # c
;
	npar=total(fitlist)
	lambdaCount=0
	repeat begin
		lambdaCount=lambdaCount+1
		if(lambdaCount gt lambdaCountMAX) then begin
			if printeach then begin
				print,'Inner loop exceeded ',lambdaCountMAX
				print,'Bailing out...'
			endif
			goto,NO_CONVERGENCE
		endif
		if(printeach eq 1) then begin
			print,'Iteration ',ITER,' repetition ',lambdaCount
			print,'Elapsed time this iteration: ',$
				systime(1)-T_iter,' Seconds'
		endif
		if log then begin
			printf,lun_log,'Iteration ',ITER,' repetition ',lambdaCount
			printf,lun_log,'Elapsed time this iteration: ',$
				systime(1)-T_iter,' Seconds'
			flush,lun_log
		endif
		array = alpha/c
		array[diag] = array[diag]*(1.d0+flambda)		
		if npar gt 1 then array=invert(array) else array=1.d0/array
		b = a+ array/c # transpose(beta) ;new params
; the order of the following two lines reversed for rgfcurvefit4
		a_all(where(fitlist))=b
		b_all=a_all
;evaluate function
		call_procedure, function_name,x,b_all,yfit,fitlist=fitlist
		chisqr = total(weights*(y-yfit)^2)/(nfree>1)  ;new chisqr
		if(printeach eq 1) then $
			print,'New chisqr=',chisqr,' previous=',chisq1
		if log then begin
			printf,lun_log,'New chisqr=',chisqr,' previous=',chisq1
			flush,lun_log
		endif
		flambda = flambda*10.d0	;assume fit got worse
		endrep until chisqr le chisq1
;
	flambda = flambda/100.d0	;decrease flambda by factor of 10
	a_prev=a
	a=b			;save new parameter estimate.
	a_all(where(fitlist))=a ; reset the array
	a_test=a
	zero=where(a_test eq 0.d0,nzero)
	if(nzero gt 0) then a_test(zero)=1.d-16 ; to avoid divide by zero
	if(printeach eq 1) then begin
	 print,'Iteration =',iter,' ,chisqr =',chisqr
	 print,a
	endif
	if log then begin
	 printf,lun_log,'Iteration =',iter,' ,chisqr =',chisqr
	 printf,lun_log,a
	 flush,lun_log
	endif
	if (default_conv eq 1 and (((chisq1-chisqr)/chisq1) le convcrit))$
		 then goto,done ;Finished?
	if (default_conv eq 0 and (max(abs(a_prev/a -1.d0)) le convcrit))$
		 then goto,done ;Finished?
	if (default_conv eq -1) then begin
		diff=abs(a-a_prev)-clist
		list=where(diff gt 0,nbad)
		if printeach eq 1 then begin
			print,'# of parameters yet to converge: ',nbad,' of ',nterms
			print,'Prev, Now, Change, Clist'
			for np=0L,nterms-1 do begin
			 print,a_prev(np),a(np),a(np)-a_prev(np),clist(np)
			endfor
		endif
		if log then begin
			printf,lun_log,'# of parameters yet to converge: ',nbad,' of ',nterms
			printf,lun_log,'Prev, Now, Change, Clist'
			for np=0L,nterms-1 do begin
			 printf,lun_log,a_prev(np),a(np),abs(a(np)-a_prev(np)),clist(np)
			endfor
	 		flush,lun_log
		endif
		if nbad le 0 then goto,done ; Finished
	endif
	endfor			;iteration loop
;
NO_CONVERGENCE:
	converged=0
	message, 'Failed to converge', /INFORMATIONAL
	if printeach then print, 'Failed to converge'
	if log then printf,lun_log, 'Failed to converge'
;
DONE:
	if printeach and converged then print,'Converged!'
	if log and converged then printf,lun_log,'Converged!'
        if keyword_set(NODERIVATIVE) then begin
; Evaluate function and estimate partial derivatives
             call_procedure, Function_name, x, a_all, yfit
	     lfit=where(fitlist)
             for term=0L, nterms-1 do begin
		iterm=lfit(term)
                p = a_all       ; Copy current parameters
; Increment size for forward difference derivative
                inc = eps[iterm] * abs(p[iterm])
                if (inc eq 0.) then inc = eps[iterm]
                p[iterm] = p[iterm] + inc
                call_procedure, function_name, x, p, yfit1
                pder[0,term] = (yfit1-yfit)/inc
             endfor
          endif else begin
; The user's procedure will return partial derivatives
		call_procedure, function_name,x,a_all,yfit,pder,fitlist=fitlist
          endelse
; the following statement added rgfcurvefit4 to make sure chisqr was 
; using the final paramters
	chisqr = total(weights*(y-yfit)^2)/(nfree>1)  ;new chisqr
	if nterms eq 1 then pder=reform(pder,n_elements(y),1)
	if nterms eq 1 then pder=reform(pder,n_elements(Y),1)
	alpha = transpose(pder) # (weights # (dblarr(nterms)+1)*pder)
	if nterms gt 1 then cmatrix=invert(alpha) else cmatrix=1.d0/alpha
	correl_matrix=dblarr(nterms,nterms)
	for i=0L,nterms-1 do begin
		for j=0L,nterms-1 do begin
			correl_matrix(i,j)=cmatrix(i,j)/$
			 sqrt(cmatrix(i,i)*cmatrix(j,j))
		endfor
	endfor
	if(printeach ge 1) then begin
	 print,'correlation matrix:'
	 print,correl_matrix
	endif
	if log then begin
	 printf,lun_log,'correlation matrix:'
	 printf,lun_log,correl_matrix
	 flush,lun_log
	endif
; This version from IDL/curvefit does NOT scale properly with
; lambda, and I assume it is because lamba is not 0 for the final
; calculation. When you use the subsequent definition, it agrees 
; exactly with the error bars that come from MRQMIN in Numerical Recipes.
;        sigma = sqrt(array(diag)/alpha(diag)) ;RETURN SIGMA'S (RSI way)
	sigma=sqrt(cmatrix(diag))	; (my way)
; if data are unweighted (all weights are 1 or 0) then scale the parameter
; errors by the sigma of the fit

; 2 April 1999 - explanation:
;	The equations used here assume that the input weights are actual
;	1/sigma^2 per data point. If we specify weight of 1 per point,
;	that clearly is not the correct sigma. We then find the post facto
;	sigma per point and assume that this properly represents the
;	1/sigma^2 weight we should have supplied to the fit in the first
;	place, if only we had known it. WE do not want to be in a
;	situation where an arbitrary scale factor applied to the weight gives us
;	different error bars. The method here assumes that if weights are 
;	all 0 or 1, then the input weights are really just a switch to
;	turn on or off a data point. In this case, use rms per point at
;	end to scale sigmas that are computed assuming that sigma per
;	point is 1.0. On this dated, in ~/ringprep/data/uranus/urall21/RECIPES
;	directory, there is a bakeoff between Numerical Recipes and 
;	rgfcurvefit4 result that illustrates this. See fit_comparisons
;	in that directory.

	weighted=where(weights ne 1.0 and weights ne 0.,n_weighted)
	if(n_weighted le 0) then sigma=sigma*sqrt(chisqr)
	if(printeach ge 1) then begin
	 print,'iteration =',iter,' ,chisqr =',chisqr
	 print,a
	 print, sigma
	endif
	niter_used=iter
	T_stop=systime(1)
	if printeach then print,'Total time for fit:',T_stop-T_start,' seconds'
	if log then begin
		close,lun_log
		free_lun,lun_log
	endif
; expand sigma into an array as long as a_all
	temp=dblarr(n_elements(a_all))
	temp(where(fitlist)) = sigma
	sigma=temp

	return,yfit		;return result
end