del-msh-cpu 0.1.52

//! methods that generate vertices connected to a vertex

use num_traits::{AsPrimitive, PrimInt};

/// point surrounding point for mesh
/// * `elem2vtx` - map element to vertex: list of vertex index for each element
/// * `num_node` - number of vertex par element (e.g., 3 for tri, 4 for quad)
/// * `num_vtx` - number of vertex
/// * `vtx2idx` - map vertex to element index: cumulative sum
/// * `idx2elem` - map vertex to element value: list of value
pub fn from_uniform_mesh_with_vtx2elem<Index, const NNODE: usize>(
    elem2vtx: &[[Index; NNODE]],
    num_vtx: usize,
    vtx2idx: &[Index],
    idx2elem: &[Index],
    is_self: bool,
) -> (Vec<Index>, Vec<Index>)
where
    Index: num_traits::PrimInt + num_traits::AsPrimitive<usize> + std::ops::AddAssign<Index>,
    usize: AsPrimitive<Index>,
{
    assert_eq!(vtx2idx.len(), num_vtx + 1);
    let mut vtx2flg = vec![usize::MAX; num_vtx];
    let mut vtx2jdx = vec![Index::zero(); num_vtx + 1];
    for i_vtx in 0..num_vtx {
        if !is_self {
            vtx2flg[i_vtx] = i_vtx;
        }
        let idx0: usize = vtx2idx[i_vtx].as_();
        let idx1: usize = vtx2idx[i_vtx + 1].as_();
        for j_elem in &idx2elem[idx0..idx1] {
            let j_elem: usize = j_elem.as_();
            for j_vtx in elem2vtx[j_elem].iter() {
                let j_vtx: usize = j_vtx.as_();
                if vtx2flg[j_vtx] != i_vtx {
                    vtx2flg[j_vtx] = i_vtx;
                    vtx2jdx[i_vtx + 1] += Index::one();
                }
            }
        }
    }
    for i_vtx in 0..num_vtx {
        let tmp = vtx2jdx[i_vtx];
        vtx2jdx[i_vtx + 1] += tmp;
    }
    let num_vtx2vtx: usize = vtx2jdx[num_vtx].as_();
    let mut jdx2vtx = vec![Index::zero(); num_vtx2vtx];
    vtx2flg.iter_mut().for_each(|v| *v = usize::MAX);
    for i_vtx in 0..num_vtx {
        if !is_self {
            vtx2flg[i_vtx] = i_vtx;
        }
        let idx0: usize = vtx2idx[i_vtx].as_();
        let idx1: usize = vtx2idx[i_vtx + 1].as_();
        for j_elem in &idx2elem[idx0..idx1] {
            let j_elem: usize = j_elem.as_();
            for j_vtx in elem2vtx[j_elem].iter() {
                let j_vtx: usize = j_vtx.as_();
                if vtx2flg[j_vtx] != i_vtx {
                    vtx2flg[j_vtx] = i_vtx;
                    let iv2v: usize = vtx2jdx[i_vtx].as_();
                    jdx2vtx[iv2v] = j_vtx.as_();
                    vtx2jdx[i_vtx] += Index::one();
                }
            }
        }
    }
    for i_vtx in (1..num_vtx).rev() {
        vtx2jdx[i_vtx] = vtx2jdx[i_vtx - 1];
    }
    vtx2jdx[0] = Index::zero();
    (vtx2jdx, jdx2vtx)
}

/// compute index of vertices adjacent to vertices for uniform mesh.
pub fn from_uniform_mesh<Index, const NNODE: usize>(
    elem2vtx: &[[Index; NNODE]],
    num_vtx: usize,
    is_self: bool,
) -> (Vec<Index>, Vec<Index>)
where
    Index: num_traits::PrimInt + std::ops::AddAssign + num_traits::AsPrimitive<usize>,
    usize: AsPrimitive<Index>,
{
    let vtx2elem = crate::vtx2elem::from_uniform_mesh(elem2vtx, num_vtx);
    let vtx2vtx =
        from_uniform_mesh_with_vtx2elem(elem2vtx, num_vtx, &vtx2elem.0, &vtx2elem.1, is_self);
    assert_eq!(vtx2vtx.0.len(), num_vtx + 1);
    vtx2vtx
}

