;-------------------------------------------------------------------------------
;
;	Function:	identity_matrix
;	Author:		Brian Vanderwende
;
;	Description
;		This function returns an integer identity matrix of requested size
;
;	Parameters
;		size - the desired number of rows and columns in the identity matrix
;
;	Usage
;		imat = identity_matrix(4)
;
;	Revision History
;		01/21/14	- Creation of the function
;
;-------------------------------------------------------------------------------

undef("identity_matrix")
function identity_matrix(size:integer)

local matrix, n

begin
	; Initialize the matrix
	matrix = new((/size,size/), integer)
	
	; Assign unity values along the diagonal
	matrix = 0
	
	do n = 0, size - 1
		matrix(n,n) = 1
	end do
	
	return(matrix)
end

;-------------------------------------------------------------------------------
;
;	Function:	lm_fit
;	Author:		Brian Vanderwende
;
;	Description
;		This function uses the Levenberg-Marquardt algorithm to estimate
;		fitted-parameters to a supplied user-defined function.
;
;	Parameters
;		y	- a 1d array of n dependent values
;		x	- a 1d or 2d array of n x m independent values
;		b	- a 1d array of initial guesses for the fit parameters
;		opt	- a logical resource variable wherein the following optional
;				settings can be applied
;			weights 	- an array of n weights (typically 1/variance)
;			iter_max	- the maximum number of iterations allowed while
;							solving for fit parameters (default = 20)
;			con_crit	- the convergence threshold for the equation
;							(r(i-1) - r(i)) / r(i - 1) where r is the residual
;							sum of squares and i is the current iteration
;
;	Usage
;		params = lm_fit(y, x, init, False)
;
;	Revision History
;		01/25/14	- Added optional settings
;		01/20/14	- Creation of the function
;
;-------------------------------------------------------------------------------

undef("lm_fit_func")
function lm_fit_func(x:numeric, b:numeric, grad:numeric)

begin
	print("ERROR: lm_fit_func has not been redefined by the user!")
	print("")
	print("In order to run the lm_fit function, you must replace this")
	print("prototype version of lm_fit_func with one containing the")
	print("function you wish to fit.")
	print("")
	print("The function should accept three arguments:")
	print("   x    - 1d/2d array of independent variable values")
	print("   b    - parameters to be fitted")
	print("   grad - storage array for gradient vector/matrix")
	print("")
	print("The function should calculate a 1d array of y-estimates using the")
	print("input fit parameters. The gradient vector (partial derivatives")
	print("with respect to each parameter) should also be calculated and")
	print("stored inside grad.")
	exit
end

undef("lm_fit")
function lm_fit(y:numeric, x:numeric, b_in:numeric, opt:logical)

local iter_max, con_crit, num_params, num_data, num_vars, x_dims, lamdba, grad,	\
		n, yf, chi_sq, delta, b, b_old, yf_old, cs_old, num_tries, temp, W

begin
	; Ensure that integer input does not cause rounding errors
	b = 1.0 * b_in
	
	; Set default values for algorithm settings if not provided
	if(isatt(opt, "iter_max")) then
		iter_max = opt@iter_max
	else
		iter_max = 20
	end if
	
	if(isatt(opt, "con_crit")) then
		con_crit = opt@con_crit
	else
		con_crit = 1e-3
	end if
	
	; Find the size of input arrays
	num_params 	= dimsizes(b)
	num_data	= dimsizes(y)
	x_dims 		= dimsizes(x)
	num_vars 	= x_dims(dimsizes(x_dims) - 1)
	
	; Check to make sure there are more data points than parameters
	if(num_params .gt. num_data) then
		print("ERROR: there must be more dependent function values "			\
				+ "than parameters!")
		exit
	end if
	
	; If weights are provided, initialize weighting matrix
	weights = new((/num_data,num_data/), float)
	weights = 0.0
	
	if(isatt(opt, "weights")) then
		if(dimsizes(opt@weights) .ne. num_data) then
			print("ERROR: number of weights must equal number of dependent"		\
					+ "data points!")
			exit
		end if
	else
		opt@weights 	= new(num_data, integer)
		opt@weights(:) 	= 1
	end if
	
	do n = 0, num_data - 1
		weights(n,n) = opt@weights(n)
	end do

	; Declare variables for iteration
	lambda 	= 1e-3
	grad	= new((/num_data,num_params/), "float")
	
	; Begin L-M iterative solve
	do iter_num = 1, iter_max - 1
		; Get y estimates and gradient values from user function
		yf = lm_fit_func(x, b, grad)
		
		beta 	= transpose(grad) # weights # (y - yf)
		alpha	= transpose(grad) # weights # grad
		
		; Compute sum of squares
		chi_sq = sum(opt@weights * (y - yf)^2)
		
		; Store old values and initialize Marquardt parameter guess counter
		b_old 		= b
		yf_old		= yf
		cs_old		= chi_sq
		num_tries	= 1
		
		do while(.not.isnan_ieee(chi_sq) .and. num_tries .le. 30 				\
				.and. chi_sq .ge. cs_old)
			; Get the next coefficients
			b = b_old + inverse_matrix(alpha + lambda * alpha 					\
					* identity_matrix(num_params)) # beta
			
			if(any(ismissing(b))) then
				break
			end if
			
			; Test fit
			yf 		= lm_fit_func(x, b, grad)
			chi_sq	= sum(opt@weights * (y - yf)^2)
			
			; Align shift vector towards steepest descent
			lambda 		= lambda * 10
			num_tries 	= num_tries + 1
		end do
		
		; Check if fit is improved and if fit is within convergence criterion
		if(chi_sq .lt. cs_old) then
			if((cs_old - chi_sq) / cs_old .le. con_crit) then
				yf@status = True
				break
			else
				; Regress lambda toward Gauss-Newton
				lambda = lambda / 100
			end if
		else			
			; Algorithm failed to converge, use last values
			b 		= b_old
			yf		= yf_old
			chi_sq 	= cs_old
			
			yf@status = False
			break
		end if
	end do
	
	; Attach metadata to result and return
	yf@chi_sq 	= chi_sq
	yf@params	= b
	yf@iter_num = iter_num
	return(yf)
end