Koch snowflake in Python with numpy and pygame

Question

I've drawn a Koch snowflake with Python 3. I'm new to the Python world and would appreciate all feedback on how to make the code more Pythonic as I'm not used to these idioms, styles and modules. All other feedback is of course also greatly appreciated.

The snowflake is constantly redrawn in preparation for the future. This main loop is based on the pygame.draw example but feels a bit clunky with done=False;while not done etc.

enter image description here

import pygame
import numpy
from math import pi, sin, cos
# some constants
FPS = 10
WINDOW_SIZE = [400,400]
LINE_WIDTH = 1
MIN_LINE_LENGTH = 1 # when to terminate recursion
# main loop
def main():
 # init pygame and window
 pygame.init()
 screen = pygame.display.set_mode(WINDOW_SIZE)
 # loop until user exits
 done = False
 clock = pygame.time.Clock()
 while not done:
 # limit frame rate
 clock.tick(FPS)
 # check if user exited
 for event in pygame.event.get():
 if event.type == pygame.QUIT:
 done=True
 # clear screen
 screen.fill([255,255,255])
 # calculate the three corners of the snowflake,
 # an equilateral triangle centered on the screen
 screen_center = numpy.array([ WINDOW_SIZE[0]/2, WINDOW_SIZE[1]/2 ])
 radius = WINDOW_SIZE[1]/2 # not really radius? but distance from center to corner
 snowflake_top = screen_center + vector_from_polar( radius, pi/2 )
 snowflake_left = screen_center + vector_from_polar( radius, pi/2 + 2*pi/3 )
 snowflake_right = screen_center + vector_from_polar( radius, pi/2 + 4*pi/3 )
 # draw the snowflake
 draw_koch_line( screen, snowflake_top, snowflake_left )
 draw_koch_line( screen, snowflake_left, snowflake_right )
 draw_koch_line( screen, snowflake_right, snowflake_top )
 # flip buffers
 pygame.display.flip()
# vector_from_polar: constructs a vector from its angle and mangitude
def vector_from_polar( magnitude, angle ):
 return magnitude * numpy.array( [ numpy.cos( angle ), numpy.sin( angle ) ] )
# draw_koch_line: (approximately) draws a Koch line between the points specified
def draw_koch_line(screen, line_start, line_end ):
 if numpy.linalg.norm( line_end - line_start ) / 3 < MIN_LINE_LENGTH:
 # last iteration: draw the line
 pygame.draw.line(screen, [0,0,0], line_start, line_end, LINE_WIDTH )
 else:
 # find the normal to this line
 line_normal = numpy.array([
 line_end[1]-line_start[1],
 line_start[0]-line_end[0] ])
 # find the three points of the "triangle" used to define the segments
 triangle_left = 2/3*line_start + 1/3*line_end # 1/3 of this line
 triangle_right = 1/3*line_start + 2/3*line_end # 2/3 of this line
 triangle_top = 1/2*line_start + 1/2*line_end + \
 numpy.sqrt(3)/2/3 * line_normal # point "above" the line
 # recurse for each segments
 draw_koch_line(screen, line_start, triangle_left)
 draw_koch_line(screen, triangle_left, triangle_top)
 draw_koch_line(screen, triangle_top, triangle_right)
 draw_koch_line(screen, triangle_right, line_end)
main()

• hpaulj: math.sqrt etc is faster than nump.sqrt when working with scalars. Small lists and tuples are also faster than arrays. Arrays are good when large or multidimensional, but have an overhead that isn't needed when small.
• hpaulj: Replace done with not_done or working. :)
• Anna: @hpaulj Thank you, I didn't know about the overhead. I used numpy to make the vector addition look neater. How do you suggest I do it with lists instead? Simply a[0]=b[0]+c[0];a[1]=b[1]+c[1]? Also, thanks for letting me know math.sqrt is faster.

Answer 1

This is good code — I only have a couple of minor points.

1. Since nothing changes or animates, there's no need to have a top-level loop or a clock. You can just draw the snowflake and then wait for the user to quit (calling event.wait instead of event.get so that the program isn't busy doing nothing).

2. It's clearer to write pygame.Color('white') and pygame.Color('black') instead of [255, 255, 255] and [0, 0, 0].

3. Instead of:

   # find the normal to this line
   line_normal = numpy.array([
    line_end[1]-line_start[1],
    line_start[0]-line_end[0] ])
   

   consider using the cross product:

   line_normal = numpy.cross(line_end - line_start, [0, 0, 1])
   

4. There's repeated code in the computation of the triangle points that could be factored out into a function:

   def interpolated(x, start, end):
    """Return vector interpolated between start and end by x."""
    return (1 - x) * start + x * end
   

   and then:

   triangle_left = interpolated(1/3, line_start, line_end)
   triangle_right = interpolated(2/3, line_start, line_end)
   

• Anna: Thank you very much for all of your comments! I didn't know pygame understood color words. And that's a clever cross product use I didn't think of!
• 200_success: Note that numpy.cross(two_dimensional_vec, three_dimensional_vec) will produce a three-dimensional result.
• Gareth Rees: @200_success: Yes, you have to pick out the first two components.
• hpaulj: github.com/numpy/numpy/blob/master/numpy/core/numeric.py - cross code is in the middle of this file.