pub fn from_specific_edges_of_uniform_mesh<Index, const NNODE: usize>(
    elem2vtx: &[[Index; NNODE]],
    edge2node: &[[usize; 2]],
    vtx2idx: &[Index],
    idx2elem: &[Index],
    is_bidirectional: bool,
) -> (Vec<Index>, Vec<Index>)
where
    Index: num_traits::PrimInt + AsPrimitive<usize>,
    usize: AsPrimitive<Index>,
{
    let num_vtx = vtx2idx.len() - 1;
    let mut vtx2jdx = vec![Index::zero(); num_vtx + 1];
    vtx2jdx[0] = Index::zero();
    let mut jdx2vtx = Vec::<Index>::new();
    let mut set_vtx: std::collections::BTreeSet<Index> = std::collections::BTreeSet::new();
    for i_vtx in 0..num_vtx {
        set_vtx.clear();
        let idx0: usize = vtx2idx[i_vtx].as_();
        let idx1: usize = vtx2idx[i_vtx + 1].as_();
        for &ielem0 in &idx2elem[idx0..idx1] {
            let i_vtx: Index = i_vtx.as_();
            let ielem0 = ielem0.as_();
            for &[inode0, inode1] in edge2node {
                let ivtx0 = elem2vtx[ielem0][inode0];
                let ivtx1 = elem2vtx[ielem0][inode1];
                if ivtx0 != i_vtx && ivtx1 != i_vtx {
                    continue;
                }
                if ivtx0 == i_vtx {
                    if is_bidirectional || ivtx1 > i_vtx {
                        set_vtx.insert(ivtx1);
                    }
                } else if is_bidirectional || ivtx0 > i_vtx {
                    set_vtx.insert(ivtx0);
                }
            }
        }
        for vtx in &set_vtx {
            jdx2vtx.push(*vtx);
        }
        vtx2jdx[i_vtx + 1] = vtx2jdx[i_vtx] + set_vtx.len().as_();
    }
    (vtx2jdx, jdx2vtx)
}

/// make vertex surrounding vertex as edges of polygon mesh.
/// A polygon mesh is a mixture of elements such as triangle, quadrilateal, pentagon.
pub fn from_polygon_mesh_edges_with_vtx2elem<INDEX>(
    elem2idx: &[INDEX],
    idx2vtx: &[INDEX],
    vtx2jdx: &[INDEX],
    jdx2elem: &[INDEX],
    is_bidirectional: bool,
) -> (Vec<INDEX>, Vec<INDEX>)
where
    INDEX: num_traits::PrimInt + num_traits::AsPrimitive<usize>,
    usize: AsPrimitive<INDEX>,
{
    let nvtx = vtx2jdx.len() - 1;

    let mut vtx2kdx = vec![INDEX::zero(); nvtx + 1];
    let mut kdx2vtx = Vec::<INDEX>::new();

    for i_vtx in 0..nvtx {
        let mut set_vtx_idx = std::collections::BTreeSet::new();
        let idx0: usize = vtx2jdx[i_vtx].as_();
        let idx1: usize = vtx2jdx[i_vtx + 1].as_();
        for &i_elem0 in &jdx2elem[idx0..idx1] {
            let ielem0: usize = i_elem0.as_();
            let num_node: usize = (elem2idx[ielem0 + 1] - elem2idx[ielem0]).as_();
            let num_edge = num_node;
            for i_edge in 0..num_edge {
                let i_node0 = i_edge;
                let i_node1 = (i_edge + 1) % num_node;
                let j_vtx0: usize = idx2vtx[elem2idx[ielem0].as_() + i_node0].as_();
                let j_vtx1: usize = idx2vtx[elem2idx[ielem0].as_() + i_node1].as_();
                if j_vtx0 != i_vtx && j_vtx1 != i_vtx {
                    continue;
                }
                if j_vtx0 == i_vtx {
                    if is_bidirectional || j_vtx1 > i_vtx {
                        set_vtx_idx.insert(j_vtx1);
                    }
                } else if is_bidirectional || j_vtx0 > i_vtx {
                    set_vtx_idx.insert(j_vtx0);
                }
            }
        }
        for &itr in &set_vtx_idx {
            kdx2vtx.push(itr.as_());
        }
        vtx2kdx[i_vtx + 1] = vtx2kdx[i_vtx] + set_vtx_idx.len().as_();
    }
    (vtx2kdx, kdx2vtx)
}

