Ray tracer

geometry.py

import math

class Point(object):
    def __init__(self, *args, **kargs):
        self.num_dimensions = kargs.get("num_dimensions", len(args))
        self.coords = [0 for i in range(self.num_dimensions)]
        for i in range(len(args)):
            self.coords[i] = args[i]

    """Gives the distance from this point to the origin."""
    def magnitude(self):
        return math.sqrt(sum(c*c for c in self.coords))

    """
    Gives the angle of this point above the x axis.
    Measured in radians.
    Ranges between -pi and pi.
    """
    def angle(self):
        assert self.num_dimensions == 2
        assert self.x != 0 or self.y != 0
        return math.atan2(self.y,self.x)
    def tuple(self):
        return tuple(self.coords)
    def map(self, func):
        new_coords = [func(a) for a in self.coords]
        return Point(*new_coords)
    def _applyVectorFunc(self, other, f):
        assert self.num_dimensions == other.num_dimensions
        new_coords = [f(a,b) for a,b in zip(self.coords, other.coords)]
        return Point(*new_coords)
    def _applyScalarFunc(self, val, f):
        return self.map(lambda a: f(a,val))
        """
        new_coords = [f(a, val) for a in self.coords]
        return Point(*new_coords)
        """

    def normalized(self):
        return self.__div__(self.magnitude())

    def __add__(self, other):
        return self._applyVectorFunc(other, lambda a,b: a+b)
    def __sub__(self, other):
        return self._applyVectorFunc(other, lambda a,b: a-b)
    def __mul__(self, a):
        return self._applyScalarFunc(a, lambda b,c: b*c)
    def __div__(self, a):
        return self._applyScalarFunc(a, lambda b,c: b/c)
    
    def __eq__(self, other):
        try:
            return self.num_dimensions == other.num_dimensions and self.coords == other.coords
        except:
            return False
    
    def __ne__(self, other):
        return not self.__eq__(other)

    #simple comparator for sorting purposes
    def __lt__(self, other):
        if not isinstance(other, Point): raise Exception("can only compare points to points")
        return self.coords < other.coords
    
    def __getattr__(self, name):
        if name == "x": return self.coords[0]
        if name == "y": return self.coords[1]
        if name == "z": return self.coords[2]
        raise AttributeError(name)
    def __setattr__(self, name, value):
        if name == "x": self.coords[0] = value
        elif name == "y": self.coords[1] = value
        elif name == "z": self.coords[2] = value
        else: object.__setattr__(self, name, value)
    def copy(self):
        return Point(*self.coords[:])
    def __hash__(self):
        return hash(self.tuple())
    def __repr__(self):
        use_nice_floats = False
        if use_nice_floats:
            return "Point(" + ", ".join("%.1f" % c for c in self.coords) + ")"
        else:
            return "Point(" + ", ".join(str(c) for c in self.coords) + ")"

#old version that is always three dimensions
"""
class Point:
    def __init__(self, x=0, y=0, z=0):
        self.x = x
        self.y = y
        self.z = z

    def magnitude(self):
        return math.sqrt(self.x*self.x + self.y*self.y + self.z*self.z)

    def tuple(self):
        return (self.x, self.y, self.z)

    def _applyVectorFunc(self, other, f):
        return Point(f(self.x, other.x), f(self.y, other.y), f(self.z, other.z))
    def _applyScalarFunc(self, a, f):
        return self._applyVectorFunc(Point(a,a,a), f)

    def __add__(self, other):
        return self._applyVectorFunc(other, lambda a,b: a+b)
    def __sub__(self, other):
        return self._applyVectorFunc(other, lambda a,b: a-b)
    def __mul__(self, a):
        return self._applyScalarFunc(a, lambda b,c: b*c)
    def __div__(self, a):
        return self._applyScalarFunc(a, lambda b,c: b/c)
    
    def __eq__(self, other):
        try:
            return self.x == other.x and self.y == other.y and self.z == other.z
        except:
            return False

    def copy(self):
        return Point(self.x, self.y, self.z)

    def __hash__(self):
        return hash(self.tuple())

    def __repr__(self):
        #return "Point({}, {}, {})".format(self.x, self.y, self.z)
        return "Point({}, {})".format(self.x, self.y)
"""    

