Merge pull request #6 from georgypv/feature/add-distance-bboxes

Feature/add distance bboxes

.github/workflows/CI.yml

@@ -33,7 +33,7 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.target }}
-          args: --release --out dist --find-interpreter
+          args: --release --out dist -i 3.10
           sccache: 'true'
           manylinux: auto
       - name: Upload wheels
@@ -57,7 +57,7 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.target }}
-          args: --release --out dist --find-interpreter
+          args: --release --out dist -i 3.10
           sccache: 'true'
       - name: Upload wheels
         uses: actions/upload-artifact@v3
@@ -79,7 +79,7 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.target }}
-          args: --release --out dist --find-interpreter
+          args: --release --out dist -i 3.10
           sccache: 'true'
       - name: Upload wheels
         uses: actions/upload-artifact@v3

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars_coord_transforms"
-version = "0.11.0"
+version = "0.12.0"
 edition = "2021"
 
 [lib]

polars_coord_transforms/__init__.py

@@ -177,6 +177,14 @@ def cosine_similarity_3d(self, other: pl.Expr) -> pl.Expr:
             is_elementwise=True
         )
 
+    def bboxes_2d(self, other: pl.Expr) -> pl.Expr:
+        return register_plugin_function(
+            plugin_path=Path(__file__).parent,
+            function_name="bboxes_2d",
+            args=[self._expr, other],
+            is_elementwise=True
+        )
+
 
 class CoordTransformExpr(pl.Expr):
     @property

pyproject.toml

@@ -13,11 +13,11 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-version = "0.11.0"
+version = "0.12.0"
 authors = [
     {name="Georgy Popov"}
 ]
 description = "Tools for coordinate transforms and S2-indexing via Polars Expression Plugins"
 
 [project.urls]
-"Homepage" = "https://github.com/georgypv/polars-coord-transforms/"
+"Homepage" = "https://github.com/georgypv/polars-coord-transforms/"

src/coord_transforms.rs

@@ -1,7 +1,7 @@
 extern crate nalgebra as na;
 
 use map_3d::{deg2rad, ecef2geodetic, geodetic2ecef, rad2deg, Ellipsoid};
-use na::{Quaternion, Rotation3, UnitQuaternion, Vector3, Vector4, Matrix4x3, Matrix3x4, Matrix1x4, Matrix4x1};
+use na::{Quaternion, Rotation3, UnitQuaternion, Vector3, Vector4, Matrix4x3, Matrix3x4, Matrix1x4};
 use nalgebra::Const;
 use utm::{lat_lon_to_zone_number, to_utm_wgs84_no_zone};
 
