Forward/Backward 2D Python Raytracer

import pygame
import pygame.gfxdraw
from pygame.locals import *
import sys, os, traceback
from math import *
import random
if sys.platform in ["win32","win64"]: os.environ["SDL_VIDEO_CENTERED"]="1"
pygame.display.init()
pygame.font.init()

fonts = [
    pygame.font.SysFont("Times New Roman",12),
    pygame.font.SysFont("Times New Roman",16)
]

screen_size = [800,400]
icon = pygame.Surface((1,1)); icon.set_alpha(0); pygame.display.set_icon(icon)
pygame.display.set_caption("Adjunct Raytracer - Ian Mallett - v.1.0.0 - 2014")
surface = pygame.display.set_mode(screen_size)

#1 == no refraction; thin lens only
#2 == refraction through thin lens
REFRACT_MODE = 2
#Focal length of lens (pixels) (only matters if REFRACT_MODE==2)
FOCAL = 50.0
#Aperture size (pixels)
APERTURE = 100

#Color of scene geometry
COLOR_GEOM = (255,128,0)
#Color of path tracing rays
COLOR_PT_HIT = (0,255,0)
COLOR_PT_MISS = (255,0,0)
#Color of light tracing rays
COLOR_LT_HIT = (0,0,255)
COLOR_LT_MISS = (255,0,255)

def rndint(x): return int(round(x))
def dot(u,v): return u[0]*v[0] + u[1]*v[1]
def perp(u,v): return u[0]*v[1] - u[1]*v[0] #perp product (2D)
def sc(sc,vec): return sc*vec[0], sc*vec[1]
def add(vec0,vec1): return vec0[0]+vec1[0], vec0[1]+vec1[1]
def sub(vec0,vec1): return vec0[0]-vec1[0], vec0[1]-vec1[1]
def lengthsq(vec): return dot(vec,vec)
def length(vec): return lengthsq(vec) ** 0.5
def normalized(vec): return sc(1.0/length(vec),vec)
def rand_semi(normal):
    angle = random.random() * 2.0*pi
    vec = (cos(angle),sin(angle))
    if dot(vec,normal)<0.0:
        return (-vec[0],-vec[1]), 1.0/(1.0*pi)
    else:
        return vec, 1.0/(1.0*pi)
def intersect_ray_circle(ray,circle):
    omc = [ray.x-circle.x,ray.y-circle.y]
    l_dot_omc = dot((ray.dx,ray.dy),omc)
    discr = l_dot_omc*l_dot_omc - lengthsq(omc) + circle.radius*circle.radius
    if discr<=0.0: return False,-1,-1
    discr_sqrt = discr ** 0.5
    d0,d1 = -l_dot_omc-discr_sqrt,-l_dot_omc+discr_sqrt
    return d0>0,d0,d1
def intersect_line_line(p0,p1, p2,p3):
    u = sub(p1,p0)
    v = sub(p3,p2)
    w = sub(p0,p2)
    D = perp(u,v)

    sI = perp(v,w) / D;
    if sI<0 or sI>1: return False,None
    tI = perp(u,w) / D;
    if tI<0 or tI>1: return False,None

    return True,add(p0,sc(sI,u))
def intersect_ray_line(ray,line):
    hit,pt = intersect_line_line(
        (ray.x,ray.y), (ray.x+10000.0*ray.dx,ray.y+10000.0*ray.dy),
        line.get_p0(), line.get_p1()
    )
    if hit:
        return True,length(sub(pt,(ray.x,ray.y)))
    else:
        return False,-1

class ObjectBase(object):
    def __init__(self, x,y):
        self.x = x
        self.y = y
    @staticmethod
    def to_screen(x,y):
        return rndint(x+screen_size[0]*0.5),rndint(screen_size[1]-(y+screen_size[1]*0.5))
    def draw(self): pass
class ObjectLineBase(ObjectBase):
    def __init__(self, x,y, size):
        ObjectBase.__init__(self, x,y)
        self.size = size
    def get_p0(self):
        return self.x, self.y-self.size*0.5
    def get_p1(self):
        return self.x, self.y+self.size*0.5
    def draw(self):
        pygame.draw.aaline(
            surface, COLOR_GEOM,
            ObjectBase.to_screen(*self.get_p0()), ObjectBase.to_screen(*self.get_p1())
        )
