Skip to content

Instantly share code, notes, and snippets.

@twirrim twirrim/rt4.py
Created Oct 22, 2017

Embed
What would you like to do?
rt4
from PIL import Image
import numpy as np
import numexpr as ne
import time
import numbers
from functools import reduce
import pdb
import timeit
def extract(cond, x):
if isinstance(x, numbers.Number):
return x
else:
return np.extract(cond, x)
class vec3():
def __init__(self, x, y, z):
(self.x, self.y, self.z) = (x, y, z)
def __mul__(self, other):
return vec3(self.x * other, self.y * other, self.z * other)
def __add__(self, other):
return vec3(self.x + other.x, self.y + other.y, self.z + other.z)
def __sub__(self, other):
return vec3(self.x - other.x, self.y - other.y, self.z - other.z)
def dot(self, other):
self_x = self.x
other_x = other.x
self_y = self.y
other_y = other.y
self_z = self.z
other_z = other.z
return ne.evaluate("(self_x * other_x) + (self_y * other_y) + (self_z * other_z)")
def __abs__(self):
return self.dot(self)
def norm(self):
mag = np.sqrt(abs(self))
return self * (1.0 / ne.evaluate("where(mag == 0, 1, mag)"))
def components(self):
return (self.x, self.y, self.z)
def extract(self, cond):
return vec3(extract(cond, self.x),
extract(cond, self.y),
extract(cond, self.z))
def place(self, cond):
r = vec3(np.zeros(cond.shape), np.zeros(cond.shape), np.zeros(cond.shape))
np.place(r.x, cond, self.x)
np.place(r.y, cond, self.y)
np.place(r.z, cond, self.z)
return r
rgb = vec3
(w, h) = (2000, 2000) # Screen size
L = vec3(5, 5., -10) # Point light position
E = vec3(0., 0.35, -1.) # Eye position
FARAWAY = 1.0e39 # an implausibly huge distance
def raytrace(O, D, scene, bounce=0):
# O is the ray origin, D is the normalized ray direction
# scene is a list of Sphere objects (see below)
# bounce is the number of the bounce, starting at zero for camera rays
distances = [s.intersect(O, D) for s in scene]
nearest = reduce(np.minimum, distances)
color = rgb(0, 0, 0)
for (s, d) in zip(scene, distances):
hit = (nearest != FARAWAY) & (d == nearest)
if np.any(hit):
dc = extract(hit, d)
Oc = O.extract(hit)
Dc = D.extract(hit)
cc = s.light(Oc, Dc, dc, scene, bounce)
color += cc.place(hit)
return color
class Sphere:
def __init__(self, center, r, diffuse, mirror=0.5):
self.c = center
self.r = r
self.diffuse = diffuse
self.mirror = mirror
def intersect(self, O, D):
b = 2 * D.dot(O - self.c)
c = abs(self.c) + abs(O) - 2 * self.c.dot(O) - (self.r * self.r)
disc = (b ** 2) - (4 * c)
np_max = np.maximum(0, disc)
sq = ne.evaluate("sqrt(np_max)")
h0 = (-b - sq) / 2
h1 = (-b + sq) / 2
h = ne.evaluate("where((h0 > 0) & (h0 < h1), h0, h1)")
pred = (disc > 0) & (h > 0)
return ne.evaluate("where(pred, h, FARAWAY)")
def diffusecolor(self, M):
return self.diffuse
def light(self, O, D, d, scene, bounce):
M = (O + D * d) # intersection point
N = (M - self.c) * (1. / self.r) # normal
toL = (L - M).norm() # direction to light
toO = (E - M).norm() # direction to ray origin
nudged = M + N * .0001 # M nudged to avoid itself
# Shadow: find if the point is shadowed or not.
# This amounts to finding out if M can see the light
light_distances = [s.intersect(nudged, toL) for s in scene]
light_nearest = reduce(np.minimum, light_distances)
seelight = light_distances[scene.index(self)] == light_nearest
# Ambient
color = rgb(0.05, 0.05, 0.05)
# Lambert shading (diffuse)
lv = np.maximum(N.dot(toL), 0)
color += self.diffusecolor(M) * lv * seelight
# Reflection
if bounce < 2:
rayD = (D - N * 2 * D.dot(N)).norm()
color += raytrace(nudged, rayD, scene, bounce + 1) * self.mirror
# Blinn-Phong shading (specular)
phong = N.dot((toL + toO).norm())
color += rgb(1, 1, 1) * np.power(np.clip(phong, 0, 1), 50) * seelight
return color
class CheckeredSphere(Sphere):
def diffusecolor(self, M):
checker = ((M.x * 2).astype(int) % 2) == ((M.z * 2).astype(int) % 2)
return self.diffuse * checker
scene = [
Sphere(vec3(.75, .1, 1.), .6, rgb(0, 0, 1)),
Sphere(vec3(-.75, .1, 2.25), .6, rgb(.5, .223, .5)),
Sphere(vec3(-2.75, .1, 3.5), .6, rgb(1., .572, .184)),
CheckeredSphere(vec3(0, -99999.5, 0), 99999, rgb(.75, .75, .75), 0.25),
]
r = float(w) / h
# Screen coordinates: x0, y0, x1, y1.
S = (-1., 1. / r + .25, 1., -1. / r + .25)
x = np.tile(np.linspace(S[0], S[2], w), h)
y = np.repeat(np.linspace(S[1], S[3], h), w)
t0 = time.time()
Q = vec3(x, y, 0)
color = raytrace(E, (Q - E).norm(), scene)
print "Took", time.time() - t0
rgb = [Image.fromarray((255 * np.clip(c, 0, 1).reshape((h, w))).astype(np.uint8), "L") for c in color.components()]
Image.merge("RGB", rgb).save("fig.png")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.