ai_pathtracer/rust/src/scene.rs
2026-09-26 08:14:00 +02:00

266 lines
9 KiB
Rust

use rand::Rng;
use rand_distr::StandardNormal;
use crate::vec3::{reflect, refract, schlick, Vec3};
pub const EPSILON: f64 = 1e-6;
pub fn random_unit_vector<R: Rng>(rng: &mut R) -> Vec3 {
loop {
let v = Vec3::new(
rng.sample(StandardNormal),
rng.sample(StandardNormal),
rng.sample(StandardNormal),
);
let n = v.length();
if n > 1e-12 {
return v / n;
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct Ray {
pub origin: Vec3,
pub direction: Vec3,
}
impl Ray {
pub fn new(origin: Vec3, direction: Vec3) -> Self {
Ray {
origin,
direction: direction.normalize(),
}
}
pub fn at(&self, t: f64) -> Vec3 {
self.origin + t * self.direction
}
}
#[derive(Clone, Copy, Debug)]
pub enum Material {
Diffuse { albedo: Vec3 },
Mirror { albedo: Vec3 },
Dielectric { ref_idx: f64, albedo: Vec3 },
Light { emission: Vec3 },
}
pub struct HitRecord {
pub point: Vec3,
pub normal: Vec3,
pub front_face: bool,
pub material: Material,
}
impl HitRecord {
/// Orients the normal against the incoming ray, as in the Python version.
fn new(r: &Ray, t: f64, outward: Vec3, material: Material) -> Self {
let front_face = r.direction.dot(outward) < 0.0;
HitRecord {
point: r.at(t),
normal: if front_face { outward } else { -outward },
front_face,
material,
}
}
}
impl Material {
/// Returns (attenuation, scattered ray), or None if the ray is absorbed.
pub fn scatter<R: Rng>(&self, ray_in: &Ray, rec: &HitRecord, rng: &mut R) -> Option<(Vec3, Ray)> {
match *self {
Material::Diffuse { albedo } => {
// normal + uniform unit vector is cosine-distributed, so the
// Lambertian BRDF * cos / pdf reduces to just the albedo.
let mut target = rec.normal + random_unit_vector(rng);
if target.length() < 1e-8 {
target = rec.normal;
}
Some((albedo, Ray::new(rec.point, target)))
}
Material::Mirror { albedo } => {
let reflected = reflect(ray_in.direction, rec.normal);
if reflected.dot(rec.normal) <= 0.0 {
return None;
}
Some((albedo, Ray::new(rec.point, reflected)))
}
Material::Dielectric { ref_idx, albedo } => {
let ratio = if rec.front_face { 1.0 / ref_idx } else { ref_idx };
let d = ray_in.direction;
let cos_theta = (-d.dot(rec.normal)).min(1.0);
let sin_theta = (1.0 - cos_theta * cos_theta).max(0.0).sqrt();
let cannot_refract = ratio * sin_theta > 1.0;
let direction = if cannot_refract || schlick(cos_theta, ratio) > rng.gen::<f64>() {
reflect(d, rec.normal)
} else {
refract(d, rec.normal, ratio).unwrap_or_else(|| reflect(d, rec.normal))
};
Some((albedo, Ray::new(rec.point, direction)))
}
Material::Light { .. } => None,
}
}
}
#[derive(Clone, Copy, Debug)]
pub enum Shape {
Sphere { center: Vec3, radius: f64 },
Plane { point: Vec3, normal: Vec3 },
Box { min: Vec3, max: Vec3 },
}
impl Shape {
fn hit(&self, r: &Ray, t_min: f64, t_max: f64) -> Option<(f64, Vec3)> {
match *self {
Shape::Sphere { center, radius } => {
let oc = r.origin - center;
let a = r.direction.dot(r.direction);
let b = oc.dot(r.direction);
let c = oc.dot(oc) - radius * radius;
let disc = b * b - a * c;
if disc < 0.0 {
return None;
}
let sq = disc.sqrt();
let mut root = (-b - sq) / a;
if root < t_min || root > t_max {
root = (-b + sq) / a;
if root < t_min || root > t_max {
return None;
}
}
let outward = (r.at(root) - center).normalize();
Some((root, outward))
}
Shape::Plane { point, normal } => {
let denom = r.direction.dot(normal);
if denom.abs() < 1e-12 {
return None;
}
let t = (point - r.origin).dot(normal) / denom;
if t < t_min || t > t_max {
return None;
}
Some((t, normal))
}
Shape::Box { min, max } => {
let mut t_lo = t_min;
let mut t_hi = t_max;
// Python indexes with -1 when no slab raises t_lo, which means the z axis.
let mut hit_axis = 2;
let mut inv_ds = [0.0; 3];
for axis in 0..3 {
let d = r.direction[axis];
let inv_d = if d.abs() > 1e-12 { 1.0 / d } else { 1e12 };
inv_ds[axis] = inv_d;
let mut t0 = (min[axis] - r.origin[axis]) * inv_d;
let mut t1 = (max[axis] - r.origin[axis]) * inv_d;
if inv_d < 0.0 {
std::mem::swap(&mut t0, &mut t1);
}
if t0 > t_lo {
t_lo = t0;
hit_axis = axis;
}
t_hi = t_hi.min(t1);
if t_hi < t_lo {
return None;
}
}
if t_lo >= t_max || t_lo <= t_min {
return None;
}
// A ray travelling +axis enters through the min face, whose outward normal is -axis.
let sign = if inv_ds[hit_axis] > 0.0 { -1.0 } else { 1.0 };
let mut n = [0.0; 3];
n[hit_axis] = sign;
Some((t_lo, Vec3::new(n[0], n[1], n[2])))
}
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct Object {
pub shape: Shape,
pub material: Material,
}
/// An axis-aligned box light, sampled explicitly for next event estimation.
#[derive(Clone, Copy, Debug)]
pub struct BoxLight {
pub min: Vec3,
pub max: Vec3,
pub emission: Vec3,
}
#[derive(Default)]
pub struct Scene {
pub objects: Vec<Object>,
pub lights: Vec<BoxLight>,
}
impl Scene {
pub fn add(&mut self, shape: Shape, material: Material) {
self.objects.push(Object { shape, material });
}
pub fn add_box_light(&mut self, min: Vec3, max: Vec3, emission: Vec3) {
self.lights.push(BoxLight { min, max, emission });
self.add(Shape::Box { min, max }, Material::Light { emission });
}
pub fn hit(&self, r: &Ray, t_min: f64, t_max: f64) -> Option<HitRecord> {
let mut closest: Option<(f64, Vec3, Material)> = None;
let mut min_t = t_max;
for obj in &self.objects {
if let Some((t, outward)) = obj.shape.hit(r, t_min, min_t) {
if t < min_t {
min_t = t;
closest = Some((t, outward, obj.material));
}
}
}
closest.map(|(t, outward, material)| HitRecord::new(r, t, outward, material))
}
pub fn blocked(&self, origin: Vec3, direction: Vec3, t_max: f64) -> bool {
let r = Ray::new(origin, direction);
self.objects
.iter()
.any(|obj| obj.shape.hit(&r, EPSILON, t_max).is_some())
}
}
pub fn cornell_box() -> Scene {
let mut scene = Scene::default();
let white = Material::Diffuse { albedo: Vec3::new(0.73, 0.73, 0.73) };
let gray = Material::Diffuse { albedo: Vec3::new(0.5, 0.5, 0.5) };
let red = Material::Diffuse { albedo: Vec3::new(0.65, 0.05, 0.05) };
let green = Material::Diffuse { albedo: Vec3::new(0.12, 0.45, 0.15) };
let mirror = Material::Mirror { albedo: Vec3::new(0.95, 0.95, 0.95) };
let glass = Material::Dielectric { ref_idx: 1.5, albedo: Vec3::ONE };
let plane = |p: Vec3, n: Vec3| Shape::Plane { point: p, normal: n.normalize() };
scene.add(plane(Vec3::new(0.0, 0.0, 0.0), Vec3::new(0.0, 1.0, 0.0)), gray);
scene.add(plane(Vec3::new(0.0, 500.0, 0.0), Vec3::new(0.0, -1.0, 0.0)), white);
scene.add(plane(Vec3::new(0.0, 0.0, 500.0), Vec3::new(0.0, 0.0, -1.0)), white);
scene.add(plane(Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 0.0, 0.0)), red);
scene.add(plane(Vec3::new(500.0, 0.0, 0.0), Vec3::new(-1.0, 0.0, 0.0)), green);
scene.add_box_light(
Vec3::new(0.0, 180.0, 180.0),
Vec3::new(25.0, 220.0, 320.0),
Vec3::splat(50.0),
);
let sphere = |c: Vec3, r: f64| Shape::Sphere { center: c, radius: r };
scene.add(sphere(Vec3::new(250.0, 180.0, 250.0), 80.0), glass);
scene.add(sphere(Vec3::new(150.0, 180.0, 380.0), 80.0), mirror);
scene.add(sphere(Vec3::new(350.0, 160.0, 300.0), 60.0), red);
scene
}