go ray tracer basic implementation

obj.go

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

func parseIndex(value string, length int) int {
	parsed, _ := strconv.ParseInt(value, 0, 0)
	n := int(parsed)
	if n < 0 {
		n += length
	}
	return n
}

func parseFloats(items []string) []float64 {
	result := make([]float64, len(items))
	for i, item := range items {
		f, _ := strconv.ParseFloat(item, 64)
		result[i] = f
	}
	return result
}

func loadOBJ(path string) []triangle {
	fmt.Printf("Loading OBJ: %s\n", path)
	file, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer file.Close()
	vs := make([]point3, 1, 1024) // 1-based indexing
	var triangles []triangle
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Fields(line)
		if len(fields) == 0 {
			continue
		}
		keyword := fields[0]
		args := fields[1:]
		switch keyword {
		case "v":
			f := parseFloats(args)
			v := point3{f[0], f[1], f[2]}
			vs = append(vs, v)
		case "f":
			fvs := make([]int, len(args))
			for i, arg := range args {
				vertex := strings.Split(arg+"//", "/")
				fvs[i] = parseIndex(vertex[0], len(vs))
			}
			for i := 1; i < len(fvs)-1; i++ {
				i1, i2, i3 := 0, i, i+1
				v1, v2, v3 := vs[fvs[i1]], vs[fvs[i2]], vs[fvs[i3]]
				triangles = append(triangles, triangle{v1, v2, v3})
			}
		}
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}
	println("loaded", len(triangles), "triangles")
	return triangles
}

raytrace.go

package main

import (
	"image"
	"image/color"
	"log"
	"math"
	"math/rand"
	"sync"
)

func randf64() float64 {
	return rand.Float64()
}

func randomInUnitSphere() vec3 {
	for {
		p := vec3{randf64(), randf64(), randf64()}.muls(2).sub(vec3{1, 1, 1})
		if p.length() < 1 {
			return p
		}
	}
}

func randomUnitVector() vec3 {
	a := randf64() * 2 * math.Pi
	z := randf64()*2 - 1
	r := math.Sqrt(1 - z*z)
	return vec3{r * math.Cos(a), r * math.Sin(a), z}
}

func randomInUnitDisk() vec3 {
	for {
		p := vec3{randf64(), randf64(), 0}.muls(2).sub(vec3{1, 1, 0})
		if p.dot(p) < 1 {
			return p
		}
	}
}

type vec3 struct {
	x float64
	y float64
	z float64
}

func (v vec3) add(v2 vec3) vec3 {
	return vec3{v.x + v2.x, v.y + v2.y, v.z + v2.z}
}

func (v vec3) sub(v2 vec3) vec3 {
	return vec3{v.x - v2.x, v.y - v2.y, v.z - v2.z}
}

func (v vec3) mul(v2 vec3) vec3 {
	return vec3{v.x * v2.x, v.y * v2.y, v.z * v2.z}
}

func (v vec3) div(v2 vec3) vec3 {
	return vec3{v.x / v2.x, v.y / v2.y, v.z / v2.z}
}

func (v vec3) muls(s float64) vec3 {
	return vec3{v.x * s, v.y * s, v.z * s}
}

func (v vec3) divs(s float64) vec3 {
	return vec3{v.x / s, v.y / s, v.z / s}
}

func (v vec3) dot(v2 vec3) float64 {
	return v.x*v2.x + v.y*v2.y + v.z*v2.z
}

func (v vec3) cross(v2 vec3) vec3 {
	return vec3{v.y*v2.z - v.z*v2.y, v.z*v2.x - v.x*v2.z, v.x*v2.y - v.y*v2.x}
}

func (v vec3) length() float64 {
	return math.Sqrt(v.dot(v))
}

func (v vec3) unit() vec3 {
	return v.divs(v.length())
}

func (v vec3) clip(min, max float64) vec3 {
	return vec3{math.Min(math.Max(v.x, min), max), math.Min(math.Max(v.y, min), max), math.Min(math.Max(v.z, min), max)}
}

