Merge pull request #339 from realpython/python-split-list

Materials for How to Split a Python List

Author: KateFinegan

python-split-list/README.md

@@ -0,0 +1,3 @@
+# How to Split a Python List or Iterable Into Chunks
+
+This folder holds sample code that supplements the Real Python tutorial on [How to Split a Python List or Iterable Into Chunks](https://realpython.com/how-to-split-a-python-list-into-chunks/).

python-split-list/parallel_demo.py

@@ -0,0 +1,126 @@
+"""
+Synthesize an image in chunks using parallel workers.
+
+Usage:
+$ python parallel_demo.py
+"""
+
+import functools
+import multiprocessing
+import time
+from dataclasses import dataclass
+from math import log
+from os import cpu_count
+from typing import Callable, Iterable, Iterator
+
+import numpy as np
+from PIL import Image
+
+from spatial_splitting import Bounds, split_multi
+
+IMAGE_WIDTH, IMAGE_HEIGHT = 1920, 1080
+CENTER = -0.7435 + 0.1314j
+SCALE = 0.0000015
+MAX_ITERATIONS = 256
+ESCAPE_RADIUS = 1000
+NUM_CHUNKS = 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: tuple[int, int]) -> int:
+        return self.pixels[self.bounds.offset(*coordinates)]
+
+    def __setitem__(self, coordinates: tuple[int, int], value: int) -> None:
+        self.pixels[self.bounds.offset(*coordinates)] = value
+
+
+@dataclass
+class MandelbrotSet:
+    max_iterations: int
+    escape_radius: float = 2.0
+
+    def __contains__(self, c):
+        return self.stability(c) == 1
+
+    def stability(self, c, smooth=False, clamp=True):
+        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, smooth=False):
+        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 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 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 combine(chunks: Iterable[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")
+
+
+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 process_sequentially(bounds_iter: Iterator[Bounds]) -> Iterator[Chunk]:
+    return map(generate_chunk, bounds_iter)
+
+
+def process_in_parallel(bounds_iter: Iterator[Bounds]) -> list[Chunk]:
+    with multiprocessing.Pool() as pool:
+        return pool.map(generate_chunk, bounds_iter)
+
+
+@timed
+def compute(worker: Callable) -> Image.Image:
+    return combine(worker(split_multi(NUM_CHUNKS, IMAGE_HEIGHT, IMAGE_WIDTH)))
+
+
+def main() -> None:
+    for worker in (process_sequentially, process_in_parallel):
+        compute(worker).show()
+
+
+if __name__ == "__main__":
+    main()

python-split-list/requirements.txt

@@ -0,0 +1,3 @@
+Pillow==9.4.0
+more-itertools==9.0.0
+numpy==1.24.1

python-split-list/spatial_splitting.py

@@ -0,0 +1,167 @@
+"""
+Split in multiple dimensions.
+"""
+
+from dataclasses import dataclass
+from functools import cached_property
+from itertools import combinations_with_replacement, product, starmap
+from math import floor, prod, sqrt
+from typing import Iterator
+
+
+@dataclass
+class Bounds:
+    """Spatial bounds defined in multiple dimensions."""
+
+    start_point: tuple[int, ...]
+    end_point: tuple[int, ...]
+
+    def __post_init__(self) -> None:
+        """Validate the bounds."""
+        assert len(self.start_point) == len(self.end_point)
+        assert self.start_point < self.end_point
+
+    def __len__(self):
+        """Get the number of points in the bounds (length, area, or volume).
+
+        Example:
+        >>> bounds = Bounds((50, 25), (150, 100))
+        >>> len(bounds)
+        7500
+        """
+        return prod(
+            map(lambda x: x[1] - x[0], zip(self.start_point, self.end_point))
+        )
+
+    def __iter__(self) -> Iterator[tuple[int, ...]]:
+        """Iterate over the bounds, yielding each point like in an odometer.
+
+        Example:
+        >>> bounds = Bounds((50, 25), (150, 100))
+        >>> for x, y in bounds:
+        ...     print(x, y)  # doctest: +ELLIPSIS
+        50 25
+        50 26
+        50 27
+        ...
+        149 97
+        149 98
+        149 99
+        """
+        return product(*starmap(range, zip(self.start_point, self.end_point)))
+
+    def slices(self) -> tuple[slice, ...]:
+        """Return the slice for each dimension.
+
+        Example:
+        >>> bounds = Bounds((50, 25), (150, 100))
+        >>> bounds.slices()
+        (slice(50, 150, None), slice(25, 100, None))
+        """
+        return tuple(
+            slice(start, end)
+            for start, end in zip(self.start_point, self.end_point)
+        )
+
+    @cached_property
+    def size(self) -> tuple[int, ...]:
+        """Return the size of the bounds in each dimension.
+
+        Example:
+        >>> bounds = Bounds((50, 25), (150, 100))
+        >>> width, height = bounds.size
+        >>> width
+        100
+        >>> height
+        75
+        """
+        return tuple(
+            self.end_point[i] - self.start_point[i]
+            for i in range(self.num_dimensions)
+        )
+
+    @cached_property
+    def num_dimensions(self) -> int:
+        """Return the number of dimensions.
+
+        Example:
+        >>> bounds = Bounds((50, 25), (150, 100))
+        >>> bounds.num_dimensions
+        2
+        """
+        return len(self.start_point)
+
+    def offset(self, *coordinates):
+        """Return the offset of the given coordinates from the start point.
+
+        Example:
+        >>> bounds = Bounds((50, 25), (150, 100))
+        >>> for x, y in bounds:
+        ...     print(bounds.offset(x, y))  # doctest: +ELLIPSIS
+        (0, 0)
+        (0, 1)
+        (0, 2)
+        ...
+        (99, 72)
+        (99, 73)
+        (99, 74)
+        """
+        return tuple(
+            coordinates[i] - self.start_point[i]
+            for i in range(self.num_dimensions)
+        )
+
+
+def split_multi(num_chunks: int, *dimensions: int) -> Iterator[Bounds]:
+    """Return a sequence of n-dimensional slices."""
+    num_chunks_along_axis = find_most_even(num_chunks, len(dimensions))
+    for slices_by_dimension in product(
+        *starmap(get_slices, zip(dimensions, num_chunks_along_axis))
+    ):
+        yield Bounds(
+            start_point=tuple(s.start for s in slices_by_dimension),
+            end_point=tuple(s.stop for s in slices_by_dimension),
+        )
+
+
+def get_slices(length: int, num_chunks: int) -> Iterator[slice]:
+    """Return a sequence of slices for the given length."""
+    chunk_size, remaining = divmod(length, 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 slice(begin, end)
+
+
+def find_most_even(number: int, num_factors: int):
+    """Return the most even tuple of integer divisors of a number."""
+    products_by_sum = {
+        sum(products): products
+        for products in find_products(number, num_factors)
+    }
+    return products_by_sum[min(products_by_sum)]
+
+
+def find_products(number: int, num_factors: int) -> Iterator[tuple[int, ...]]:
+    """Return all possible products of a number."""
+    divisors = find_divisors(number)
+    for factors in combinations_with_replacement(divisors, num_factors):
+        if prod(factors) == number:
+            yield factors
+
+
+def find_divisors(number: int) -> set[int]:
+    """Return unique integer divisors of a number."""
+    divisors = {1, number}
+    for divisor in range(2, floor(sqrt(number)) + 1):
+        factor, remainder = divmod(number, divisor)
+        if remainder == 0:
+            divisors.add(divisor)
+            divisors.add(factor)
+    return divisors
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()