del-msh-cpu 0.1.49

pub trait ScalarRender<T> {
    fn fwd(
        &self,
        bc: &[T; 3],
        i_tri: u32,
        tri2vtx: &[[u32; 3]],
        vtx2xyz: &[[T; 3]],
        transform_world2ndc: &[T; 16],
    ) -> T;

    #[allow(clippy::too_many_arguments)]
    fn bwd(
        &self,
        dldw_val: T,
        p0: &[T; 3],
        p1: &[T; 3],
        p2: &[T; 3],
        ray_org: &[T; 3],
        ray_dir: &[T; 3],
        transform_world2ndc: &[T; 16],
    ) -> ([T; 3], [T; 3], [T; 3]);
}

#[allow(clippy::too_many_arguments)]
pub fn bwd_continuous<T: ScalarRender<f32>>(
    pix2tri: &[u32],
    tri2vtx: &[[u32; 3]],
    vtx2xyz: &[[f32; 3]],
    dldw_pix2val: &[f32],
    transform_ndc2world: &[f32; 16],
    img_shape: (usize, usize),
    dldw_vtx2xyz: &mut [[f32; 3]],
    mode: &T,
) {
    let transform_world2ndc =
        del_geo_core::mat4_col_major::try_inverse_with_pivot(transform_ndc2world).unwrap();
    for (i_pix, &i_tri) in pix2tri.iter().enumerate() {
        if i_tri == u32::MAX {
            continue;
        }
        let i_w = i_pix % img_shape.0;
        let i_h = i_pix / img_shape.0;
        let (ray_org, ray_dir) =
            del_geo_core::mat4_col_major::ray_from_transform_ndc2world_and_pixel_coordinates(
                (i_w as f32 + 0.5, i_h as f32 + 0.5),
                &(img_shape.0 as f32, img_shape.1 as f32),
                transform_ndc2world,
            );
        let i_tri = i_tri as usize;
        let i0 = tri2vtx[i_tri][0] as usize;
        let i1 = tri2vtx[i_tri][1] as usize;
        let i2 = tri2vtx[i_tri][2] as usize;
        let p0 = &vtx2xyz[i0];
        let p1 = &vtx2xyz[i1];
        let p2 = &vtx2xyz[i2];
        let dldw_val = dldw_pix2val[i_pix];
        let (dp0, dp1, dp2) = mode.bwd(
            dldw_val,
            p0,
            p1,
            p2,
            &ray_org,
            &ray_dir,
            &transform_world2ndc,
        );
        use del_geo_core::vec3::Vec3;
        dldw_vtx2xyz[i0].add_in_place(&dp0);
        dldw_vtx2xyz[i1].add_in_place(&dp1);
        dldw_vtx2xyz[i2].add_in_place(&dp2);
    }
}

pub fn fwd_continuous<T: ScalarRender<f32>>(
    pix2tri: &[u32],
    img_shape: (usize, usize),
    tri2vtx: &[[u32; 3]],
    vtx2xyz: &[[f32; 3]],
    transform_ndc2world: &[f32; 16],
    model: &T,
) -> Vec<f32> {
    let transform_world2ndc =
        del_geo_core::mat4_col_major::try_inverse_with_pivot(transform_ndc2world).unwrap();
    let mut pix2vin = vec![0.; pix2tri.len()];
    for (i_pix, &i_tri) in pix2tri.iter().enumerate() {
        if i_tri == u32::MAX {
            continue;
        }
        let i_w = i_pix % img_shape.0;
        let i_h = i_pix / img_shape.0;
        let (ray_org, ray_dir) =
            del_geo_core::mat4_col_major::ray_from_transform_ndc2world_and_pixel_coordinates(
                (i_w as f32 + 0.5, i_h as f32 + 0.5),
                &(img_shape.0 as f32, img_shape.1 as f32),
                transform_ndc2world,
            );
        let Some((_t, bc)) = crate::trimesh3::to_tri3(tri2vtx, vtx2xyz, i_tri as usize)
            .intersection_against_ray(&ray_org, &ray_dir)
        else {
            unreachable!()
        };
        pix2vin[i_pix] = model.fwd(&bc, i_tri, tri2vtx, vtx2xyz, &transform_world2ndc);
    }
    pix2vin
}

/// Generates sub-pixel sample positions for Monte Carlo integration over a pixel.
///
/// Each call to `sample` returns a 2-D offset in [0,1)^2 (relative to the pixel's
/// top-left corner) and the PDF at that point.  The returned value accumulator in
/// `multi_sample` divides each radiance contribution by its PDF, so returning
/// `pdf = 1.0` (uniform) leaves the averaging unchanged.
pub trait PixelSampler: Sized {
    /// Create a sampler for pixel `i_pix`, seeded deterministically from its index.
    fn initialize(i_pix: usize) -> Self;
    fn sample(&mut self) -> ([f32; 2], f32);
}

