Issue410

and also some clean up (removing code that was not being used and required a lot of if's).

@josenavas, can you check if this solves the issue? Thanks!
Can this and the next line be reduced by using `np.ones`?