micropython/tests/perf_bench/misc_raytrace.py

259 lines
7.1 KiB
Python
Raw Normal View History

# A simple ray tracer
# MIT license; Copyright (c) 2019 Damien P. George
INF = 1e30
EPS = 1e-6
class Vec:
def __init__(self, x, y, z):
self.x, self.y, self.z = x, y, z
def __neg__(self):
return Vec(-self.x, -self.y, -self.z)
def __add__(self, rhs):
return Vec(self.x + rhs.x, self.y + rhs.y, self.z + rhs.z)
def __sub__(self, rhs):
return Vec(self.x - rhs.x, self.y - rhs.y, self.z - rhs.z)
def __mul__(self, rhs):
return Vec(self.x * rhs, self.y * rhs, self.z * rhs)
def length(self):
return (self.x ** 2 + self.y ** 2 + self.z ** 2) ** 0.5
def normalise(self):
l = self.length()
return Vec(self.x / l, self.y / l, self.z / l)
def dot(self, rhs):
return self.x * rhs.x + self.y * rhs.y + self.z * rhs.z
RGB = Vec
class Ray:
def __init__(self, p, d):
self.p, self.d = p, d
class View:
def __init__(self, width, height, depth, pos, xdir, ydir, zdir):
self.width = width
self.height = height
self.depth = depth
self.pos = pos
self.xdir = xdir
self.ydir = ydir
self.zdir = zdir
def calc_dir(self, dx, dy):
return (self.xdir * dx + self.ydir * dy + self.zdir * self.depth).normalise()
class Light:
def __init__(self, pos, colour, casts_shadows):
self.pos = pos
self.colour = colour
self.casts_shadows = casts_shadows
class Surface:
def __init__(self, diffuse, specular, spec_idx, reflect, transp, colour):
self.diffuse = diffuse
self.specular = specular
self.spec_idx = spec_idx
self.reflect = reflect
self.transp = transp
self.colour = colour
@staticmethod
def dull(colour):
return Surface(0.7, 0.0, 1, 0.0, 0.0, colour * 0.6)
@staticmethod
def shiny(colour):
return Surface(0.2, 0.9, 32, 0.8, 0.0, colour * 0.3)
@staticmethod
def transparent(colour):
return Surface(0.2, 0.9, 32, 0.0, 0.8, colour * 0.3)
class Sphere:
def __init__(self, surface, centre, radius):
self.surface = surface
self.centre = centre
self.radsq = radius ** 2
def intersect(self, ray):
v = self.centre - ray.p
b = v.dot(ray.d)
det = b ** 2 - v.dot(v) + self.radsq
if det > 0:
det **= 0.5
t1 = b - det
if t1 > EPS:
return t1
t2 = b + det
if t2 > EPS:
return t2
return INF
def surface_at(self, v):
return self.surface, (v - self.centre).normalise()
class Plane:
def __init__(self, surface, centre, normal):
self.surface = surface
self.normal = normal.normalise()
self.cdotn = centre.dot(normal)
def intersect(self, ray):
ddotn = ray.d.dot(self.normal)
if abs(ddotn) > EPS:
t = (self.cdotn - ray.p.dot(self.normal)) / ddotn
if t > 0:
return t
return INF
def surface_at(self, p):
return self.surface, self.normal
class Scene:
def __init__(self, ambient, light, objs):
self.ambient = ambient
self.light = light
self.objs = objs
def trace_scene(canvas, view, scene, max_depth):
for v in range(canvas.height):
y = (-v + 0.5 * (canvas.height - 1)) * view.height / canvas.height
for u in range(canvas.width):
x = (u - 0.5 * (canvas.width - 1)) * view.width / canvas.width
ray = Ray(view.pos, view.calc_dir(x, y))
c = trace_ray(scene, ray, max_depth)
canvas.put_pix(u, v, c)
def trace_ray(scene, ray, depth):
# Find closest intersecting object
hit_t = INF
hit_obj = None
for obj in scene.objs:
t = obj.intersect(ray)
if t < hit_t:
hit_t = t
hit_obj = obj
# Check if any objects hit
if hit_obj is None:
return RGB(0, 0, 0)
# Compute location of ray intersection
point = ray.p + ray.d * hit_t
surf, surf_norm = hit_obj.surface_at(point)
if ray.d.dot(surf_norm) > 0:
surf_norm = -surf_norm
# Compute reflected ray
reflected = ray.d - surf_norm * (surf_norm.dot(ray.d) * 2)
# Ambient light
col = surf.colour * scene.ambient
# Diffuse, specular and shadow from light source
light_vec = scene.light.pos - point
light_dist = light_vec.length()
light_vec = light_vec.normalise()
ndotl = surf_norm.dot(light_vec)
ldotv = light_vec.dot(reflected)
if ndotl > 0 or ldotv > 0:
light_ray = Ray(point + light_vec * EPS, light_vec)
light_col = trace_to_light(scene, light_ray, light_dist)
if ndotl > 0:
col += light_col * surf.diffuse * ndotl
if ldotv > 0:
col += light_col * surf.specular * ldotv ** surf.spec_idx
# Reflections
if depth > 0 and surf.reflect > 0:
col += trace_ray(scene, Ray(point + reflected * EPS, reflected), depth - 1) * surf.reflect
# Transparency
if depth > 0 and surf.transp > 0:
col += trace_ray(scene, Ray(point + ray.d * EPS, ray.d), depth - 1) * surf.transp
return col
def trace_to_light(scene, ray, light_dist):
col = scene.light.colour
for obj in scene.objs:
t = obj.intersect(ray)
if t < light_dist:
col *= obj.surface.transp
return col
class Canvas:
def __init__(self, width, height):
self.width = width
self.height = height
self.data = bytearray(3 * width * height)
def put_pix(self, x, y, c):
off = 3 * (y * self.width + x)
self.data[off] = min(255, max(0, int(255 * c.x)))
self.data[off + 1] = min(255, max(0, int(255 * c.y)))
self.data[off + 2] = min(255, max(0, int(255 * c.z)))
def write_ppm(self, filename):
with open(filename, "wb") as f:
f.write(bytes("P6 %d %d 255\n" % (self.width, self.height), "ascii"))
f.write(self.data)
def main(w, h, d):
canvas = Canvas(w, h)
view = View(32, 32, 64, Vec(0, 0, 50), Vec(1, 0, 0), Vec(0, 1, 0), Vec(0, 0, -1))
scene = Scene(
0.5,
Light(Vec(0, 8, 0), RGB(1, 1, 1), True),
[
Plane(Surface.dull(RGB(1, 0, 0)), Vec(-10, 0, 0), Vec(1, 0, 0)),
Plane(Surface.dull(RGB(0, 1, 0)), Vec(10, 0, 0), Vec(-1, 0, 0)),
Plane(Surface.dull(RGB(1, 1, 1)), Vec(0, 0, -10), Vec(0, 0, 1)),
Plane(Surface.dull(RGB(1, 1, 1)), Vec(0, -10, 0), Vec(0, 1, 0)),
Plane(Surface.dull(RGB(1, 1, 1)), Vec(0, 10, 0), Vec(0, -1, 0)),
Sphere(Surface.shiny(RGB(1, 1, 1)), Vec(-5, -4, 3), 4),
Sphere(Surface.dull(RGB(0, 0, 1)), Vec(4, -5, 0), 4),
Sphere(Surface.transparent(RGB(0.2, 0.2, 0.2)), Vec(6, -1, 8), 4),
],
)
trace_scene(canvas, view, scene, d)
return canvas
# For testing
# main(256, 256, 4).write_ppm('rt.ppm')
###########################################################################
# Benchmark interface
bm_params = {
(100, 100): (5, 5, 2),
(1000, 100): (18, 18, 3),
(5000, 100): (40, 40, 3),
}
def bm_setup(params):
return lambda: main(*params), lambda: (params[0] * params[1] * params[2], None)