RustでRayTracing(チャプター11-最後まで)

チャプター11はカメラの場所を動かしたときに、どのようにrayへ影響させるかでした。
ベクトルの演算がわかれば、すごく素直な実装だなと思いました。

チャプター12では焦点を合わせるというお話で、非常に勉強になりました。
カメラとかをやってると、理解しやすいのかなと思います。
光をどのように集めるのかで、画像に影響が出るんですね。

チャプター13では、最後に多数の球を置いてのレンダリングでした。
速度を意識せずにコードを作ってたので、レンダリングの時間が43分ぐらいでした。

次は、コードを整理して、どこまで高速化できるのかをやってみたいと思います。

最後のコードをのせておきますね。

use std::fmt;
use rand::{thread_rng, Rng};
use rand::distributions::{Distribution, Standard};
use rand::distributions::uniform::SampleUniform;
use cgmath::prelude::*;
use cgmath::{Vector3, vec3, dot};

type Color = Vector3<i32>;

pub struct HitRecord<'a> {
    pub p: Vector3<f64>,
    pub normal: Vector3<f64>,
    pub mat_ptr: Option<&'a Box<dyn Material>>,
    pub t: f64,
    pub front_face: bool
}

impl HitRecord<'_> {
    pub fn set_face_normal(&mut self, ray: &Ray, outward_normal: Vector3<f64>) {
        self.front_face = dot(ray.direction, outward_normal) < 0.0;
        self.normal = if self.front_face {
            outward_normal
        } else {
            -outward_normal
        }
    }
}

pub trait HitTable {
    fn hit(&self, ray: &Ray, t_min: f64, t_max: f64) -> Option<HitRecord>;
}

pub struct HitTableList {
    list: Vec<Box<dyn HitTable>>
}

impl HitTableList {
    pub fn new() -> Self {
        HitTableList {
            list: vec![],
        }
    }

    pub fn push(&mut self, hitable: impl HitTable + 'static) {
        self.list.push(Box::new(hitable))
    }

    pub fn clear(&mut self) {
        self.list.clear()
    }
}

impl HitTable for HitTableList {
    fn hit(&self, ray: &Ray, t_min: f64, t_max: f64) -> Option<HitRecord> {
        let mut closest_so_far = t_max;
        let mut hit_anything: Option<HitRecord> = None;
        for h in self.list.iter() {
            if let Some(hit) = h.hit(ray, t_min, closest_so_far) {
                closest_so_far = hit.t;
                hit_anything = Some(hit);
            }
        }
        hit_anything
    }
}

pub struct Sphere {
    center: Vector3<f64>,
    radius: f64,
    mat_ptr: Option<Box<dyn Material>>,
}

impl Sphere {
    pub fn new(center: Vector3<f64>, radius: f64, mat_ptr: Option<Box<dyn Material>>) -> Self {
        Sphere {center, radius, mat_ptr}
    }
}

impl HitTable for Sphere {
    fn hit(&self, ray: &Ray, t_min: f64, t_max: f64) -> Option<HitRecord> {
        let oc = ray.origin - self.center;
        let a = length_squared(ray.direction);
        let half_b = dot(oc, ray.direction);
        let c = length_squared(oc) - self.radius * self.radius;
    
        let discriminant = half_b.powi(2) - a * c;

        if discriminant > 0.0 {
            let sqrtd = discriminant.sqrt();

            // Find the nearest root that lies in the acceptable range.
            let mut root = (- half_b - sqrtd) / a;
            if root < t_min || t_max < root {
                root = (- half_b + sqrtd) / a;
                if root < t_min || t_max < root {
                    return None
                }
            }

            let p = ray.at(root);
            let outward_normal = (p - self.center) / self.radius;

            // set face normal 
            let front_face = dot(ray.direction, outward_normal) < 0.0;
            let normal = if front_face {outward_normal} else {-outward_normal};

            return Some(HitRecord{t: root, p: p, normal: normal, front_face: front_face, mat_ptr: self.mat_ptr.as_ref()});
        }

        None
    }
}

