ray-tracing2/ray-tracing-scene/src/examples/mis_test.rs

use std::cell::OnceCell;

use crate::triangle_bvh::{BVHMaterial, Triangle, TriangleBVH};

use super::ExampleScene;

use ray_tracing_core::{light::AreaLight, prelude::*, scene::Scene};
use ray_tracing_material::{
    diffuse::DiffuseMaterial,
    microfacet::{BeckmannDistribution, Microfacet},
    mirror::Mirror,
};

pub struct MISTest<R: Rng> {
    scene: OnceCell<TriangleBVH<R>>,
    dist: Float,
    light_height: Float,
    camera_height: Float,
    camera_dist: Float,
    num_lights: usize,
    plane_width: Float,
    plane_depth: Float,
    light_min: Float,
    light_max: Float,
    materials: fn() -> Vec<BVHMaterial<R>>,
}

impl<R: Rng> MISTest<R> {
    pub fn new() -> Self {
        MISTest {
            scene: OnceCell::new(),
            dist: 2.0,
            light_height: 6.0,
            camera_height: 10.0,
            camera_dist: 6.0,
            num_lights: 4,
            plane_width: 10.0,
            plane_depth: 0.5,
            light_min: 0.05,
            light_max: 0.7,
            materials: || {
                vec![
                    BVHMaterial::new_material(Mirror::new(Color::white())),
                    BVHMaterial::new_material(Microfacet::new(
                        BeckmannDistribution::new(0.005175),
                        Color::white(),
                    )),
                    BVHMaterial::new_material(Microfacet::new(
                        BeckmannDistribution::new(0.0125),
                        Color::white(),
                    )),
                    BVHMaterial::new_material(Microfacet::new(
                        BeckmannDistribution::new(0.025),
                        Color::white(),
                    )),
                    BVHMaterial::new_material(Microfacet::new(
                        BeckmannDistribution::new(0.05),
                        Color::white(),
                    )),
                ]
            },
        }
    }

    fn camera_x(&self) -> Float {
        -((((self.materials)().len() - 1) as Float) * 0.5 + self.camera_dist)
    }
}

impl<R: Rng + 'static> ExampleScene<R> for MISTest<R> {
    fn get_scene(&self) -> Box<dyn Scene<R> + Sync> {
        let s = self.scene.get_or_init(move || {
            let mut vertices = Vec::new();
            let mut triangles = Vec::new();
            let mut materials = Vec::new();

            for i in 0..self.num_lights {
                let center_z = self.dist * ((i as Float) - 0.5 * ((self.num_lights - 1) as Float));
                let light_offset = 0.5
                    * (self.light_min
                        + ((i as Float) / ((self.num_lights - 1) as Float))
                            * (self.light_max - self.light_min));

                let vertex_offset = vertices.len() as u32;
                vertices.extend_from_slice(&[
                    Pos3::new(-light_offset, self.light_height, -light_offset + center_z),
                    Pos3::new(light_offset, self.light_height, -light_offset + center_z),
                    Pos3::new(-light_offset, self.light_height, light_offset + center_z),
                    Pos3::new(light_offset, self.light_height, light_offset + center_z),
                ]);

                let area = light_offset * light_offset * 4.0;
                let material_offset = materials.len() as u32;

                let color = match i % 3 {
                    0 => Color::new(1.0, 0.0, 0.0),
                    1 => Color::new(0.0, 1.0, 0.0),
                    2 => Color::new(0.0, 0.0, 1.0),
                    _ => unreachable!(),
                };
                materials.push(BVHMaterial::new_light(AreaLight::new(color * 10.0 / area)));

                triangles.extend_from_slice(&[
                    Triangle::new(
                        [vertex_offset, vertex_offset + 1, vertex_offset + 3],
                        material_offset,
                    ),
                    Triangle::new(
                        [vertex_offset, vertex_offset + 3, vertex_offset + 2],
                        material_offset,
                    ),
                ]);
            }

            let mats = (self.materials)();
            let num_materials = mats.len();

            for (i, mat) in mats.into_iter().enumerate() {
                let center_x = self.dist * ((i as Float) - 0.5 * ((num_materials - 1) as Float));

                let theta_light = Float::atan2(center_x, self.light_height);
                let theta_camera = Float::atan2(center_x - self.camera_x(), self.camera_height);

                let theta = (theta_camera + theta_light) * 0.5;

                let offset_x = self.plane_depth * Float::cos(theta);
                let offset_y = self.plane_depth * Float::sin(theta);

                let vertex_offset = vertices.len() as u32;
                vertices.extend_from_slice(&[
                    Pos3::new(center_x + offset_x, offset_y, -0.5 * self.plane_width),
                    Pos3::new(center_x - offset_x, -offset_y, -0.5 * self.plane_width),
                    Pos3::new(center_x + offset_x, offset_y, 0.5 * self.plane_width),
                    Pos3::new(center_x - offset_x, -offset_y, 0.5 * self.plane_width),
                ]);

                let material_offset = materials.len() as u32;
                materials.push(mat);

                triangles.extend_from_slice(&[
                    Triangle::new(
                        [vertex_offset, vertex_offset + 1, vertex_offset + 3],
                        material_offset,
                    ),
                    Triangle::new(
                        [vertex_offset, vertex_offset + 3, vertex_offset + 2],
                        material_offset,
                    ),
                ]);
            }

            TriangleBVH::new(vertices, triangles, materials)
        });

        Box::new(s.clone()) as Box<dyn Scene<R> + Sync>
    }

    fn get_camera_pos(&self) -> Pos3 {
        Pos3::new(self.camera_x(), self.camera_height, 0.0)
    }

    fn get_camera_look_at(&self) -> Pos3 {
        Pos3::new(0.0, 0.0, 0.0)
    }

    fn get_camera_up(&self) -> Dir3 {
        Dir3::up()
    }

    fn get_horizontal_fov(&self) -> Float {
        Float::to_radians(90.0)
    }
}