@@ -127,7 +127,7 @@ mod transform_tests {
         let ecef_coords: (f64, f64, f64) =
             (2830593.6327610738, 2062375.5703225536, 5312896.0721501345);
         let expected_result: (f64, f64, f64) =
-            (36.077147686805766, 56.783927007002866, 165.8986865637805);
+            (36.077147686805766, 56.783927007002845, 165.8986865637805);
 
         assert_eq!(
             ecef_to_lla_elementwise(ecef_coords.0, ecef_coords.1, ecef_coords.2),

src/distance.rs

@@ -1,3 +1,9 @@
+#[derive(Debug, Clone, Copy)]
+pub struct PointCoords {
+    pub x: f64,
+    pub y: f64,
+}
+
 pub fn euclidean_3d_elementwise(x1: f64, y1: f64, z1: f64, x2: f64, y2: f64, z2: f64) -> f64 {
     (((x2 - x1).powi(2)) + ((y2 - y1).powi(2)) + ((z2 - z1).powi(2))).sqrt()
 }
@@ -38,3 +44,96 @@ pub fn cosine_similarity_3d_elementwise(
     };
     res
 }
+
+fn round(num: f64, precision: u8) -> f64 {
+    let multiplier = 10f64.powi(precision as i32);
+    (num * multiplier).round() / multiplier
+}
+
+pub fn dist_to_segment(point: PointCoords, start: PointCoords, end: PointCoords) -> f64 {
+    let l2 = euclidean_2d_elementwise(start.x, start.y, end.x, end.y).powi(2);
+    if l2 == 0.0 {
+        return euclidean_2d_elementwise(point.x, point.y,  start.x, start.y).powi(2);
+    }
+    let t = ((point.x - start.x) * (end.x - start.x) + (point.y - start.y) * (end.y - start.y)) / l2;
+    let t = t.max(0.0).min(1.0);
+    let projection = PointCoords{x: start.x + t * (end.x - start.x), y: start.y + t * (end.y - start.y)};
+    let dist_to_segment = euclidean_2d_elementwise(point.x, point.y, projection.x, projection.y);
+    dist_to_segment
+}
+
+pub fn bboxes_2d_elementwise(box1: [PointCoords; 4], box2: [PointCoords; 4]) -> f64 {
+    let mut min_distance = f64::MAX;
+
+    // Sides of the first box
+    for i in 0..4 {
+        let start = box1[i];
+        let end = box1[(i + 1) % 4];
+
+        // Distance from each apex of the second box to the side of the first box
+        for &point in &box2 {
+            let distance = dist_to_segment(point, start, end);
+            min_distance = min_distance.min(distance);
+        }
+    }
+
+    // Sides of the second box
+    for i in 0..4 {
+        let start = box2[i];
+        let end = box2[(i + 1) % 4];
+
+        // Distance from each apex of the first box to the side of the second box
+        for &point in &box1 {
+            let distance = dist_to_segment(point, start, end);
+            min_distance = min_distance.min(distance);
+        }
+    }
+
+    round(min_distance, 5)
+}
+
+
+#[cfg(test)]
+mod distance_tests {
+    use crate::distance::{dist_to_segment, bboxes_2d_elementwise, PointCoords};
+
+    #[test]
+    fn test_dist_to_segment() {
+
+        let point  =  PointCoords { x: 5.0, y: 2.0};
+        let start  = PointCoords { x: 1.0, y: 0.0};
+        let end = PointCoords { x: 3.0, y: 0.0};
+
+        let expected_distance = 2.8284271247461903;
+
+        assert_eq!(
+            dist_to_segment(point, start, end),
+                    expected_distance
+                )
+    }
+
+    #[test]
+    fn test_bboxes_2d_elementwise() {
+
+        let bbox1 = [
+            PointCoords {x: 2.0, y: 0.0,},
+            PointCoords {x: 0.0, y: 3.0,},
+            PointCoords {x: 2.0, y: 4.0,},
+            PointCoords {x: 4.0, y: 1.0,},
+        ];
+
+        let bbox2 = [
+            PointCoords {x: 3.0, y: 5.0,},
+            PointCoords {x: 3.0, y: 8.0,},
+            PointCoords {x: 5.0, y: 8.0,},
+            PointCoords {x: 5.0, y: 5.0,},
+        ];
+
+        let expected_distance = 1.41421;
+
+        assert_eq!(
+            bboxes_2d_elementwise(bbox1, bbox2),
+                    expected_distance
+                )
+    }
+}

src/expressions.rs

@@ -769,3 +769,100 @@ fn cosine_similarity_3d(inputs: &[Series]) -> PolarsResult<Series> {
         iter.collect_ca_with_dtype("cosine_similarity", DataType::Float64);
     Ok(out_ca.into_series())
 }
+
+#[polars_expr(output_type = Float64)]
+fn bboxes_2d(inputs: &[Series]) -> PolarsResult<Series> {
+    let box_a = inputs[0].struct_()?;
+    let box_b = inputs[1].struct_()?;
+
+
+    let out_chunked = get_distance_between_bboxes_2d(&box_a, &box_b);
+
+    Ok(out_chunked?.into_series())
+}
+
+fn get_distance_between_bboxes_2d(
+    box1: &StructChunked,
+    box2: &StructChunked,
+) -> Result<Series, PolarsError> {
+
+    let bbox1_1_x_ser = &box1.field_by_name("bbox_1_x")?;
+    let bbox1_1_y_ser = &box1.field_by_name("bbox_1_y")?;
+    let bbox1_2_x_ser = &box1.field_by_name("bbox_2_x")?;
+    let bbox1_2_y_ser = &box1.field_by_name("bbox_2_y")?;
+    let bbox1_3_x_ser = &box1.field_by_name("bbox_3_x")?;
+    let bbox1_3_y_ser = &box1.field_by_name("bbox_3_y")?;
+    let bbox1_4_x_ser = &box1.field_by_name("bbox_4_x")?;
+    let bbox1_4_y_ser = &box1.field_by_name("bbox_4_y")?;
+
+    let bbox2_1_x_ser = &box2.field_by_name("bbox_1_x")?;
+    let bbox2_1_y_ser = &box2.field_by_name("bbox_1_y")?;
+    let bbox2_2_x_ser = &box2.field_by_name("bbox_2_x")?;
+    let bbox2_2_y_ser = &box2.field_by_name("bbox_2_y")?;
+    let bbox2_3_x_ser = &box2.field_by_name("bbox_3_x")?;
+    let bbox2_3_y_ser = &box2.field_by_name("bbox_3_y")?;
+    let bbox2_4_x_ser = &box2.field_by_name("bbox_4_x")?;
+    let bbox2_4_y_ser = &box2.field_by_name("bbox_4_y")?;
+
+
+    let mut distances: PrimitiveChunkedBuilder<Float64Type> =
+        PrimitiveChunkedBuilder::new("distance".into(), box1.len());
+
+    for (
+        bbox1_1_x,
+        bbox1_1_y,
+        bbox1_2_x,
+        bbox1_2_y,
+        bbox1_3_x,
+        bbox1_3_y,
+        bbox1_4_x,
+        bbox1_4_y,
+        bbox2_1_x,
+        bbox2_1_y,
+        bbox2_2_x,
+        bbox2_2_y,
+        bbox2_3_x,
+        bbox2_3_y,
+        bbox2_4_x,
+        bbox2_4_y
+    ) in izip!(
+        bbox1_1_x_ser.f64().unwrap(),
+        bbox1_1_y_ser.f64().unwrap(),
+        bbox1_2_x_ser.f64().unwrap(),
+        bbox1_2_y_ser.f64().unwrap(),
+        bbox1_3_x_ser.f64().unwrap(),
+        bbox1_3_y_ser.f64().unwrap(),
+        bbox1_4_x_ser.f64().unwrap(),
+        bbox1_4_y_ser.f64().unwrap(),
+        bbox2_1_x_ser.f64().unwrap(),
+        bbox2_1_y_ser.f64().unwrap(),
+        bbox2_2_x_ser.f64().unwrap(),
+        bbox2_2_y_ser.f64().unwrap(),
+        bbox2_3_x_ser.f64().unwrap(),
+        bbox2_3_y_ser.f64().unwrap(),
+        bbox2_4_x_ser.f64().unwrap(),
+        bbox2_4_y_ser.f64().unwrap()
+    ) {
+        let bbox1 = [
+            PointCoords { x: bbox1_1_x.unwrap(), y: bbox1_1_y.unwrap()},
+            PointCoords { x: bbox1_2_x.unwrap(), y: bbox1_2_y.unwrap()},
+            PointCoords { x: bbox1_3_x.unwrap(), y: bbox1_3_y.unwrap()},
+            PointCoords { x: bbox1_4_x.unwrap(), y: bbox1_4_y.unwrap()},
+        ];
+
+        let bbox2 = [
+            PointCoords { x: bbox2_1_x.unwrap(), y: bbox2_1_y.unwrap()},
+            PointCoords { x: bbox2_2_x.unwrap(), y: bbox2_2_y.unwrap()},
+            PointCoords { x: bbox2_3_x.unwrap(), y: bbox2_3_y.unwrap()},
+            PointCoords { x: bbox2_4_x.unwrap(), y: bbox2_4_y.unwrap()},
+        ];
+
+        let distance = bboxes_2d_elementwise(bbox1, bbox2);
+
+        distances.append_value(distance);
+    }
+
+    let ser_out_distances = distances.finish().into_series();
+
+    Ok(ser_out_distances)
+}