Skip to content

Instantly share code, notes, and snippets.

@exhuma
Last active Feb 2, 2020
Embed
What would you like to do?
Mandelbrot
# 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()
# 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment