STY: Fix flake8 errors

Most of them were very minor, fixes #547.
While we're looking at this code, we might as well change this to the simpler/faster
```python
counts = np.ones(1001)
counts[0] = 9000
```