ray tracer in a weekend

tracer.go

package main

import (
	"fmt"
	"image"
	"image/jpeg"
	"log"
	"math"
	"math/rand"
	"os"

	"github.com/larspensjo/Go-simplex-noise/simplexnoise"
)

func main() {
	Save(trace(), "image.jpg")
}

const SamplesPerPixel = 10
const MaxDepth = 20

func trace() image.Image {
	cam := Camera{}
	cam.SetImageFormat(1600, 16.0/9.0)
	cam.SetPosition()
	bounds := image.Rect(0, 0, cam.ImageWidth, cam.ImageHeight)
	img := image.NewNRGBA(bounds)

	world := make(HittableList, 0, 2)
	// world = append(world, Sphere{Vec3{0, 0, -1}, 0.5})
	// world = append(world, Sphere{Vec3{0, -100.5, -1}, 100})
	for i := 0; i < 500; i++ {
		world = append(world, Sphere{RandomInUnitSphere().Scaled(10), 0.5})
	}
	for y := 0; y < cam.ImageHeight; y++ {
		fmt.Printf("\rProgress: %d%%", int(y*100/cam.ImageHeight))
		for x := 0; x < cam.ImageWidth; x++ {
			col := Vec3{0, 0, 0}
			for s := 0; s < SamplesPerPixel; s++ {
				v := (float64(cam.ImageHeight-y-1) + rand.Float64()) / float64(cam.ImageHeight-1)
				u := (float64(x) + rand.Float64()) / float64(cam.ImageWidth-1)
				ray := cam.GetRay(u, v)
				col = col.Add(RayColor(ray, world, MaxDepth))
			}
			col = col.Scaled(1.0 / SamplesPerPixel).Gamma2()
			pos := (y-bounds.Min.Y)*img.Stride + (x-bounds.Min.X)*4
			col.WriteAsColor(img.Pix, pos)
		}
	}

	return img
}

type Camera struct {
	AspectRatio     float64
	ImageWidth      int
	ImageHeight     int
	ViewportWidth   float64
	ViewportHeight  float64
	FocalLenght     float64
	Origin          Vec3
	Horizontal      Vec3
	Vertical        Vec3
	lowerLeftCorner Vec3
}

func (c *Camera) SetImageFormat(width int, aspectRatio float64) {
	c.AspectRatio = aspectRatio
	c.ImageWidth = width
	c.ImageHeight = int(float64(width) / aspectRatio)

	c.ViewportHeight = 2.0
	c.ViewportWidth = aspectRatio * c.ViewportHeight
	c.FocalLenght = 1.0
}

func (c *Camera) SetPosition() {
	c.Origin = Vec3{0, 0, 0}
	c.Horizontal = Vec3{c.ViewportWidth, 0, 0}
	c.Vertical = Vec3{0, c.ViewportHeight, 0}
	c.lowerLeftCorner = c.Origin.Add(c.Horizontal.Scaled(-0.5)).Add(c.Vertical.Scaled(-0.5)).Add(Vec3{0, 0, -c.FocalLenght})
}

func (c Camera) GetRay(u, v float64) Ray {
	return Ray{c.Origin, c.lowerLeftCorner.Add(c.Horizontal.Scaled(u)).Add(c.Vertical.Scaled(v)).Subtract(c.Origin)}
}

const maximum = 1.0/1.0 + 1.0/2.0 + 1.0/3.0 + 1.0/4.0 + 1.0/5.0 + 1.0/6.0 + 1.0/7.0 + 1.0/8.0

func cloud(x, y, z float64) float64 {
	sum := 0.0
	for i := 0; i < 8; i++ {
		f := (float64(i) + 1.0) * 4
		sum += simplexnoise.Noise3(-5*f+x*f, y*f, z*f) / f * 4
	}
	return sum/maximum/2.0 + 0.5
}

// Return gradient from white to blue, depending on Y coordinate of ray
func RayColor(ray Ray, world Hittable, depth int) Vec3 {
	if depth <= 0 {
		return Vec3{0, 0, 0}
	}
	hit := world.Hit(ray, 0.001, 10000)
	unitDirection := ray.Direction.Unit()
	if hit != nil {
		// target := hit.Point.Add(hit.Normal).Add(RandomOnUnitSphere())
		// reflected := target.Subtract(hit.Point)
		reflected := reflect(unitDirection, hit.Normal)
		return Hadamar(Vec3{0.8, 0.8, 0.8}, RayColor(Ray{hit.Point, reflected}, world, depth-1))
	}
	// t := 0.5 * (ray.Direction.Unit().Y + 1.0)
	t := cloud(unitDirection.X, unitDirection.Y, unitDirection.Z)
	return Vec3{1.0, 1.0, 1.0}.Scaled(1.0 - t).Add(Vec3{0.5, 0.7, 1.0}.Scaled(t))
}

type Hittable interface {
	Hit(r Ray, t_min, t_max float64) *HitRecord
}