class ObjectSensor(ObjectLineBase): #a line
    def __init__(self, x,y, size):
        ObjectLineBase.__init__(self, x,y, size)
    def get_random(self):
        x,y = self.x,self.y+(random.random()-0.5)*self.size
        pdf_pos = 1.0 / self.size
        vec,pdf_dir = rand_semi((1,0))
        pdf = pdf_pos * pdf_dir
        return Ray(x,y,vec[0],vec[1]),(1,0),pdf
    def draw(self):
        ObjectLineBase.draw(self)
        surface.blit(fonts[0].render("sensor",True,(0,0,0)),ObjectBase.to_screen(self.x-10,self.y+self.size//2+25))
class ObjectLens(ObjectLineBase): #a line
    def __init__(self, x,y, size, f):
        ObjectLineBase.__init__(self, x,y, size)
        self.size = size
        self.f = f
    def get_refracted(self, incoming_ray,hit_distance):
        lensx,lensy = incoming_ray.at(hit_distance)
        if   REFRACT_MODE == 1:
            #Do nothing (no refraction; just an aperture)
            return Ray(lensx,lensy, incoming_ray.dx,incoming_ray.dy)
        elif REFRACT_MODE == 2:
            #Thin lens
            #   Solve by considering the ray's origin to be an "object", finding the "image" point that focuses to it,
            #       then the refracted ray will start at the hit position on the lens and pass through the image point.
            o = abs(self.x - incoming_ray.x)
            lens_to_obj = sub((incoming_ray.x,incoming_ray.y),(self.x,self.y))
            lens_to_objn = normalized(lens_to_obj)
            if o == self.f:
                #Object is exactly at the focal length; rays are parallel and no finite image is formed.
                d = (-lens_to_objn[0],-lens_to_objn[1])
            else:
                #Thin lens equation
                i = 1.0 / (1.0/self.f - 1.0/o)
                if o < self.f:
                    #Object is inside the focal length; rays diverge and virtual image is behind object on the same
                    #   side of the lens.
                    image = add( sc(-i/o,lens_to_obj), (self.x,self.y) )
                    d = normalized(sub( (lensx,lensy), image ))
                else:
                    #Object is outside the focal length; rays converge and real image is on the other side of the lens.
                    image = add( sc(i/-o,lens_to_obj), (self.x,self.y) )
                    d = normalized(sub( image, (lensx,lensy) ))
            return Ray(lensx,lensy, d[0],d[1])
    def draw(self):
        ObjectLineBase.draw(self)
        if REFRACT_MODE == 2:
            pygame.draw.circle(surface,COLOR_GEOM,ObjectBase.to_screen(self.x-self.f,self.y),2)
            pygame.draw.circle(surface,COLOR_GEOM,ObjectBase.to_screen(self.x+self.f,self.y),2)
        surface.blit(fonts[0].render("lens",True,(0,0,0)),ObjectBase.to_screen(self.x-10,self.y+self.size//2+20))
class ObjectLight(ObjectBase): #a circle
    def __init__(self, x,y,radius, radiance):
        ObjectBase.__init__(self, x,y)
        self.radius = radius
        self.radiance = radiance
    def get_random(self):
        angle = random.random() * 2*pi
        nx,ny = cos(angle),sin(angle)
        pdf_pos = 1.0 / (2.0*pi*self.radius)
        vec,pdf_dir = rand_semi((nx,ny))
        pdf = pdf_pos * pdf_dir
        return Ray(self.x+self.radius*nx,self.y+self.radius*ny,vec[0],vec[1]),(nx,ny),pdf
    def draw(self):
        sx,sy = ObjectBase.to_screen(self.x,self.y)
        #pygame.draw.circle(surface, COLOR_GEOM, (sx,sy), self.radius,1)
        pygame.gfxdraw.aacircle(surface, sx,sy, self.radius, COLOR_GEOM)
        surface.blit(fonts[0].render("light",True,(0,0,0)),(sx-10,sy-6))
class Ray(ObjectBase):
    def __init__(self, x,y,dx,dy):
        ObjectBase.__init__(self, x,y)
        self.dx = dx
        self.dy = dy
    def at(self, t):
        return self.x+self.dx*t, self.y+self.dy*t
    def draw(self,t,color):
        pygame.draw.aaline(
            surface, color,
            ObjectBase.to_screen(self.x,self.y), ObjectBase.to_screen(*self.at(t))
        )

img = ObjectSensor(-300,0,100)
lns = ObjectLens(-200,0,APERTURE, FOCAL)
lgt = ObjectLight(200,0,50, 1.0)

def path_trace():
    #Get random ray from image plane
    sensor_ray,normal,pdf = img.get_random()
    hit,d_1 = intersect_ray_line(sensor_ray,lns)
    if not hit: #didn't hit the lens
        sensor_ray.draw(10.0, COLOR_PT_MISS)
        return 0.0 #monte carlo estimate for flux is zero since radiance is zero
    #If the sensor ray hit the lens, refract it
    lens_ray = lns.get_refracted(sensor_ray,d_1)
    hit,d_2,d_3 = intersect_ray_circle(lens_ray,lgt)
    if not hit:
        sensor_ray.draw(d_1, COLOR_PT_HIT)
        lens_ray.draw(10.0, COLOR_PT_MISS)
        return 0.0 #monte carlo estimate for flux is zero since radiance is zero
    #The refracted sensor ray hit the light
    sensor_ray.draw(d_1, COLOR_PT_HIT)
    lens_ray.draw(d_2, COLOR_PT_HIT)
    #   Note: sensor_ray.dx is img's normal==(1,0) dotted with sensor_ray.dir==(sensor_ray.dx,sensor_ray.dy)
    #   Note: 1.0 is the importance (self-importance of sensor)
    return lgt.radiance * 1.0 * sensor_ray.dx / pdf
def light_trace():
    #Get a random ray from the light
    light_ray,normal,pdf = lgt.get_random()
    hit,d_1 = intersect_ray_line(light_ray,lns)
    if not hit:
        light_ray.draw(10.0, COLOR_LT_MISS)
        return 0.0 #monte carlo estimate for flux is zero since importance is zero
    #If the light ray hit the lens, refract it
    lens_ray = lns.get_refracted(light_ray,d_1)
    hit,d_2 = intersect_ray_line(lens_ray,img)
    if not hit:
        light_ray.draw(d_1, COLOR_LT_HIT)
        lens_ray.draw(10.0, COLOR_LT_MISS)
        return 0.0 #monte carlo estimate for flux is zero since importance is zero
    #The refracted light ray hit the sensor
    light_ray.draw(d_1, COLOR_LT_HIT)
    lens_ray.draw(d_2, COLOR_LT_HIT)
    return lgt.radiance * 1.0 * dot(normal,(light_ray.dx,light_ray.dy)) / pdf

accum = [0.0,0.0] #estimate from eye, estimate from light
n = 0
def update_and_draw():
    global accum, n
    surface.fill((255,255,255))

    N = 100
    #Trace from eye, drawing paths: path tracing
    for i in range(N):
        mc_flux = path_trace()
        accum[0] += mc_flux
    #Trace from light, drawing paths: light tracing
    for i in range(N):
        mc_flux = light_trace()
        accum[1] += mc_flux
    n += N

    #Draw scene
    img.draw()
    lns.draw()
    lgt.draw()

    #Draw readout(s)
    surface.blit(fonts[1].render(u"\u03a6 (eye)  = " +str(accum[0]/n),True,(0,0,0)),(20,20))
    surface.blit(fonts[1].render(u"\u03a6 (light) = "+str(accum[1]/n),True,(0,0,0)),(20,40))
    surface.blit(fonts[1].render("n = "+str(n),True,(0,0,0)),(20,60))

    pygame.display.flip()

def main():
    clock = pygame.time.Clock()
    continuing = True
    while continuing:
        for event in pygame.event.get():
            if   event.type == QUIT: continuing=False
            elif event.type == KEYDOWN:
                if   event.key == K_ESCAPE: continuing=False
        update_and_draw()
        clock.tick(60)
    pygame.quit()
if __name__ == "__main__":
    try:
        main()
    except:
        traceback.print_exc()
        pygame.quit()
        input()