Mandelbrot

mandelbrot.py

# pylint: disable=no-member, no-name-in-module
"""
This program is inspired on exercise 4.4 from the book for 1CB "CNESC
Informatique" (2019-2020): The Mandelbrot set. The exercise demands the
following features:

* Display the Mandelbrot set on the complex plane. Initially between -2.2-1.2j
  and 1.0+1.2j and later by coordinates selected by the user (see below).
* Points within the set should be drawn in black, points outside of the set
  should be colored based on the number of iterations necessary to determine
  that the point was outside the set.
* The complex plane should be shown in a window of size 600x450 pixels
* The user should be able to draw a rectangle to select a new region to
  display. This zoom should draw a white rectangle while the mouse button is
  pressed down (just like in graphical programs).


NOTE: This file makes use of "type-hinting" which can be removed. This was
      added to the code to help the IDE give more helpful code-completion.
"""
import sys
from colorsys import hls_to_rgb
from typing import Generic, Optional, Tuple, TypeVar

import pygame
from pygame import MOUSEBUTTONDOWN, MOUSEBUTTONUP, MOUSEMOTION, QUIT, Color

T = TypeVar("T")

#: Values inside the mandelbrot set will fall into an endless loop. Stop after
#: this many iterations to consider the loop to be "infinite"
MAX_ITERATIONS = 255

#: If |z| >= this value, we consider the point to be *outside* of the
#: Mandelbrot set.
MAX_RADIUS = 10


def percentage_color(percentage: float) -> Color:
    """
    This function returns a color-value based on a percentage.

    The percentage value should be a float value ranging from 0 to 1.
    """

    # We use the HSL/HLS color system and map the given percentage 1:1 to the
    # "hue" value. This gives us a nice "color-wheel" gradient looking a bit
    # like a rainbow. Using 50% lightness (midpoint between black and white) is
    # a good value for colors. The same goes for a "saturation" value of 100%.
    #
    # The values need to be converted back to RGB for pygame color instances
    # which use components ranging from 0-255.
    r, g, b = hls_to_rgb(percentage, 0.5, 1.0)
    return Color(round(r * 255), round(g * 255), round(b * 255))


def mandelbrot_distance(c: complex, max_radius: int, max_iterations: int) -> int:
    """
    Calculates the "distance" of the value *c* from the Mandelbrot set. If
    this returns "0", the value *c* is part of the set. Otherwise it will
    return the number of iterations needed to determine that it was not part
    of it (the "distance").

    The distance has *max_iterations* as upper bound.

    The value for *max_iterations* has a direct impact on performance and
    maximum resolution. The higher this value, the slower the allication.
    """
    z = complex(0, 0)
    iterations = 0
    while abs(z) < max_radius and iterations < max_iterations:
        z = (z * z) + c
        iterations += 1
    if iterations == max_iterations:
        return 0
    else:
        return iterations


class Range(Generic[T]):
    """
    This is a helper class to make the source-code more readable. It
    represents a range of values between an upper and lower bound. Because
    these values should always be used in pairs (upper and lower), we wrap
    them in a class, reducing the number of variables needed in our code (and
    thus reducing mental load).
    """
    def __init__(self, lower: T, upper: T) -> None:
        self.lower = lower
        self.upper = upper

    def __repr__(self) -> str:
        return "Range(%r, %r)" % (self.lower, self.upper)


class Point:
    """
    A simple helper class to make the source code more readable.

    Similarly to the "Range" class above, this wraps two values (the X and Y
    components of a point) into one value because they usually need to be
    passed around as "one" in the code. Which again reduces the number of
    variables that need to be passed around. This again makes code more
    readable and reduces mental load.
    """
    def __init__(self, x: int, y: int) -> None:
        self.x = x
        self.y = y

    def __repr__(self) -> str:
        """
        The "__repr__" function is automatically used when using "print" on an instance of this class.
        """
        return "Point(%r, %r)" % (self.x, self.y)


