#!/usr/bin/env python
from numpy import *
from numpy.linalg import solve
from matplotlib import pyplot as plt

def main():
	n=100 #spatial samples
	dev = .1 #deviation from test signal
	x = random.random(n) #spatial sampling points
	z = random.normal(loc=0.,scale=dev,size=n) #gaussian noise
	y = sin(2 + 3*x) + z  #samples

	def Y(a):
		return sin(a[0]+a[1]*x)
	def dY(a):
		return array([cos(a[0]+a[1]*x), x*cos(a[0]+a[1]*x)])
	def ddY(a):
		return array([ [-sin(a[0]+a[1]*x), -x*sin(a[0]+a[1]*x)], [-x*sin(a[0]+a[1]*x), -x*x*sin(a[0]+a[1]*x)]])
	def M(a,lamb):
		ddchi = sum(2*dY(a)*dY(a).reshape(2,1,-1) - 2*(y-Y(a))*ddY(a),axis=-1)
		return .5*ddchi*array([[1+lamb,0],[0,1+lamb]])
	def M_inv(a,lamb):
		m = M(a,lamb)
		c = 1./(m[0,0]*m[1,1]-m[0,1]*m[1,0])
		return c*array([[m[1,1],-m[1,0]],[-m[0,1],m[0,0]]])
	def chi_sq(a):
		return sum((y-Y(a))**2,axis=-1) #assume sample variances are identical
	def d_chi_sq(a):
		return sum(-2*(y-Y(a))*dY(a), axis=-1)

	a = array([1.,1.]) #decision variable
	a_history = array([a])
	chi_history = array([chi_sq(a)])
	lamb = 10.
	l_history = array([lamb])
	beta = 1.01
	n_steps = 10000
	n_rec = 10
	for i in range(n_steps):
		error = chi_sq(a)
		#a += -solve(M(a,lamb),d_chi_sq(a))
		a += -dot(M_inv(a,lamb),d_chi_sq(a))
		if chi_sq(a)>error:
			lamb *= beta #increase lambda, use more gradient descent
		else:
			lamb /= beta #decrease lambda, use more newton
		if i%(n_steps/n_rec)==0:
			a_history = vstack((a_history,a))
		chi_history = vstack((chi_history,chi_sq(a)))
		l_history = vstack((l_history,lamb))

	print a, 'should be', [2,3]
	print "lambda=%.3f"%lamb
	print "chi_sq=%.3f"%chi_sq(a)

	plt.figure()
	plt.plot(x,y,marker='.',linestyle='',label='samples')
	t = linspace(0,1,50)
	for ai in a_history:
		plt.plot(t,sin(ai[0]+ai[1]*t))
	plt.title('Nonlinear least squares fitting')

	plt.figure()
	plt.plot(chi_history)
	plt.title('Error history')
	plt.xlabel('iteration')
	plt.ylabel('$\chi^2$')

	plt.figure()
	plt.plot(l_history)
	plt.title('$\lambda$ history')
	plt.xlabel('iteration')
	plt.ylabel('$\lambda$')

	plt.show()

if __name__ == '__main__':
	main()