Examples

Every example below is executed as part of building the docs (via sphinx-exec-code).

SerialPool

"""Minimal working example of SerialPool."""

import time

from richpool import SerialPool

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 20

    start = time.perf_counter()
    with SerialPool() as pool:
        results = pool.map(square_slowly, range(n), desc="squaring")
    elapsed = time.perf_counter() - start

    print(results)
    print(f"serial: {elapsed:.2f}s (no speedup, SerialPool runs one item at a time)")

  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 20/20 0:00:20 0:00:00
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361]
serial: 20.01s (no speedup, SerialPool runs one item at a time)

MultiPool

"""Minimal working example of MultiPool."""

import time

from richpool import MultiPool

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 20
    serial_estimate = n * SLEEP_TIME

    start = time.perf_counter()
    with MultiPool(processes=4) as pool:
        results = pool.map(square_slowly, range(n), desc="squaring")
    elapsed = time.perf_counter() - start

    print(results)
    print(f"parallel: {elapsed:.2f}s (serial would take ~{serial_estimate:.2f}s)")

  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 20/20 0:00:05 0:00:00
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361]
parallel: 5.02s (serial would take ~20.00s)

MultiPool with emcee

Adapted from schwimmbad’s Using MultiPool with emcee example:

emcee’s EnsembleSampler accepts any pool object exposing a .map() method, so richpool’s MultiPool can be passed directly to parallelize likelihood evaluations across worker processes. log_probability sleeps briefly to stand in for a realistic, non-trivial likelihood (e.g. one involving a simulation or numerical integration); without that, each evaluation is a handful of floating-point ops that finishes in microseconds.

emcee calls pool.map() several times per MCMC step internally, so letting MultiPool render its usual per-call progress bar produces a handful of tiny, near-instant, low-information bars, not real progress bars. Passing disable=True to MultiPool’s constructor turns off those per-call bars (the default for every map() call unless a call overrides it), and a single rich progress bar below tracks real progress across the n_steps MCMC steps instead, while the likelihood evaluations for each step still run in parallel underneath it.

Needs the extra emcee, numpy, and matplotlib dependencies:

uv pip install emcee numpy matplotlib

"""Using MultiPool with emcee for parallel MCMC sampling.

Adapted from schwimmbad's "Using MultiPool with emcee" example:
https://schwimmbad.readthedocs.io/en/latest/examples/index.html#using-multipool-with-emcee

emcee's EnsembleSampler accepts any pool object exposing a .map() method, so
richpool's MultiPool can be passed directly to parallelize likelihood
evaluations across worker processes. log_probability sleeps briefly to stand
in for a realistic, non-trivial likelihood (e.g. one involving a simulation
or numerical integration); without that, each evaluation is a handful of
floating-point ops that finishes in microseconds.

emcee calls pool.map() several times per MCMC step internally (its default
move splits the ensemble in half and updates each half in turn), so letting
MultiPool render its usual per-call progress bar produces a handful of tiny,
near-instant, low-information bars, not real progress bars. Passing disable=True to
MultiPool's constructor turns off those per-call bars (the default for every
map() call unless a call overrides it), and a single rich progress bar below
tracks real progress across the n_steps MCMC steps instead, while the
likelihood evaluations for each step still run in parallel underneath it.

Needs the extra emcee, numpy, and matplotlib dependencies:
pip install emcee numpy matplotlib.
"""

import time

import matplotlib.pyplot as plt  # ty: ignore[unresolved-import]
import numpy as np  # ty: ignore[unresolved-import]
from emcee import EnsembleSampler  # ty: ignore[unresolved-import]
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)

from richpool import MultiPool

LIKELIHOOD_SLEEP = 0.01


def log_probability(p, mean, var):
    time.sleep(LIKELIHOOD_SLEEP)
    x = p[0]
    return -0.5 * (x - mean) ** 2 / var


def main():
    n_walkers = 20
    n_dim = 1
    n_steps = 50
    processes = 4
    p0 = np.random.uniform(0, 10, (n_walkers, n_dim))
    serial_estimate = n_steps * n_walkers * LIKELIHOOD_SLEEP

    start = time.perf_counter()
    with MultiPool(processes, disable=True) as pool:
        sampler = EnsembleSampler(n_walkers, n_dim, log_probability, args=(5, 1.2), pool=pool)

        columns = (
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            MofNCompleteColumn(),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
        )
        with Progress(*columns) as progress:
            task_id = progress.add_task("sampling", total=n_steps)
            for _ in sampler.sample(p0, iterations=n_steps):
                progress.advance(task_id)
    elapsed = time.perf_counter() - start

    print(f"mean of flatchain: {sampler.flatchain.mean():.3f}")
    print(f"parallel: {elapsed:.2f}s (serial would take ~{serial_estimate:.2f}s)")

    plt.figure()
    plt.hist(sampler.flatchain, bins=30)
    plt.xlabel("x")
    plt.ylabel("count")
    plt.title("emcee samples drawn via richpool's MultiPool")


if __name__ == "__main__":
    main()

  sampling ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 50/50 0:00:03 0:00:00
mean of flatchain: 5.031
parallel: 3.38s (serial would take ~10.00s)

