python ray tracer

Ray.py

from math import sqrt
import pdb

DELTA = 1.5e-8
INFINITY = 1e5000


def Vec(x=0.0,y=0.0,z=0.0):
    return [float(x),float(y),float(z)]


def add(a,b):
    return [a[0]+b[0],a[1]+b[1],a[2]+b[2]]


def sub(a,b):
    return [a[0]-b[0],a[1]-b[1],a[2]-b[2]]


def scale(s,a):
    return [a[0]*s,a[1]*s,a[2]*s]


def dot(a,b):
    return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]


def unitize(a):
    return scale( 1.0 / ( sqrt( dot( a, a) ) ), a)


class Ray:
    def __init__(self, orig, dir ):
        self.orig = orig 
        self.dir = dir

class Hit:
    def __init__(self, lamb, normal ):
        self.lamb = float(lamb)
        self.normal = normal

## all objects need to expose 
## .intersect( hit, ray )

class Sphere:
    def __init__(self, center, radius):
        self.center = center
        self.radius = radius

    def __str__(self):
        return "%s,%f" % ( str(self.center), self.radius)

    def ray_sphere(self, ray):
        v = sub( self.center, ray.orig)
        b = dot( v, ray.dir)
        disc = b*b - dot(v,v) + self.radius*self.radius
        if disc < 0.0:
            return INFINITY
        d = sqrt(disc)
        t2 = b + d
        if t2  < 0.0:
            return INFINITY
        t1 = b - d
        if t1 > 0.0:
            return t1
        else:
            return t2

    def intersect(self, hit, ray):
        l = self.ray_sphere( ray)
        if l >= hit[0]:
            return hit
        n = add( ray.orig, sub( scale( l, ray.dir), self.center))
        return [l, unitize( n)]


class Group:
    def __init__(self, sphere):
        self.bound = sphere
        self.objs = []

    def __str__(self):
        s = "bound: %s\n" % (self.bound)
        for t in self.objs:
            s += "\t" + str(t) + "\n"
        return s

    def intersect(self, hit, ray):
        l = self.bound.ray_sphere( ray)
        if l >= hit[0]:
            return hit
        for s in self.objs:
            hit = s.intersect( hit, ray)
        return hit

def ray_trace( light, ray, scene):
    _ = (INFINITY, Vec( 0,0,0))
    hit = scene.intersect( _, ray)
    if hit[0] == INFINITY:
        return 0.0
    o = add( ray.orig, add( scale( hit[0] , ray.dir), scale( DELTA, hit[1]) ) )
    g = dot( hit[1], light)
    if g >= 0.0:
        return 0.0
    sray = Ray( o, scale( -1, light))
    si = scene.intersect( (INFINITY, Vec( 0,0,0)), sray)
    if si[0] == INFINITY:
        return -g
    else:
        return 0.0

def create( level, c, r):
    sphere = Sphere( c, r)
    if level == 1:
        return sphere
    group = Group( Sphere( c, 3*r))
    group.objs.append(sphere)
    rn = 3*r/sqrt(12)
    for dz in range(-1,2,2):
        for dx in range(-1,2,2):
            c2 = Vec(c[0]+dx*rn, c[1]+rn, c[2]+dz*rn)
            group.objs.append( create( level-1, c2, r/2.0))


    return group
            



def run( n, level):
    scene = create( level, Vec(0, -1, 0), 1)
    out = open("image.pgm","w")
    out.write( "P5\n" + str(n) + " " + str(n) + "\n255\n")
    for y in range( n-1, -1, -1):
        for x in range( 0, n, 1):
            d = Vec( x - (n/2.0), y - (n/2.0), n)
            ray = Ray( Vec( 0, 0, -4), unitize( d))
            g = ray_trace( unitize( Vec( -1, -3, 2)), ray, scene)
            out.write( chr( int(0.5 + (255.0 * g))))
    out.close()

if __name__ == "__main__":
    import sys
    if "--run" in sys.argv:
        run( int(sys.argv[3]), int(sys.argv[2]))

run with

pypy Ray.py --run 5 512
convert image.pgm image.png