def distance(a,b):
    return (a-b).magnitude()

def dot_product(a,b):
    return sum(a._applyVectorFunc(b, lambda x,y: x*y).coords)

def cross_product(u,v):
    #todo: support more dimensions than three, if it is possible to do so
    x = u.y*v.z - u.z*v.y
    y = u.z*v.x - u.x*v.z
    z = u.x*v.y - u.y*v.x
    return Point(x,y,z)

def midpoint(a,b, frac=0.5):
    return a*(1-frac) + b*frac

main.py

from geometry import Point
import raytrace

bodies = [
    {
        "type": "sphere",
        "center": Point(0,1,2),
        "radius": 0.5,
        "material": {
            "type": "mirror"
        }
    },
    {
        "type": "sphere",
        "center": Point(1,0,2.5),
        "radius": 0.5,
        "color": (0,255,0)
    },
    {
        "type": "sphere",
        "center": Point(-1,0,1.5),
        "radius": 0.5,
        "color": (0,0,255)
    },
    {
        "type": "plane",
        "point": Point(0,-10.5,0),
        "normal": Point(0,1,0),
        "material":{
            "type": "checkerboard",
            "colors": [(0x88,0x88,0x88), (0xFF, 0xFF, 0xFF)],
            "forward_vector": Point(1,0,1),
            "right_vector": Point(1,0,-1)
        }
    }
]

lighting = [
    {
        "type": "diffuse",
        "position": Point(0, 100, 0),
        "color": (255,255,255),
        "intensity": 1.0
    },
    {
        "type": "diffuse",
        "position": Point(50, -100, -10),
        "color": (255,255,255),
        "intensity": 0.3
    }
]

img = raytrace.render(bodies, lighting, verbose=True, width=200, height=200)
img = img.resize((400,400))
img.save("output/output.png")

progress.py

class ProgressIndicator:
    def __init__(self, max_steps):
        self.max_steps = max_steps
        self.cur = 0
    def tick(self):
        prev_percentage = int(100 * self.cur / self.max_steps)
        self.cur += 1
        cur_percentage = int(100 * self.cur / self.max_steps)
        if prev_percentage != cur_percentage:
            print "\r{}%".format(cur_percentage),
        if cur_percentage == 100:
            print ""

raytrace.py

import math
from PIL import Image
from geometry import Point, distance, dot_product, cross_product
from progress import ProgressIndicator

def angle(a,b):
    A = cross_product(a,b).magnitude()
    return math.asin(A / (a.magnitude() / b.magnitude()))

def rebase(num, a,b, x,y):
    num = (num-a) / float(b-a)
    num = (num * (y-x)) + x
    return num

def solve_quad(a,b,c):
    r = b**2 - 4*a*c
    if r < 0:
        return []
    vals = [
        (-b + math.sqrt(r)) / (2*a),
        (-b - math.sqrt(r)) / (2*a)
    ]
    if vals[0] == vals[1]:
        return [vals[0]]
    return vals


def line_sphere_intersect(v,d, center, radius):
    d = d - center
    return solve_quad(v.x**2 + v.y**2 + v.z**2, 2 * (v.x*d.x + v.y*d.y + v.z*d.z), -radius**2 + d.x**2 + d.y**2 + d.z**2)


def line_plane_intersect(v,d, c,n):
    numerator = dot_product(c - d, n)
    denominator = dot_product(v,n)
    if denominator == 0:
        if numerator == 0:
            return (float("inf"), None)
        else:
            return (0, [])
    else:
        return (1, [float(numerator)/denominator])

def get_normal(body, p):
    if body["type"] == "plane":
        return body["normal"]
    elif body["type"] == "sphere":
        return p - body["center"]

