533 lines
18 KiB
Rust
533 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();
|
|
let spheres: Vec<Option<(Vec3, f64)>> = self
|
|
.objects
|
|
.iter()
|
|
.map(|o| match o.shape {
|
|
Shape::Sphere { center, radius } => Some((center, radius)),
|
|
_ => None,
|
|
})
|
|
.collect();
|
|
Scene {
|
|
bvh: Bvh::build(&bounds, &spheres),
|
|
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 bvh_depth(&self) -> usize {
|
|
self.bvh.depth()
|
|
}
|
|
|
|
pub fn uses_avx512(&self) -> bool {
|
|
self.bvh.uses_avx512()
|
|
}
|
|
|
|
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));
|
|
}
|
|
}
|
|
// Spheres are intersected inside the BVH (8 at a time); other shapes
|
|
// come back through this callback, which remembers their normal.
|
|
let mut other_normal = Vec3::ZERO;
|
|
let bvh_hit = self.bvh.closest(r, t_min, min_t, |i, t_max| {
|
|
let (t, outward) = self.objects[i].shape.hit(r, t_min, t_max)?;
|
|
other_normal = outward;
|
|
Some(t)
|
|
});
|
|
if let Some((i, t, is_sphere)) = bvh_hit {
|
|
let obj = &self.objects[i];
|
|
let outward = match obj.shape {
|
|
Shape::Sphere { center, .. } if is_sphere => (r.at(t) - center).normalize(),
|
|
_ => other_normal,
|
|
};
|
|
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::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())
|
|
}
|
|
}
|