/// \[I + lambda * L\] {vtx2lhs} = {vtx2rhs}
/// L = \[..-1,..,valence, ..,-1 \]
#[allow(clippy::too_many_arguments)]
pub fn graph_screend_poisson<IDX>(
    vtx2idx_offset: &[IDX],
    idx2vtx: &[IDX],
    lambda: f32,
    num_vdim: usize,
    vtx2lhs: &mut [f32],
    vtx2rhs: &[f32],
    num_iter: usize,
    vtx2lhstmp: &mut [f32],
) where
    IDX: num_traits::PrimInt + AsPrimitive<usize> + AsPrimitive<f32> + std::marker::Sync,
{
    let num_vtx = vtx2idx_offset.len() - 1;
    assert_eq!(vtx2lhs.len(), num_vtx * num_vdim);
    assert_eq!(vtx2rhs.len(), num_vtx * num_vdim);
    assert_eq!(vtx2lhstmp.len(), num_vtx * num_vdim);
    let func_upd = |i_vtx: usize, lhs_next: &mut [f32], vtx2lhs_prev: &[f32]| {
        let mut buff = vec![0f32; num_vdim];
        buff.copy_from_slice(&vtx2rhs[i_vtx * num_vdim..(i_vtx + 1) * num_vdim]);
        for &j_vtx in &idx2vtx[vtx2idx_offset[i_vtx].as_()..vtx2idx_offset[i_vtx + 1].as_()] {
            let j_vtx: usize = j_vtx.as_();
            for i in 0..num_vdim {
                buff[i] += lambda * vtx2lhs_prev[j_vtx * num_vdim + i];
            }
        }
        let valence: f32 = (vtx2idx_offset[i_vtx + 1] - vtx2idx_offset[i_vtx]).as_();
        let inv_dia = 1f32 / (1f32 + lambda * valence);
        for i in 0..num_vdim {
            lhs_next[i] = buff[i] * inv_dia;
        }
    };
    use rayon::prelude::*;
    for _iter in 0..num_iter {
        vtx2lhstmp
            .par_chunks_mut(num_vdim)
            .enumerate()
            .for_each(|(i_vtx, lhs1)| func_upd(i_vtx, lhs1, vtx2lhs));
        vtx2lhs
            .par_chunks_mut(num_vdim)
            .enumerate()
            .for_each(|(i_vtx, lhs)| func_upd(i_vtx, lhs, vtx2lhstmp));
    }
}

/// \[I + lambda * L\] {vtx2lhs} = {vtx2rhs}
/// L = \[..-1,..,valence, ..,-1 \]
#[allow(clippy::too_many_arguments)]
pub fn laplacian_smoothing<IDX>(
    vtx2idx_offset: &[IDX],
    idx2vtx: &[IDX],
    lambda: f32,
    num_vdim: usize,
    vtx2val: &mut [f32],
    vtx2ave: &mut [f32],
    num_iter: usize,
) where
    IDX: num_traits::PrimInt + AsPrimitive<usize> + AsPrimitive<f32> + std::marker::Sync,
{
    let num_vtx = vtx2idx_offset.len() - 1;
    assert_eq!(vtx2ave.len(), num_vtx * num_vdim);
    assert_eq!(vtx2val.len(), num_vtx * num_vdim);
    let func_upd = |i_vtx: usize, ave: &mut [f32], vtx2val: &[f32]| {
        ave.iter_mut().for_each(|x| *x *= lambda);
        for &j_vtx in &idx2vtx[vtx2idx_offset[i_vtx].as_()..vtx2idx_offset[i_vtx + 1].as_()] {
            let j_vtx: usize = j_vtx.as_();
            for i in 0..num_vdim {
                ave[i] += vtx2val[j_vtx * num_vdim + i];
            }
        }
        let valence: f32 = (vtx2idx_offset[i_vtx + 1] - vtx2idx_offset[i_vtx]).as_();
        let inv_dia = 1f32 / (lambda + valence);
        ave.iter_mut().for_each(|x| *x *= inv_dia);
    };
    use rayon::prelude::*;
    for _iter in 0..num_iter {
        vtx2ave
            .par_chunks_mut(num_vdim)
            .enumerate()
            .for_each(|(i_vtx, ave)| func_upd(i_vtx, ave, vtx2val));
        vtx2val.copy_from_slice(vtx2ave);
    }
}