pub struct Camera {
    pub origin: Vector3<f64>,
    pub horizontal: Vector3<f64>,
    pub vertical: Vector3<f64>,
    pub lower_left_corner: Vector3<f64>,
    pub lens_radius: f64,
    pub w: Vector3<f64>,
    pub u: Vector3<f64>,
    pub v: Vector3<f64>,
}

impl Camera {
    pub fn new(lookfrom: Vector3<f64>, lookat: Vector3<f64>, vup: Vector3<f64>, vfov: f64, aspect_ratio: f64, aperture: f64, focus_dist: f64) -> Self {
        let theta = vfov.to_radians();
        let h = (theta / 2.0).tan();
        let viewport_height: f64 = 2.0 * h;
        let viewport_width = aspect_ratio * viewport_height;

        let w = (lookfrom - lookat).normalize();
        let u = vup.cross(w).normalize();
        let v = w.cross(u);


        let origin: Vector3<f64> = lookfrom;
        let horizontal: Vector3<f64> = focus_dist * viewport_width * u;
        let vertical: Vector3<f64> = focus_dist * viewport_height * v;
        let lower_left_corner = origin - horizontal / 2.0 - vertical / 2.0 - focus_dist * w;
        let lens_radius = aperture / 2.0;

        Camera {
            origin,
            horizontal,
            vertical,
            lower_left_corner,
            lens_radius,
            w, u, v,
        }
    }

    pub fn get_ray(&self, s: f64, t: f64) -> Ray {
        let rd = self.lens_radius * random_in_unit_disk();
        let offset = self.u * rd.x + self.v * rd.y;

        Ray::new(self.origin + offset, self.lower_left_corner + s * self.horizontal + t * self.vertical - self.origin - offset)
    }
}

pub trait Material {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> (bool, Vector3<f64>, Ray);
}

pub struct Lambertian {
    albedo: Vector3<f64>,
}

impl Lambertian {
    pub fn new(albedo: Vector3<f64>) -> Self {
        Lambertian{albedo}
    }
}

pub struct Metal {
    albedo: Vector3<f64>,
    fuzz: f64,
}

impl Metal {
    pub fn new(albedo: Vector3<f64>, f: f64) -> Self {
        Metal{albedo, fuzz: f.min(1.0)}
    }
}

impl Material for Metal {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> (bool, Vector3<f64>, Ray) {
        let reflected = reflect(r_in.direction.normalize(), rec.normal);
        let scattered = Ray::new(rec.p, reflected + self.fuzz * random_in_unit_sphere());

        (dot(scattered.direction, rec.normal) > 0.0, self.albedo, scattered)
    }
}

pub fn reflect(v:Vector3<f64>, n: Vector3<f64>) -> Vector3<f64> {
    v - 2.0 * dot(v, n) * n
}

impl Material for Lambertian {
    fn scatter(&self, _r_in: &Ray, rec: &HitRecord) -> (bool, Vector3<f64>, Ray) {
        let scatter_direction = rec.normal + random_unit_vector();

        let scattered = if near_zero(scatter_direction) {
            Ray::new(rec.p, rec.normal)
        }
        else {
            Ray::new(rec.p, scatter_direction)
        };

        (true, self.albedo, scattered)
    }
}

pub fn near_zero(v: Vector3<f64>) -> bool {
    const s:f64 = 1e-8;
    v.x.abs() < s && v.y.abs() < s && v.z.abs() < s

}

pub fn write_color(pixel_color: Vector3<f64>, samples_per_pixel: i32) {
    let scale = 1.0 / samples_per_pixel as f64;

    let color:Color = pixel_color.map(|e| 256.0 * (e * scale).sqrt().min(0.999).max(0.0)).cast().unwrap();

    println!("{} {} {}", color.x, color.y, color.z);
}

pub struct Ray {
    pub origin: Vector3<f64>,
    pub direction: Vector3<f64>
}

impl Ray {
    pub fn new(o:Vector3<f64>, d:Vector3<f64>) -> Self {
        Ray{origin: o, direction: d}
    }

    pub fn at(&self, t: f64) -> Vector3<f64> {
        self.origin + self.direction * t
    }
}

