pbrt/shared/src/shapes/disk.rs

use crate::core::geometry::{
    Bounds3f, DirectionCone, Normal3f, Point2f, Point3f, Point3fi, Ray, Vector2f, Vector3f,
    Vector3fi, VectorLike,
};
use crate::core::interaction::{Interaction, InteractionTrait, SurfaceInteraction};
use crate::core::shape::{
    QuadricIntersection, ShapeIntersection, ShapeSample, ShapeSampleContext, ShapeTrait,
};
use crate::utils::interval::Interval;
use crate::utils::math::{clamp, radians, square};
use crate::utils::sampling::sample_uniform_disk_concentric;
use crate::utils::Transform;
use crate::{Float, PI};
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct DiskShape {
    pub radius: Float,
    pub inner_radius: Float,
    pub height: Float,
    pub phi_max: Float,
    pub render_from_object: Transform,
    pub object_from_render: Transform,
    pub reverse_orientation: bool,
    pub transform_swap_handedness: bool,
}

impl DiskShape {
    pub fn new(
        radius: Float,
        inner_radius: Float,
        height: Float,
        phi_max: Float,
        render_from_object: Transform,
        object_from_render: Transform,
        reverse_orientation: bool,
    ) -> Self {
        Self {
            radius,
            inner_radius,
            height,
            // pbrt: `phiMax(Radians(Clamp(phiMax, 0, 360)))`. The parameter arrives in
            // DEGREES (default 360); storing it raw made `area()` 360/2pi = 57.3x too
            // large, so the sampling pdf was 57.3x too small and every direct-lighting
            // contribution from a disk area light was 57.3x too bright.
            phi_max: radians(clamp(phi_max, 0., 360.)),
            render_from_object: render_from_object.clone(),
            object_from_render,
            reverse_orientation,
            transform_swap_handedness: render_from_object.swaps_handedness(),
        }
    }

    fn basic_intersect(&self, r: &Ray, t_max: Float) -> Option<QuadricIntersection> {
        let oi = self
            .object_from_render
            .apply_to_interval(&Point3fi::new_from_point(r.o));
        let di = self
            .object_from_render
            .apply_to_vector_interval(&Vector3fi::new_from_vector(r.d));

        if Float::from(di.z()) == 0. {
            return None;
        }
        // pbrt computes the plane hit as a SCALAR:
        //   Float tShapeHit = (height - Float(oi.z)) / Float(di.z);
        //   if (tShapeHit <= 0 || tShapeHit >= tMax) return {};
        // Interval arithmetic here produced NaN, and neither `high <= 0` nor
        // `low >= t_max` rejects a NaN (every NaN comparison is false). The NaN then
        // flowed into p_hit -- so the `dist2 > radius^2` test could not reject it either,
        // and the disk swallowed every ray it was tested against -- and on into dpdu/dpdv
        // and the shading normal, poisoning the area-light MIS pdf.
        let t_shape_hit: Float = (self.height - Float::from(oi.z())) / Float::from(di.z());
        if t_shape_hit <= 0. || t_shape_hit >= t_max {
            return None;
        }

        let oi_f = Point3f::from(oi);
        let di_f = Vector3f::from(di);
        let t = t_shape_hit;
        let p_hit: Point3f = oi_f + di_f * t;

        let dist2 = square(p_hit.x()) + square(p_hit.y());
        if dist2 > square(self.radius) || dist2 < square(self.inner_radius) {
            return None;
        }
        let mut phi = p_hit.y().atan2(p_hit.x());
        if phi < 0. {
            phi += 2. * PI;
        }
        if phi > self.phi_max {
            return None;
        }

        Some(QuadricIntersection {
            t_hit: t,
            p_obj: p_hit,
            phi,
        })
    }

