Do we want to use `np.arange` instead of `range`? The line above use one or the other.  I would opt for `numpy`.  In fact, could not we make a variable that contains that range and reuse it? Just to avoid repetition everywhere.