impl fmt::Debug for Ray {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({:?}, {:?})", self.origin, self.direction)
    }
}

fn ray_color(ray: &Ray, world: &dyn HitTable, depth: i32) -> Vector3<f64>
{
    const white: Vector3<f64> = vec3(1.0, 1.0, 1.0);
    const blue: Vector3<f64> = vec3(0.5, 0.7, 1.0);

    if depth <= 0 {
        return Vector3::zero();
    }

    if let Some(hit) = world.hit(ray, 0.001, f64::INFINITY) {

        let (cond, a, scattered) = hit.mat_ptr.as_ref().unwrap().scatter(ray, &hit);
        if cond {
            let n = ray_color(&scattered, world, depth - 1);
            return vec3(a.x * n.x, a.y * n.y, a.z *n.z)
        }
        else {
            return Vector3::zero();
        }
    }

    let unit_direction = ray.direction.normalize();
    let t = 0.5 * (unit_direction.y + 1f64);
    let rgb = white.lerp(blue, t);

    return rgb
}

fn length_squared(v: Vector3<f64>) -> f64 {
    v.x * v.x + v.y * v.y + v.z * v.z
}

fn hit_sphere(center: Vector3<f64>, radius: f64, ray:&Ray) -> f64 {
    let oc = ray.origin - center;
    let a = length_squared(ray.direction);
    let half_b = dot(oc, ray.direction);
    let c = length_squared(oc) - radius * radius;

    let discriminant = half_b * half_b - a * c;

    if discriminant < 0.0 {
        return -1.0;
    }
    else {
        return (- half_b - discriminant.sqrt()) / a;
    }
}

fn random_vector3<T>() -> Vector3<T>
where Standard: Distribution<T>
{
    vec3(rand::random::<T>(), rand::random::<T>(), rand::random::<T>())
}

fn random_vector3_range<T>(min: T, max: T) -> Vector3<T>
where Standard: Distribution<T>, T: SampleUniform + PartialOrd + Copy
{
    let mut rng = thread_rng();
    vec3(rng.gen_range(min..max), rng.gen_range(min..max), rng.gen_range(min..max))
}

pub fn random_unit_vector() -> Vector3<f64> {
    random_in_unit_sphere().normalize()
}

fn random_in_range(min: f64, max: f64) -> f64 {
    let mut rng = thread_rng();
    rng.gen_range(min..max)
}

pub fn random_in_hemisphere(normal: Vector3<f64>) -> Vector3<f64> {
    let in_unit_sphere = random_in_unit_sphere();
    if dot(in_unit_sphere, normal) > 0.0 {
        return in_unit_sphere;
    }
    else {
        return - in_unit_sphere;
    }
}

fn random_in_unit_sphere() -> Vector3<f64> {
    loop {
        let p = random_vector3_range(-1.0, 1.0);

        if length_squared(p) < 1.0 {
            return p;
        }
    }
}

fn random_in_unit_disk() -> Vector3<f64> {
    loop {
        let p0 = random_vector3_range(-1.0, 1.0);
        let p = vec3(p0.x, p0.y, 0.0);

        if length_squared(p) < 1.0 {
            return p;
        }
    }
}

pub fn refract(uv: Vector3<f64>, n: Vector3<f64>, etai_over_etat: f64) -> Vector3<f64> {
    let cos_theta = dot(-uv, n);
    let r_out_perp = etai_over_etat * (uv + cos_theta * n);
    let r_out_parallel = - (1.0 - length_squared(r_out_perp)).abs().sqrt() * n;
    return r_out_perp + r_out_parallel;
}

pub struct Dielectric {
    ir: f64,
}

impl Dielectric {
    pub fn new(ir: f64) -> Self {
        Dielectric {ir}
    }

    fn reflectance(cosine: f64, ref_idx: f64) -> f64 {
        let r0 = ((1.0 - ref_idx) / (1.0 + ref_idx)).powi(2);
        r0 + (1.0 - r0) * (1.0 - cosine).powi(5)
    }

}

