#!/usr/bin/env python

import vector
import math, types, operator

"""
A sparse matrix class based on a dictionary, supporting matrix (dot)
product and a conjugate gradient solver. 

In this version, the sparse class inherits from the dictionary; this
requires Python 2.2 or later.
"""

class sparse(dict):
	"""
	A complex sparse matrix 
	A. Pletzer 5 Jan 00/12 April 2002

	Dictionary storage format { (i,j): value, ... }
	where (i,j) are the matrix indices
       	"""

	# no c'tor

	def size(self):
		" returns # of rows and columns "
		nrow = 0
		ncol = 0
		for key in self.keys():
			nrow = max([nrow, key[0]+1])
			ncol = max([ncol, key[1]+1])
		return (nrow, ncol)

	def __add__(self, other):
		res = sparse(self.copy())
		for ij in other:
			res[ij] = self.get(ij,0.) + other[ij]
		return res
		
	def __neg__(self):
		return sparse(zip(self.keys(), map(operator.neg, self.values())))

	def __sub__(self, other):
		res = sparse(self.copy())
		for ij in other:
			res[ij] = self.get(ij,0.) - other[ij]
		return res
		
	def __mul__(self, other):
		" element by element multiplication: other can be scalar or sparse "
		try:
			# other is sparse
			nval = len(other)
			res = sparse()
			if nval < len(self):
				for ij in other:
					res[ij] = self.get(ij,0.)*other[ij]
			else:
				for ij in self:
					res[ij] = self[ij]*other.get(ij,0j)
			return res
		except:
			# other is scalar
			return sparse(zip(self.keys(), map(lambda x: x*other, self.values())))


	def __rmul__(self, other): return self.__mul__(other)

	def __div__(self, other):
		" element by element division self/other: other is scalar"
		return sparse(zip(self.keys(), map(lambda x: x/other, self.values())))
		
	def __rdiv__(self, other):
		" element by element division other/self: other is scalar"
		return sparse(zip(self.keys(), map(lambda x: other/x, self.values())))

	def abs(self):
		return sparse(zip(self.keys(), map(operator.abs, self.values())))

	def out(self):
		print '# (i, j) -- value'
		for k in self.keys():
			print k, self[k]

	def plot(self, width_in=400, height_in=400):

		import colormap
		import Tkinter

		cmax = max(self.values())
		cmin = min(self.values())
		
		offset =  0.05*min(width_in, height_in)
		xmin, ymin, xmax, ymax = 0,0,self.size()[0], self.size()[1]
		scale =  min(0.9*width_in, 0.9*height_in)/max(xmax-xmin, ymax-ymin)

		root = Tkinter.Tk()
		frame = Tkinter.Frame(root)
		frame.pack()
		
		text = Tkinter.Label(width=20, height=10, text='matrix sparsity')
		text.pack()
		

		canvas = Tkinter.Canvas(bg="black", width=width_in, height=height_in)
		canvas.pack()

		button = Tkinter.Button(frame, text="OK?", fg="red", command=frame.quit)
		button.pack()

		for index in self.keys():
			ix, iy = index[0], ymax-index[1]-1
			ya, xa = offset+scale*(ix  ), height_in -offset-scale*(iy  )
			yb, xb = offset+scale*(ix+1), height_in -offset-scale*(iy  )
			yc, xc = offset+scale*(ix+1), height_in -offset-scale*(iy+1)
			yd, xd = offset+scale*(ix  ), height_in -offset-scale*(iy+1)
			color = colormap.strRgb(self[index], cmin, cmax)
			canvas.create_polygon(xa, ya, xb, yb, xc, yc, xd, yd, fill=color)
		
		root.mainloop()

	def CGsolve(self, x0, b, tol=1.0e-10, nmax = 1000, verbose=1):
		"""
		Solve self*x = b and return x using the conjugate gradient method
		"""
		if not vector.isVector(b):
			raise TypeError, self.__class__,' in solve '
		else:
			if self.size()[0] != len(b) or self.size()[1] != len(b):
				print '**Incompatible sizes in solve'
				print '**size()=', self.size()[0], self.size()[1]
				print '**len=', len(b)
			else:
				kvec = diag(self) # preconditionner
				n = len(b)
				x = x0 # initial guess
				r =  b - dot(self, x)
				try:
					w = r/kvec
				except: print '***singular kvec'
				p = vector.zeros(n);
				beta = 0.0;
				rho = vector.dot(r, w);
				err = vector.norm(dot(self,x) - b);
				k = 0
				if verbose: print " conjugate gradient convergence (log error)"
				while abs(err) > tol and k < nmax:
					p = w + beta*p;
					z = dot(self, p);
					alpha = rho/vector.dot(p, z);
					r = r - alpha*z;
					w = r/kvec;
					rhoold = rho;
					rho = vector.dot(r, w);
					x = x + alpha*p;
					beta = rho/rhoold;
					err = vector.norm(dot(self, x) - b);
					if verbose: print k,' %5.1f ' % math.log10(err)
					k = k+1
				return x
				
		
	    		
	def biCGsolve(self,x0, b, tol=1.0e-10, nmax = 1000):
		
		"""
		Solve self*x = b and return x using the bi-conjugate gradient method
		"""

		try:
			if not vector.isVector(b):
				raise TypeError, self.__class__,' in solve '
			else:
				if self.size()[0] != len(b) or self.size()[1] != len(b):
					print '**Incompatible sizes in solve'
					print '**size()=', self.size()[0], self.size()[1]
					print '**len=', len(b)
				else:
					kvec = diag(self) # preconditionner 
					n = len(b)
					x = x0 # initial guess
					r =  b - dot(self, x)
					rbar =  r
					w = r/kvec;
					wbar = rbar/kvec;
					p = vector.zeros(n);
					pbar = vector.zeros(n);
					beta = 0.0;
					rho = vector.dot(rbar, w);
					err = vector.norm(dot(self,x) - b);
					k = 0
					print " bi-conjugate gradient convergence (log error)"
					while abs(err) > tol and k < nmax:
						p = w + beta*p;
						pbar = wbar + beta*pbar;
						z = dot(self, p);
						alpha = rho/vector.dot(pbar, z);
						r = r - alpha*z;
						rbar = rbar - alpha* dot(pbar, self);
						w = r/kvec;
						wbar = rbar/kvec;
						rhoold = rho;
						rho = vector.dot(rbar, w);
						x = x + alpha*p;
						beta = rho/rhoold;
						err = vector.norm(dot(self, x) - b);
						print k,' %5.1f ' % math.log10(err)
						k = k+1
					return x
			
		except: print 'ERROR ',self.__class__,'::biCGsolve'


	def save(self, filename, OneBased=0):
		"""
		Save matrix in file <filaname> using format:
		OneBased, nrow, ncol, nnonzeros
		[ii, jj, data]

		"""
		m = n = 0
		nnz = len(self)
		for ij in self.keys():
			m = max(ij[0], m)
			n = max(ij[1], n)

		f = open(filename,'w')
		f.write('%d %d %d %d\n' % (OneBased, m+1,n+1,nnz))
		for ij in self.keys():
			i,j = ij
			f.write('%d %d %20.17f \n'% \
				(i+OneBased,j+OneBased,self[ij]))
		f.close()
				
