mod linear;
mod quadratic;
use ::burn::tensor::{Bool, Tensor, backend::Backend};
use crate::base::geometry::{Command, Indices};
use super::{Contour, contour};
pub type Mask<B> = Tensor<B, 1, Bool>;
type Values<B> = Tensor<B, 1>;
pub fn contains<B, I>(contours: I, x: Tensor<B, 1>, y: Tensor<B, 1>) -> Mask<B>
where
B: Backend,
I: IntoIterator<Item = Contour<B>>,
{
evaluate(contours, x, y).not_equal_elem(0.0)
}
fn evaluate<B, I>(contours: I, x: Values<B>, y: Values<B>) -> Values<B>
where
B: Backend,
I: IntoIterator<Item = Contour<B>>,
{
assert!(
x.dims() == y.dims(),
"sample coordinate columns must have matching shapes"
);
let mut values = x.zeros_like();
for contour in contours {
values = values + evaluate_contour(contour, x.clone(), y.clone());
}
values
}
fn evaluate_contour<B: Backend>(contour: Contour<B>, x: Values<B>, y: Values<B>) -> Values<B> {
let [segment_count, point_count, coordinate_count] = contour.dims();
assert!(
segment_count > 0
&& point_count == 2
&& coordinate_count == 2
&& contour.commands().len() == segment_count,
"contours must have shape [segments, 2, 2]"
);
let mut values = x.zeros_like();
let arguments = contour.arguments();
for segment in 0..segment_count {
let indices = Indices::new(segment, segment_count);
values = values
+ match contour.command(segment) {
Command::Linear => {
let (x_coefficients, y_coefficients) =
linear::coefficients(arguments.clone(), indices);
linear::evaluate(x_coefficients, y_coefficients, x.clone(), y.clone())
}
Command::Quadratic => {
let (x_coefficients, y_coefficients) =
quadratic::coefficients(arguments.clone(), indices);
quadratic::evaluate(x_coefficients, y_coefficients, x.clone(), y.clone())
}
};
}
values
}
fn segment_coordinate<B: Backend>(
arguments: contour::Arguments<B>,
point_index: [usize; 2],
coordinate: usize,
) -> Values<B> {
arguments
.slice([
point_index[0]..point_index[0] + 1,
point_index[1]..point_index[1] + 1,
coordinate..coordinate + 1,
])
.squeeze_dim::<2>(0)
.squeeze_dim::<1>(0)
}
#[cfg(test)]
mod tests {
use ::burn::tensor::{Tensor, TensorData};
use crate::{base::geometry::Command, burn::tests::Backend};
#[test]
fn classifies_samples_inside_curved_quadratic_contour() {
let inside = contour_contains(
[curved_square()],
sample_coordinates([[1.25, 0.5], [1.8, 0.5]]),
);
assert_bool(inside, [true, false]);
}
#[test]
fn classifies_samples_inside_multiple_contours() {
let inside = contour_contains(
[
square([0.0, 0.0], [4.0, 4.0], Direction::CounterClockwise),
square([1.0, 1.0], [3.0, 3.0], Direction::Clockwise),
],
sample_coordinates([[0.5, 0.5], [2.0, 2.0], [5.0, 2.0]]),
);
assert_bool(inside, [true, false, false]);
}
#[test]
fn classifies_samples_inside_one_contour() {
let inside = contour_contains(
[square([0.0, 0.0], [1.0, 1.0], Direction::CounterClockwise)],
sample_coordinates([[0.5, 0.5], [1.5, 0.5], [-0.5, 0.5]]),
);
assert_bool(inside, [true, false, false]);
}
#[test]
#[should_panic]
fn rejects_invalid_contours() {
let arguments = Tensor::<Backend, 3>::from_data(
TensorData::new(vec![0.0; 6], [1, 3, 2]),
&Default::default(),
);
let _ = evaluate_contour_values(
super::Contour::new(vec![Command::Quadratic], arguments),
sample_coordinates([[0.5, 0.5]]),
);
}
#[test]
fn winding_sign_depends_on_direction() {
let ccw = evaluate_contour_values(
square([0.0, 0.0], [1.0, 1.0], Direction::CounterClockwise),
sample_coordinates([[0.5, 0.5]]),
);
let cw = evaluate_contour_values(
square([0.0, 0.0], [1.0, 1.0], Direction::Clockwise),
sample_coordinates([[0.5, 0.5]]),
);
assert_close(scalar(ccw), 1.0);
assert_close(scalar(cw), -1.0);
}
enum Direction {
Clockwise,
CounterClockwise,
}
fn assert_bool<const N: usize>(
tensor: Tensor<Backend, 1, ::burn::tensor::Bool>,
expected: [bool; N],
) {
let actual = tensor.into_data().to_vec::<bool>().unwrap();
assert_eq!(actual, expected);
}
fn assert_close(actual: f32, expected: f32) {
assert!((actual - expected).abs() < 1e-6);
}
fn contour_contains(
contours: impl IntoIterator<Item = super::Contour<Backend>>,
(x, y): (Tensor<Backend, 1>, Tensor<Backend, 1>),
) -> Tensor<Backend, 1, ::burn::tensor::Bool> {
super::contains(contours, x, y)
}
fn curved_square() -> super::Contour<Backend> {
super::Contour::new(
vec![Command::Quadratic; 4],
Tensor::<Backend, 3>::from_data(
TensorData::from([
[[0.0, 0.0], [0.5, 0.0]],
[[1.0, 0.0], [2.0, 0.5]],
[[1.0, 1.0], [0.5, 1.0]],
[[0.0, 1.0], [0.0, 0.5]],
]),
&Default::default(),
),
)
}
fn evaluate_contour_values(
contour: super::Contour<Backend>,
(x, y): (Tensor<Backend, 1>, Tensor<Backend, 1>),
) -> Tensor<Backend, 1> {
super::evaluate([contour], x, y)
}
fn interpolate(start: [f32; 2], end: [f32; 2], t: f32) -> [f32; 2] {
[
start[0] + (end[0] - start[0]) * t,
start[1] + (end[1] - start[1]) * t,
]
}
fn sample_coordinates<const N: usize>(
values: [[f32; 2]; N],
) -> (Tensor<Backend, 1>, Tensor<Backend, 1>) {
let x = values.iter().map(|value| value[0]).collect::<Vec<_>>();
let y = values.iter().map(|value| value[1]).collect::<Vec<_>>();
(
Tensor::<Backend, 1>::from_data(TensorData::new(x, [N]), &Default::default()),
Tensor::<Backend, 1>::from_data(TensorData::new(y, [N]), &Default::default()),
)
}
fn scalar(tensor: Tensor<Backend, 1>) -> f32 {
tensor.into_scalar()
}
fn square(min: [f32; 2], max: [f32; 2], direction: Direction) -> super::Contour<Backend> {
let segments = match direction {
Direction::Clockwise => square_segments(
[min[0], min[1]],
[min[0], max[1]],
[max[0], max[1]],
[max[0], min[1]],
),
Direction::CounterClockwise => square_segments(
[min[0], min[1]],
[max[0], min[1]],
[max[0], max[1]],
[min[0], max[1]],
),
};
super::Contour::new(
vec![Command::Linear; 4],
Tensor::<Backend, 3>::from_data(TensorData::from(segments), &Default::default()),
)
}
fn square_segments(a: [f32; 2], b: [f32; 2], c: [f32; 2], d: [f32; 2]) -> [[[f32; 2]; 2]; 4] {
[
[a, interpolate(a, b, 0.5)],
[b, interpolate(b, c, 0.5)],
[c, interpolate(c, d, 0.5)],
[d, interpolate(d, a, 0.5)],
]
}
}