leibniz 0.2.0

//! Boundary loss signal sampling.

use ::burn::tensor::{Bool, Int, Tensor, backend::Backend};

use crate::burn::{filter::Weights, geometry::Coordinates};

const EPSILON: f32 = 1e-6;

/// Loss signal image with shape `[height, width]`.
pub type Image<B> = Tensor<B, 2>;

/// Sampled loss signal values with shape `[samples]`.
pub type Values<B> = Tensor<B, 1>;

/// Sample a grayscale loss signal at boundary points.
///
/// Mirrors DiffVG's `gather_d_color`: every pixel within `radius` of a boundary
/// point contributes its loss signal divided by that pixel's `weight_image`,
/// the per-pixel forward filter weight. For the box filter the constant filter
/// weight cancels, so the gather reduces to a sum of `image / weights` over the
/// footprint.
///
/// # Panics
///
/// Panics if the image is empty, if it does not match the weights, or if the
/// boundary point columns do not have matching non-empty shapes.
pub fn sample<B: Backend>(
    image: Image<B>,
    points: Coordinates<B>,
    radius: f32,
    weights: Weights<B>,
) -> Values<B> {
    let [height, width] = image.dims();
    let (x, y) = points;
    let [sample_count] = x.dims();

    assert!(
        height > 0
            && width > 0
            && sample_count > 0
            && y.dims() == [sample_count]
            && weights.dims() == [height, width],
        "image must be non-empty, match the weights, and point columns must have matching non-empty shapes"
    );

    // Precompute the weighted image once; each footprint cell then only gathers
    // and masks samples instead of repeating the image/weight division.
    let device = image.device();
    let weighted = (image / weights.clamp_min(EPSILON)).reshape([height * width]);
    let x_floor = x.clone().floor();
    let y_floor = y.clone().floor();
    let zeros = Tensor::<B, 1>::zeros([sample_count], &device);
    let mut values = zeros.clone();

    // Box-filter footprint of ceil(radius) cells on each side, mirroring
    // DiffVG's gather_d_color. The inclusive `radius` cutoff matches DiffVG's
    // exclusive `|d| > radius`, so a point on a pixel edge lands in both cells,
    // whose contributions are summed. Radius 0.5 can only touch the current
    // floor cell and the previous cell on each axis; wider radii use the
    // general clipped footprint.
    let (minimum_offset, maximum_offset) = if radius == 0.5 {
        (-1, 0)
    } else {
        let offset = (radius.ceil() as i32).min(width.max(height) as i32);

        (-offset, offset)
    };

    // Hoist the filter predicate, the bounds checks, and the integer indices
    // out of the footprint loop: a whole row or column of cells shares its
    // axis, so each is computed once per offset instead of once per cell.
    let columns = (minimum_offset..=maximum_offset)
        .map(|offset| axis(&x, &x_floor, offset, width, radius))
        .collect::<Vec<_>>();
    let rows = (minimum_offset..=maximum_offset)
        .map(|offset| {
            let (inside, index) = axis(&y, &y_floor, offset, height, radius);

            (inside, index * width as i64)
        })
        .collect::<Vec<_>>();

    for (row_inside, row_index) in &rows {
        for (column_inside, column_index) in &columns {
            let inside_filter = row_inside.clone().bool_and(column_inside.clone());
            let index = row_index.clone() + column_index.clone();
            // Gather instead of select: the backend select funnels the index
            // tensor through a per-element iterator into a vector before the
            // lookup, measuring several times slower than gather for
            // sample-sized lookups.
            let sample = weighted.clone().gather(0, index);

            values = values + zeros.clone().mask_where(inside_filter, sample);
        }
    }

    values
}

fn axis<B: Backend>(
    coordinate: &Tensor<B, 1>,
    floor: &Tensor<B, 1>,
    offset: i32,
    extent: usize,
    radius: f32,
) -> (Tensor<B, 1, Bool>, Tensor<B, 1, Int>) {
    let candidate = floor.clone() + offset as f32;
    let center = candidate.clone() + 0.5;
    let inside = (center - coordinate.clone())
        .abs()
        .lower_equal_elem(radius)
        .bool_and(candidate.clone().greater_equal_elem(0.0))
        .bool_and(candidate.clone().lower_elem(extent as f32))
        .bool_and(coordinate.clone().greater_equal_elem(0.0))
        .bool_and(coordinate.clone().lower_elem(extent as f32));
    let index = candidate.int().clamp(0, extent as i64 - 1);

    (inside, index)
}

#[cfg(test)]
mod tests {
    use super::super::tests::{assert_floats, image, samples};
    use super::sample;

    #[test]
    fn divides_by_weights() {
        let values = sample(
            image([[2.0, 4.0], [6.0, 8.0]]),
            (samples([0.5, 1.5, 0.5, 1.5]), samples([0.5, 0.5, 1.5, 1.5])),
            0.5,
            image([[2.0, 2.0], [2.0, 2.0]]),
        );

        assert_floats(values, [1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn gathers_within_wider_radius() {
        let values = sample(
            image([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
            (samples([1.5]), samples([0.5])),
            1.0,
            image([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]),
        );

        assert_floats(values, [21.0]);
    }

    #[test]
    fn samples_grayscale_loss_signal() {
        let values = sample(
            image([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
            (
                samples([0.5, 1.5, 0.5, 1.0, 3.0]),
                samples([0.5, 0.5, 1.5, 1.0, 1.0]),
            ),
            0.5,
            image([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]),
        );

        assert_floats(values, [1.0, 2.0, 4.0, 12.0, 0.0]);
    }
}