Session outline

• Presentation on optimisation
  • Profiling
  • Testing
  • Speeding up Python code
• Practical
  • Open ended session with a short code to try some of these techniques out on.
  • Competition : Who can achieve the greatest speed-up?

The optimisation cycle

cProfile

import cProfile
import pstats

profile = cProfile.Profile()

profile.enable()

dn_dlogm_list = []
for z in z_list:
    dn_dlogm_list.append(press_schecter(M_list, z, A_std_func()))

profile.disable()

ps = pstats.Stats(profile).strip_dirs().sort_stats(pstats.SortKey.TIME)
ps.print_stats(10)

cProfile

cProfile is part of the Python standard library (you don't need to manually install it).
It gives statistics for function calls and is often the best place to start when profiling Python code.

import cProfile
import pstats

profile = cProfile.Profile()

profile.enable()
dn_dlogm_list = []
for z in z_list:
    dn_dlogm_list.append(press_schecter(M_list, z, A_std_func()))

profile.disable()
ps = pstats.Stats(profile).strip_dirs().sort_stats(pstats.SortKey.TIME)
ps.print_stats(10)

cProfile

cProfile is part of the Python standard library (you don't need to manually install it).
It gives statistics for function calls and is often the best place to start when profiling Python code.

import cProfile
import pstats

profile = cProfile.Profile()

profile.enable()

dn_dlogm_list = []
for z in z_list:
    dn_dlogm_list.append(press_schecter(M_list, z, A_std_func()))

profile.disable()

ps = pstats.Stats(profile).strip_dirs().sort_stats(pstats.SortKey.TIME)ps.print_stats(10)

cProfile

• tottime: time spent inside this function
• cumtime: time spent inside this function and all functions it calls
• percall: cumulative time per call

line_profiler

What if we want to go deeper though?
Where in the function are we spending all of our time? Calculating logarithms, multiplying numbers?

line_profiler

• Hits: Number of times this line is called.
• Total: Cumulative time spent on this line.
• Per Hit: Time spent on line each time it's called.
• % Time: Percentage of time spent on this line relative to whole function.

Pytest

Covering pytest and all it has to offer would take a whole session itself!
There are a lot of great examples and articles online.

Some useful topics to check out include:

• Fixtures
• Parameterizing tests
• Setting up continuous integration (CI) for automated testing
• Coverage measurements

For regression testing, the regtest plugin is popular¹...

Regression tests with pytest

Naming of directories, files and functions all matter with pytest. You can configure this, but standard practice is:

• Place all of your tests in a directory called tests.
• You need to have a __init__.py file in tests (to make them part of a module) but it can be empty.
• Each file should be called test_*.py.
• Each test function should be called test_*.

Regression tests with pytest

We need to initialise and store the expected output the first time around.

The results will go in tests/_regtest_outputs/ which should be committed to version control.

Regression tests with pytest

Now, after trying to optimise the code, we can check we still get the right result using:

Regression tests

If the output of the press_schechter function changes:

Things to know about Python and speed

Python is an interpreted language and (almost) everything is an object

This means there is a cost with looking up and accessing values.

Translated to C this is actually something more like (but in reality way more complicated!):

The most important scientific python optimisation rule

For loops should be avoided if possible. Take advantage of Numpy's ufuncs (which will vectorize it and do it more efficiently).

An aside...

A common misconception:

Numpy vectorize does not "vectorize" your code.

From the docs:

"The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop."
-- Numpy docs

I've removed all the loops I can, but I still need more speed!!!

There are a number of techniques and tools at our disposal. Which ones will work best is heavily dependent on the problem. e.g.:

• Use 3rd-party libraries which are already optimised!
  (I won't cover this, but it should always be your first stop.)
• Algorithmic changes (e.g. identifying redundant calculations).

Memoisation

The idea

Sometimes we have a function that may be called many times with the same input. In these situations it may be faster to store the results and reuse them rather than recalculating them each time.

Works well when

• You have an expensive function being called a number of times with the same input.
• You have any non-trivial function being called many, many times with the same small input.

Memoisation

A (very contrived) example¹

5.02 s ± 242 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Memoisation

A (very contrived) example¹

from functools import cache
import numpy as np

points = np.random.rand(10_000, 2)

@cachedef myfunc(x, y):
    return (np.log10(np.sqrt(x**2 + y**2) * 5) / 0.2)**6

result = 0.0
for _ in range(100):
    result += np.array([myfunc(x, y) for x, y in points])

1.28 s ± 423 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Memoisation

However...

functools.cache stores the result in RAM and needs to compare each input argument against the cached results.
This works well when arguments and output are small e.g floats ints etc.

joblib.Memory

These limitations can be overcome in some cases with joblib.Memory.
It caches to disk and removes much of the overhead with large input and output arrays.

Going lower level...

Sometimes the best way to improve the speed of our code is to use another language!...

But fear not! This does not mean having to re-write your code in C/C++/Rust etc.

Numba

Using Numba can be incredibly easy:

• njit means compile in "no-python" mode. This requires the function to be simple and use only support functions and types (most of Numpy) but not Python classes or dicts.
• Loops are fine! They will be optimised by LLVM (the compiler) for your CPU.

• 15.7 ms with decorator vs. 42.5 ms without = ~2.5x speed-up

Where Numba might not be enough
(WARNING: very subjective!)

• It only supports a subset of Python + Numpy (although a very large one).
• It can be hard to debug if things don't work as expected.
• There isn't much scope for workarounds. If it doesn't "just work" then often it won't work.

Despite these, I definitely recommend Numba as the first-stop for speeding up numerical code!

Parallelisation

Almost all computers have multiple cores. In some cases we can take advantage of this to speed up our calculations. Numpy & Scipy actually do this for some functions without you necessarily realising.

This technique works well when:

• you have an expensive function
• being called multiple times with different arguments
• independently of each other.

Extra
Native parallelisation in Python is expensive

The GIL

The Global Interpreter Lock (GIL) is a mutex which ensures only one thread can use the interpreter at any one time. It's a feature that makes Python simple and easy to use.
But it's a hurdle to efficient parallelisation.

To get around this natively we can use the multiprocessing standard library.
It uses multiple processes, each running their own Python interpreter, with their own private memory space.
But this is expensive, so we need to have enough work for each process to make it worthwhile.

Parallelisation with Numba

As it turns out, Numba (does) and Cython (can) release the GIL!

Numba can parallelise code automatically in some cases (see the docs), but we can also manually request it:

import numba
import numpy as np

@numba.njit(parallel=True)def myfunc(arr, n_loops):
    result = np.zeros(arr.shape[0])
    for _ in numba.prange(n_loops):        result += (np.log10(np.sqrt(arr[:, 0]**2 + arr[:, 1]**2) * 5) / 0.2)**6
    return result

points = np.random.rand(10_000, 2)
result = myfunc(points, 10_000)

Even with Numba, the best performance gains only come when the units of work are large enough.

For more information and a good, flexible library for multiple different parallelisation backends, see joblib.

Practical

Main suggestion

• Optimize code_prac/press_schechter.py and see how fast you can make it¹.
  • Time your attempts using
    python -m timeit -s 'import code_prac' 'code_prac.press_schechter.main()'
• If you followed the setup instructions in the repo then you should have all of the libraries and packages you may need installed.
• Compete for the title of "Optimizer Prime" !
  • There are no rules, except that we must be able to call code_prac.press_schechter.mass_function with the same arguments as the starting code and get the same result.

Practical

and/or

Take advantage of the knowledge of your peers and practice your skills with the press_schechter code.
Some ideas include:

• Add inline documentation and build docs using Sphinx
  • Tip: use the Napolean extension and Numpy or Googledoc style docstrings.
  • Bonus: Fork of the code_prac_hwsa2021 repo on Github and serve the docs online using Github Pages.
• Develop unit tests for the code
  • Bonus: Use coverage.py tool to measure your unit test coverage and aim for >90% coverage.
  • Bonus bonus: Fork the code_prac_hwsa2021 repo and use Github Actions to automatically test the code on every push.
• Add type annotations to the code then use mypy to check these.
• If you manage to speed up the press_schechter function enough (or ask me for a faster version), try making an interactive tool for exploring the PS mass function using Jupyter Notebooks and ipywidgets (or a tool of your choice).

Remember

1. Profile
2. Develop regression tests
3. Baseline timing
   
   
4. Optimise identified bottlenecks
5. Time
6. Test
7. Profile then go back to 4
   or call it a day