How to Split a Python List - Initial commit (materials)

Author: bzaczynski

python-split-list/README.md

@@ -0,0 +1,3 @@
+# How to Split a Python List
+
+This folder holds the code for the Real Python How to Split a Python List tutorial.

python-split-list/parallelization.py

@@ -0,0 +1,132 @@
+"""
+Plot the Mandelbrot set using parallel workers.
+
+Usage:
+$ python parallelization.py
+"""
+
+import functools
+import multiprocessing
+import os
+import time
+from dataclasses import dataclass
+from math import log
+from typing import Callable, Iterable
+
+import numpy as np
+from PIL import Image
+
+from splitting.multidimensional import Bounds, split_multi
+
+IMAGE_WIDTH, IMAGE_HEIGHT = 1280, 720
+CENTER = -0.7435 + 0.1314j
+SCALE = 0.0000025
+MAX_ITERATIONS = 256
+ESCAPE_RADIUS = 1000
+NUM_CHUNKS = os.cpu_count() or 4
+
+
+class Chunk:
+    """A chunk of the image to be computed and rendered."""
+
+    def __init__(self, bounds: Bounds) -> None:
+        self.bounds = bounds
+        self.height = bounds.size[0]
+        self.width = bounds.size[1]
+        self.pixels = np.zeros((self.height, self.width), dtype=np.uint8)
+
+    def __getitem__(self, coordinates) -> int:
+        """Get the value of a pixel at the given absolute coordinates."""
+        return self.pixels[self.bounds.offset(*coordinates)]
+
+    def __setitem__(self, coordinates, value: int) -> None:
+        """Set the value of a pixel at the given absolute coordinates."""
+        self.pixels[self.bounds.offset(*coordinates)] = value
+
+
+@dataclass
+class MandelbrotSet:
+    max_iterations: int
+    escape_radius: float = 2.0
+
+    def __contains__(self, c: complex) -> bool:
+        return self.stability(c) == 1
+
+    def stability(self, c: complex, smooth=False, clamp=True) -> float:
+        value = self.escape_count(c, smooth) / self.max_iterations
+        return max(0.0, min(value, 1.0)) if clamp else value
+
+    def escape_count(self, c: complex, smooth=False) -> int | float:
+        z = 0 + 0j
+        for iteration in range(self.max_iterations):
+            z = z**2 + c
+            if abs(z) > self.escape_radius:
+                if smooth:
+                    return iteration + 1 - log(log(abs(z))) / log(2)
+                return iteration
+        return self.max_iterations
+
+
+def timed(function: Callable) -> Callable:
+    @functools.wraps(function)
+    def wrapper(*args, **kwargs):
+        start = time.perf_counter()
+        result = function(*args, **kwargs)
+        end = time.perf_counter()
+        print(f"{function.__name__}() took {end - start:.2f} seconds")
+        return result
+
+    return wrapper
+
+
+def main() -> None:
+    for worker in (process_chunks_parallel, process_chunks_sequential):
+        image = compute(worker)
+        image.show()
+
+
+@timed
+def process_chunks_sequential(chunked_bounds: Iterable[Bounds]) -> list[Chunk]:
+    return list(map(generate_chunk, chunked_bounds))
+
+
+@timed
+def process_chunks_parallel(chunked_bounds: Iterable[Bounds]) -> list[Chunk]:
+    with multiprocessing.Pool() as pool:
+        return pool.map(generate_chunk, chunked_bounds)
+
+
+def generate_chunk(bounds: Bounds) -> Chunk:
+    """Generate a chunk of pixels for the given bounds."""
+    chunk = Chunk(bounds)
+    mandelbrot_set = MandelbrotSet(MAX_ITERATIONS, ESCAPE_RADIUS)
+    for y, x in bounds:
+        c = transform(y, x)
+        instability = 1 - mandelbrot_set.stability(c, smooth=True)
+        chunk[y, x] = int(instability * 255)
+    return chunk
+
+
+def transform(y: int, x: int) -> complex:
+    """Transform the given pixel coordinates to the complex plane."""
+    im = SCALE * (IMAGE_HEIGHT / 2 - y)
+    re = SCALE * (x - IMAGE_WIDTH / 2)
+    return complex(re, im) + CENTER
+
+
+def compute(worker: Callable) -> Image.Image:
+    """Render the image using the given worker function."""
+    chunks = worker(split_multi(NUM_CHUNKS, IMAGE_HEIGHT, IMAGE_WIDTH))
+    return combine(chunks)
+
+
+def combine(chunks: list[Chunk]) -> Image.Image:
+    """Combine the chunks into a single image."""
+    pixels = np.zeros((IMAGE_HEIGHT, IMAGE_WIDTH), dtype=np.uint8)
+    for chunk in chunks:
+        pixels[chunk.bounds.slices()] = chunk.pixels
+    return Image.fromarray(pixels, mode="L")
+
+
+if __name__ == "__main__":
+    main()

