ai_pathtracer/rust/src/scene.rs
2026-09-26 10:46:01 +02:00

507 lines
18 KiB
Rust

use rand::Rng;
use crate::bvh::{Aabb, Bvh};
use crate::vec3::{reflect, refract, schlick, Vec3};
pub const EPSILON: f64 = 1e-6;
/// Cosine-weighted direction in the hemisphere around the unit vector `n`.
/// Uses the branchless orthonormal basis of Duff et al. 2017 ("Building an
/// Orthonormal Basis, Revisited"); the result is already unit length.
pub fn cosine_hemisphere<R: Rng>(n: Vec3, rng: &mut R) -> Vec3 {
let r1: f64 = rng.gen();
let r2: f64 = rng.gen();
let phi = 2.0 * std::f64::consts::PI * r1;
let r = r2.sqrt();
let (x, y, z) = (r * phi.cos(), r * phi.sin(), (1.0 - r2).max(0.0).sqrt());
let sign = 1.0_f64.copysign(n.z);
let a = -1.0 / (sign + n.z);
let b = n.x * n.y * a;
let t = Vec3::new(1.0 + sign * n.x * n.x * a, sign * b, -sign * n.x);
let bt = Vec3::new(b, sign + n.y * n.y * a, -n.y);
t * x + bt * y + n * z
}
#[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(),
}
}
/// For directions already known to be unit length.
pub fn from_unit(origin: Vec3, direction: Vec3) -> Self {
Ray { origin, direction }
}
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 } => {
// Cosine-weighted sampling: the Lambertian BRDF * cos / pdf
// reduces to just the albedo.
let direction = cosine_hemisphere(rec.normal, rng);
Some((albedo, Ray::from_unit(rec.point, direction)))
}
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 },
/// Axis-aligned rectangle spanning min..max, flat along the axis of
/// `normal` (which must be ±x, ±y or ±z). As a light it emits only on the
/// `normal` side.
Rect { min: Vec3, max: Vec3, normal: 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])))
}
Shape::Rect { min, max, normal } => {
let denom = r.direction.dot(normal);
if denom.abs() < 1e-12 {
return None;
}
let t = (min - r.origin).dot(normal) / denom;
if t < t_min || t > t_max {
return None;
}
let p = r.at(t);
let tol = 1e-9;
let inside = (0..3).all(|a| p[a] >= min[a] - tol && p[a] <= max[a] + tol);
inside.then_some((t, normal))
}
}
}
}
impl Shape {
/// None for unbounded shapes (planes), which are tested outside the BVH.
fn bounds(&self) -> Option<Aabb> {
match *self {
Shape::Sphere { center, radius } => Some(Aabb {
min: center - Vec3::splat(radius),
max: center + Vec3::splat(radius),
}),
Shape::Box { min, max } => Some(Aabb { min, max }),
Shape::Rect { min, max, .. } => Some(Aabb {
min: min - Vec3::splat(1e-9),
max: max + Vec3::splat(1e-9),
}),
Shape::Plane { .. } => None,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct Object {
pub shape: Shape,
pub material: Material,
}
/// An emitter that next event estimation can sample: a box, sphere or rectangle.
#[derive(Clone, Copy, Debug)]
pub struct Light {
shape: Shape,
pub emission: Vec3,
/// Total emitted power (up to a constant): luminance * area.
power: f64,
center: Vec3,
/// Half the bounding box diagonal: how close a point can be to the centre
/// while still outside the light.
extent: f64,
}
impl Light {
fn new(shape: Shape, emission: Vec3) -> Self {
let bounds = shape.bounds().expect("planes cannot be lights");
let mut light = Light {
shape,
emission,
power: 0.0,
center: bounds.centroid(),
extent: 0.5 * (bounds.max - bounds.min).length(),
};
light.power = emission.luminance() * light.area();
light
}
pub fn area(&self) -> f64 {
match self.shape {
Shape::Box { min, max } => {
let d = max - min;
2.0 * (d.x * d.y + d.y * d.z + d.x * d.z)
}
Shape::Sphere { radius, .. } => 4.0 * std::f64::consts::PI * radius * radius,
Shape::Rect { min, max, normal } => {
let d = max - min;
if normal.x != 0.0 {
d.y * d.z
} else if normal.y != 0.0 {
d.x * d.z
} else {
d.x * d.y
}
}
Shape::Plane { .. } => unreachable!("planes cannot be lights"),
}
}
/// Rough estimate of how much this light contributes at `point`: its
/// power over the squared distance, clamped so points near or inside the
/// light's bounds don't blow up.
fn importance(&self, point: Vec3) -> f64 {
let d = point - self.center;
self.power / d.dot(d).max(self.extent * self.extent)
}
/// A uniformly distributed point on the surface (pdf 1 / area) and the
/// outward normal there.
pub fn sample<R: Rng>(&self, rng: &mut R) -> (Vec3, Vec3) {
match self.shape {
Shape::Box { min, max } => {
// Pick a face with probability proportional to its area, then
// a uniform point on it.
let size = max - min;
let face_areas = [size.y * size.z, size.x * size.z, size.x * size.y];
let u = rng.gen::<f64>() * (face_areas[0] + face_areas[1] + face_areas[2]);
let axis = if u < face_areas[0] {
0
} else if u < face_areas[0] + face_areas[1] {
1
} else {
2
};
let positive_side = rng.gen::<f64>() < 0.5;
let mut p = [
min.x + size.x * rng.gen::<f64>(),
min.y + size.y * rng.gen::<f64>(),
min.z + size.z * rng.gen::<f64>(),
];
p[axis] = if positive_side { max[axis] } else { min[axis] };
let mut n = [0.0; 3];
n[axis] = if positive_side { 1.0 } else { -1.0 };
(Vec3::new(p[0], p[1], p[2]), Vec3::new(n[0], n[1], n[2]))
}
Shape::Sphere { center, radius } => {
let z = 1.0 - 2.0 * rng.gen::<f64>();
let r = (1.0 - z * z).max(0.0).sqrt();
let phi = 2.0 * std::f64::consts::PI * rng.gen::<f64>();
let n = Vec3::new(r * phi.cos(), r * phi.sin(), z);
(center + radius * n, n)
}
Shape::Rect { min, max, normal } => {
let size = max - min;
let p = Vec3::new(
min.x + size.x * rng.gen::<f64>(),
min.y + size.y * rng.gen::<f64>(),
min.z + size.z * rng.gen::<f64>(),
);
(p, normal)
}
Shape::Plane { .. } => unreachable!("planes cannot be lights"),
}
}
}
/// What a ray that leaves the scene sees.
#[derive(Clone, Copy, Debug)]
pub enum Background {
/// White-to-blue gradient, as in raytracer.py.
Sky,
Color(Vec3),
}
impl Background {
pub fn radiance(&self, direction: Vec3) -> Vec3 {
match *self {
Background::Sky => {
let t = 0.5 * (direction.y + 1.0);
Vec3::ONE * (1.0 - t) + Vec3::new(0.5, 0.7, 1.0) * t
}
Background::Color(c) => c,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct View {
pub lookfrom: Vec3,
pub lookat: Vec3,
pub vfov: f64,
}
pub struct SceneBuilder {
objects: Vec<Object>,
planes: Vec<Object>,
lights: Vec<Light>,
background: Background,
view: View,
}
impl SceneBuilder {
pub fn new(view: View, background: Background) -> Self {
SceneBuilder {
objects: Vec::new(),
planes: Vec::new(),
lights: Vec::new(),
background,
view,
}
}
pub fn add(&mut self, shape: Shape, material: Material) {
let obj = Object { shape, material };
match shape {
Shape::Plane { .. } => self.planes.push(obj),
_ => self.objects.push(obj),
}
}
pub fn plane(&mut self, point: Vec3, normal: Vec3, material: Material) {
self.add(Shape::Plane { point, normal: normal.normalize() }, material);
}
pub fn sphere(&mut self, center: Vec3, radius: f64, material: Material) {
self.add(Shape::Sphere { center, radius }, material);
}
pub fn cuboid(&mut self, min: Vec3, max: Vec3, material: Material) {
self.add(Shape::Box { min, max }, material);
}
/// Every emitter must be added through a *_light method: after a diffuse
/// bounce, trace() only counts light found by sampling these.
pub fn box_light(&mut self, min: Vec3, max: Vec3, emission: Vec3) {
self.add_light(Shape::Box { min, max }, emission);
}
pub fn sphere_light(&mut self, center: Vec3, radius: f64, emission: Vec3) {
self.add_light(Shape::Sphere { center, radius }, emission);
}
/// One-sided rectangular light emitting towards `normal` (±x, ±y or ±z).
pub fn rect_light(&mut self, min: Vec3, max: Vec3, normal: Vec3, emission: Vec3) {
self.add_light(Shape::Rect { min, max, normal }, emission);
}
fn add_light(&mut self, shape: Shape, emission: Vec3) {
self.lights.push(Light::new(shape, emission));
self.add(shape, Material::Light { emission });
}
pub fn build(self) -> Scene {
let bounds: Vec<Aabb> = self.objects.iter().map(|o| o.shape.bounds().unwrap()).collect();
Scene {
bvh: Bvh::build(&bounds),
objects: self.objects,
planes: self.planes,
lights: self.lights,
background: self.background,
view: self.view,
}
}
}
pub struct Scene {
objects: Vec<Object>,
planes: Vec<Object>,
bvh: Bvh,
lights: Vec<Light>,
pub background: Background,
pub view: View,
}
impl Scene {
pub fn object_count(&self) -> usize {
self.objects.len() + self.planes.len()
}
pub fn light_count(&self) -> usize {
self.lights.len()
}
/// Picks one light for shading `point`, with probability proportional to
/// its estimated contribution there (power / distance^2), so a scene with
/// many lights costs one shadow ray per shading point and mostly samples
/// nearby lights. Returns the light and the probability of picking it.
/// This is O(lights) per call; thousands of lights would want a light tree.
pub fn pick_light<R: Rng>(&self, point: Vec3, rng: &mut R) -> Option<(&Light, f64)> {
match self.lights.len() {
0 => None,
1 => Some((&self.lights[0], 1.0)),
_ => {
let total: f64 = self.lights.iter().map(|l| l.importance(point)).sum();
let u = rng.gen::<f64>() * total;
let mut acc = 0.0;
for light in &self.lights {
let w = light.importance(point);
acc += w;
if u < acc {
return Some((light, w / total));
}
}
// Rounding left u just past the last bucket.
let last = self.lights.last().unwrap();
Some((last, last.importance(point) / total))
}
}
}
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.planes {
if let Some((t, outward)) = obj.shape.hit(r, t_min, min_t) {
min_t = t;
closest = Some((t, outward, obj.material));
}
}
self.bvh.closest(r, t_min, min_t, |i, t_max| {
let obj = &self.objects[i];
let (t, outward) = obj.shape.hit(r, t_min, t_max)?;
closest = Some((t, outward, obj.material));
Some(t)
});
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::from_unit(origin, direction);
self.planes.iter().any(|obj| obj.shape.hit(&r, EPSILON, t_max).is_some())
|| self
.bvh
.any(&r, EPSILON, t_max, |i| self.objects[i].shape.hit(&r, EPSILON, t_max).is_some())
}
}