/// Box-filter sampler: uniform random position over the pixel area; PDF is always 1.0.
pub struct BoxPixelSampler<R: rand::Rng> {
    rng: R,
}

impl<R: rand::Rng + rand::SeedableRng> PixelSampler for BoxPixelSampler<R> {
    fn initialize(i_pix: usize) -> Self {
        Self {
            rng: R::seed_from_u64(i_pix as u64),
        }
    }
    fn sample(&mut self) -> ([f32; 2], f32) {
        use rand::RngExt;
        let x = self.rng.random_range(0.0..1.0);
        let y = self.rng.random_range(0.0..1.0);
        ([x, y], 1.0)
    }
}

pub fn multi_sample<T, S>(
    tri2vtx: &[[u32; 3]],
    vtx2xyz: &[[f32; 3]],
    transform_world2ndc: &[f32; 16],
    img_shape: (usize, usize),
    num_sample: usize,
    model: &T,
) -> Vec<f32>
where
    T: ScalarRender<f32> + Sync,
    S: PixelSampler + Send,
{
    let transform_ndc2world =
        del_geo_core::mat4_col_major::try_inverse_with_pivot(transform_world2ndc).unwrap();
    let bvhnodes = crate::bvhnodes_morton::from_triangle_mesh(tri2vtx, vtx2xyz);
    let bvhnode2aabb =
        crate::bvhnode2aabb3::from_uniform_mesh_with_bvh(0, &bvhnodes, tri2vtx, vtx2xyz, None);
    let fn_pix2val = |i_pix: usize| -> f32 {
        let mut sampler = S::initialize(i_pix);
        let i_h = i_pix / img_shape.0;
        let i_w = i_pix - i_h * img_shape.0;
        let mut sum = 0.0f32;
        for _itr in 0..num_sample {
            let ([x_offset, y_offset], pdf) = sampler.sample();
            let (ray_org, ray_dir) =
                del_geo_core::mat4_col_major::ray_from_transform_ndc2world_and_pixel_coordinates(
                    (i_w as f32 + x_offset, i_h as f32 + y_offset),
                    &(img_shape.0 as f32, img_shape.1 as f32),
                    &transform_ndc2world,
                );
            if let Some((_t, i_tri, bc)) = crate::search_bvh3::first_intersection_ray(
                &ray_org,
                &ray_dir,
                &crate::search_bvh3::TriMesh3WithBvhRef {
                    tri2vtx,
                    vtx2xyz,
                    bvhnode2tri_btree: &bvhnodes,
                    bvhnode2aabb: &bvhnode2aabb,
                },
                0,
                f32::INFINITY,
            ) {
                sum += model.fwd(&bc, i_tri as u32, tri2vtx, vtx2xyz, transform_world2ndc) / pdf;
            }
        }
        sum / num_sample as f32
    };
    let mut pix2val = vec![0f32; img_shape.0 * img_shape.1];
    use rayon::prelude::*;
    pix2val
        .par_iter_mut()
        .enumerate()
        .for_each(|(i_pix, val)| *val = fn_pix2val(i_pix));
    pix2val
}

#[allow(clippy::too_many_arguments)]
pub fn finite_difference<T, S>(
    tri2vtx: &[[u32; 3]],
    vtx2xyz: &[[f32; 3]],
    vtx2uvw: &[[f32; 3]],
    transform_world2ndc: &[f32; 16],
    img_shape: (usize, usize),
    mode: &T,
    num_sample: usize,
    eps: f32,
) -> (Vec<f32>, Vec<f32>)
where
    T: ScalarRender<f32> + Sync,
    S: PixelSampler + Send,
{
    let pix2val0 = {
        multi_sample::<T, S>(
            tri2vtx,
            vtx2xyz,
            transform_world2ndc,
            img_shape,
            num_sample,
            mode,
        )
    };
    let pix2val1 = {
        let vtx2xyz = vtx2xyz
            .iter()
            .zip(vtx2uvw.iter())
            .map(|(x, u)| [x[0] + eps * u[0], x[1] + eps * u[1], x[2] + eps * u[2]])
            .collect::<Vec<_>>();
        multi_sample::<T, S>(
            tri2vtx,
            &vtx2xyz,
            transform_world2ndc,
            img_shape,
            num_sample,
            mode,
        )
    };
    let pix2diff = pix2val1
        .iter()
        .zip(pix2val0.iter())
        .map(|(u, v)| (u - v) / eps)
        .collect::<Vec<_>>();
    (pix2diff, pix2val0)
}