python-split-list/requirements.txt

@@ -0,0 +1,3 @@
+Pillow==9.3.0
+more-itertools==9.0.0
+numpy==1.23.5

python-split-list/splitting/__init__.py

@@ -0,0 +1,31 @@
+"""
+Split a Python List Into a Fixed Number of Chunks With Roughly Equal Sizes
+"""
+
+from itertools import zip_longest
+from typing import Any, Iterable, Iterator, Sequence
+
+import numpy as np
+
+
+def split_with_numpy(array: np.array, num_chunks: int) -> Iterable[np.array]:
+    return np.array_split(array, num_chunks)
+
+
+def split_transposed(
+    sequence: Sequence[Any], num_chunks: int
+) -> Iterator[tuple[Any]]:
+    for i in range(num_chunks):
+        yield sequence[i::num_chunks]
+
+
+def transpose(chunks, padding: Any = None) -> Iterator[tuple[Any]]:
+    return zip_longest(*chunks, value=padding)
+
+
+def split(sequence: Sequence[Any], num_chunks: int) -> Iterator[Sequence[Any]]:
+    chunk_size, remaining = divmod(len(sequence), num_chunks)
+    for i in range(num_chunks):
+        begin = i * chunk_size + min(i, remaining)
+        end = (i + 1) * chunk_size + min(i + 1, remaining)
+        yield sequence[begin:end]

python-split-list/splitting/fixednumber.py

@@ -0,0 +1,73 @@
+"""
+Split a Python List Into Fixed-Sized Chunks
+"""
+
+import math
+import sys
+from itertools import islice, zip_longest
+from typing import Any, Iterable, Iterator, Sequence
+
+import numpy as np
+
+
+def batched(
+    iterable: Iterable[Any], max_chunk_size: int
+) -> Iterator[list[Any]]:
+    if sys.version_info >= (3, 12):
+        from itertools import batched
+
+        return batched(iterable, max_chunk_size)
+    try:
+        from more_itertools import batched
+
+        return batched(iterable, max_chunk_size)
+    except ImportError:
+        iterator = iter(iterable)
+        while chunk := list(islice(iterator, max_chunk_size)):
+            yield chunk
+
+
+def batched_functional(
+    iterable: Iterable[Any], max_chunk_size: int
+) -> Iterator[list[Any]]:
+    iterator = iter(iterable)
+    return iter(lambda: list(islice(iterator, max_chunk_size)), [])
+
+
+def split_with_padding(
+    sequence: Iterable[Any], chunk_size: int, padding: Any = None
+) -> Iterable[tuple[Any]]:
+    return zip_longest(*[iter(sequence)] * chunk_size, fillvalue=padding)
+
+
+def split_drop_last(
+    sequence: Sequence[Any], chunk_size: int
+) -> Iterator[Sequence[tuple[Any]]]:
+    return zip(*[sequence[i::chunk_size] for i in range(chunk_size)])
+
+
+def split_strings(
+    sequence: Sequence[Any], max_chunk_size: int
+) -> Iterator[Sequence[Any]]:
+    for i in range(0, len(sequence), max_chunk_size):
+        yield sequence[i : i + max_chunk_size]
+
+
+def split_into_slices(
+    sequence: Sequence[Any], max_chunk_size: int
+) -> Iterator[slice]:
+    for i in range(0, len(sequence), max_chunk_size):
+        yield slice(i, i + max_chunk_size)
+
+
+def split_list_with_padding(
+    sequence: Sequence[Any], chunk_size: int, fill_with: Any
+) -> Iterable[Sequence[Any]]:
+    for i in range(0, len(sequence), chunk_size):
+        chunk = list(sequence[i : i + chunk_size])
+        padding = [fill_with] * (chunk_size - len(chunk))
+        yield chunk + padding
+
+
+def split_with_numpy(array: np.array, chunk_size: int) -> Iterable[np.array]:
+    return np.array_split(array, math.ceil(len(array) / chunk_size))