type point3 struct {
	x float64
	y float64
	z float64
}

func (p point3) sub(p2 point3) vec3 {
	return vec3{p.x - p2.x, p.y - p2.y, p.z - p2.z}
}

func (p point3) add(v vec3) point3 {
	return point3{p.x + v.x, p.y + v.y, p.z + v.z}
}

func (p point3) subv(v vec3) point3 {
	return point3{p.x - v.x, p.y - v.y, p.z - v.z}
}

type ray struct {
	origin       point3
	direction    vec3
	scatterCount int
}

func (r ray) at(t float64) point3 {
	return r.origin.add(r.direction.muls(t))
}

type collider interface {
	collide(r ray, tmin float64, tmax float64) (bool, collision)
}

type collision struct {
	t      float64
	at     point3
	normal vec3
}

type sphere struct {
	center point3
	radius float64
}

func (s sphere) collide(r ray, tmin, tmax float64) (bool, collision) {
	oc := r.origin.sub(s.center)
	a := r.direction.dot(r.direction)
	b := oc.dot(r.direction)
	c := oc.dot(oc) - s.radius*s.radius
	discriminant := b*b - a*c
	if discriminant < 0 {
		return false, collision{}
	}
	sqrtD := math.Sqrt(discriminant)
	t := (-b - sqrtD) / a
	if t < tmin || t > tmax {
		t = (-b + sqrtD) / a
		if t < tmin || t > tmax {
			return false, collision{}
		}
	}
	at := r.at(t)
	return true, collision{t, at, at.sub(s.center).divs(s.radius)}
}

type rectangle struct {
	x0 float64
	x1 float64
	y0 float64
	y1 float64
	z  float64
}

func (rect rectangle) collide(r ray, tmin, tmax float64) (bool, collision) {
	t := (rect.z - r.origin.z) / r.direction.z
	if t < tmin || t > tmax {
		return false, collision{}
	}
	x := r.origin.x + t*r.direction.x
	y := r.origin.y + t*r.direction.y
	if x < rect.x0 || x > rect.x1 || y < rect.y0 || y > rect.y1 {
		return false, collision{}
	}
	return true, collision{t, r.at(t), vec3{0, 0, 1}}
}

type triangle struct {
	p0 point3
	p1 point3
	p2 point3
}

func (tri triangle) collide(r ray, tmin, tmax float64) (bool, collision) {
	edge1 := tri.p1.sub(tri.p0)
	edge2 := tri.p2.sub(tri.p0)
	h := r.direction.cross(edge2)
	a := edge1.dot(h)
	if a > -1e-8 && a < 1e-8 {
		return false, collision{}
	}
	f := 1 / a
	s := r.origin.sub(tri.p0)
	u := f * s.dot(h)
	if u < 0 || u > 1 {
		return false, collision{}
	}
	q := s.cross(edge1)
	v := f * r.direction.dot(q)
	if v < 0 || u+v > 1 {
		return false, collision{}
	}
	t := f * edge2.dot(q)
	if t < tmin || t > tmax {
		return false, collision{}
	}
	at := r.at(t)
	return true, collision{t, at, edge1.cross(edge2).unit()}
}

func (t triangle) normal() vec3 {
	return t.p1.sub(t.p0).cross(t.p2.sub(t.p0)).unit()
}

type mesh struct {
	triangles []triangle
}

func (m mesh) collide(r ray, tmin, tmax float64) (bool, collision) {
	hit := false
	var closest collision
	for _, tri := range m.triangles {
		if h, c := tri.collide(r, tmin, tmax); h {
			hit = true
			tmax = c.t
			closest = c
		}
	}
	return hit, closest
}

func (m *mesh) translate(v vec3) {
	for i := range m.triangles {
		m.triangles[i].p0 = m.triangles[i].p0.add(v)
		m.triangles[i].p1 = m.triangles[i].p1.add(v)
		m.triangles[i].p2 = m.triangles[i].p2.add(v)
	}
}