pub fn multiply_graph_laplacian<IDX>(
    vtx2idx_offset: &[IDX],
    idx2vtx: &[IDX],
    num_dim: usize,
    vtx2rhs: &[f32],
    vtx2lhs: &mut [f32],
) where
    IDX: PrimInt + AsPrimitive<usize> + AsPrimitive<f32> + std::marker::Sync,
{
    let num_vtx = vtx2idx_offset.len() - 1;
    assert_eq!(vtx2rhs.len(), num_vtx * num_dim);
    assert_eq!(vtx2lhs.len(), num_vtx * num_dim);
    let func_upd = |i_vtx: usize, lhs: &mut [f32]| {
        let valence: f32 = (vtx2idx_offset[i_vtx + 1] - vtx2idx_offset[i_vtx]).as_();
        for i in 0..num_dim {
            lhs[i] = valence * vtx2rhs[i_vtx * num_dim + i];
        }
        for &j_vtx in &idx2vtx[vtx2idx_offset[i_vtx].as_()..vtx2idx_offset[i_vtx + 1].as_()] {
            let j_vtx: usize = j_vtx.as_();
            for i in 0..num_dim {
                lhs[i] -= vtx2rhs[j_vtx * num_dim + i];
            }
        }
    };
    use rayon::prelude::*;
    vtx2lhs
        .par_chunks_mut(num_dim)
        .enumerate()
        .for_each(|(i_vtx, lhs)| func_upd(i_vtx, lhs));
}

#[test]
fn test_laplacian_smoothing() {
    let (tri2vtx, vtx2xyz) = crate::trimesh3_primitive::torus_zup::<usize, f32>(1.0, 0.3, 32, 32);
    let (vtx2idx, idx2vtx) = crate::vtx2vtx::from_uniform_mesh(&tri2vtx, vtx2xyz.len(), false);
    let num_vdim = 3;
    let vtx2rhs = {
        use rand::RngExt;
        use rand::SeedableRng;
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42);
        (0..vtx2xyz.len() * num_vdim)
            .map(|_| rng.random())
            .collect::<Vec<f32>>()
    };
    let lambda = 1f32;
    let mut vtx2lhs = vec![0f32; vtx2xyz.len() * num_vdim];
    let res0: f32 = {
        let mut vtx2tmp = vec![0f32; vtx2xyz.len() * num_vdim];
        multiply_graph_laplacian::<usize>(&vtx2idx, &idx2vtx, num_vdim, &vtx2lhs, &mut vtx2tmp);
        vtx2tmp
            .iter()
            .zip(&vtx2lhs)
            .zip(&vtx2rhs)
            .map(|((t, l), r)| {
                let v = t * lambda + l - r;
                v * v
            })
            .sum()
    };
    dbg!(res0);
    assert!(res0 > 1000.);
    {
        let mut vtx2lhs_tmp = vtx2lhs.clone();
        graph_screend_poisson::<usize>(
            &vtx2idx,
            &idx2vtx,
            lambda,
            num_vdim,
            &mut vtx2lhs,
            &vtx2rhs,
            100,
            &mut vtx2lhs_tmp,
        );
    }
    let res1: f32 = {
        let mut vtx2tmp = vec![0f32; vtx2xyz.len() * num_vdim];
        multiply_graph_laplacian::<usize>(&vtx2idx, &idx2vtx, num_vdim, &vtx2lhs, &mut vtx2tmp);
        vtx2tmp
            .iter()
            .zip(&vtx2lhs)
            .zip(&vtx2rhs)
            .map(|((t, l), r)| {
                let v = t * lambda + l - r;
                v * v
            })
            .sum()
    };
    dbg!(res1);
    assert!(res1 < 1.0e-9);
}

