Simulates refraction of light inside a raindrop.

raindrop.py

#!/usr/bin/python
#

from math import sin,cos,tan,sqrt,pi,asin, atan2 as atan
from GUI import View, Application, Window, Task
from GUI.Colors import Color

pale_blue = Color(.5,.5,1)
black = Color(0,0,0)
white = Color(1,1,1)
red = Color(1,0,0)

MIN_WAVELENGTH = 400
MAX_WAVELENGTH = 700

class RaindropView( View ):
    hpadding = 100
    lpadding = 100
    rpadding = 300
    r = 100.
    width = 2*r + lpadding + rpadding
    height = 2*r + 2*hpadding
    circle_x = lpadding+r
    circle_y = hpadding+r

    b_speed = 2
    n_speed = .01

    def __init__( self ):
        self.keys_held = set([])
        self.b = 0
        self.base_n = 1.333
        View.__init__( self, size=(RaindropView.width,
                                       RaindropView.height) )
        self.updater = Task( self.update, .03, repeat=True, start=True )

    def main_angle( self, w ):
        n = refractive_index( self.base_n, w )
        result = 2*asin(sqrt((4-n**2)/3)) - 4*asin(sqrt((4-n**2)/(3*n**2)))
        if self.b < 0:
            return -result
        else:
            return result

    def stroke_wavelength( self, canvas, w ):
        """Draw the path of a single monochromatic ray."""
        if 400 <= w <= 700:
            n = refractive_index( self.base_n, w )
            # Draw the light ray. The starting point:
            start_x = self.width
            start_y = self.circle_y + self.b
            # The place it first hits the drop:
            in_x = self.circle_x + sqrt(self.r**2 - self.b**2)
            in_y = start_y
            # The place it bounces:
            theta_in = atan(in_y-self.circle_y, in_x-self.circle_x)
            theta_bounce = theta_in + (pi - 2*asin(self.b/(self.r*n)))
            bounce_x = self.circle_x + self.r*cos(theta_bounce)
            bounce_y = self.circle_y + self.r*sin(theta_bounce)
            # The place it exits:
            theta_out = theta_bounce + (pi - 2*asin(self.b/(self.r*n)))
            out_x = self.circle_x + self.r*cos(theta_out)
            out_y = self.circle_y + self.r*sin(theta_out)
            # And the place it goes off the screen:
            outgoing_angle = theta_out + asin(self.b/self.r)
            finish_x = out_x + (self.width+self.height)*cos(outgoing_angle)
            finish_y = out_y + (self.width+self.height)*sin(outgoing_angle)

            # For comparison, draw the max-refraction angle's exit trajectory too.
            canvas.pencolor = black
            if n <= 2:
                main_x = out_x + (self.width+self.height)*cos(self.main_angle(w))
                main_y = out_y + (self.width+self.height)*sin(self.main_angle(w))
            else:
                main_x = self.width
                main_y = out_y
            canvas.stroke_lines( ((out_x,out_y), (main_x,main_y)) )

            # Draw the colored line on the canvas. (I'd put this first, but I want the colors to be on top of the black lines.)
            canvas.pencolor = wavelength_to_color(w)
            canvas.stroke_lines( ((start_x,start_y), (in_x,in_y),
                                  (bounce_x,bounce_y), (out_x,out_y),
                                  (finish_x,finish_y)) )

        else:
            raise ValueError( "illegal wavelength: %d" %w )

    def draw( self, canvas, update_rect ):
        canvas.fillcolor = white
        canvas.fill_rect(update_rect)

        # First, fill the raindrop.
        oval = (self.circle_x-self.r, self.circle_y-self.r,
                self.circle_x+self.r, self.circle_y+self.r)
        canvas.fillcolor = pale_blue
        canvas.pencolor = black
        canvas.fill_oval( oval )
        canvas.stroke_oval( oval )

        # Now stroke lines for different wavelengths.
        for i in range(10):
            w = 400 + 300*(i/10.)
            self.stroke_wavelength( canvas, w )

        # Finally, show the refractive index.
        canvas.moveto( 2, 2+canvas.font.height )
        canvas.show_text( "%.2f" %self.base_n )

    def key_down( self, event ):
        self.keys_held.add(event.key)
    def key_up( self, event ):
        self.keys_held.remove(event.key)

    def update( self ):
        changed = False
        if ('down_arrow' in self.keys_held) and (self.b+self.b_speed) < self.r:
            changed = True
            self.b += self.b_speed
        if ('up_arrow' in self.keys_held) and (self.b-self.b_speed) > -self.r:
            changed = True
            self.b -= self.b_speed
        if ('left_arrow' in self.keys_held) and (self.base_n-self.n_speed) > 1:
            changed = True
            self.base_n -= self.n_speed
        if ('right_arrow' in self.keys_held) and (self.base_n+self.n_speed) < 4:
            changed = True
            self.base_n += self.n_speed

        if changed:
            self.invalidate()


def refractive_index( n, w ):
    """Returns the refractive index of water for a given wavelength and base n."""
    return (1.35 - (w-400)*0.022/300) * n/1.333

def wavelength_to_color( w ):
    """Turns a wavelength (in nm) into a color we can display on the screen.

    Algorithm stolen from... somewhere on the Internet.
    """
    if (w >= 400) and (w <= 439):
        r= -(w - 440) / (440 - 350)
        g = 0.0
        b= 1.0
    elif(w >= 440) and (w <= 489):
        r= 0.0
        g = (w - 440) / (490 - 440)
        b= 1.0
    elif (w >= 490) and (w <= 509):
        r = 0.0
        g = 1.0
        b = -(w - 510) / (510 - 490)
    elif (w >= 510) and (w <= 579):
        r = (w - 510) / (580 - 510)
        g = 1.0
        b = 0.0
    elif (w >= 580) and (w <= 644):
        r = 1.0
        g = -(w - 645) / (645 - 580)
        b = 0.0
    elif (w >= 645) and (w <= 700):
        r = 1.0
        g = 0.0
        b = 0.0
    else:
        raise ValueError("illegal wavelength: %d" %w)

    return Color(r,g,b)

class RaindropApp( Application ):
    def open_app( self ):
        self.make_window()

    def make_window( self ):
        win = Window()
        view = RaindropView()
        win.place(view, left=0, top=0)
        win.shrink_wrap()
        win.show()
        view.become_target()

if __name__ == '__main__':
    print "Up/down to move incoming beam; left/right to change refractive index."
    print "Black lines indicate the maximum possible exit angle for each wavelength, for comparison."
    print "(Interesting things happen near n=1 and n=2.)"
    RaindropApp().run()