impl Material for Dielectric {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> (bool, Vector3<f64>, Ray) {
        let refraction_ratio = if rec.front_face {
            1.0 / self.ir
        } else {
            self.ir
        };

        let unit_direction = r_in.direction.normalize();
        let cos_theta = dot(-unit_direction, rec.normal).min(1.0);
        let sin_theta = (1.0 - cos_theta * cos_theta).sqrt();

        let direction = if refraction_ratio * sin_theta > 1.0 || Dielectric::reflectance(cos_theta, refraction_ratio) > rand::random() {
            reflect(unit_direction, rec.normal)
        } else {
            refract(unit_direction, rec.normal, refraction_ratio)
        };

        let scattered = Ray::new(rec.p, direction);
        let attenuation = vec3(1.0, 1.0, 1.0);

        (true, attenuation, scattered)
    }
}

fn random_scene() -> HitTableList {
    let mut world = HitTableList::new();

    for a in -11..11 {
        for b in -11..11 {
            let choose_mat: f64 = rand::random();
            let center = vec3(a as f64 + 0.9 * rand::random::<f64>(), 0.2, b as f64 + 0.9 * rand::random::<f64>());

            if length_squared(center - vec3(4.0, 0.2, 0.0)).sqrt() > 0.9 {
                let sphere_material: Box<dyn Material> = if choose_mat < 0.8 {
                    let albedo = random_vector3();
                    Box::new(Lambertian::new(albedo))
                } else if choose_mat < 0.95 {
                    let albedo = random_vector3_range(0.5, 1.0);
                    let fuzz = random_in_range(0.0, 0.5);
                    Box::new(Metal::new(albedo, fuzz))
                } else {
                    Box::new(Dielectric::new(1.5))
                };
                world.push(Sphere::new(center, 0.2, Some(sphere_material)));
            }
        }
    }

    let ground_material = Box::new(Lambertian::new(vec3(0.5, 0.5, 0.5)));
    world.push(Sphere::new(vec3(0.0, -1000.0, 0.0), 1000.0, Some(ground_material)));

    let material1 = Box::new(Dielectric::new(1.5));
    world.push(Sphere::new(vec3(0.0, 1.0, 0.0), 1.0, Some(material1)));

    let material2 = Box::new(Lambertian::new(vec3(0.4, 0.2, 0.1)));
    world.push(Sphere::new(vec3(-4.0, 1.0, 0.0), 1.0, Some(material2)));

    let material3 = Box::new(Metal::new(vec3(0.7, 0.6, 0.5), 0.0));
    world.push(Sphere::new(vec3(4.0, 1.0, 0.0), 1.0, Some(material3)));

    world
}

fn main() {

    // Image

    const aspect_ratio: f64 = 3.0 / 2.0;
    const image_width: i32 = 1200;
    const image_height: i32 = (image_width as f64 / aspect_ratio) as i32;
    const samples_per_pixel: i32 = 500;
    const max_depth: i32 = 50;

    // World

    let world = random_scene();

    // Camera

    let lookfrom: Vector3<f64> = vec3(13.0, 2.0, 3.0);
    let lookat: Vector3<f64> = vec3(0.0, 0.0, 0.0);
    let vup: Vector3<f64> = vec3(0.0, 1.0, 0.0);
    let dist_to_focus: f64 = 10.0;
    let aperture: f64 = 0.1;

    let camera = Camera::new(lookfrom, lookat, vup, 20.0, aspect_ratio, aperture, dist_to_focus);

    // Render

    println!("P3\n {} {} \n255", image_width, image_height);

    for j in (0..image_height).rev() {
        for i in 0..image_width {
            let mut pixel_color: Vector3<f64> = Vector3::zero();
            for s in 0..samples_per_pixel {
                let u = (i as f64 + rand::random::<f64>()) / (image_width - 1) as f64;
                let v = (j as f64 + rand::random::<f64>()) / (image_height - 1) as f64;
                let r = camera.get_ray(u, v);
                pixel_color += ray_color(&r, &world, max_depth);
            }
            write_color(pixel_color, samples_per_pixel);
        }
    }
}