    fn interaction_from_intersection(
        &self,
        isect: QuadricIntersection,
        wo: Vector3f,
        time: Float,
    ) -> SurfaceInteraction {
        let mut p_hit = isect.p_obj;
        let phi = isect.phi;
        // Find parametric representation of disk hit
        let u = phi / self.phi_max;
        let r_hit = (square(p_hit.x()) + square(p_hit.y())).sqrt();
        let v = (self.radius - r_hit) / (self.radius - self.inner_radius);
        let dpdu = Vector3f::new(-self.phi_max * p_hit.y(), self.phi_max * p_hit.x(), 0.);
        let dpdv =
            Vector3f::new(p_hit.x(), p_hit.y(), 0.) * (self.inner_radius - self.radius) / r_hit;
        let dndu = Normal3f::zero();
        let dndv = Normal3f::zero();

        p_hit[2] = self.height;

        let p_error = Vector3f::zero();
        let flip_normal = self.reverse_orientation ^ self.transform_swap_handedness;
        let wo_object = self.object_from_render.apply_to_vector(wo);
        let intr = SurfaceInteraction::new(
            Point3fi::new_with_error(p_hit, p_error),
            Point2f::new(u, v),
            wo_object,
            dpdu,
            dpdv,
            dndu,
            dndv,
            time,
            flip_normal,
        );

        match self
            .render_from_object
            .apply_to_interaction(&Interaction::Surface(intr))
        {
            Interaction::Surface(si) => si,
            _ => unreachable!("Only surfaces need apply"),
        }
    }
}

impl ShapeTrait for DiskShape {
    fn area(&self) -> Float {
        self.phi_max * 0.5 * (square(self.radius) - square(self.inner_radius))
    }

    fn bounds(&self) -> Bounds3f {
        self.render_from_object
            .apply_to_bounds(Bounds3f::from_points(
                Point3f::new(-self.radius, -self.radius, self.height),
                Point3f::new(self.radius, self.radius, self.height),
            ))
    }
    fn normal_bounds(&self) -> DirectionCone {
        let mut n = self
            .render_from_object
            .apply_to_normal(Normal3f::new(0., 0., 1.));
        if self.reverse_orientation {
            n = -n;
        }
        DirectionCone::new_from_vector(Vector3f::from(n))
    }

    fn intersect(&self, ray: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        let t = t_max.unwrap_or(Float::INFINITY);
        if let Some(isect) = self.basic_intersect(ray, t) {
            let intr = self.interaction_from_intersection(isect.clone(), -ray.d, ray.time);
            Some(ShapeIntersection::new(intr, isect.t_hit))
        } else {
            None
        }
    }

    fn sample(&self, u: Point2f) -> Option<ShapeSample> {
        let pd = sample_uniform_disk_concentric(u);
        let p_obj = Point3f::new(pd.x() * self.radius, pd.y() * self.radius, self.height);
        let pi = self
            .render_from_object
            .apply_to_interval(&Point3fi::new_from_point(p_obj));
        let mut n: Normal3f = self
            .render_from_object
            .apply_to_normal(Normal3f::new(0., 0., 1.))
            .normalize();
        if self.reverse_orientation {
            n = -n;
        }
        let mut phi = pd.y().atan2(pd.x());
        if phi < 0. {
            phi += 2. * PI;
        }
        let radius_sample = (square(p_obj.x()) + square(p_obj.y())).sqrt();
        let uv = Point2f::new(
            phi / self.phi_max,
            (self.radius - radius_sample) / (self.radius - self.inner_radius),
        );

        Some(ShapeSample {
            intr: Interaction::Surface(SurfaceInteraction::new_simple(pi, n, uv)),
            pdf: 1. / self.area(),
        })
    }

    fn intersect_p(&self, ray: &Ray, t_max: Option<Float>) -> bool {
        if let Some(t) = t_max {
            self.basic_intersect(ray, t).is_some()
        } else {
            self.basic_intersect(ray, Float::INFINITY).is_some()
        }
    }

    fn sample_from_context(&self, ctx: &ShapeSampleContext, u: Point2f) -> Option<ShapeSample> {
        let mut ss = self.sample(u)?;
        ss.intr.get_common_mut().time = ctx.time;
        let mut wi = ss.intr.p() - ctx.p();
        if wi.norm_squared() == 0. {
            return None;
        }
        wi = wi.normalize();

        ss.pdf /= Vector3f::from(ss.intr.n()).abs_dot(-wi) / ctx.p().distance_squared(ss.intr.p());
        if ss.pdf.is_infinite() {
            return None;
        }

        Some(ss)
    }

    fn pdf(&self, _interaction: &Interaction) -> Float {
        1. / self.area()
    }

    fn pdf_from_context(&self, ctx: &ShapeSampleContext, wi: Vector3f) -> Float {
        let ray = ctx.spawn_ray(wi);
        if let Some(isect) = self.intersect(&ray, None) {
            let n = isect.intr.n();
            let absdot = Vector3f::from(n).dot(-wi).abs();
            let pdf = (1. / self.area()) / (absdot / ctx.p().distance_squared(isect.intr.p()));
            if pdf.is_infinite() {
                return 0.;
            }
            pdf
        } else {
            0.
        }
    }
}