Cleaned up the code
Added documentation; return correct numpy matrix
Finished polish; merging to master