From 77d63f96bb2f5fb8760a1e4473da503535e211f6 Mon Sep 17 00:00:00 2001 From: GeorgyPopov Date: Thu, 8 Feb 2024 17:15:34 +0300 Subject: [PATCH] fmt code --- src/coord_transforms.rs | 97 +++++---- src/distance.rs | 25 ++- src/expressions.rs | 472 ++++++++++++++++++++-------------------- src/lib.rs | 4 +- src/s2_functions.rs | 23 +- 5 files changed, 322 insertions(+), 299 deletions(-) diff --git a/src/coord_transforms.rs b/src/coord_transforms.rs index 1cb5aa8..ab1db95 100644 --- a/src/coord_transforms.rs +++ b/src/coord_transforms.rs @@ -1,28 +1,37 @@ extern crate nalgebra as na; -use utm::{to_utm_wgs84_no_zone, lat_lon_to_zone_number}; -use na::{Vector3, Vector4, Quaternion, UnitQuaternion, Rotation3}; -use map_3d::{ecef2geodetic, geodetic2ecef, rad2deg, deg2rad, Ellipsoid}; +use map_3d::{deg2rad, ecef2geodetic, geodetic2ecef, rad2deg, Ellipsoid}; +use na::{Quaternion, Rotation3, UnitQuaternion, Vector3, Vector4}; +use utm::{lat_lon_to_zone_number, to_utm_wgs84_no_zone}; fn rotation_from_quat(rotation: Vec) -> na::Rotation { - let quat = UnitQuaternion::from_quaternion(Quaternion::from_vector(Vector4::from_vec(rotation))); + let quat = + UnitQuaternion::from_quaternion(Quaternion::from_vector(Vector4::from_vec(rotation))); Rotation3::from(quat) -} - -pub fn map_to_ecef_elementwise(map_coords: Vec, rotation: Vec, offset: Vec) -> (f64, f64, f64) { +} +pub fn map_to_ecef_elementwise( + map_coords: Vec, + rotation: Vec, + offset: Vec, +) -> (f64, f64, f64) { let r: na::Rotation = rotation_from_quat(rotation); - let ecef_vector = r.transform_vector(&Vector3::from_vec(map_coords)) + Vector3::from_vec(offset); + let ecef_vector = + r.transform_vector(&Vector3::from_vec(map_coords)) + Vector3::from_vec(offset); (ecef_vector.x, ecef_vector.y, ecef_vector.z) } -pub fn ecef_to_map_elementwise(ecef_coords: Vec, rotation: Vec, offset: Vec) -> (f64, f64, f64) { +pub fn ecef_to_map_elementwise( + ecef_coords: Vec, + rotation: Vec, + offset: Vec, +) -> (f64, f64, f64) { let r_inverse: na::Rotation = rotation_from_quat(rotation).inverse(); - - let map_vector = r_inverse.transform_vector(&(Vector3::from_vec(ecef_coords) - Vector3::from_vec(offset))); - (map_vector.x, map_vector.y, map_vector.z) + let map_vector = + r_inverse.transform_vector(&(Vector3::from_vec(ecef_coords) - Vector3::from_vec(offset))); + (map_vector.x, map_vector.y, map_vector.z) } pub fn ecef_to_lla_elementwise(x: f64, y: f64, z: f64) -> (f64, f64, f64) { @@ -35,60 +44,70 @@ pub fn lla_to_ecef_elementwise(lon: f64, lat: f64, alt: f64) -> (f64, f64, f64) (x, y, z) } - - pub fn lla_to_utm_zone_number_elementwise(lon: f64, lat: f64) -> u8 { let zone_number = lat_lon_to_zone_number(lat, lon); zone_number -} - +} pub fn lla_to_utm_elementwise(lon: f64, lat: f64, alt: f64) -> (f64, f64, f64) { let (northing, easting, _meridian_convergence) = to_utm_wgs84_no_zone(lat, lon); (easting, northing, alt) } - -pub fn rotate_map_coords_elementwise(map_coords: Vec, rotation: Vec, scale: Vec) -> (f64, f64, f64) { +pub fn rotate_map_coords_elementwise( + map_coords: Vec, + rotation: Vec, + scale: Vec, +) -> (f64, f64, f64) { let r: na::Rotation = rotation_from_quat(rotation); let scale_rotated = r.transform_vector(&Vector3::from_vec(scale)); let map_coords_rotated = Vector3::from_vec(map_coords) + scale_rotated; - (map_coords_rotated.x, map_coords_rotated.y, map_coords_rotated.z) - + ( + map_coords_rotated.x, + map_coords_rotated.y, + map_coords_rotated.z, + ) } - -pub fn interpolate_linear_elementwise(coords: Vec, other: Vec, coef: f64) -> (f64, f64, f64) { - let interpolated = (Vector3::from_vec(coords) * coef) + (Vector3::from_vec(other) * (1.0 - coef)); +pub fn interpolate_linear_elementwise( + coords: Vec, + other: Vec, + coef: f64, +) -> (f64, f64, f64) { + let interpolated = + (Vector3::from_vec(coords) * coef) + (Vector3::from_vec(other) * (1.0 - coef)); (interpolated.x, interpolated.y, interpolated.z) -} - +} #[cfg(test)] -mod transform_tests { - use crate::coord_transforms::{map_to_ecef_elementwise, ecef_to_lla_elementwise}; +mod transform_tests { + use crate::coord_transforms::{ecef_to_lla_elementwise, map_to_ecef_elementwise}; #[test] fn test_map_to_ecef() { - let map_coords: Vec = vec![-97066.730132, 122807.787398, -1888.737721]; + let map_coords: Vec = vec![-97066.730132, 122807.787398, -1888.737721]; let rotation: Vec = vec![0.13007119, 0.26472049, 0.85758219, 0.42137553]; let offset: Vec = vec![2852423.40536658, 2201848.41975346, 5245234.74365368]; - let expected_result: (f64, f64, f64) = (2830593.6327610738, 2062375.5703225536, 5312896.0721501345); - - assert_eq!(map_to_ecef_elementwise(map_coords, rotation, offset), expected_result) + let expected_result: (f64, f64, f64) = + (2830593.6327610738, 2062375.5703225536, 5312896.0721501345); + assert_eq!( + map_to_ecef_elementwise(map_coords, rotation, offset), + expected_result + ) } #[test] fn test_ecef_to_lla() { - - let ecef_coords: (f64, f64, f64) = (2830593.6327610738, 2062375.5703225536, 5312896.0721501345); - let expected_result: (f64, f64, f64) = (36.077147686805766, 56.783927007002866, 165.8986865637805); - - assert_eq!(ecef_to_lla_elementwise(ecef_coords.0, ecef_coords.1, ecef_coords.2), expected_result) + let ecef_coords: (f64, f64, f64) = + (2830593.6327610738, 2062375.5703225536, 5312896.0721501345); + let expected_result: (f64, f64, f64) = + (36.077147686805766, 56.783927007002866, 165.8986865637805); + + assert_eq!( + ecef_to_lla_elementwise(ecef_coords.0, ecef_coords.1, ecef_coords.2), + expected_result + ) } - } - - diff --git a/src/distance.rs b/src/distance.rs index 662f8fc..2090b48 100644 --- a/src/distance.rs +++ b/src/distance.rs @@ -1,35 +1,40 @@ 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() + (((x2 - x1).powi(2)) + ((y2 - y1).powi(2)) + ((z2 - z1).powi(2))).sqrt() } - pub fn euclidean_2d_elementwise(x1: f64, y1: f64, x2: f64, y2: f64) -> f64 { - (((x2 - x1).powi(2)) + ((y2 - y1).powi(2))).sqrt() + (((x2 - x1).powi(2)) + ((y2 - y1).powi(2))).sqrt() } - pub fn cosine_similarity_2d_elementwise(x1: f64, y1: f64, x2: f64, y2: f64) -> f64 { - let dot_product = (x1*x2) + (y1*y2); + let dot_product = (x1 * x2) + (y1 * y2); let magnitude1 = (x1.powi(2) + y1.powi(2)).powf(0.5); let magnitude2 = (x2.powi(2) + y2.powi(2)).powf(0.5); let res = if magnitude1 == 0.0 || magnitude2 == 0.0 { 0.0 } else { - dot_product / (magnitude1*magnitude2) + dot_product / (magnitude1 * magnitude2) }; res } -pub fn cosine_similarity_3d_elementwise(x1: f64, y1: f64, z1: f64, x2: f64, y2: f64, z2: f64) -> f64 { - let dot_product = (x1*x2) + (y1*y2) + (z1*z2); +pub fn cosine_similarity_3d_elementwise( + x1: f64, + y1: f64, + z1: f64, + x2: f64, + y2: f64, + z2: f64, +) -> f64 { + let dot_product = (x1 * x2) + (y1 * y2) + (z1 * z2); let magnitude1 = (x1.powi(2) + y1.powi(2) + z1.powi(2)).powf(0.5); let magnitude2 = (x2.powi(2) + y2.powi(2) + z2.powi(2)).powf(0.5); let res = if magnitude1 == 0.0 || magnitude2 == 0.0 { 0.0 } else { - dot_product / (magnitude1*magnitude2) + dot_product / (magnitude1 * magnitude2) }; res -} \ No newline at end of file +} diff --git a/src/expressions.rs b/src/expressions.rs index a909fdf..d985a1c 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -1,21 +1,16 @@ use itertools::Itertools; -use polars::prelude::*; use polars::datatypes::DataType; +use polars::prelude::*; use pyo3_polars::derive::polars_expr; use itertools::izip; use serde::Deserialize; -use crate::s2_functions::*; use crate::coord_transforms::*; use crate::distance::*; +use crate::s2_functions::*; -fn unpack_xyz(ca: &StructChunked, lonlat: bool) -> ( - Series, - Series, - Series -) { - +fn unpack_xyz(ca: &StructChunked, lonlat: bool) -> (Series, Series, Series) { let (field_x, field_y, field_z) = if lonlat { ("lon", "lat", "alt") } else { @@ -24,83 +19,99 @@ fn unpack_xyz(ca: &StructChunked, lonlat: bool) -> ( let x: Series = match ca.field_by_name(field_x) { Ok(series) => series, - Err(_) => panic!("Field `x` not found in `{}`!", &ca.name()) + Err(_) => panic!("Field `x` not found in `{}`!", &ca.name()), }; let y: Series = match ca.field_by_name(field_y) { Ok(field) => field, - Err(_) => panic!("Field `y` not found in `{}`!", &ca.name()) + Err(_) => panic!("Field `y` not found in `{}`!", &ca.name()), }; let z: Series = match ca.field_by_name(field_z) { Ok(field) => field, - Err(_) => panic!("Field `z` not found in `{}`!", &ca.name()) + Err(_) => panic!("Field `z` not found in `{}`!", &ca.name()), }; (x, y, z) -} +} fn apply_rotation_to_map( coords_ca: &StructChunked, rotation_ca: &StructChunked, offset_ca: &StructChunked, result_struct_name: &str, - func_elementwise: impl Fn(Vec, Vec, Vec) -> (f64, f64, f64) - ) -> Result { - - let (x_ser, y_ser, z_ser) = unpack_xyz(coords_ca, false); - let (rotation_x, rotation_y, rotation_z) = unpack_xyz(rotation_ca, false); - let rotation_w = rotation_ca.field_by_name("w").expect("Unable to find `w` field for rotation!"); - - let (offset_x, offset_y, offset_z) = unpack_xyz(offset_ca, false); - - let mut x_cb: PrimitiveChunkedBuilder = - PrimitiveChunkedBuilder::new("x", coords_ca.len()); - let mut y_cb: PrimitiveChunkedBuilder = - PrimitiveChunkedBuilder::new("y", coords_ca.len()); - let mut z_cb: PrimitiveChunkedBuilder = - PrimitiveChunkedBuilder::new("z", coords_ca.len()); - for (x_val, y_val, z_val, rotation_x_val, rotation_y_val, rotation_z_val, rotation_w_val, offset_x_val, offset_y_val, offset_z_val) - in izip!( - x_ser.f64().unwrap(), - y_ser.f64().unwrap(), - z_ser.f64().unwrap(), - rotation_x.f64().unwrap(), - rotation_y.f64().unwrap(), - rotation_z.f64().unwrap(), - rotation_w.f64().unwrap(), - offset_x.f64().unwrap(), - offset_y.f64().unwrap(), - offset_z.f64().unwrap() - ) { - let map_vec = vec![x_val.unwrap(), y_val.unwrap(), z_val.unwrap()]; - let rotation_vec = vec![rotation_x_val.unwrap(), rotation_y_val.unwrap(), rotation_z_val.unwrap(),rotation_w_val.unwrap()]; - let offset_vec = vec![offset_x_val.unwrap(), offset_y_val.unwrap(), offset_z_val.unwrap()]; - - let (x, y, z) = func_elementwise(map_vec, rotation_vec, offset_vec); - - x_cb.append_value(x); - y_cb.append_value(y); - z_cb.append_value(z); - } - - let ser_out_x = x_cb.finish().into_series(); - let ser_out_y = y_cb.finish().into_series(); - let ser_out_z = z_cb.finish().into_series(); - - let out_chunked = StructChunked::new(result_struct_name, &[ser_out_x, ser_out_y, ser_out_z]); - out_chunked - -} + func_elementwise: impl Fn(Vec, Vec, Vec) -> (f64, f64, f64), +) -> Result { + let (x_ser, y_ser, z_ser) = unpack_xyz(coords_ca, false); + let (rotation_x, rotation_y, rotation_z) = unpack_xyz(rotation_ca, false); + let rotation_w = rotation_ca + .field_by_name("w") + .expect("Unable to find `w` field for rotation!"); + + let (offset_x, offset_y, offset_z) = unpack_xyz(offset_ca, false); + + let mut x_cb: PrimitiveChunkedBuilder = + PrimitiveChunkedBuilder::new("x", coords_ca.len()); + let mut y_cb: PrimitiveChunkedBuilder = + PrimitiveChunkedBuilder::new("y", coords_ca.len()); + let mut z_cb: PrimitiveChunkedBuilder = + PrimitiveChunkedBuilder::new("z", coords_ca.len()); + for ( + x_val, + y_val, + z_val, + rotation_x_val, + rotation_y_val, + rotation_z_val, + rotation_w_val, + offset_x_val, + offset_y_val, + offset_z_val, + ) in izip!( + x_ser.f64().unwrap(), + y_ser.f64().unwrap(), + z_ser.f64().unwrap(), + rotation_x.f64().unwrap(), + rotation_y.f64().unwrap(), + rotation_z.f64().unwrap(), + rotation_w.f64().unwrap(), + offset_x.f64().unwrap(), + offset_y.f64().unwrap(), + offset_z.f64().unwrap() + ) { + let map_vec = vec![x_val.unwrap(), y_val.unwrap(), z_val.unwrap()]; + let rotation_vec = vec![ + rotation_x_val.unwrap(), + rotation_y_val.unwrap(), + rotation_z_val.unwrap(), + rotation_w_val.unwrap(), + ]; + let offset_vec = vec![ + offset_x_val.unwrap(), + offset_y_val.unwrap(), + offset_z_val.unwrap(), + ]; + + let (x, y, z) = func_elementwise(map_vec, rotation_vec, offset_vec); + + x_cb.append_value(x); + y_cb.append_value(y); + z_cb.append_value(z); + } + let ser_out_x = x_cb.finish().into_series(); + let ser_out_y = y_cb.finish().into_series(); + let ser_out_z = z_cb.finish().into_series(); + + let out_chunked = StructChunked::new(result_struct_name, &[ser_out_x, ser_out_y, ser_out_z]); + out_chunked +} // SSNameSpace #[derive(Deserialize)] struct S2Kwargs { - level: u64 + level: u64, } - #[polars_expr(output_type=UInt64)] fn lonlat_to_cellid(inputs: &[Series], kwargs: S2Kwargs) -> PolarsResult { - let lonlat_ca = inputs[0].struct_()?; let lon = lonlat_ca.field_by_name("lon")?; @@ -121,18 +132,18 @@ fn lonlat_to_cellid(inputs: &[Series], kwargs: S2Kwargs) -> PolarsResult let lon_ca = lon.f64()?; let lat_ca = lat.f64()?; - let iter = lon_ca.into_iter() + let iter = lon_ca + .into_iter() .zip(lat_ca.into_iter()) .map(|(lon_opt, lat_opt)| match (lon_opt, lat_opt) { (Some(lon), Some(lat)) => lonlat_to_cellid_elementwise(lon, lat, kwargs.level), - _ => 0 + _ => 0, }); - let out_ca: ChunkedArray = iter.collect_ca_with_dtype("s2_cellid", DataType::UInt64); + let out_ca: ChunkedArray = + iter.collect_ca_with_dtype("s2_cellid", DataType::UInt64); Ok(out_ca.into_series()) - } - fn cellid_to_lonlat_output(_: &[Field]) -> PolarsResult { let v: Vec = vec![ Field::new("lon", DataType::Float64), @@ -143,7 +154,6 @@ fn cellid_to_lonlat_output(_: &[Field]) -> PolarsResult { #[polars_expr(output_type_func=cellid_to_lonlat_output)] fn cellid_to_lonlat(inputs: &[Series]) -> PolarsResult { - let ca: &ChunkedArray = inputs[0].u64()?; let mut longitude: PrimitiveChunkedBuilder = @@ -151,27 +161,25 @@ fn cellid_to_lonlat(inputs: &[Series]) -> PolarsResult { let mut latitude: PrimitiveChunkedBuilder = PrimitiveChunkedBuilder::new("lat", ca.len()); - - for cellid_op in ca.into_iter() { - match cellid_op { - Some(cellid) => { - let (lon, lat) = cellid_to_lonlat_elementwise(cellid); - longitude.append_value(lon); - latitude.append_value(lat) - }, - _ => { - longitude.append_null(); - latitude.append_null(); - } + for cellid_op in ca.into_iter() { + match cellid_op { + Some(cellid) => { + let (lon, lat) = cellid_to_lonlat_elementwise(cellid); + longitude.append_value(lon); + latitude.append_value(lat) + } + _ => { + longitude.append_null(); + latitude.append_null(); } } - + } + let ser_lon = longitude.finish().into_series(); let ser_lat = latitude.finish().into_series(); let out_chunked = StructChunked::new("coordinates", &[ser_lon, ser_lat])?; Ok(out_chunked.into_series()) - -} +} #[polars_expr(output_type=Boolean)] fn cell_contains_point(inputs: &[Series]) -> PolarsResult { @@ -181,16 +189,21 @@ fn cell_contains_point(inputs: &[Series]) -> PolarsResult { let lon_ser = lonlat_ca.field_by_name("lon")?; let lat_ser = lonlat_ca.field_by_name("lat")?; - + let lon_ca = lon_ser.f64()?; let lat_ca = lat_ser.f64()?; - let out_ca: ChunkedArray = izip!(cell_ca.into_iter(), lon_ca.into_iter(), lat_ca.into_iter()).map( - | (cellid_op, lon_op, lat_op) | match (cellid_op, lon_op, lat_op) { - (Some(cellid), Some(lon), Some(lat)) => cell_contains_point_elementwise(cellid, lon, lat), - _ => false - } - ).collect_ca(cell_ca.name()); + let out_ca: ChunkedArray = + izip!(cell_ca.into_iter(), lon_ca.into_iter(), lat_ca.into_iter()) + .map( + |(cellid_op, lon_op, lat_op)| match (cellid_op, lon_op, lat_op) { + (Some(cellid), Some(lon), Some(lat)) => { + cell_contains_point_elementwise(cellid, lon, lat) + } + _ => false, + }, + ) + .collect_ca(cell_ca.name()); Ok(out_ca.into_series()) } @@ -199,13 +212,8 @@ fn cellid_to_vertices_output(_: &[Field]) -> PolarsResult { let mut v: Vec = vec![]; for i in 0..4 { - v.push( - Field::new(&format!("v{i}_lon"), DataType::Float64), - ); - v.push( - Field::new(&format!("v{i}_lat"), DataType::Float64), - ); - + v.push(Field::new(&format!("v{i}_lon"), DataType::Float64)); + v.push(Field::new(&format!("v{i}_lat"), DataType::Float64)); } Ok(Field::new("vertices", DataType::Struct(v))) } @@ -218,13 +226,15 @@ fn cellid_to_vertices(inputs: &[Series]) -> PolarsResult { let mut v_lat: Vec> = Vec::with_capacity(4); for i in 0..4 { - v_lon.push( - PrimitiveChunkedBuilder::new(&format!("v{i}_lon"), cell_ca.len()) - ); - v_lat.push( - PrimitiveChunkedBuilder::new(&format!("v{i}_lat"), cell_ca.len()) - ); - } + v_lon.push(PrimitiveChunkedBuilder::new( + &format!("v{i}_lon"), + cell_ca.len(), + )); + v_lat.push(PrimitiveChunkedBuilder::new( + &format!("v{i}_lat"), + cell_ca.len(), + )); + } for cell_op in cell_ca.into_iter() { match cell_op { @@ -234,30 +244,33 @@ fn cellid_to_vertices(inputs: &[Series]) -> PolarsResult { let (lon, lat) = vertices[i]; v_lon[i].append_value(lon); v_lat[i].append_value(lat); - } - }, + } + } _ => { for i in 0..4 { - v_lon[i].append_null(); - v_lat[i].append_null(); - } + v_lon[i].append_null(); + v_lat[i].append_null(); } + } } } - let v_coords_ser: Vec = v_lon.into_iter().zip(v_lat.into_iter()).flat_map( - |(cb_lon, cb_lat)| [cb_lon.finish().into_series(), cb_lat.finish().into_series()].into_iter() - ).collect_vec(); + let v_coords_ser: Vec = v_lon + .into_iter() + .zip(v_lat.into_iter()) + .flat_map(|(cb_lon, cb_lat)| { + [cb_lon.finish().into_series(), cb_lat.finish().into_series()].into_iter() + }) + .collect_vec(); let out_chunked = StructChunked::new("vertices", &v_coords_ser[..])?; Ok(out_chunked.into_series()) } - //TransfromNameSpace #[derive(Deserialize)] struct TransformInterpolateKwargs { - coef: f64 + coef: f64, } fn output_3d(_: &[Field]) -> PolarsResult { @@ -270,12 +283,13 @@ fn output_3d(_: &[Field]) -> PolarsResult { } #[polars_expr(output_type_func=output_3d)] -fn interpolate_linear(inputs: &[Series], kwargs: TransformInterpolateKwargs) -> PolarsResult { - +fn interpolate_linear( + inputs: &[Series], + kwargs: TransformInterpolateKwargs, +) -> PolarsResult { let ca = inputs[0].struct_()?; let ca_other = inputs[1].struct_()?; - let (x_ser, y_ser, z_ser) = unpack_xyz(ca, false); let (x_other_ser, y_other_ser, z_other_ser) = unpack_xyz(ca_other, false); @@ -294,20 +308,24 @@ fn interpolate_linear(inputs: &[Series], kwargs: TransformInterpolateKwargs) -> y_other_ser.f64()?.into_iter(), z_other_ser.f64()?.into_iter() ) { - match (x_op, y_op, z_op, x_other_op, y_other_op, z_other_op) { - (Some(x), Some(y), Some(z), Some(x_other), Some(y_other), Some(z_other)) => { - - let (x_interpolated, y_interpolated, z_interpolated) = interpolate_linear_elementwise(vec![x, y, z], vec![x_other, y_other, z_other], kwargs.coef); - x_cb.append_value(x_interpolated); - y_cb.append_value(y_interpolated); - z_cb.append_value(z_interpolated); - }, - _ => { - x_cb.append_null(); - y_cb.append_null(); - z_cb.append_null(); - } + match (x_op, y_op, z_op, x_other_op, y_other_op, z_other_op) { + (Some(x), Some(y), Some(z), Some(x_other), Some(y_other), Some(z_other)) => { + let (x_interpolated, y_interpolated, z_interpolated) = + interpolate_linear_elementwise( + vec![x, y, z], + vec![x_other, y_other, z_other], + kwargs.coef, + ); + x_cb.append_value(x_interpolated); + y_cb.append_value(y_interpolated); + z_cb.append_value(z_interpolated); } + _ => { + x_cb.append_null(); + y_cb.append_null(); + z_cb.append_null(); + } + } } let ser_x = x_cb.finish().into_series(); @@ -316,11 +334,8 @@ fn interpolate_linear(inputs: &[Series], kwargs: TransformInterpolateKwargs) -> let out_chunked = StructChunked::new("coords", &[ser_x, ser_y, ser_z])?; Ok(out_chunked.into_series()) - } - - fn ecef_output(_: &[Field]) -> PolarsResult { let v: Vec = vec![ Field::new("x", DataType::Float64), @@ -330,18 +345,21 @@ fn ecef_output(_: &[Field]) -> PolarsResult { Ok(Field::new("ecef", DataType::Struct(v))) } - #[polars_expr(output_type_func=ecef_output)] fn map_to_ecef(inputs: &[Series]) -> PolarsResult { - let map_ca = inputs[0].struct_()?; let rotation_ca = inputs[1].struct_()?; let offset_ca = inputs[2].struct_()?; - let out_chunked = apply_rotation_to_map(map_ca, rotation_ca, offset_ca, "ecef", map_to_ecef_elementwise); + let out_chunked = apply_rotation_to_map( + map_ca, + rotation_ca, + offset_ca, + "ecef", + map_to_ecef_elementwise, + ); Ok(out_chunked?.into_series()) - } fn map_output(_: &[Field]) -> PolarsResult { @@ -353,20 +371,22 @@ fn map_output(_: &[Field]) -> PolarsResult { Ok(Field::new("map", DataType::Struct(v))) } - #[polars_expr(output_type_func=map_output)] fn ecef_to_map(inputs: &[Series]) -> PolarsResult { - let ecef_ca = inputs[0].struct_()?; let rotation_ca = inputs[1].struct_()?; let offset_ca = inputs[2].struct_()?; - let out_chunked = apply_rotation_to_map(ecef_ca, rotation_ca, offset_ca, "map", ecef_to_map_elementwise); + let out_chunked = apply_rotation_to_map( + ecef_ca, + rotation_ca, + offset_ca, + "map", + ecef_to_map_elementwise, + ); Ok(out_chunked?.into_series()) - } - fn lla_output(_: &[Field]) -> PolarsResult { let v: Vec = vec![ Field::new("lon", DataType::Float64), @@ -378,7 +398,6 @@ fn lla_output(_: &[Field]) -> PolarsResult { #[polars_expr(output_type_func=lla_output)] fn ecef_to_lla(inputs: &[Series]) -> PolarsResult { - let ca = inputs[0].struct_()?; let (ecef_x_ser, ecef_y_ser, ecef_z_ser) = unpack_xyz(ca, false); @@ -387,7 +406,6 @@ fn ecef_to_lla(inputs: &[Series]) -> PolarsResult { let ecef_y = ecef_y_ser.f64()?; let ecef_z = ecef_z_ser.f64()?; - let mut longitude: PrimitiveChunkedBuilder = PrimitiveChunkedBuilder::new("lon", ca.len()); let mut latitude: PrimitiveChunkedBuilder = @@ -402,7 +420,7 @@ fn ecef_to_lla(inputs: &[Series]) -> PolarsResult { longitude.append_value(lon); latitude.append_value(lat); altitude.append_value(alt); - }, + } _ => { longitude.append_null(); latitude.append_null(); @@ -417,14 +435,10 @@ fn ecef_to_lla(inputs: &[Series]) -> PolarsResult { let out_chunked = StructChunked::new("coordinates", &[ser_lon, ser_lat, ser_alt])?; Ok(out_chunked.into_series()) - - } - #[polars_expr(output_type_func=ecef_output)] fn lla_to_ecef(inputs: &[Series]) -> PolarsResult { - let ca = inputs[0].struct_()?; let (lon_ser, lat_ser, alt_ser) = unpack_xyz(ca, true); @@ -433,7 +447,6 @@ fn lla_to_ecef(inputs: &[Series]) -> PolarsResult { let lat = lat_ser.f64()?; let alt = alt_ser.f64()?; - let mut ecef_x: PrimitiveChunkedBuilder = PrimitiveChunkedBuilder::new("x", ca.len()); let mut ecef_y: PrimitiveChunkedBuilder = @@ -448,7 +461,7 @@ fn lla_to_ecef(inputs: &[Series]) -> PolarsResult { ecef_x.append_value(x); ecef_y.append_value(y); ecef_z.append_value(z); - }, + } _ => { ecef_x.append_null(); ecef_y.append_null(); @@ -463,10 +476,8 @@ fn lla_to_ecef(inputs: &[Series]) -> PolarsResult { let out_chunked = StructChunked::new("coordinates", &[ser_x, ser_y, ser_z])?; Ok(out_chunked.into_series()) - } - fn utm_output(_: &[Field]) -> PolarsResult { let v: Vec = vec![ Field::new("x", DataType::Float64), @@ -478,7 +489,6 @@ fn utm_output(_: &[Field]) -> PolarsResult { #[polars_expr(output_type_func=utm_output)] fn lla_to_utm(inputs: &[Series]) -> PolarsResult { - let coords_ca = inputs[0].struct_()?; let (lon_ser, lat_ser, alt_ser) = unpack_xyz(coords_ca, true); @@ -489,80 +499,67 @@ fn lla_to_utm(inputs: &[Series]) -> PolarsResult { let mut utm_z: PrimitiveChunkedBuilder = PrimitiveChunkedBuilder::new("z", coords_ca.len()); - for (lon_op, lat_op, alt_op) in izip!( - lon_ser.f64()?, - lat_ser.f64()?, - alt_ser.f64()? - ) { + for (lon_op, lat_op, alt_op) in izip!(lon_ser.f64()?, lat_ser.f64()?, alt_ser.f64()?) { match (lon_op, lat_op, alt_op) { (Some(lon), Some(lat), Some(alt)) => { let (easting, northing, alt) = lla_to_utm_elementwise(lon, lat, alt); utm_x.append_value(easting); utm_y.append_value(northing); utm_z.append_value(alt); - - }, + } _ => { utm_x.append_null(); utm_y.append_null(); utm_z.append_null(); - } } } - let ser_x = utm_x.finish().into_series(); - let ser_y = utm_y.finish().into_series(); - let ser_z = utm_z.finish().into_series(); - + let ser_x = utm_x.finish().into_series(); + let ser_y = utm_y.finish().into_series(); + let ser_z = utm_z.finish().into_series(); + let out_chunked: StructChunked = StructChunked::new("utm", &[ser_x, ser_y, ser_z])?; Ok(out_chunked.into_series()) - -} - - +} #[polars_expr(output_type=UInt8)] fn lla_to_utm_zone_number(inputs: &[Series]) -> PolarsResult { - let coords_ca = inputs[0].struct_()?; let (lon_ser, lat_ser, _alt_ser) = unpack_xyz(coords_ca, true); - let mut utm_number_cb: PrimitiveChunkedBuilder = PrimitiveChunkedBuilder::new("utm_zone_number", coords_ca.len()); - for (lon_op, lat_op) in izip!( - lon_ser.f64()?, - lat_ser.f64()?, - ) { + for (lon_op, lat_op) in izip!(lon_ser.f64()?, lat_ser.f64()?,) { match (lon_op, lat_op) { (Some(lon), Some(lat)) => { let utm_number = lla_to_utm_zone_number_elementwise(lon, lat); utm_number_cb.append_value(utm_number); - }, + } _ => { utm_number_cb.append_null(); - } } } - let ser_utm_number = utm_number_cb.finish().into_series(); + let ser_utm_number = utm_number_cb.finish().into_series(); Ok(ser_utm_number) - -} - +} #[polars_expr(output_type_func=map_output)] fn rotate_map_coords(inputs: &[Series]) -> PolarsResult { - let map_ca = inputs[0].struct_()?; let rotation_ca = inputs[1].struct_()?; let scale_ca = inputs[2].struct_()?; - let out_chunked = apply_rotation_to_map(map_ca, rotation_ca, scale_ca, "map", rotate_map_coords_elementwise); - - Ok(out_chunked?.into_series()) + let out_chunked = apply_rotation_to_map( + map_ca, + rotation_ca, + scale_ca, + "map", + rotate_map_coords_elementwise, + ); + Ok(out_chunked?.into_series()) } //distance @@ -574,25 +571,22 @@ fn euclidean_2d(inputs: &[Series]) -> PolarsResult { let (x1, y1, _z1) = unpack_xyz(ca1, false); let (x2, y2, _z2) = unpack_xyz(ca2, false); - let iter = izip!( - x1.f64()?, - y1.f64()?, - x2.f64()?, - y2.f64()? - ).into_iter().map( - |(x1_op, y1_op, x2_op, y2_op)| { - match (x1_op, y1_op, x2_op, y2_op) { - (Some(x1), Some(y1), Some(x2), Some(y2)) => euclidean_2d_elementwise(x1, y1, x2, y2), - _ => panic!("Unable to find euclidean distance!") - } - }); + let iter = izip!(x1.f64()?, y1.f64()?, x2.f64()?, y2.f64()?) + .into_iter() + .map( + |(x1_op, y1_op, x2_op, y2_op)| match (x1_op, y1_op, x2_op, y2_op) { + (Some(x1), Some(y1), Some(x2), Some(y2)) => { + euclidean_2d_elementwise(x1, y1, x2, y2) + } + _ => panic!("Unable to find euclidean distance!"), + }, + ); - let out_ca: ChunkedArray = iter.collect_ca_with_dtype("distance", DataType::Float64); + let out_ca: ChunkedArray = + iter.collect_ca_with_dtype("distance", DataType::Float64); Ok(out_ca.into_series()) - } - #[polars_expr(output_type=Float64)] fn euclidean_3d(inputs: &[Series]) -> PolarsResult { let ca1: &StructChunked = inputs[0].struct_()?; @@ -604,23 +598,26 @@ fn euclidean_3d(inputs: &[Series]) -> PolarsResult { let iter = izip!( x1.f64()?, y1.f64()?, - z1.f64()?, - x2.f64()?, - y2.f64()?, + z1.f64()?, + x2.f64()?, + y2.f64()?, z2.f64()? - ).into_iter().map( - |(x1_op, y1_op, z1_op, x2_op, y2_op, z2_op)| { - match (x1_op, y1_op, z1_op, x2_op, y2_op, z2_op) { - (Some(x1), Some(y1), Some(z1), Some(x2), Some(y2), Some(z2),) => euclidean_3d_elementwise(x1, y1, z1, x2, y2, z2), - _ => panic!("Unable to find euclidean distance!") + ) + .into_iter() + .map(|(x1_op, y1_op, z1_op, x2_op, y2_op, z2_op)| { + match (x1_op, y1_op, z1_op, x2_op, y2_op, z2_op) { + (Some(x1), Some(y1), Some(z1), Some(x2), Some(y2), Some(z2)) => { + euclidean_3d_elementwise(x1, y1, z1, x2, y2, z2) + } + _ => panic!("Unable to find euclidean distance!"), } }); - let out_ca: ChunkedArray = iter.collect_ca_with_dtype("distance", DataType::Float64); + let out_ca: ChunkedArray = + iter.collect_ca_with_dtype("distance", DataType::Float64); Ok(out_ca.into_series()) } - #[polars_expr(output_type=Float64)] fn cosine_similarity_2d(inputs: &[Series]) -> PolarsResult { let ca1: &StructChunked = inputs[0].struct_()?; @@ -629,26 +626,22 @@ fn cosine_similarity_2d(inputs: &[Series]) -> PolarsResult { let (x1, y1, _z1) = unpack_xyz(ca1, false); let (x2, y2, _z2) = unpack_xyz(ca2, false); - let iter = izip!( - x1.f64()?, - y1.f64()?, - x2.f64()?, - y2.f64()? - ).into_iter().map( - |(x1_op, y1_op, x2_op, y2_op)| { - match (x1_op, y1_op, x2_op, y2_op) { - (Some(x1), Some(y1), Some(x2), Some(y2)) => cosine_similarity_2d_elementwise(x1, y1, x2, y2), - _ => panic!("Unable to find cosine similarity!") - } - }); + let iter = izip!(x1.f64()?, y1.f64()?, x2.f64()?, y2.f64()?) + .into_iter() + .map( + |(x1_op, y1_op, x2_op, y2_op)| match (x1_op, y1_op, x2_op, y2_op) { + (Some(x1), Some(y1), Some(x2), Some(y2)) => { + cosine_similarity_2d_elementwise(x1, y1, x2, y2) + } + _ => panic!("Unable to find cosine similarity!"), + }, + ); - let out_ca: ChunkedArray = iter.collect_ca_with_dtype("cosine_similarity", DataType::Float64); + let out_ca: ChunkedArray = + iter.collect_ca_with_dtype("cosine_similarity", DataType::Float64); Ok(out_ca.into_series()) - } - - #[polars_expr(output_type=Float64)] fn cosine_similarity_3d(inputs: &[Series]) -> PolarsResult { let ca1: &StructChunked = inputs[0].struct_()?; @@ -659,20 +652,23 @@ fn cosine_similarity_3d(inputs: &[Series]) -> PolarsResult { let iter = izip!( x1.f64()?, - y1.f64()?, - z1.f64()?, - x2.f64()?, + y1.f64()?, + z1.f64()?, + x2.f64()?, y2.f64()?, z2.f64()?, - ).into_iter().map( - |(x1_op, y1_op, z1_op, x2_op, y2_op, z2_op)| { - match (x1_op, y1_op, z1_op, x2_op, y2_op, z2_op) { - (Some(x1), Some(y1), Some(z1), Some(x2), Some(y2), Some(z2)) => cosine_similarity_3d_elementwise(x1, y1, z1, x2, y2, z2), - _ => panic!("Unable to find cosine similarity!") + ) + .into_iter() + .map(|(x1_op, y1_op, z1_op, x2_op, y2_op, z2_op)| { + match (x1_op, y1_op, z1_op, x2_op, y2_op, z2_op) { + (Some(x1), Some(y1), Some(z1), Some(x2), Some(y2), Some(z2)) => { + cosine_similarity_3d_elementwise(x1, y1, z1, x2, y2, z2) + } + _ => panic!("Unable to find cosine similarity!"), } }); - let out_ca: ChunkedArray = iter.collect_ca_with_dtype("cosine_similarity", DataType::Float64); + let out_ca: ChunkedArray = + iter.collect_ca_with_dtype("cosine_similarity", DataType::Float64); Ok(out_ca.into_series()) - -} \ No newline at end of file +} diff --git a/src/lib.rs b/src/lib.rs index 995b89c..3e57b96 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,11 +1,11 @@ -mod s2_functions; mod coord_transforms; mod distance; mod expressions; +mod s2_functions; #[cfg(target_os = "linux")] use jemallocator::Jemalloc; #[global_allocator] #[cfg(target_os = "linux")] -static ALLOC: Jemalloc = Jemalloc; \ No newline at end of file +static ALLOC: Jemalloc = Jemalloc; diff --git a/src/s2_functions.rs b/src/s2_functions.rs index d3ce0f3..0b8f88b 100644 --- a/src/s2_functions.rs +++ b/src/s2_functions.rs @@ -1,6 +1,6 @@ extern crate s2; -use s2::cellid::CellID; use s2::cell::Cell; +use s2::cellid::CellID; use s2::latlng::LatLng; pub fn lonlat_to_cellid_elementwise(lng: f64, lat: f64, level: u64) -> u64 { @@ -8,7 +8,6 @@ pub fn lonlat_to_cellid_elementwise(lng: f64, lat: f64, level: u64) -> u64 { cell_id.parent(level).0 } - pub fn cellid_to_lonlat_elementwise(cellid: u64) -> (f64, f64) { let lnglat = LatLng::from(CellID(cellid)); let lat = lnglat.lat.deg(); @@ -18,7 +17,7 @@ pub fn cellid_to_lonlat_elementwise(cellid: u64) -> (f64, f64) { #[allow(dead_code)] fn cell_area_elementwise(cellid: u64) -> f64 { - let cell = Cell::from(CellID(cellid)); + let cell = Cell::from(CellID(cellid)); cell.approx_area() } @@ -30,14 +29,16 @@ pub fn cell_contains_point_elementwise(cellid: u64, point_lon: f64, point_lat: f pub fn cellid_to_vertices_elementwise(cellid: u64) -> Vec<(f64, f64)> { let cell = Cell::from(CellID(cellid)); - cell.vertices().into_iter().map(|point: s2::point::Point| (point.longitude().deg(), point.latitude().deg())).collect::>() + cell.vertices() + .into_iter() + .map(|point: s2::point::Point| (point.longitude().deg(), point.latitude().deg())) + .collect::>() } - #[cfg(test)] mod s2_tests { - use crate::s2_functions::{lonlat_to_cellid_elementwise, cellid_to_lonlat_elementwise}; + use crate::s2_functions::{cellid_to_lonlat_elementwise, lonlat_to_cellid_elementwise}; #[test] fn test_lonlat_to_cellid() { @@ -46,14 +47,17 @@ mod s2_tests { let level: u64 = 30; let expected_cellid: u64 = 5095400969591719543; - assert_eq!(lonlat_to_cellid_elementwise(lon, lat, level), expected_cellid); + assert_eq!( + lonlat_to_cellid_elementwise(lon, lat, level), + expected_cellid + ); } #[test] fn test_cellid_to_lonlat() { let cellid: u64 = 5095400969591719540; let expected_lonlat: (f64, f64) = (36.077147799420544, 56.78392700893653); - + assert_eq!(cellid_to_lonlat_elementwise(cellid), expected_lonlat) } @@ -64,5 +68,4 @@ mod s2_tests { assert_eq!(lonlat_to_cellid_elementwise(lon, lat, 30), cellid) } - -} \ No newline at end of file +}