#!/usr/bin/env python
from __future__ import division
from numpy import *
from numpy.random import rand
import numpy.linalg as la
from matplotlib import pyplot as plt

def C(x):
	return sum(x*x.reshape(-1,1,N),axis=-1)/N
def expect(x):
	return sum(x,axis=-1)[...,None]/N
def plot(x):
	plt.figure()
	plt.plot(x[0,::dplot],x[1,::dplot],linestyle='',marker='.')
	plt.gca().set_aspect(1.)

N = 20000
dplot = 5 #how many points to skip for plotting
s = rand(2,N)
plot(s)

A = array([[1,2],[3,1]])
print "A_inv=",la.inv(A)
x = dot(A,s)
plot(x)

x -= expect(x) #make zero mean
w,v = la.eig(C(x))
print w
M = (v/sqrt(w)).T
y = dot(M,x) #diagonalize with unit variance
print C(y)
plot(y)

#ICA
df = lambda x: tanh(x)
ddf = lambda x: 1/cosh(x)**2
def ica(w):
	i=0
	while True:
		i+=1
		w_new = (expect(df(dot(w,y))*y) - w[...,None]*expect(ddf(dot(w,y))) )[...,0]
		w_new /= la.norm(w_new) #misprint is squared, I think.
		if allclose(w_new+w,0): w_new *= -1 #avoid oscillation from normalization
		if la.norm(w_new-w)<1e-8 or i>1e3: break
		w = w_new
	return w

w1 = ica(rand(2)) #random seed for weights
#second ica component must be orthogonal to first, since we're in 2d, this
#is unique, up to sign.
w2 = array([-w1[1],w1[0]])

plt.plot([0,w1[0]],[0,w1[1]],marker='.',lw=4,c='r')
plt.plot([0,w2[0]],[0,w2[1]],marker='.',lw=4,c='m')
plt.gca().set_aspect(1.)
plt.show()