Plot timings for a range of inputs MkII

Question

To celebrate questions getting more reputation points per upvote, I revisited my previous question, Plot timings for a range of inputs.

Since the time of that original question (more than two years!), this code has undergone quite a few changes. I incorporated some of the advice given in the answer by @MaartenFabré. I also added multiprocessing to the mix to perform the timings in parallel. I added some convenience flags as well as using the minimum time recorded as the value, instead of the mean.

Code:

from __future__ import print_function
from functools import partial
from itertools import product, count
import matplotlib.pyplot as plt
import multiprocessing
import numpy as np
import pandas as pd
import timeit
from uncertainties import ufloat, unumpy as unp

N_CPU = multiprocessing.cpu_count()
DEBUG = False

def get_time(func, *x):
    """Run a timer for `func` five times with the given input."""
    timer = timeit.Timer(partial(func, *x))
    t = timer.repeat(repeat=5, number=1)
    if DEBUG:
        print(func.__name__, np.min(t))
    return ufloat(np.min(t), np.std(t) / np.sqrt(len(t)))

def flatten_once(it):
    """For passing multiple arguments to functions, flatten each inner list."""
    for x in it:
        yield [x[0], *x[1]]

def identity(x):
    """The identity function."""
    return x

def get_timings_df(funcs, inputs, key=identity, star=False, doc=False):
    """Use multiprocessing to time all `funcs` using the `inputs`.

       key:    Function that is called once on each input to determine the
               x-value. Default: identity
       star:   If the functions take multiple argument, and `inputs` contains
               tuples of arguments, splat them out. Default: False
       doc:    Use `func.__doc__` as plotting label instead of `func.__name__`.
               Default: False
    """
    df = pd.DataFrame(list(map(key, inputs)), columns=["x"])
    labels = [func.__name__ for func in funcs]
    if doc:
        labels = [func.__doc__ for func in funcs]
    y = product(funcs, inputs)
    if star:
        y = flatten_once(y)
    with multiprocessing.Pool(processes=N_CPU-1) as pool:
        times = np.array(pool.starmap(get_time, y)).reshape(len(funcs), -1)
    for label, t in zip(labels, times):
        df[label] = t
    return df

def plot_times(funcs, inputs, key=identity, xlabel="x", ylabel="Time [s]", 
               logx=False, logy=False, star=False, ratio=False, doc=False):
    """Plot timings of `funcs` using `inputs`.

       key:    Function that is called once on each input to determine the
               x-value. Default: identity.
       xlabel: Label of x-axis. Default: 'x'.
       ylabel: Label of y-axis (may be overwritten for some flags). 
               Default: 'Time [s]'.
       logx:   Make x-axis logarithmic. Default: False.
       logy:   Make y-axis logarithmic. Default: False.   
       star:   If the functions take multiple argument, and `inputs` contains
               tuples of arguments, splat them out. Default: False
       ratio:  Plot timings relative to time of first function. Default: False.
       doc:    Use `func.__doc__` as plotting label instead of `func.__name__`.
               Default: False
    """
    df = get_timings_df(funcs, inputs, key, star, doc)
    for label in df.columns[1:]:
        x, y = df["x"], df[label]
        if ratio:
            y = y / df.T.iloc[1]
        plt.errorbar(x, unp.nominal_values(y), unp.std_devs(y),
                     fmt='o-', label=label)
    plt.xlabel(xlabel)
    if ratio:
        ylabel = "{} / Time of {} [s]".format(ylabel, df.columns[1])
    plt.ylabel(ylabel)
    if logx:
        plt.xscale("log")
    if logy:
        plt.yscale("log")
    plt.legend()
    plt.show()

For now, I want to keep Python 2 compatibility, so, unfortunately, no f-strings. I am also not a big fan of type-hints, and I don't think they would help a lot here, but I'm open to being convinced otherwise.

If you can find a way to simplify any of the functions (or get effectively rid of them), I would love to hear it. The same is true for simplifying the interface of plot_times. All other recommendations are welcome as well.

Example usages:

This function is quite flexible. It allows the typical use-case of timing a bunch of single-argument functions with a bunch of inputs:

import time 
from timing import plot_times

def linear(x):
    """$\mathcal{O}(n)$"""
    time.sleep(x)

def quadratic(x):
    """$\mathcal{O}(n^2)$"""
    time.sleep(x**2)

if __name__ == "__main__":
    x = np.arange(0, 1, 0.1)
    plot_times([linear, quadratic], x, doc=True)

To timing multi-argument functions and plotting them relative to the first function:

if __name__ == "__main__":
    x = np.arange(0.1, 1, 0.1)
    plot_times([linear, quadratic], x, doc=True, ratio=True)

To use a different function to map values to the x-axis and semi- or log-log plotting:

import string
import random
import numpy as np
from timing import plot_times

def count_digits(x):
    return sum(c.isdigit() for c in x)

def count_digits2(x):
    return sum(1 for c in x if c.isdigit())

if __name__ == "__main__":
    alpha_num = string.ascii_letters + string.digits
    x = ["".join(random.choices(alpha_num, k=n))
         for n in np.logspace(1, 5, dtype=int)]
    plot_times([count_digits, count_digits2], x, key=len, xlabel="len(s)",
               logx=True, logy=True)

Answer 1

These plots are a great visual aid to support algorithmic arguments. I've been guilty of creating some of them myself, although more on the go and not with a dedicated "framework" like yours.

The code itself is clean and to the point. It's a great example of the power and beauty of Python. In under 100 lines of code you can create a usable tool, including documentation and multiprocessing. I'd support your opinion on that type hints would not really make it better here.

In essence there are only a few points that I personally would change:

get_timings_df

• I prefer list comprehensions over map, so I would likely use df = pd.DataFrame([key(input_) for input_ in inputs], columns=["x"]), mainly because I find it clearer.
• Since we are at it, I don't like key as a parameter name here. I'd suggest to use something like preprocess, input_transform or another more telling name instead. I also understand that maybe one could argue that key is "more in-line" with general purpose functions like sorted(...), but I reckon both versions of the data frame creation would become easier to understand with a more straightforward name.
• labels = [func.__name__ for func in funcs] could be moved into an else branch of if doc: to avoid iterating over the functions twice.
• I'm not entirely sure how I feel about mandatory multiprocessing with a number of workers I cannot control. I think this should be optional since it may or may not influence the timing, especially if the code that is timed already uses multiprocessing, although I have no hard facts to back that up.

plot_times

• The note on key above applies here to
• Maybe it would be a good idea to make the final plt.show() optional. IMHO that would make it easier to use this functionality in scripts without user interaction, e.g. to automatically save several plots before showing them.
• The documentation should explain that the function always uses plt.gcf()/plt.gca() as plotting target. This would make it more obvious for the user that they need to call plt.figure() manually when doing separate tests in one go in order to avoid messing up the previous plot.

"Auto review"

Just for the sake of completeness, some minor complaints from flake8:

• itertools.count is imported but not used
• there are quite a few instances of trailing whitespace in the docstrings