Some Code improvements
- prefer ravel of flatten when possible (https://stackoverflow.com/questions/28930465/what-is-the-difference-between-flatten-and-ravel-functions-in-numpy)
- prefer enumerate over range(len(...)) (https://www.reddit.com/r/learnpython/comments/nn0il2/rangelens_vs_enumerate/)