/// Merge two CSR adjacency arrays into one, deduplicating shared neighbors.
pub fn merge<IDX>(
    vtx2idx_offset: &[IDX],
    idx2vtx: &[IDX],
    vtx2jdx_offset: &[IDX],
    jdx2vtx: &[IDX],
    num_vtx: usize,
) -> (Vec<IDX>, Vec<IDX>)
where
    IDX: num_traits::PrimInt + AsPrimitive<usize>,
    usize: AsPrimitive<IDX>,
{
    assert_eq!(vtx2idx_offset.len(), num_vtx + 1);
    assert_eq!(vtx2jdx_offset.len(), num_vtx + 1);
    let mut vtx2flg = vec![usize::MAX; num_vtx];
    let mut vtx2kdx_offset = vec![IDX::zero(); num_vtx + 1];
    // pass 1: count unique neighbors per vertex
    for i_vtx in 0..num_vtx {
        let i0: usize = vtx2idx_offset[i_vtx].as_();
        let i1: usize = vtx2idx_offset[i_vtx + 1].as_();
        for &j in &idx2vtx[i0..i1] {
            let j: usize = j.as_();
            if vtx2flg[j] != i_vtx {
                vtx2flg[j] = i_vtx;
                vtx2kdx_offset[i_vtx + 1] = vtx2kdx_offset[i_vtx + 1] + IDX::one();
            }
        }
        let j0: usize = vtx2jdx_offset[i_vtx].as_();
        let j1: usize = vtx2jdx_offset[i_vtx + 1].as_();
        for &j in &jdx2vtx[j0..j1] {
            let j: usize = j.as_();
            if vtx2flg[j] != i_vtx {
                vtx2flg[j] = i_vtx;
                vtx2kdx_offset[i_vtx + 1] = vtx2kdx_offset[i_vtx + 1] + IDX::one();
            }
        }
    }
    // prefix sum
    for i_vtx in 0..num_vtx {
        vtx2kdx_offset[i_vtx + 1] = vtx2kdx_offset[i_vtx + 1] + vtx2kdx_offset[i_vtx];
    }
    let num_kdx: usize = vtx2kdx_offset[num_vtx].as_();
    let mut kdx2vtx = vec![IDX::zero(); num_kdx];
    // pass 2: fill neighbor list
    vtx2flg.iter_mut().for_each(|v| *v = usize::MAX);
    let mut vtx2cur = vtx2kdx_offset[..num_vtx].to_vec();
    for i_vtx in 0..num_vtx {
        let i0: usize = vtx2idx_offset[i_vtx].as_();
        let i1: usize = vtx2idx_offset[i_vtx + 1].as_();
        for &j in &idx2vtx[i0..i1] {
            let j: usize = j.as_();
            if vtx2flg[j] != i_vtx {
                vtx2flg[j] = i_vtx;
                let pos: usize = vtx2cur[i_vtx].as_();
                kdx2vtx[pos] = j.as_();
                vtx2cur[i_vtx] = vtx2cur[i_vtx] + IDX::one();
            }
        }
        let j0: usize = vtx2jdx_offset[i_vtx].as_();
        let j1: usize = vtx2jdx_offset[i_vtx + 1].as_();
        for &j in &jdx2vtx[j0..j1] {
            let j: usize = j.as_();
            if vtx2flg[j] != i_vtx {
                vtx2flg[j] = i_vtx;
                let pos: usize = vtx2cur[i_vtx].as_();
                kdx2vtx[pos] = j.as_();
                vtx2cur[i_vtx] = vtx2cur[i_vtx] + IDX::one();
            }
        }
    }
    (vtx2kdx_offset, kdx2vtx)
}