type emitter interface {
	emit(u, v float64) ray
}

type perspectiveCamera struct {
	lowerLeftCorner point3
	origin          point3
	horizontal      vec3
	vertical        vec3
}

func (c perspectiveCamera) emit(u, v float64) ray {
	return ray{c.origin, c.lowerLeftCorner.add(c.horizontal.muls(u)).add(c.vertical.muls(v)).sub(c.origin), 0}
}

type positionableCamera struct {
	lookfrom        point3
	lookat          point3
	vup             vec3
	fovHeight       float64
	fovWidth        float64
	aperture        float64
	workingDistance float64
}

func (cam positionableCamera) emit(s, t float64) ray {
	w := cam.lookfrom.sub(cam.lookat).unit()
	u := cam.vup.cross(w).unit()
	v := w.cross(u)
	horizontal := u.muls(cam.fovWidth * cam.workingDistance)
	vertical := v.muls(cam.fovHeight * cam.workingDistance)
	lowerLeftCorner := cam.lookfrom.subv(horizontal.divs(2)).subv(vertical.divs(2)).subv(w.muls(cam.workingDistance))
	lensRadius := cam.aperture / 2
	rd := randomInUnitDisk().muls(lensRadius)
	offset := u.muls(rd.x).add(v.muls(rd.y))
	origin := cam.lookfrom.add(offset)
	direction := lowerLeftCorner.add(horizontal.muls(s)).add(vertical.muls(t)).sub(origin)
	return ray{origin, direction, 0}
}

type parallelCamera struct {
	lookfrom  point3
	lookat    point3
	vup       vec3
	fovHeight float64
	fovWidth  float64
}

func (cam parallelCamera) emit(s, t float64) ray {
	w := cam.lookfrom.sub(cam.lookat).unit()
	u := cam.vup.cross(w).unit()
	v := w.cross(u)
	origin := cam.lookfrom.add(u.muls(cam.fovWidth * (s - 0.5))).add(v.muls(cam.fovHeight * (t - 0.5)))
	direction := cam.lookat.sub(cam.lookfrom)
	return ray{origin, direction, 0}
}

type scatterer interface {
	scatter(old ray, c collision) (new ray, color vec3)
}

type lambertian struct {
	albedo vec3
}

func (m lambertian) scatter(old ray, c collision) (ray, vec3) {
	p := old.at(c.t)
	target := p.add(c.normal).add(randomUnitVector())
	return ray{p, target.sub(p), 0}, m.albedo
}

type metal struct {
	albedo vec3
	fuzz   float64
}

func (m metal) scatter(old ray, c collision) (ray, vec3) {
	reflected := old.direction.unit().sub(c.normal.muls(old.direction.unit().dot(c.normal) * 2))
	if reflected.dot(c.normal) > 0 {
		return ray{old.at(c.t), reflected.add(randomUnitVector().muls(m.fuzz)), 0}, m.albedo
	}
	return ray{}, vec3{}
}

type dielectric struct {
	refractionIndex float64
}

func (m dielectric) scatter(old ray, c collision) (new ray, color vec3) {
	var outwardNormal vec3
	var niOverNt float64
	var cosine float64
	if old.direction.dot(c.normal) > 0 {
		outwardNormal = c.normal.muls(-1)
		niOverNt = m.refractionIndex
		cosine = m.refractionIndex * old.direction.dot(c.normal) / old.direction.length()
	} else {
		outwardNormal = c.normal
		niOverNt = 1.35 / m.refractionIndex // NOTE SCOTT 1.35 HACK
		cosine = -old.direction.dot(c.normal) / old.direction.length()
	}
	refracted, ok := refract(old.direction, outwardNormal, niOverNt)
	if ok {
		reflectProbability := schlick(cosine, m.refractionIndex)
		if randf64() < reflectProbability {
			return ray{c.at, reflect(old.direction, c.normal), 0}, vec3{1, 1, 1}
		} else {
			return ray{c.at, refracted, 0}, vec3{1, 1, 1}
		}
	} else {
		// total internal reflection
		return ray{c.at, reflect(old.direction, c.normal), 0}, vec3{1, 1, 1}
	}
}

