aobench.go

package main

import (
  "image"
	"image/color"
	"image/png"
	"math"
	"math/rand"
	"os"
)

const (
	WIDTH       = 256
	HEIGHT      = 256
	NSUBSAMPLES = 2
	NAO_SAMPLES = 8
)

type Vec struct {
	x float64
	y float64
	z float64
}

type Isect struct {
	t   float64
	p   Vec
	n   Vec
	hit int
}

type Sphere struct {
	center Vec
	radius float64
}

type Plane struct {
	p Vec
	n Vec
}

type Ray struct {
	org Vec
	dir Vec
}

var spheres [3]Sphere
var plane Plane

func dot(v0 Vec, v1 Vec) float64 {
	return v0.x*v1.x + v0.y*v1.y + v0.z*v1.z
}

func cross(v0 Vec, v1 Vec) (c Vec) {
	c.x = v0.y*v1.z - v0.z*v1.y
	c.y = v0.z*v1.x - v0.x*v1.z
	c.z = v0.x*v1.y - v0.y*v1.x
	return c
}

func normalize(v Vec) (c Vec) {
	c = v
	length := math.Sqrt(dot(c, c))

	if math.Abs(length) > 1.0e-17 {
		c.x /= length
		c.y /= length
		c.z /= length
	}
	return c
}

func (isect *Isect) raySphereIntersect(ray Ray, sphere Sphere) {
	var rs Vec
	rs.x = ray.org.x - sphere.center.x
	rs.y = ray.org.y - sphere.center.y
	rs.z = ray.org.z - sphere.center.z

	B := dot(rs, ray.dir)
	C := dot(rs, rs) - sphere.radius*sphere.radius
	D := B*B - C

	if D > 0.0 {
		t := -B - math.Sqrt(D)

		if (t > 0.0) && (t < isect.t) {
			isect.t = t
			isect.hit = 1

			isect.p.x = ray.org.x + ray.dir.x*t
			isect.p.y = ray.org.y + ray.dir.y*t
			isect.p.z = ray.org.z + ray.dir.z*t

			isect.n.x = isect.p.x - sphere.center.x
			isect.n.y = isect.p.y - sphere.center.y
			isect.n.z = isect.p.z - sphere.center.z

			isect.n = normalize(isect.n)
		}
	}
}

func (isect *Isect) rayPlaneIntersect(ray Ray, plane Plane) {
	d := -dot(plane.p, plane.n)
	v := dot(ray.dir, plane.n)

	if math.Abs(v) < 1.0e-17 {
		return
	}

	t := -(dot(ray.org, plane.n) + d) / v

	if (t > 0.0) && (t < isect.t) {
		isect.t = t
		isect.hit = 1

		isect.p.x = ray.org.x + ray.dir.x*t
		isect.p.y = ray.org.y + ray.dir.y*t
		isect.p.z = ray.org.z + ray.dir.z*t

		isect.n = plane.n
	}
}

func orthoBasis(n Vec) (basis [3]Vec) {
	basis[2] = n
	basis[1].x = 0.0
	basis[1].y = 0.0
	basis[1].z = 0.0

	if (n.x < 0.6) && (n.x > -0.6) {
		basis[1].x = 1.0
	} else if (n.y < 0.6) && (n.y > -0.6) {
		basis[1].y = 1.0
	} else if (n.z < 0.6) && (n.z > -0.6) {
		basis[1].z = 1.0
	} else {
		basis[1].x = 1.0
	}

	basis[0] = cross(basis[1], basis[2])
	basis[0] = normalize(basis[0])

	basis[1] = cross(basis[2], basis[0])
	basis[1] = normalize(basis[1])

	return basis
}