class Canvas:
    """
    This class represents the drawable area (in screen pixels) for our
    Mandelbrot application.

    It abstracts the conversion between X/Y coordinates for screen pixels and
    complex numbers. The x-axis is used for the real part of the imaginary
    numbers, and the y-axis is used for the imaginary parts.

    In order to allow zooming and panning around, the canvas accepts a range
    for the displayed complex numbers where the lower bound of that range
    corresponds to the top-left corner and the upper bound corresponds to the
    lower-right corner.
    """

    def __init__(self, base_layer, complex_layer, selection_layer, range: Range[complex]) -> None:
        self.width = complex_layer.get_width()
        self.height = complex_layer.get_height()
        self.base_layer = base_layer
        self.complex_layer = complex_layer
        self.selection_layer = selection_layer
        self.set_range(range)

    def set_range(self, range: Range) -> None:
        """
        This function updates the range of complex values that can be shown
        on this canvas. It also updaes the internal "zoom" value which help
        projecting a screen pixel with X/Y components onto are of the complex
        plane visible on this canvas.
        """
        self.range = range

        # Our canvas does not represent the complex-plane on a 1:1 scale. We
        # need to determine how "big" or "small" one pixel on screen is in the
        # complex space. We can determine this using the given range of complex
        # values.
        self.zoom_x = (abs(range.lower.real) + abs(range.upper.real)) / self.width
        self.zoom_y = (abs(range.lower.imag) + abs(range.upper.imag)) / self.height

    def complex_at(self, point: Point) -> complex:
        """
        Given an X/Y coordinate on this canvas, returns the corresponding
        complex value.
        """
        output = complex(
            point.x * self.zoom_x + self.range.lower.real,
            point.y * self.zoom_y + self.range.lower.imag,
        )
        return output

    def draw_selection(self, area_select: "AreaSelect") -> None:
        """
        This function draws a transparent rectangle onto the screen. We use
        an instance of the class "AreaSelect" because it gives us the helper
        function "bounds()".
        """
        if area_select.p2 is None:
            return

        # We set the transparent color here for one main reason: Only this
        # function knows (and cares about) which color should be seen as
        # transparent (Demeter's Law).
        #
        # If we wanted to change the transparency behaviour, we only need to
        # look at the lines in this functoin and nowhere else.
        self.selection_layer.set_colorkey(Color("black"))

        # First, we reset this layer. The "fill" method makes this easy and
        # efficient because we work on a separate layer (surface) so we don't
        # need to recalculate either the mandelbrot pixels or the regions we
        # need to "update"
        self.selection_layer.fill(Color("black"))

        # The AreaSelect class gives us the helper method "bounds()" which
        # always ensures that p1 is the upper-left corner, no matter in which
        # direction the user moved the mouse.
        p1, p2 = area_select.bounds()

        # Draw the rectangle to the layer
        width = p2.x - p1.x
        height = p2.y - p1.y
        rectangle = (p1.x, p1.y, width, height)
        pygame.draw.rect(
            self.selection_layer,
            Color("white"),
            rectangle,
            1
        )

        # Now we overwrite the base layer with our existing mandelbrot pixels
        # which are on the "complex layer". This erases any existing selection
        # rectangles
        self.base_layer.blit(self.complex_layer, (0, 0))

        # And next we put the pixels of the selection layer on top. Because we
        # have a "transparency" color using the "colorkey", we don't need to do
        # anything special here.
        self.base_layer.blit(self.selection_layer, (0, 0))

    def update_complex_layer(self) -> None:
        """
        This function is the main workhorse. It recalculates the colors of
        every pixel on our complex layer for the Mandelbrot representation.

        This is slow and should only be called when really necessary.
        """
        for y in range(self.height):
            for x in range(self.width):
                c = self.complex_at(Point(x, y))
                iterations = mandelbrot_distance(c, MAX_RADIUS, MAX_ITERATIONS)
                if iterations == 0:
                    color = Color("black")
                else:
                    color = percentage_color(iterations / MAX_ITERATIONS)
                self.complex_layer.set_at((x, y), color)
            self.base_layer.blit(self.complex_layer, (0, 0))

            # Using "pygame.display.update()" inside this function violates
            # "Demeter's Law" making this function harder to test and also more
            # fragile. But for the sake of this exercise it is kept here to
            # keep the overall program code easier to understand.

            # This update is called on each line and slows down the application
            # a bit. But without it, we would not see any progress.
            pygame.display.update()
        
        # This update is needed to ensure the final image is shown
        pygame.display.update()