JoblibPool

"""Minimal working example of JoblibPool."""

import time

from richpool import JoblibPool

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 20
    serial_estimate = n * SLEEP_TIME

    start = time.perf_counter()
    with JoblibPool(processes=4) as pool:
        results = pool.map(square_slowly, range(n), desc="squaring")
    elapsed = time.perf_counter() - start

    print(results)
    print(f"parallel: {elapsed:.2f}s (serial would take ~{serial_estimate:.2f}s)")

  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 20/20 0:00:05 0:00:00
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361]
parallel: 5.43s (serial would take ~20.00s)

MPIPool

MPIPool needs to be launched with mpiexec, so it can’t run inline like the examples above. Run it with:

mpiexec -n 4 python examples/mpi_pool_example.py

"""Minimal working example of MPIPool.

Run with: mpiexec -n 4 python mpi_example.py
Only the master rank (rank 0) prints output and shows the progress bar.
"""

import time

from richpool import MPIPool

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 20
    serial_estimate = n * SLEEP_TIME

    with MPIPool() as pool:
        if pool.is_master():
            start = time.perf_counter()
            results = pool.map(square_slowly, range(n), desc="squaring")
            elapsed = time.perf_counter() - start

            print(results)
            print(f"parallel: {elapsed:.2f}s (serial would take ~{serial_estimate:.2f}s)")

p_map

"""Minimal working example of richpool's p_map: a parallel ordered map."""

import time

from richpool import p_map

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 20
    serial_estimate = n * SLEEP_TIME

    start = time.perf_counter()
    results = p_map(square_slowly, range(n), desc="squaring")
    elapsed = time.perf_counter() - start

    print(results)
    print(f"parallel: {elapsed:.2f}s (serial would take ~{serial_estimate:.2f}s)")

  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 20/20 0:00:10 0:00:00
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361]
parallel: 10.02s (serial would take ~20.00s)

p_imap

"""Minimal working example of richpool's p_imap: a parallel ordered iterator.

Unlike p_map, p_imap doesn't wait for the whole batch, it yields each result
as soon as it's ready, in the same order as the input.
"""

import time

from richpool import p_imap

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 10
    serial_estimate = n * SLEEP_TIME

    start = time.perf_counter()
    for result in p_imap(square_slowly, range(n), desc="squaring"):
        print("got:", result)
    elapsed = time.perf_counter() - start

    print(f"parallel: {elapsed:.2f}s (serial would take ~{serial_estimate:.2f}s)")

got: 0
got: 1
got: 4
got: 9
got: 16
got: 25
got: 36
got: 49
got: 64
got: 81
  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 10/10 0:00:05 0:00:00
parallel: 5.02s (serial would take ~10.00s)

p_umap

"""Minimal working example of richpool's p_umap: a parallel unordered map.

Unlike p_map, results come back in completion order, not input order. To make
that visible, earlier items are made to sleep longer than later ones, so later
items finish first.
"""

import time

from richpool import p_umap


def square_reverse_delay(x):
    time.sleep((9 - x) * 0.3)
    return x * x


if __name__ == "__main__":
    results = p_umap(square_reverse_delay, range(10), desc="squaring")
    print("input order: ", list(range(10)))
    print("result order:", results)

  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 10/10 0:00:06 0:00:00
input order:  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
result order: [1, 0, 9, 4, 25, 16, 49, 36, 81, 64]

p_uimap

"""Minimal working example of richpool's p_uimap: a parallel unordered iterator.

Combines p_imap's lazy, as-ready yielding with p_umap's unordered delivery.
Earlier items sleep longer than later ones, so results arrive out of order.
"""

import time

from richpool import p_uimap


def square_reverse_delay(x):
    time.sleep((9 - x) * 0.3)
    return x * x


if __name__ == "__main__":
    for result in p_uimap(square_reverse_delay, range(10), desc="squaring"):
        print("got:", result)

got: 1
got: 0
got: 9
got: 4
got: 25
got: 16
got: 49
got: 36
got: 81
got: 64
  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 10/10 0:00:06 0:00:00

t_map

"""Minimal working example of richpool's t_map: a sequential map with a progress bar.

No parallelism, t_map just wraps the builtin map() with a rich progress bar.
"""

import time

from richpool import t_map

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    n = 10

    start = time.perf_counter()
    results = t_map(square_slowly, range(n), desc="squaring")
    elapsed = time.perf_counter() - start

    print(results)
    print(f"sequential: {elapsed:.2f}s (no speedup, t_map runs one item at a time)")

  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 10/10 0:00:10 0:00:00
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
sequential: 10.01s (no speedup, t_map runs one item at a time)

t_imap

"""Minimal working example of richpool's t_imap: a sequential iterator.

Same as t_map, but yields each result immediately instead of waiting to
return a full list.
"""

import time

from richpool import t_imap

SLEEP_TIME = 1


def square_slowly(x):
    time.sleep(SLEEP_TIME)
    return x * x


if __name__ == "__main__":
    for result in t_imap(square_slowly, range(5), desc="squaring"):
        print("got:", result)

got: 0
got: 1
got: 4
got: 9
got: 16
  squaring ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 5/5 0:00:05 0:00:00