These changes drastically improve the speed of the calculations in the main loop by replacing nested for-loops with numpy matrix operations. In general, I replaced many slow for-loops with corresponding matrix operations, always keeping the idea of the original algorithm. This sometimes introduces small numerical differences, which I believe are insignificant. In total, calculations (without plotting) are ~5x faster. Plotting now takes by far the most amount of time in each iteration.
Detailed comparison of performance improvements:
Part
Improvement
Radiation
~35x
Velocity
~20x
Advection
~120x
Projection*
~4x
* grid_velocities_[north|south] should also be optimized, but I could not wrap my head around the side-effects in the nested for-loops of these functions
These changes drastically improve the speed of the calculations in the main loop by replacing nested for-loops with numpy matrix operations. In general, I replaced many slow for-loops with corresponding matrix operations, always keeping the idea of the original algorithm. This sometimes introduces small numerical differences, which I believe are insignificant. In total, calculations (without plotting) are ~5x faster. Plotting now takes by far the most amount of time in each iteration.
*
grid_velocities_[north|south]
should also be optimized, but I could not wrap my head around the side-effects in the nested for-loops of these functions