class AreaSelect:
    """
    Another simple helper class representing a rectangular selection with two
    points (p1 and p2).

    It provides the helper method "bounds()" which returns two points, and
    ensuring that the first of the two points is *always* the top-left point.
    """

    def __init__(self, p1: Point) -> None:
        self.p1 = p1
        self.p2 = None
    
    def bounds(self) -> Optional[Tuple[Point, Point]]:
        """
        Return the top-left and bottom-right point of this selection. It may
        return "None" if it only has one point set.
        """
        if self.p2 is None:
            return None
        x1 = min(self.p1.x, self.p2.x)
        x2 = max(self.p1.x, self.p2.x)
        y1 = min(self.p1.y, self.p2.y)
        y2 = max(self.p1.y, self.p2.y)
        return Point(x1, y1), Point(x2, y2)


def main():
    """
    The main function of the code, which deals with setting up the pygame window and handling events.
    """
    pygame.init()
    w, h = 600, 450
    screen = pygame.display.set_mode((w, h))
    mouselayer = pygame.Surface(screen.get_size())
    mandel_layer = pygame.Surface(screen.get_size())
    pygame.display.set_caption("Mandelbrot")

    canvas = Canvas(screen, mandel_layer, mouselayer, Range(complex(-2.2, -1.2), complex(1.0, 1.2)))
    canvas.update_complex_layer()

    done = False
    area_select = None
    while not done:
        for event in pygame.event.get():
            if event.type == QUIT:
                done = True
            elif event.type == MOUSEBUTTONDOWN:
                area_select = AreaSelect(Point(event.pos[0], event.pos[1]))
            elif event.type == MOUSEBUTTONUP and area_select is not None:
                area_select.p2 = Point(event.pos[0], event.pos[1])
                p1, p2 = area_select.bounds()
                canvas.set_range(Range(
                    canvas.complex_at(p1),
                    canvas.complex_at(p2)
                ))
                canvas.update_complex_layer()
                area_select = None
            elif area_select is not None and event.type == MOUSEMOTION:
                area_select.p2 = Point(event.pos[0], event.pos[1])
                canvas.draw_selection(area_select)
                pygame.display.update()

    pygame.quit()


if __name__ == "__main__":
    main()

mandelbrot_multiprocessing.py

# pylint: disable=no-member, no-name-in-module
"""
******************************************************************************
  This is modified for multiprocessing! All changes made for this change are
  marked with a "[MP]" comment so you can search for it in the code
******************************************************************************

This program is inspired on exercise 4.4 from the book for 1CB "CNESC
Informatique" (2019-2020): The Mandelbrot set. The exercise demands the
following features:

* Display the Mandelbrot set on the complex plane. Initially between -2.2-1.2j
  and 1.0+1.2j and later by coordinates selected by the user (see below).
* Points within the set should be drawn in black, points outside of the set
  should be colored based on the number of iterations necessary to determine
  that the point was outside the set.
* The complex plane should be shown in a window of size 600x450 pixels
* The user should be able to draw a rectangle to select a new region to
  display. This zoom should draw a white rectangle while the mouse button is
  pressed down (just like in graphical programs).


NOTE: This file makes use of "type-hinting" which can be removed. This was
      added to the code to help the IDE give more helpful code-completion.
"""
import sys
from colorsys import hls_to_rgb
from concurrent.futures import ProcessPoolExecutor, as_completed
from random import shuffle
from typing import Generic, Optional, Tuple, TypeVar

import pygame
from pygame import MOUSEBUTTONDOWN, MOUSEBUTTONUP, MOUSEMOTION, QUIT, Color

T = TypeVar("T")

#: Values inside the mandelbrot set will fall into an endless loop. Stop after
#: this many iterations to consider the loop to be "infinite"
MAX_ITERATIONS = 255

#: If |z| >= this value, we consider the point to be *outside* of the
#: Mandelbrot set.
MAX_RADIUS = 10