###############################################################################

def isSparse(x):
    return hasattr(x,'__class__') and x.__class__ is sparse

def transp(a):
	" transpose "
	new = sparse({})
	for ij in a:
		new[(ij[1], ij[0])] = a[ij]
	return new

def dotDot(y,a,x):
	" double dot product y^+ *A*x "
	if vector.isVector(y) and isSparse(a) and vector.isVector(x):
		res = 0.
		for ij in a.keys():
			i,j = ij
			res += y[i]*a[ij]*x[j]
		return res
	else:
		print 'sparse::Error: dotDot takes vector, sparse , vector as args'

def dot(a, b):
	" vector-matrix, matrix-vector or matrix-matrix product "
	if isSparse(a) and vector.isVector(b):
		new = vector.zeros(a.size()[0])
		for ij in a.keys():
			new[ij[0]] += a[ij]* b[ij[1]]
		return new
	elif vector.isVector(a) and isSparse(b):
		new = vector.zeros(b.size()[1])
		for ij in b.keys():
			new[ij[1]] += a[ij[0]]* b[ij]
		return new
	elif isSparse(a) and isSparse(b):
		if a.size()[1] != b.size()[0]:
			print '**Warning shapes do not match in dot(sparse, sparse)'
		new = sparse({})
		n = min([a.size()[1], b.size()[0]])
		for i in range(a.size()[0]):
			for j in range(b.size()[1]):
				sum = 0.
				for k in range(n):
					sum += a.get((i,k),0.)*b.get((k,j),0.)
				if sum != 0.:
					new[(i,j)] = sum
		return new
	else:
		raise TypeError, 'in dot'