func ambientOcclusion(isect Isect) (col Vec) {
	ntheta := NAO_SAMPLES
	nphi := NAO_SAMPLES
	eps := 0.0001

	var p Vec
	p.x = isect.p.x + eps*isect.n.x
	p.y = isect.p.y + eps*isect.n.y
	p.z = isect.p.z + eps*isect.n.z

	basis := orthoBasis(isect.n)

	occlusion := 0.0

	for j := 0; j < ntheta; j++ {
		for i := 0; i < nphi; i++ {
			theta := math.Sqrt(rand.Float64())
			phi := 2.0 * math.Pi * rand.Float64()

			x := math.Cos(phi) * theta
			y := math.Sin(phi) * theta
			z := math.Sqrt(1.0 - theta*theta)

			rx := x*basis[0].x + y*basis[1].x + z*basis[2].x
			ry := x*basis[0].y + y*basis[1].y + z*basis[2].y
			rz := x*basis[0].z + y*basis[1].z + z*basis[2].z

			var ray Ray
			ray.org = p
			ray.dir.x = rx
			ray.dir.y = ry
			ray.dir.z = rz

			var occIsect Isect
			occIsect.t = 1.0e+17
			occIsect.hit = 0

			occIsect.raySphereIntersect(ray, spheres[0])
			occIsect.raySphereIntersect(ray, spheres[1])
			occIsect.raySphereIntersect(ray, spheres[2])
			occIsect.rayPlaneIntersect(ray, plane)

			if occIsect.hit > 0.0 {
				occlusion += 1.0
			}
		}
	}

	occlusion = (float64(ntheta)*float64(nphi) - occlusion) / float64(ntheta*nphi)

	col.x = occlusion
	col.y = occlusion
	col.z = occlusion

	return col
}

func clamp(f float64) uint8 {
	i := int(f * 255)

	if i < 0 {
		i = 0
	}
	if i > 255 {
		i = 255
	}

	return uint8(i)
}

func render(img []uint8, w int, h int, nsubsamples int) {
	fimg := make([]float64, w*h*3)

	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {

			for v := 0; v < nsubsamples; v++ {
				for u := 0; u < nsubsamples; u++ {
					px := (float64(x) + (float64(u) / float64(nsubsamples)) - (float64(w) / 2.0)) / (float64(w) / 2.0)
					py := -(float64(y) + (float64(v) / float64(nsubsamples)) - (float64(h) / 2.0)) / (float64(h) / 2.0)

					var ray Ray
					ray.dir.x = px
					ray.dir.y = py
					ray.dir.z = -1.0
					ray.dir = normalize(ray.dir)

					var isect Isect
					isect.t = 1.0e+17
					isect.hit = 0

					isect.raySphereIntersect(ray, spheres[0])
					isect.raySphereIntersect(ray, spheres[1])
					isect.raySphereIntersect(ray, spheres[2])
					isect.rayPlaneIntersect(ray, plane)

					if isect.hit > 0 {
						col := ambientOcclusion(isect)

						fimg[3*(y*w+x)+0] += col.x
						fimg[3*(y*w+x)+1] += col.y
						fimg[3*(y*w+x)+2] += col.z
					}
				}
			}

			fimg[3*(y*w+x)+0] /= float64(nsubsamples * nsubsamples)
			fimg[3*(y*w+x)+1] /= float64(nsubsamples * nsubsamples)
			fimg[3*(y*w+x)+2] /= float64(nsubsamples * nsubsamples)

			img[3*(y*w+x)+0] = clamp(fimg[3*(y*w+x)+0])
			img[3*(y*w+x)+1] = clamp(fimg[3*(y*w+x)+1])
			img[3*(y*w+x)+2] = clamp(fimg[3*(y*w+x)+2])
		}
	}
}

func savePng(fname string, w int, h int, img []uint8) {
	i := image.NewRGBA(image.Rect(0, 0, w, h))

	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			r := img[3*(y*w+x)+0]
			g := img[3*(y*w+x)+1]
			b := img[3*(y*w+x)+2]
			i.Set(x, y, color.RGBA{r, g, b, 255})
		}
	}

	f, e := os.OpenFile(fname, os.O_CREATE|os.O_WRONLY, 0666)
	if e != nil {
		panic(e)
	}
	defer f.Close()

	png.Encode(f, i)
}

func initScene() {
	spheres[0].center.x = -2.0
	spheres[0].center.y = 0.0
	spheres[0].center.z = -3.5
	spheres[0].radius = 0.5

	spheres[1].center.x = -0.5
	spheres[1].center.y = 0.0
	spheres[1].center.z = -3.0
	spheres[1].radius = 0.5

	spheres[2].center.x = 1.0
	spheres[2].center.y = 0.0
	spheres[2].center.z = -2.2
	spheres[2].radius = 0.5

	plane.p.x = 0.0
	plane.p.y = -0.5
	plane.p.z = 0.0

	plane.n.x = 0.0
	plane.n.y = 1.0
	plane.n.z = 0.0
}

func main() {
	img := make([]uint8, WIDTH*HEIGHT*3)
	initScene()
	render(img, WIDTH, HEIGHT, NSUBSAMPLES)
	savePng("ao.png", WIDTH, HEIGHT, img)
}