def percentage_color(percentage: float) -> Color:
    """
    This function returns a color-value based on a percentage.

    The percentage value should be a float value ranging from 0 to 1.
    """

    # We use the HSL/HLS color system and map the given percentage 1:1 to the
    # "hue" value. This gives us a nice "color-wheel" gradient looking a bit
    # like a rainbow. Using 50% lightness (midpoint between black and white) is
    # a good value for colors. The same goes for a "saturation" value of 100%.
    #
    # The values need to be converted back to RGB for pygame color instances
    # which use components ranging from 0-255.
    r, g, b = hls_to_rgb(percentage, 0.5, 1.0)
    return Color(round(r * 255), round(g * 255), round(b * 255))


def process_fragment(fragment, zoom_x, zoom_y, range):
    """
    [MP] This is the "slow" function which is spread out over multiple
    processes. It takes one "fragment" (see below) and calculates the pixel
    colors based on the Mandelbrot set. It returns a list of X/Y coordinates
    with their corresponding colors.

    As return value of this function we chose to a list of X/Y coordinates
    with the color. This makes the result independent on the shapes of our
    fragments.
    """
    output = []
    for x, y in fragment:
        c = complex_at(Point(x, y), zoom_x, zoom_y, range)
        iterations = mandelbrot_distance(c, MAX_RADIUS, MAX_ITERATIONS)
        if iterations == 0:
            color = Color("black")
        else:
            color = percentage_color(iterations / MAX_ITERATIONS)
        output.append((x, y, color))
    return output


def complex_at(point: "Point", zoom_x, zoom_y, range) -> complex:
    """
    Given an X/Y coordinate on this canvas, returns the corresponding
    complex value.

    [MP] For multiprocessing this had to be extracted from the "Canvas"
    class. See the class documentation for the reason.
    """
    output = complex(
        point.x * zoom_x + range.lower.real, point.y * zoom_y + range.lower.imag,
    )
    return output


def mandelbrot_distance(c: complex, max_radius: int, max_iterations: int) -> int:
    """
    Calculates the "distance" of the value *c* from the Mandelbrot set. If
    this returns "0", the value *c* is part of the set. Otherwise it will
    return the number of iterations needed to determine that it was not part
    of it (the "distance").

    The distance has *max_iterations* as upper bound.

    The value for *max_iterations* has a direct impact on performance and
    maximum resolution. The higher this value, the slower the allication.
    """
    z = complex(0, 0)
    iterations = 0
    while abs(z) < max_radius and iterations < max_iterations:
        z = (z * z) + c
        iterations += 1
    if iterations == max_iterations:
        return 0
    else:
        return iterations


class Range(Generic[T]):
    """
    This is a helper class to make the source-code more readable. It
    represents a range of values between an upper and lower bound. Because
    these values should always be used in pairs (upper and lower), we wrap
    them in a class, reducing the number of variables needed in our code (and
    thus reducing mental load).
    """

    def __init__(self, lower: T, upper: T) -> None:
        self.lower = lower
        self.upper = upper

    def __repr__(self) -> str:
        return "Range(%r, %r)" % (self.lower, self.upper)


class Point:
    """
    A simple helper class to make the source code more readable.

    Similarly to the "Range" class above, this wraps two values (the X and Y
    components of a point) into one value because they usually need to be
    passed around as "one" in the code. Which again reduces the number of
    variables that need to be passed around. This again makes code more
    readable and reduces mental load.
    """

    def __init__(self, x: int, y: int) -> None:
        self.x = x
        self.y = y

    def __repr__(self) -> str:
        """
        The "__repr__" function is automatically used when using "print" on an instance of this class.
        """
        return "Point(%r, %r)" % (self.x, self.y)