def diag(b):
	# given a sparse matrix b return its diagonal
	res = vector.zeros(b.size()[0])
	for i in range(b.size()[0]):
		res[i] = b.get((i,i), 0.)
	return res
		
def identity(n):
	if type(n) != types.IntType:
		raise TypeError, ' in identity: # must be integer'
	else:
		new = sparse({})
		for i in range(n):
			new[(i,i)] = 1+0.
		return new

###############################################################################
if __name__ == "__main__":

	print 'a = sparse()'
	a = sparse()

	print 'a.__doc__=',a.__doc__

	print 'a[(0,0)] = 1.0'
	a[(0,0)] = 1.0
	a.out()

	print 'a[(2,3)] = 3.0'
	a[(2,3)] = 3.0
	a.out()

	print 'len(a)=',len(a)
	print 'a.size()=', a.size()
			
	b = sparse({(0,0):2.0, (0,1):1.0, (1,0):1.0, (1,1):2.0, (1,2):1.0, (2,1):1.0, (2,2):2.0})
	print 'a=', a
	print 'b=', b
	b.out()

	print 'a+b'
	c = a + b
	c.out()

	print '-a'
	c = -a
	c.out()
	a.out()

	print 'a-b'
	c = a - b
	c.out()

	print 'a*1.2'
	c = a*1.2
	c.out()


	print '1.2*a'
	c = 1.2*a
	c.out()
	print 'a=', a

	print 'dot(a, b)'
	print 'a.size()[1]=',a.size()[1],' b.size()[0]=', b.size()[0]
	c = dot(a, b)
	c.out()

	print 'dot(b, a)'
	print 'b.size()[1]=',b.size()[1],' a.size()[0]=', a.size()[0]
	c = dot(b, a)
	c.out()

	try:
		print 'dot(b, vector.vector([1,2,3]))'
		c = dot(b, vector.vector([1,2,3]))
		c.out()
	
		print 'dot(vector.vector([1,2,3]), b)'
		c = dot(vector.vector([1,2,3]), b)
		c.out()

		print 'b.size()=', b.size()
	except: pass
	
	print 'a*b -> element by element product'
	c = a*b
	c.out()

	print 'b*a -> element by element product'
	c = b*a
	c.out()
	
	print 'a/1.2'
	c = a/1.2
	c.out()

	print 'c = identity(4)'
	c = identity(4)
	c.out()

	print 'c = transp(a)'
	c = transp(a)
	c.out()


	b[(2,2)]=-10.0
	b[(2,0)]=+10.0

	try:
		import vector
		print 'Check conjugate gradient solver'
		s = vector.vector([1, 0, 0])
		print 's'
		s.out()
		x0 = s 
		print 'x = b.biCGsolve(x0, s, 1.0e-10, len(b)+1)'
		x = b.biCGsolve(x0, s, 1.0e-10,  len(b)+1)
		x.out()

		print 'check validity of CG'
		c = dot(b, x) - s
		c.out()
	except: pass

	print 'plot b matrix'
	b.out()
	b.plot()

	print 'del b[(2,2)]'
	del b[(2,2)]

	print 'del a'
	del a
	#a.out()