def get_reflection(v, normal):
    normal = normal.normalized()
    projection = normal * dot_product(v, normal)
    delta = projection - v
    result = v + delta*2
    if result.magnitude() == 0: #does this ever actually happen?
        result = v
    return result * -1

def get_material(body, p):
    if "material" not in body:
        return body["color"]
    material = body["material"]
    if material["type"] == "checkerboard":
        if body["type"] == "plane":
            a,b = material["forward_vector"], material["right_vector"]
            if (int(math.floor(dot_product(p, a))) + int(math.floor(dot_product(p,b)))) % 2:
                return material["colors"][0]
            else:
                return material["colors"][1]
        else:
            raise Exception("Checkerboard material not implemented for body {}".format(body["type"]))
    elif material["type"] == "mirror":
        raise Exception("can't get surface data for mirrors")
    else:
        raise Exception("Unrecognized material type {}".format(material["type"]))

def get_color(body, p, lighting):
    normal = get_normal(body, p).normalized()
    intensities = []
    for light in lighting:
        assert light["type"] == "diffuse", "not implemented yet for other lighting types"
        if light["color"] != (255,255,255):
            raise Exception("lightning not implemented for colored lights")

        light_vector = light["position"] - p
        m = max(0, dot_product(light_vector.normalized(), normal)) * light["intensity"]
        intensities.append(m)
    m = sum(intensities)
    r,g,b = get_material(body, p)
    color = (r*m, g*m, b*m)
    return tuple(map(int, color))

def get_ray(v,d, bodies, max_depth=5, emitting_body = None):
    void = {"type": "sphere", "center": Point(0,0,0), "radius": 999, "color": (0,0,0)} #don't want this to be a plane because what if the camera is facing away from it?
    candidates = []
    for body in bodies + [void]:
        if body == emitting_body: continue
        if body["type"] == "sphere":
            dv = v.copy() - body["center"]
            ts = line_sphere_intersect(v,d,body["center"], body["radius"])
            candidates.extend((t, body) for t in ts)
        elif body["type"] == "plane":
            count, solutions = line_plane_intersect(v,d, body["point"], body["normal"])
            if count == float("inf"):
                candidates.append((0, body))
            elif count == 1:
                candidates.append((solutions[0], body))
        else:
            raise Exception("Unknown body type {}".format(body["type"]))
    candidates = [c for c in candidates if c[0] > 0]
    if max_depth == 0: candidates = [c for c in candidates if "material" not in c[1] or c[1]["material"]["type"] != "color"]
    winner = min(candidates, key=lambda pair: pair[0])
    winner = (v*winner[0] + d, winner[1])
    if "material" in winner[1] and winner[1]["material"]["type"] == "mirror":
        p, body = winner
        n = get_normal(body, p)
        return get_ray(get_reflection(v, n), p, bodies, max_depth-1, body)
    else:
        return winner

#the viewer is at (0,0,0) and is facing towards (0,0,1)
#the viewing quad lies on the plane z=1. Its left and right edges are at x=-1, x=1 respectively.
#The top and bottom edges have whatever y coordinates preserve the aspect ratio at 1:1.

def render(bodies, lighting, **kargs):
    width, height = kargs.get("width", 100), kargs.get("height", 100)
    verbose = kargs.get("verbose", False)
    img = Image.new("RGB", (width, height), "white")
    pix = img.load()

    indicator = ProgressIndicator(width*height)

    d = Point(0,0,0)
    for i in range(width):
        for j in range(height):
            if verbose: indicator.tick()
            z = 1
            x = rebase(i, 0, width, -1, 1)
            y = rebase(j, 0, width, 1, -1) #yes we want to rebase by width. This preserves aspect ratio for rectangular viewing quads.
            p = Point(x,y,z)
            collision_point, body = get_ray(p, d, bodies, max_depth=float("inf"))
            color = get_color(body, collision_point, lighting)
            pix[i,j] = color

    return img