class Canvas:
    """
    [MP] The different pygame layers needed to be removed from this class for
         multiprocessing. The reason for this is that the main process needs
         to send data from the main process to the subprocces. This does only
         work for data which can be converted to bytes (via the Python
         "pickle" protocol). And this is not possible for pygame surfaces.

         This puts the whole class structure into question, but in order to
         keep the overall code comparable to the original code it was kept
         as-is. This is also a good demonstration of "pure vs impure" code.
         The original implementation of this class was "impure" because it
         causes side-effects by drawing pixels onto the screen. If it were
         pure, we would not have needed to make so many changes in the
         multiprocessing code (those changes would have been already applied
         to the old code instead).

    This class represents the drawable area (in screen pixels) for our
    Mandelbrot application.

    It abstracts the conversion between X/Y coordinates for screen pixels and
    complex numbers. The x-axis is used for the real part of the imaginary
    numbers, and the y-axis is used for the imaginary parts.

    In order to allow zooming and panning around, the canvas accepts a range
    for the displayed complex numbers where the lower bound of that range
    corresponds to the top-left corner and the upper bound corresponds to the
    lower-right corner.
    """

    def __init__(self, executor, width: int, height: int, range: Range[complex]) -> None:
        self.executor = executor
        self.width = width
        self.height = height
        self.set_range(range)

    def set_range(self, range: Range) -> None:
        """
        This function updates the range of complex values that can be shown
        on this canvas. It also updaes the internal "zoom" value which help
        projecting a screen pixel with X/Y components onto are of the complex
        plane visible on this canvas.
        """
        self.range = range

        # Our canvas does not represent the complex-plane on a 1:1 scale. We
        # need to determine how "big" or "small" one pixel on screen is in the
        # complex space. We can determine this using the given range of complex
        # values.
        self.zoom_x = (abs(range.lower.real) + abs(range.upper.real)) / self.width
        self.zoom_y = (abs(range.lower.imag) + abs(range.upper.imag)) / self.height

    def draw_selection(self, base_layer, complex_layer, selection_layer, area_select: "AreaSelect") -> None:
        """
        [MP] Because the layers needed to be removed for multiprocessing,
             they now need to be passed as arguments.

        This function draws a transparent rectangle onto the screen. We use
        an instance of the class "AreaSelect" because it gives us the helper
        function "bounds()".
        """
        if area_select.p2 is None:
            return

        # We set the transparent color here for one main reason: Only this
        # function knows (and cares about) which color should be seen as
        # transparent (Demeter's Law).
        #
        # If we wanted to change the transparency behaviour, we only need to
        # look at the lines in this functoin and nowhere else.
        selection_layer.set_colorkey(Color("black"))

        # First, we reset this layer. The "fill" method makes this easy and
        # efficient because we work on a separate layer (surface) so we don't
        # need to recalculate either the mandelbrot pixels or the regions we
        # need to "update"
        selection_layer.fill(Color("black"))

        # The AreaSelect class gives us the helper method "bounds()" which
        # always ensures that p1 is the upper-left corner, no matter in which
        # direction the user moved the mouse.
        p1, p2 = area_select.bounds()

        # Draw the rectangle to the layer
        width = p2.x - p1.x
        height = p2.y - p1.y
        rectangle = (p1.x, p1.y, width, height)
        pygame.draw.rect(selection_layer, Color("white"), rectangle, 1)

        # Now we overwrite the base layer with our existing mandelbrot pixels
        # which are on the "complex layer". This erases any existing selection
        # rectangles
        base_layer.blit(complex_layer, (0, 0))

        # And next we put the pixels of the selection layer on top. Because we
        # have a "transparency" color using the "colorkey", we don't need to do
        # anything special here.
        base_layer.blit(selection_layer, (0, 0))

    def explode(self):
        """
        [MP] A helper function for multiprocessing.

        For multiprocessing we must decide how to spread the work over the
        available processes. So the input data needs to be "split" into
        multiples parts (we "explode" it into "fragments").

        For demonstration, multiple implementations are available.

        As expected return value of these functions we simply chose a list of
        X/Y coordinates.
        """
        implementation = self.explode_as_lines
        return implementation()

    def explode_as_columns(self):
        """
        [MP] One of the implementations for the multiprocessing "split".

        In this implementation each "fragment" is a single-pixel column of the data.
        """
        fragments = []
        for x in range(self.width):
            column = []
            for y in range(self.height):
                column.append((x, y))
            fragments.append(column)
        return fragments

    def explode_as_lines(self):
        """
        [MP] One of the implementations for the multiprocessing "split".

        In this implementation each "fragment" is a single-pixel row of the data.
        """
        fragments = []
        for y in range(self.height):
            row = []
            for x in range(self.width):
                row.append((x, y))
            fragments.append(row)
        return fragments

    def update_complex_layer(self, base_layer, complex_layer) -> None:
        """
        [MP] Because the layers needed to be removed for multiprocessing,
             they now need to be passed as arguments.

        This function will create sub-processes (the "workers"), split the
        data into "fragments" and submit them to the workers.

        This is slow and should only be called when really necessary.
        """
        complex_layer.fill(Color("black"))
        base_layer.fill(Color("black"))
        pygame.display.update()
        fragments = self.explode()

        # [MP] We keep a list of all our jobs
        jobs = []
        for fragment in fragments:
            job = self.executor.submit(process_fragment, fragment, self.zoom_x, self.zoom_y, self.range)
            jobs.append(job)

        # [MP] Once they are submitted, we will wait for them to complete and
        #      will process them in order of completion. Note that the order here is
        #      not guaranteed.
        for job in as_completed(jobs):
            result = job.result()
            # The result is the return value of "process_fragment". In our case
            # this is a simple list of X/Y coordinates with their calculated
            # color.
            for x, y, color in result:
                complex_layer.set_at((x, y), color)
            base_layer.blit(complex_layer, (0, 0))
            # We can update the screen for each completed job
            pygame.display.update()
        pygame.display.update()


