many of the calculations could be optimized by separating out into functions and using the `@jit`decorator