lorenz_waterwheel.py

#!/usr/bin/python


import colorsys
import math
import threading
import turtle


class async_wait_for_input(threading.Thread):
    """A Thread that waits for any user input."""
    quit = False
    def __init__(self):
        threading.Thread.__init__(self)
        self.quit = False
    def run(self):
        foo = input('Please enter anything to quit.')
        self.quit = True


def norm_angle(angle):
    """Convert angle to the 0-360 degree range."""
    return angle % 360


def deg_to_rad(angle):
    """Convert degree to radians."""
    return math.pi * angle / 180


def lever(angle):
    """Calculate the lever for the momentum of a
    bucket with a specific angle"""
    return math.sin(deg_to_rad(angle))


def bucket_force_sum(buckets, angle):
    """Calculate the sum off the forces all buckets generate."""
    return sum([b*lever(angle-360*n/len(buckets))
               for n, b in enumerate(buckets)])


def fill_upper_bucket(buckets, bucket_sizeInDeg, angle, flow, bucket_size):
    """Fill the upper bucket with new water."""
    newBuckets = buckets
    for n,b in enumerate(buckets):
        left = n*360/len(buckets)-bucket_sizeInDeg/2
        right = n*360/len(buckets)+bucket_sizeInDeg/2
        if angle > left and angle < right:
            newBuckets[n] = min(newBuckets[n]+flow, bucket_size)
    return newBuckets


def bucket_leakage(buckets, bucket_hole):
    """Let buckets leak some water."""
    return [max(0, bucket-bucket_hole) for bucket in buckets]


def calc_color(i):
    """Return a RGB color on the edge of the HSV cylinder (rainbow)."""
    return colorsys.hsv_to_rgb(i/256 % 1, 1, 1)


def lorenz_waterwheel():
    """Simulate the chaotic behavior of a lorenz waterwheel."""
    angle = 1
    speed = 0
    empty_mass = 1000
    flow = 41
    bucket_hole = 1
    bucket_size = 100000
    bucket_start_fill = 0
    bucket_count = 8
    bucket_sizeInDeg = 360/(bucket_count*2)
    brake = 0.99
    buckets = [bucket_start_fill for c in range(bucket_count)]
    print("angle speed norm_angle speed acceleration buckets")
    turtle.speed(0)
    turtle.up()
    turtle.hideturtle()
    turtle.bgcolor('black')
    counter = 0
    background = async_wait_for_input()
    background.start()
    while(not background.quit):
        counter += 1
        turtle.pencolor(calc_color(counter))
        speed *= brake
        angle += speed
        acceleration = (bucket_force_sum(buckets, norm_angle(angle))/
                        (empty_mass+sum(buckets)))
        speed += acceleration
        buckets = fill_upper_bucket(buckets, bucket_sizeInDeg,
                                    norm_angle(angle), flow, bucket_size)
        buckets = bucket_leakage(buckets, bucket_hole)
        print(angle, speed, norm_angle(angle), speed, acceleration, buckets)
        turtle.setpos(0.5*(180-angle), 20 * speed)
        turtle.down()
    background.join()
    turtle_screen = turtle.getscreen()
    turtle_screen.getcanvas().postscript(file="LorenzWaterwheel.eps")
    turtle.bye()


if __name__ == "__main__":
    exit(lorenz_waterwheel())