class AreaSelect:
    """
    Another simple helper class representing a rectangular selection with two
    points (p1 and p2).

    It provides the helper method "bounds()" which returns two points, and
    ensuring that the first of the two points is *always* the top-left point.
    """

    def __init__(self, p1: Point) -> None:
        self.p1 = p1
        self.p2 = None

    def bounds(self) -> Optional[Tuple[Point, Point]]:
        """
        Return the top-left and bottom-right point of this selection. It may
        return "None" if it only has one point set.
        """
        if self.p2 is None:
            return None
        x1 = min(self.p1.x, self.p2.x)
        x2 = max(self.p1.x, self.p2.x)
        y1 = min(self.p1.y, self.p2.y)
        y2 = max(self.p1.y, self.p2.y)
        return Point(x1, y1), Point(x2, y2)


def main(executor):
    """
    The main function of the code, which deals with setting up the pygame window and handling events.
    """
    pygame.init()
    w, h = 600, 450
    screen = pygame.display.set_mode((w, h))
    mouselayer = pygame.Surface(screen.get_size())
    mandel_layer = pygame.Surface(screen.get_size())
    pygame.display.set_caption("Mandelbrot")

    # [MP] The layers have been removed from this class for multiprocessing and
    #      will be passed below as arguments to the methods
    canvas = Canvas(executor, w, h, Range(complex(-2.2, -1.2), complex(1.0, 1.2)))
    canvas.update_complex_layer(screen, mandel_layer)

    done = False
    area_select = None
    while not done:
        for event in pygame.event.get():
            if event.type == QUIT:
                done = True
            elif event.type == MOUSEBUTTONDOWN:
                area_select = AreaSelect(Point(event.pos[0], event.pos[1]))
            elif event.type == MOUSEBUTTONUP and area_select is not None:
                area_select.p2 = Point(event.pos[0], event.pos[1])
                p1, p2 = area_select.bounds()
                canvas.set_range(Range(
                    complex_at(p1, canvas.zoom_x, canvas.zoom_y, canvas.range),
                    complex_at(p2, canvas.zoom_x, canvas.zoom_y, canvas.range),
                ))
                canvas.update_complex_layer(screen, mandel_layer)
                area_select = None
            elif area_select is not None and event.type == MOUSEMOTION:
                area_select.p2 = Point(event.pos[0], event.pos[1])
                canvas.draw_selection(screen, mandel_layer, mouselayer, area_select)
                pygame.display.update()

    pygame.quit()


if __name__ == "__main__":
    # [MP] The ProcessPoolExecutor provides us with an easy way to run jobs in
    #      parallel
    with ProcessPoolExecutor() as executor:
        main(executor)