func reflect(v, n vec3) vec3 {
	return v.sub(n.muls(v.dot(n) * 2))
}

// refract returns the refraction vector and a boolean indicating whether
// refraction is possible.
func refract(v, n vec3, niOverNt float64) (refracted vec3, ok bool) {
	uv := v.unit()
	dt := uv.dot(n)
	discriminant := 1 - niOverNt*niOverNt*(1-dt*dt)
	if discriminant > 0 {
		return uv.sub(n.muls(dt)).muls(niOverNt).sub(n.muls(math.Sqrt(discriminant))), true
	}
	return vec3{}, false
}

func schlick(cosine, refractionIndex float64) float64 {
	r0 := (1 - refractionIndex) / (1 + refractionIndex)
	r0 = r0 * r0
	return r0 + (1-r0)*math.Pow(1-cosine, 5)
}

type light struct {
	shape collider
	color vec3
}

type entity struct {
	shape    collider
	material scatterer
}

type scene struct {
	lights   []light
	samples  int
	camera   emitter
	entities []entity
	seed     int64
}

func (s scene) paint(r ray) vec3 {
	closest := math.MaxFloat64
	var closestCollision collision
	var closestEntity entity
	for _, e := range s.entities {
		var hit bool
		hit, c := e.shape.collide(r, 0.001, math.MaxFloat64)
		if hit && c.t < closest {
			closest = c.t
			closestEntity = e
			closestCollision = c
		}
	}
	var closestLight light
	for _, light := range s.lights {
		var hit bool
		hit, c := light.shape.collide(r, 0.001, math.MaxFloat64)
		if hit && c.t < closest {
			closest = c.t
			closestLight = light
			closestCollision = c
		}
	}
	if closestLight != (light{}) {
		unitDirection := r.direction.unit()
		t := 0.5 * (unitDirection.x + 1)
		return vec3{1, 1, 1}.muls(1 - t).add(vec3{0.5, 0.7, 1}.muls(t))
		// return closestLight.color
	}

	if closestEntity != (entity{}) {
		new, color := closestEntity.material.scatter(r, closestCollision)
		new.scatterCount = r.scatterCount + 1
		if new.scatterCount > 50 {
			println("scattered", new.scatterCount)
			return vec3{}
		}
		return color.mul(s.paint(new))
	}
	return vec3{}
}

func (s scene) render(dst *image.RGBA) {
	rand.Seed(s.seed)
	dx := dst.Bounds().Dx()
	dy := dst.Bounds().Dy()
	dsty := 0 // y coordinate of dst image which is flipped vertically compared to the scene.
	wg := sync.WaitGroup{}
	wg.Add(dy)
	for y := dy - 1; y >= 0; y-- {
		go func(y int, dsty int) {
			defer wg.Done()
			for x := 0; x < dx; x++ {
				col := vec3{0, 0, 0}
				for n := 0; n < s.samples; n++ {
					u := (float64(x) + randf64()) / float64(dx)
					v := (float64(y) + randf64()) / float64(dy)
					r := s.camera.emit(u, v)
					col = col.add(s.paint(r))
				}
				col = col.divs(float64(s.samples))
				// col = vec3{math.Sqrt(col.x), math.Sqrt(col.y), math.Sqrt(col.z)} // Gamma 2 correction.
				dst.Set(x, dsty, color.RGBA{uint8(255.99 * col.x), uint8(255.99 * col.y), uint8(255.99 * col.z), 255})
			}
			log.Printf("rendered %d/%d", dsty, dy)
		}(y, dsty)
		dsty++

	}
	wg.Wait()
}