type Material interface {
	Scatter(ray Ray, hit HitRecord) float64 // TODO
}

type HitRecord struct {
	Point     Vec3
	Normal    Vec3
	t         float64
	FrontFace bool
	Material  Material
}

func (hr *HitRecord) SetFaceNormal(r Ray, outwardNormal Vec3) {
	hr.FrontFace = dot(r.Direction, outwardNormal) < 0
	if hr.FrontFace {
		hr.Normal = outwardNormal
	} else {
		hr.Normal = outwardNormal.Scaled(-1)
	}
}

type Sphere struct {
	Center Vec3
	Radius float64
}

func (s Sphere) Hit(r Ray, t_min, t_max float64) *HitRecord {
	oc := r.Origin.Subtract(s.Center)
	a := r.Direction.LenghtSqr()
	half_b := dot(oc, r.Direction)
	c := oc.LenghtSqr() - s.Radius*s.Radius
	discriminant := half_b*half_b - a*c

	rec := &HitRecord{}
	if discriminant > 0 {
		root := math.Sqrt(discriminant)
		temp := (-half_b - root) / a
		if temp < t_max && temp > t_min {
			rec.t = temp
			rec.Point = r.At(rec.t)
			rec.SetFaceNormal(r, rec.Point.Subtract(s.Center).Scaled(1/s.Radius))
			return rec
		}
		temp = (-half_b + root) / a
		if temp < t_max && temp > t_min {
			rec.t = temp
			rec.Point = r.At(rec.t)
			rec.SetFaceNormal(r, rec.Point.Subtract(s.Center).Scaled(1/s.Radius))
			return rec
		}
	}
	return nil
}

type HittableList []Hittable

func (hl HittableList) Hit(r Ray, t_min, t_max float64) *HitRecord {
	closest_so_far := t_max
	var res *HitRecord

	for _, object := range hl {
		hit := object.Hit(r, t_min, closest_so_far)
		if hit != nil {
			closest_so_far = hit.t
			res = hit
		}
	}

	return res
}

type Vec3 struct {
	X float64
	Y float64
	Z float64
}

// WriteAsColor sets value of pixel of NRGBA image in given position
func (v Vec3) WriteAsColor(pixels []uint8, pos int) {
	pixels[pos] = uint8(v.X * 255.0)
	pixels[pos+1] = uint8(v.Y * 255.0)
	pixels[pos+2] = uint8(v.Z * 255.0)
	pixels[pos+3] = 255 // Full opacity
}

func RandomVec() Vec3 {
	return Vec3{rand.Float64(), rand.Float64(), rand.Float64()}.Scaled(2.0).Subtract(Vec3{1, 1, 1})
}

func RandomInUnitSphere() Vec3 {
	for {
		v := RandomVec()
		if v.LenghtSqr() <= 1.0 {
			return v
		}
	}
}

func RandomOnUnitSphere() Vec3 {
	return RandomInUnitSphere().Unit()
}

func (v Vec3) LenghtSqr() float64 {
	return v.X*v.X + v.Y*v.Y + v.Z*v.Z
}

func (v Vec3) Lenght() float64 {
	return math.Sqrt(v.LenghtSqr())
}

func (v Vec3) Scaled(s float64) Vec3 {
	return Vec3{v.X * s, v.Y * s, v.Z * s}
}

func (v Vec3) Unit() Vec3 {
	return v.Scaled(1 / v.Lenght())
}

func (v Vec3) Add(v2 Vec3) Vec3 {
	return Vec3{v.X + v2.X, v.Y + v2.Y, v.Z + v2.Z}
}

func (v Vec3) Subtract(v2 Vec3) Vec3 {
	return Vec3{v.X - v2.X, v.Y - v2.Y, v.Z - v2.Z}
}

func reflect(v, n Vec3) Vec3 {
	return v.Subtract(n.Scaled(-2 * dot(v, n)))
}

func Hadamar(u, v Vec3) Vec3 {
	return Vec3{u.X * v.X, u.Y * v.Y, u.Z * v.Z}
}

func dot(u, v Vec3) float64 {
	return u.X*v.X + u.Y*v.Y + u.Z*v.Z
}

func (v Vec3) Gamma2() Vec3 {
	return Vec3{
		math.Sqrt(v.X),
		math.Sqrt(v.Y),
		math.Sqrt(v.Z),
	}
}

type Ray struct {
	Origin    Vec3
	Direction Vec3
}

// At return point at a distance from ray
func (r Ray) At(d float64) Vec3 {
	return r.Origin.Add(r.Direction.Scaled(d))
}

func Save(img image.Image, filename string) {
	f, err := os.Create(filename)
	last(err)

	last(jpeg.Encode(f, img, &jpeg.Options{Quality: 90}))

	defer f.Close()
}

func last(err error) {
	if err != nil {
		log.Fatal(err)
	}
}