From 5cf2cc72288a7413aa40ad921ac269bbd6b1f719 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 15:27:43 -0700 Subject: [PATCH 01/12] Restructure augmentation code. --- Cargo.lock | 2 ++ rust/Cargo.toml | 2 ++ rust/src/geometry/augmentations.rs | 31 +++++++++++++------------ rust/src/geometry/se3.rs | 36 +++++++++++++++++++++++++++++- rust/src/lib.rs | 2 +- 5 files changed, 57 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4c380fab..696e73f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -213,6 +213,8 @@ dependencies = [ "polars", "pyo3", "pyo3-polars", + "rand", + "rand_distr", "rayon", "serde", "strum", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index ac74c1fc..80976673 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -44,6 +44,8 @@ pyo3 = { version = "0.18.3", features = ["extension-module"] } pyo3-polars = { git = "https://github.com/benjaminrwilson/pyo3-polars", rev = "993d22a6bf54ceb8f93c5ed9082621330b186a52", features = [ "serde", ] } +rand = "0.8.5" +rand_distr = "0.4.3" rayon = "1.7.0" serde = "1.0.160" strum = "0.24.1" diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 30297422..7513e797 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -2,22 +2,25 @@ //! //! Geometric augmentations. -use std::f32::consts::PI; +use ndarray::{Array, Ix2}; +use rand_distr::{Bernoulli, Distribution}; -use ndarray::{Array, ArrayView, Ix2}; +use super::se3::{reflect_pose_x, reflect_pose_y}; -use crate::geometry::so3::{quat_to_yaw, yaw_to_quat}; - -/// Reflect pose across the x-axis. -pub fn reflect_pose_x(quat_wxyz: &ArrayView) -> Array { - let yaw_rad = quat_to_yaw(quat_wxyz); - let reflected_yaw_rad = -yaw_rad; - yaw_to_quat(&reflected_yaw_rad.view()) +/// Reflect all poses across the x-axis with probability `p` (sampled from a Bernoulli distribution). +pub fn sample_global_reflect_pose_x(xyzlwh_qwxyz: Array, p: f64) -> Array { + let d = Bernoulli::new(p).unwrap(); + if d.sample(&mut rand::thread_rng()) { + return reflect_pose_x(&xyzlwh_qwxyz.view()); + } + xyzlwh_qwxyz } -/// Reflect pose across the y-axis. -pub fn reflect_pose_y(quat_wxyz: &ArrayView) -> Array { - let yaw_rad = quat_to_yaw(quat_wxyz); - let reflected_yaw_rad = PI - yaw_rad; - yaw_to_quat(&reflected_yaw_rad.view()) +/// Reflect all poses across the y-axis with probability `p` (sampled from a Bernoulli distribution). +pub fn sample_global_reflect_pose_y(xyzlwh_qwxyz: Array, p: f64) -> Array { + let d = Bernoulli::new(p).unwrap(); + if d.sample(&mut rand::thread_rng()) { + return reflect_pose_y(&xyzlwh_qwxyz.view()); + } + xyzlwh_qwxyz } diff --git a/rust/src/geometry/se3.rs b/rust/src/geometry/se3.rs index bc3a0207..663f642f 100644 --- a/rust/src/geometry/se3.rs +++ b/rust/src/geometry/se3.rs @@ -2,7 +2,11 @@ //! //! Special Euclidean Group 3. -use ndarray::{s, Array1, Array2, ArrayView2}; +use std::{f32::consts::PI, ops::Neg}; + +use ndarray::{concatenate, s, Array, Array1, Array2, ArrayView, ArrayView2, Axis, Ix2}; + +use super::so3::{quat_to_yaw, yaw_to_quat}; /// Special Euclidean Group 3 (SE(3)). /// Rigid transformation parameterized by a rotation and translation in $R^3$. @@ -57,3 +61,33 @@ impl SE3 { } } } + +/// Reflect pose across the x-axis. +/// (N,10) `xyzlwh_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. +pub fn reflect_pose_x(xyzlwh_qwxyz: &ArrayView) -> Array { + let xyz_m = xyzlwh_qwxyz.slice(s![.., ..3]); + let quat_wxyz = xyzlwh_qwxyz.slice(s![.., -4..]); + + let mut reflected_xyz_m = xyz_m.clone().to_owned(); + reflected_xyz_m.assign(&reflected_xyz_m.slice(s![.., 1]).neg()); + + let yaw_rad = quat_to_yaw(&quat_wxyz); + let reflected_yaw_rad = -yaw_rad; + let reflected_quat_wxyz = yaw_to_quat(&reflected_yaw_rad.view()); + concatenate![Axis(1), reflected_xyz_m, reflected_quat_wxyz] +} + +/// Reflect pose across the y-axis. +/// (N,10) `xyzlwh_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. +pub fn reflect_pose_y(xyzlwh_qwxyz: &ArrayView) -> Array { + let xyz_m = xyzlwh_qwxyz.slice(s![.., ..3]); + let quat_wxyz = xyzlwh_qwxyz.slice(s![.., -4..]); + + let mut reflected_xyz_m = xyz_m.clone().to_owned(); + reflected_xyz_m.assign(&reflected_xyz_m.slice(s![.., 0]).neg()); + + let yaw_rad = quat_to_yaw(&quat_wxyz); + let reflected_yaw_rad = PI - yaw_rad; + let reflected_quat_wxyz = yaw_to_quat(&reflected_yaw_rad.view()); + concatenate![Axis(1), reflected_xyz_m, reflected_quat_wxyz] +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0e6b0f5c..d0344c4b 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -15,7 +15,7 @@ pub mod share; pub mod structures; use data_loader::{DataLoader, Sweep}; -use geometry::augmentations::{reflect_pose_x, reflect_pose_y}; +use geometry::se3::{reflect_pose_x, reflect_pose_y}; use ndarray::{Dim, Ix1, Ix2}; use numpy::PyReadonlyArray; use numpy::{IntoPyArray, PyArray}; From 89de9d6914c36681eb59bd8b4e50ec86c4d93e5e Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 15:32:31 -0700 Subject: [PATCH 02/12] Add python bindings. --- rust/src/geometry/augmentations.rs | 10 ++++----- rust/src/lib.rs | 33 ++++++++++++++++++++++++++---- 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 7513e797..766c9a15 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -2,25 +2,25 @@ //! //! Geometric augmentations. -use ndarray::{Array, Ix2}; +use ndarray::{Array, ArrayView, Ix2}; use rand_distr::{Bernoulli, Distribution}; use super::se3::{reflect_pose_x, reflect_pose_y}; /// Reflect all poses across the x-axis with probability `p` (sampled from a Bernoulli distribution). -pub fn sample_global_reflect_pose_x(xyzlwh_qwxyz: Array, p: f64) -> Array { +pub fn sample_global_reflect_pose_x(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { let d = Bernoulli::new(p).unwrap(); if d.sample(&mut rand::thread_rng()) { return reflect_pose_x(&xyzlwh_qwxyz.view()); } - xyzlwh_qwxyz + xyzlwh_qwxyz.to_owned() } /// Reflect all poses across the y-axis with probability `p` (sampled from a Bernoulli distribution). -pub fn sample_global_reflect_pose_y(xyzlwh_qwxyz: Array, p: f64) -> Array { +pub fn sample_global_reflect_pose_y(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { let d = Bernoulli::new(p).unwrap(); if d.sample(&mut rand::thread_rng()) { return reflect_pose_y(&xyzlwh_qwxyz.view()); } - xyzlwh_qwxyz + xyzlwh_qwxyz.to_owned() } diff --git a/rust/src/lib.rs b/rust/src/lib.rs index d0344c4b..9135b098 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -15,6 +15,7 @@ pub mod share; pub mod structures; use data_loader::{DataLoader, Sweep}; +use geometry::augmentations::{sample_global_reflect_pose_x, sample_global_reflect_pose_y}; use geometry::se3::{reflect_pose_x, reflect_pose_y}; use ndarray::{Dim, Ix1, Ix2}; use numpy::PyReadonlyArray; @@ -90,9 +91,9 @@ fn py_yaw_to_quat<'py>( #[allow(clippy::type_complexity)] fn py_reflect_pose_x<'py>( py: Python<'py>, - quat_wxyz: PyReadonlyArray, + xyzlwh_qwxyz: PyReadonlyArray, ) -> &'py PyArray { - reflect_pose_x(&quat_wxyz.as_array().view()).into_pyarray(py) + reflect_pose_x(&xyzlwh_qwxyz.as_array().view()).into_pyarray(py) } #[pyfunction] @@ -100,9 +101,31 @@ fn py_reflect_pose_x<'py>( #[allow(clippy::type_complexity)] fn py_reflect_pose_y<'py>( py: Python<'py>, - quat_wxyz: PyReadonlyArray, + xyzlwh_qwxyz: PyReadonlyArray, +) -> &'py PyArray { + reflect_pose_y(&xyzlwh_qwxyz.as_array().view()).into_pyarray(py) +} + +#[pyfunction] +#[pyo3(name = "sample_global_reflect_pose_x")] +#[allow(clippy::type_complexity)] +fn py_sample_reflect_pose_x<'py>( + py: Python<'py>, + xyzlwh_qwxyz: PyReadonlyArray, + p: f64, +) -> &'py PyArray { + sample_global_reflect_pose_x(xyzlwh_qwxyz.as_array().view(), p).into_pyarray(py) +} + +#[pyfunction] +#[pyo3(name = "sample_global_reflect_pose_y")] +#[allow(clippy::type_complexity)] +fn py_sample_reflect_pose_y<'py>( + py: Python<'py>, + xyzlwh_qwxyz: PyReadonlyArray, + p: f64, ) -> &'py PyArray { - reflect_pose_y(&quat_wxyz.as_array().view()).into_pyarray(py) + sample_global_reflect_pose_y(xyzlwh_qwxyz.as_array(), p).into_pyarray(py) } /// A Python module implemented in Rust. @@ -114,6 +137,8 @@ fn _r(_py: Python, m: &PyModule) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_quat_to_yaw, m)?)?; m.add_function(wrap_pyfunction!(py_reflect_pose_x, m)?)?; m.add_function(wrap_pyfunction!(py_reflect_pose_y, m)?)?; + m.add_function(wrap_pyfunction!(py_sample_reflect_pose_x, m)?)?; + m.add_function(wrap_pyfunction!(py_sample_reflect_pose_y, m)?)?; m.add_function(wrap_pyfunction!(py_voxelize, m)?)?; m.add_function(wrap_pyfunction!(py_yaw_to_quat, m)?)?; Ok(()) From 819e22ee500cf51c29552a95c2890c52d8776e1d Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 17:51:48 -0700 Subject: [PATCH 03/12] Clean up augmentations. --- rust/src/geometry/augmentations.rs | 91 +++++++++++++++++++++++++++++- rust/src/geometry/se3.rs | 26 +++++---- rust/src/lib.rs | 48 ---------------- rust/src/share.rs | 25 +++++++- 4 files changed, 127 insertions(+), 63 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 766c9a15..8caed95f 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -3,12 +3,19 @@ //! Geometric augmentations. use ndarray::{Array, ArrayView, Ix2}; +use polars::{ + lazy::dsl::{col, GetOutput}, + prelude::{DataFrame, DataType, IntoLazy}, + series::Series, +}; use rand_distr::{Bernoulli, Distribution}; +use crate::share::{data_frame_to_ndarray_f32, ndarray_to_series_vec}; + use super::se3::{reflect_pose_x, reflect_pose_y}; /// Reflect all poses across the x-axis with probability `p` (sampled from a Bernoulli distribution). -pub fn sample_global_reflect_pose_x(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { +pub fn _sample_global_reflect_pose_x(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { let d = Bernoulli::new(p).unwrap(); if d.sample(&mut rand::thread_rng()) { return reflect_pose_x(&xyzlwh_qwxyz.view()); @@ -17,10 +24,90 @@ pub fn sample_global_reflect_pose_x(xyzlwh_qwxyz: ArrayView, p: f64) - } /// Reflect all poses across the y-axis with probability `p` (sampled from a Bernoulli distribution). -pub fn sample_global_reflect_pose_y(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { +pub fn _sample_global_reflect_pose_y(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { let d = Bernoulli::new(p).unwrap(); if d.sample(&mut rand::thread_rng()) { return reflect_pose_y(&xyzlwh_qwxyz.view()); } xyzlwh_qwxyz.to_owned() } + +/// Sample a scene reflection. +/// This reflects both a point cloud and cuboids across the x-axis. +pub fn sample_scene_reflection_x( + lidar: DataFrame, + cuboids: DataFrame, + p: f64, +) -> (DataFrame, DataFrame) { + let d = Bernoulli::new(p).unwrap(); + let is_augmented = d.sample(&mut rand::thread_rng()); + if is_augmented { + let augmented_lidar = lidar + .clone() + .lazy() + .with_column(col("y").map( + move |x| { + Ok(Some( + x.f32() + .unwrap() + .into_no_null_iter() + .map(|y| -y) + .collect::(), + )) + }, + GetOutput::from_type(DataType::Float32), + )) + .collect() + .unwrap(); + + // println!("{}", augmented_lidar); + + let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; + let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); + let augmented_poses = reflect_pose_x(&poses.view()); + let series_vec = ndarray_to_series_vec(augmented_poses, column_names); + let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); + (augmented_lidar, augmented_cuboids) + } else { + (lidar, cuboids) + } +} + +/// Sample a scene reflection. +/// This reflects both a point cloud and cuboids across the y-axis. +pub fn sample_scene_reflection_y( + lidar: DataFrame, + cuboids: DataFrame, + p: f64, +) -> (DataFrame, DataFrame) { + let d = Bernoulli::new(p).unwrap(); + let is_augmented = d.sample(&mut rand::thread_rng()); + if is_augmented { + let augmented_lidar = lidar + .clone() + .lazy() + .with_column(col("x").map( + move |x| { + Ok(Some( + x.f32() + .unwrap() + .into_no_null_iter() + .map(|x| -x) + .collect::(), + )) + }, + GetOutput::from_type(DataType::Float32), + )) + .collect() + .unwrap(); + + let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; + let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); + let augmented_poses = reflect_pose_y(&poses.view()); + let series_vec = ndarray_to_series_vec(augmented_poses, column_names); + let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); + (augmented_lidar, augmented_cuboids) + } else { + (lidar, cuboids) + } +} diff --git a/rust/src/geometry/se3.rs b/rust/src/geometry/se3.rs index 663f642f..a95b2318 100644 --- a/rust/src/geometry/se3.rs +++ b/rust/src/geometry/se3.rs @@ -2,7 +2,7 @@ //! //! Special Euclidean Group 3. -use std::{f32::consts::PI, ops::Neg}; +use std::f32::consts::PI; use ndarray::{concatenate, s, Array, Array1, Array2, ArrayView, ArrayView2, Axis, Ix2}; @@ -64,12 +64,14 @@ impl SE3 { /// Reflect pose across the x-axis. /// (N,10) `xyzlwh_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. -pub fn reflect_pose_x(xyzlwh_qwxyz: &ArrayView) -> Array { - let xyz_m = xyzlwh_qwxyz.slice(s![.., ..3]); - let quat_wxyz = xyzlwh_qwxyz.slice(s![.., -4..]); +pub fn reflect_pose_x(xyz_qwxyz: &ArrayView) -> Array { + let xyz_m = xyz_qwxyz.slice(s![.., ..3]); + let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); let mut reflected_xyz_m = xyz_m.clone().to_owned(); - reflected_xyz_m.assign(&reflected_xyz_m.slice(s![.., 1]).neg()); + reflected_xyz_m + .slice_mut(s![.., 1]) + .par_mapv_inplace(|y| -y); let yaw_rad = quat_to_yaw(&quat_wxyz); let reflected_yaw_rad = -yaw_rad; @@ -78,13 +80,17 @@ pub fn reflect_pose_x(xyzlwh_qwxyz: &ArrayView) -> Array { } /// Reflect pose across the y-axis. -/// (N,10) `xyzlwh_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. -pub fn reflect_pose_y(xyzlwh_qwxyz: &ArrayView) -> Array { - let xyz_m = xyzlwh_qwxyz.slice(s![.., ..3]); - let quat_wxyz = xyzlwh_qwxyz.slice(s![.., -4..]); +/// (N,10) `xyz_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. +pub fn reflect_pose_y(xyz_qwxyz: &ArrayView) -> Array { + let xyz_m = xyz_qwxyz.slice(s![.., ..3]); + let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); let mut reflected_xyz_m = xyz_m.clone().to_owned(); - reflected_xyz_m.assign(&reflected_xyz_m.slice(s![.., 0]).neg()); + reflected_xyz_m + .slice_mut(s![.., 0]) + .par_mapv_inplace(|x| -x); + + println!("{}", reflected_xyz_m); let yaw_rad = quat_to_yaw(&quat_wxyz); let reflected_yaw_rad = PI - yaw_rad; diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 9135b098..f8540610 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -15,8 +15,6 @@ pub mod share; pub mod structures; use data_loader::{DataLoader, Sweep}; -use geometry::augmentations::{sample_global_reflect_pose_x, sample_global_reflect_pose_y}; -use geometry::se3::{reflect_pose_x, reflect_pose_y}; use ndarray::{Dim, Ix1, Ix2}; use numpy::PyReadonlyArray; use numpy::{IntoPyArray, PyArray}; @@ -86,48 +84,6 @@ fn py_yaw_to_quat<'py>( yaw_to_quat(&quat_wxyz.as_array().view()).into_pyarray(py) } -#[pyfunction] -#[pyo3(name = "reflect_pose_x")] -#[allow(clippy::type_complexity)] -fn py_reflect_pose_x<'py>( - py: Python<'py>, - xyzlwh_qwxyz: PyReadonlyArray, -) -> &'py PyArray { - reflect_pose_x(&xyzlwh_qwxyz.as_array().view()).into_pyarray(py) -} - -#[pyfunction] -#[pyo3(name = "reflect_pose_y")] -#[allow(clippy::type_complexity)] -fn py_reflect_pose_y<'py>( - py: Python<'py>, - xyzlwh_qwxyz: PyReadonlyArray, -) -> &'py PyArray { - reflect_pose_y(&xyzlwh_qwxyz.as_array().view()).into_pyarray(py) -} - -#[pyfunction] -#[pyo3(name = "sample_global_reflect_pose_x")] -#[allow(clippy::type_complexity)] -fn py_sample_reflect_pose_x<'py>( - py: Python<'py>, - xyzlwh_qwxyz: PyReadonlyArray, - p: f64, -) -> &'py PyArray { - sample_global_reflect_pose_x(xyzlwh_qwxyz.as_array().view(), p).into_pyarray(py) -} - -#[pyfunction] -#[pyo3(name = "sample_global_reflect_pose_y")] -#[allow(clippy::type_complexity)] -fn py_sample_reflect_pose_y<'py>( - py: Python<'py>, - xyzlwh_qwxyz: PyReadonlyArray, - p: f64, -) -> &'py PyArray { - sample_global_reflect_pose_y(xyzlwh_qwxyz.as_array(), p).into_pyarray(py) -} - /// A Python module implemented in Rust. #[pymodule] fn _r(_py: Python, m: &PyModule) -> PyResult<()> { @@ -135,10 +91,6 @@ fn _r(_py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; m.add_function(wrap_pyfunction!(py_quat_to_mat3, m)?)?; m.add_function(wrap_pyfunction!(py_quat_to_yaw, m)?)?; - m.add_function(wrap_pyfunction!(py_reflect_pose_x, m)?)?; - m.add_function(wrap_pyfunction!(py_reflect_pose_y, m)?)?; - m.add_function(wrap_pyfunction!(py_sample_reflect_pose_x, m)?)?; - m.add_function(wrap_pyfunction!(py_sample_reflect_pose_y, m)?)?; m.add_function(wrap_pyfunction!(py_voxelize, m)?)?; m.add_function(wrap_pyfunction!(py_yaw_to_quat, m)?)?; Ok(()) diff --git a/rust/src/share.rs b/rust/src/share.rs index db5cf66a..a6c753b4 100644 --- a/rust/src/share.rs +++ b/rust/src/share.rs @@ -3,10 +3,14 @@ //! Conversion methods between different libraries. use ndarray::{Array, Ix2}; -use polars::{prelude::NamedFrom, series::Series}; +use polars::{ + lazy::dsl::{cols, lit, Expr}, + prelude::{DataFrame, Float32Type, IntoLazy, NamedFrom}, + series::Series, +}; /// Convert the columns of an `ndarray::Array` into a vector of `polars::series::Series`. -pub fn ndarray_to_series_vec(arr: Array, column_names: Vec<&str>) -> Vec { +pub fn ndarray_to_series_vec(arr: Array, column_names: Vec<&str>) -> Vec { let num_dims = arr.shape()[1]; if num_dims != column_names.len() { panic!("Number of array columns and column names must match."); @@ -18,7 +22,22 @@ pub fn ndarray_to_series_vec(arr: Array, column_names: Vec<&str>) -> V column_name, column.as_standard_layout().to_owned().into_raw_vec(), ); - series_vec.push(series); + series_vec.push(lit(series)); } series_vec } + +/// Convert a data frame to an `ndarray::Array::`. +pub fn data_frame_to_ndarray_f32( + data_frame: DataFrame, + column_names: Vec<&str>, +) -> Array { + data_frame + .clone() + .lazy() + .select(&[cols(column_names)]) + .collect() + .unwrap() + .to_ndarray::() + .unwrap() +} From 9f1bc643f4e3c04e7e9d22b4e325d0a4dcb041e7 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 17:52:17 -0700 Subject: [PATCH 04/12] Remove comment. --- rust/src/geometry/augmentations.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 8caed95f..ed48df01 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -60,8 +60,6 @@ pub fn sample_scene_reflection_x( .collect() .unwrap(); - // println!("{}", augmented_lidar); - let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); let augmented_poses = reflect_pose_x(&poses.view()); From 14ac85f5c725e4e135983bc8bd237eea211ab67c Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 17:55:29 -0700 Subject: [PATCH 05/12] Remove dead code. --- rust/src/geometry/augmentations.rs | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index ed48df01..51973c64 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -2,7 +2,6 @@ //! //! Geometric augmentations. -use ndarray::{Array, ArrayView, Ix2}; use polars::{ lazy::dsl::{col, GetOutput}, prelude::{DataFrame, DataType, IntoLazy}, @@ -14,24 +13,6 @@ use crate::share::{data_frame_to_ndarray_f32, ndarray_to_series_vec}; use super::se3::{reflect_pose_x, reflect_pose_y}; -/// Reflect all poses across the x-axis with probability `p` (sampled from a Bernoulli distribution). -pub fn _sample_global_reflect_pose_x(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { - let d = Bernoulli::new(p).unwrap(); - if d.sample(&mut rand::thread_rng()) { - return reflect_pose_x(&xyzlwh_qwxyz.view()); - } - xyzlwh_qwxyz.to_owned() -} - -/// Reflect all poses across the y-axis with probability `p` (sampled from a Bernoulli distribution). -pub fn _sample_global_reflect_pose_y(xyzlwh_qwxyz: ArrayView, p: f64) -> Array { - let d = Bernoulli::new(p).unwrap(); - if d.sample(&mut rand::thread_rng()) { - return reflect_pose_y(&xyzlwh_qwxyz.view()); - } - xyzlwh_qwxyz.to_owned() -} - /// Sample a scene reflection. /// This reflects both a point cloud and cuboids across the x-axis. pub fn sample_scene_reflection_x( From b5f910ea1a1204b94979a78fa0e62e5afc1d2556 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 17:56:27 -0700 Subject: [PATCH 06/12] Update var name. --- rust/src/geometry/augmentations.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 51973c64..9c05c4f2 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -20,8 +20,8 @@ pub fn sample_scene_reflection_x( cuboids: DataFrame, p: f64, ) -> (DataFrame, DataFrame) { - let d = Bernoulli::new(p).unwrap(); - let is_augmented = d.sample(&mut rand::thread_rng()); + let distribution = Bernoulli::new(p).unwrap(); + let is_augmented = distribution.sample(&mut rand::thread_rng()); if is_augmented { let augmented_lidar = lidar .clone() @@ -59,8 +59,8 @@ pub fn sample_scene_reflection_y( cuboids: DataFrame, p: f64, ) -> (DataFrame, DataFrame) { - let d = Bernoulli::new(p).unwrap(); - let is_augmented = d.sample(&mut rand::thread_rng()); + let distribution: Bernoulli = Bernoulli::new(p).unwrap(); + let is_augmented = distribution.sample(&mut rand::thread_rng()); if is_augmented { let augmented_lidar = lidar .clone() From 3f18590680bce1c7eb9a14c011269df7c187391e Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 18:00:02 -0700 Subject: [PATCH 07/12] Fix clippy. --- rust/src/geometry/augmentations.rs | 2 -- rust/src/share.rs | 1 - 2 files changed, 3 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 9c05c4f2..25cb3913 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -24,7 +24,6 @@ pub fn sample_scene_reflection_x( let is_augmented = distribution.sample(&mut rand::thread_rng()); if is_augmented { let augmented_lidar = lidar - .clone() .lazy() .with_column(col("y").map( move |x| { @@ -63,7 +62,6 @@ pub fn sample_scene_reflection_y( let is_augmented = distribution.sample(&mut rand::thread_rng()); if is_augmented { let augmented_lidar = lidar - .clone() .lazy() .with_column(col("x").map( move |x| { diff --git a/rust/src/share.rs b/rust/src/share.rs index a6c753b4..8416d273 100644 --- a/rust/src/share.rs +++ b/rust/src/share.rs @@ -33,7 +33,6 @@ pub fn data_frame_to_ndarray_f32( column_names: Vec<&str>, ) -> Array { data_frame - .clone() .lazy() .select(&[cols(column_names)]) .collect() From 6ab49c069fcb0a3b585e5a5b82a1fc20db0d5d80 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 18:31:59 -0700 Subject: [PATCH 08/12] Update comments. --- rust/src/geometry/se3.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/src/geometry/se3.rs b/rust/src/geometry/se3.rs index a95b2318..12ca501d 100644 --- a/rust/src/geometry/se3.rs +++ b/rust/src/geometry/se3.rs @@ -63,7 +63,7 @@ impl SE3 { } /// Reflect pose across the x-axis. -/// (N,10) `xyzlwh_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. +/// (N,7) `xyz_qwxyz` represent the translation and orientation of `N` rigid objects. pub fn reflect_pose_x(xyz_qwxyz: &ArrayView) -> Array { let xyz_m = xyz_qwxyz.slice(s![.., ..3]); let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); @@ -80,7 +80,7 @@ pub fn reflect_pose_x(xyz_qwxyz: &ArrayView) -> Array { } /// Reflect pose across the y-axis. -/// (N,10) `xyz_qwxyz` represent the translation, dimensions, and orientation of `N` rigid objects. +/// (N,7) `xyz_qwxyz` represent the translation and orientation of `N` rigid objects. pub fn reflect_pose_y(xyz_qwxyz: &ArrayView) -> Array { let xyz_m = xyz_qwxyz.slice(s![.., ..3]); let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); From 6f73ea37dcf3475225f83ee67733e74542e2e812 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Fri, 5 May 2023 23:43:53 -0700 Subject: [PATCH 09/12] Remove print statement. --- rust/src/geometry/se3.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/rust/src/geometry/se3.rs b/rust/src/geometry/se3.rs index 12ca501d..3792cd1f 100644 --- a/rust/src/geometry/se3.rs +++ b/rust/src/geometry/se3.rs @@ -90,8 +90,6 @@ pub fn reflect_pose_y(xyz_qwxyz: &ArrayView) -> Array { .slice_mut(s![.., 0]) .par_mapv_inplace(|x| -x); - println!("{}", reflected_xyz_m); - let yaw_rad = quat_to_yaw(&quat_wxyz); let reflected_yaw_rad = PI - yaw_rad; let reflected_quat_wxyz = yaw_to_quat(&reflected_yaw_rad.view()); From e2f9150c25eca2dd7da8d35096450eab7c289fa2 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Sat, 6 May 2023 01:17:06 -0700 Subject: [PATCH 10/12] Fix duplicate transform. --- rust/src/geometry/augmentations.rs | 6 ++--- rust/src/geometry/se3.rs | 41 ++---------------------------- rust/src/geometry/so3.rs | 22 +++++++++++++++- 3 files changed, 26 insertions(+), 43 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 25cb3913..7b9dc7bc 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -11,7 +11,7 @@ use rand_distr::{Bernoulli, Distribution}; use crate::share::{data_frame_to_ndarray_f32, ndarray_to_series_vec}; -use super::se3::{reflect_pose_x, reflect_pose_y}; +use super::so3::{reflect_orientation_x, reflect_orientation_y}; /// Sample a scene reflection. /// This reflects both a point cloud and cuboids across the x-axis. @@ -42,7 +42,7 @@ pub fn sample_scene_reflection_x( let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); - let augmented_poses = reflect_pose_x(&poses.view()); + let augmented_poses = reflect_orientation_x(&poses.view()); let series_vec = ndarray_to_series_vec(augmented_poses, column_names); let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); (augmented_lidar, augmented_cuboids) @@ -80,7 +80,7 @@ pub fn sample_scene_reflection_y( let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); - let augmented_poses = reflect_pose_y(&poses.view()); + let augmented_poses = reflect_orientation_y(&poses.view()); let series_vec = ndarray_to_series_vec(augmented_poses, column_names); let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); (augmented_lidar, augmented_cuboids) diff --git a/rust/src/geometry/se3.rs b/rust/src/geometry/se3.rs index 3792cd1f..9fc95d54 100644 --- a/rust/src/geometry/se3.rs +++ b/rust/src/geometry/se3.rs @@ -2,11 +2,8 @@ //! //! Special Euclidean Group 3. -use std::f32::consts::PI; - -use ndarray::{concatenate, s, Array, Array1, Array2, ArrayView, ArrayView2, Axis, Ix2}; - -use super::so3::{quat_to_yaw, yaw_to_quat}; +use ndarray::ArrayView2; +use ndarray::{s, Array1, Array2}; /// Special Euclidean Group 3 (SE(3)). /// Rigid transformation parameterized by a rotation and translation in $R^3$. @@ -61,37 +58,3 @@ impl SE3 { } } } - -/// Reflect pose across the x-axis. -/// (N,7) `xyz_qwxyz` represent the translation and orientation of `N` rigid objects. -pub fn reflect_pose_x(xyz_qwxyz: &ArrayView) -> Array { - let xyz_m = xyz_qwxyz.slice(s![.., ..3]); - let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); - - let mut reflected_xyz_m = xyz_m.clone().to_owned(); - reflected_xyz_m - .slice_mut(s![.., 1]) - .par_mapv_inplace(|y| -y); - - let yaw_rad = quat_to_yaw(&quat_wxyz); - let reflected_yaw_rad = -yaw_rad; - let reflected_quat_wxyz = yaw_to_quat(&reflected_yaw_rad.view()); - concatenate![Axis(1), reflected_xyz_m, reflected_quat_wxyz] -} - -/// Reflect pose across the y-axis. -/// (N,7) `xyz_qwxyz` represent the translation and orientation of `N` rigid objects. -pub fn reflect_pose_y(xyz_qwxyz: &ArrayView) -> Array { - let xyz_m = xyz_qwxyz.slice(s![.., ..3]); - let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); - - let mut reflected_xyz_m = xyz_m.clone().to_owned(); - reflected_xyz_m - .slice_mut(s![.., 0]) - .par_mapv_inplace(|x| -x); - - let yaw_rad = quat_to_yaw(&quat_wxyz); - let reflected_yaw_rad = PI - yaw_rad; - let reflected_quat_wxyz = yaw_to_quat(&reflected_yaw_rad.view()); - concatenate![Axis(1), reflected_xyz_m, reflected_quat_wxyz] -} diff --git a/rust/src/geometry/so3.rs b/rust/src/geometry/so3.rs index c4f52b96..9b43db12 100644 --- a/rust/src/geometry/so3.rs +++ b/rust/src/geometry/so3.rs @@ -2,7 +2,9 @@ //! //! Special Orthogonal Group 3 (SO(3)). -use ndarray::{par_azip, Array, Array2, ArrayView, Ix1, Ix2}; +use std::f32::consts::PI; + +use ndarray::{par_azip, s, Array, Array2, ArrayView, Ix1, Ix2}; /// Convert a quaternion in scalar-first format to a 3x3 rotation matrix. pub fn quat_to_mat3(quat_wxyz: &ArrayView) -> Array { @@ -71,3 +73,21 @@ pub fn _yaw_to_quat(yaw_rad: f32) -> Array { let qz = f32::sin(0.5 * yaw_rad); Array::::from_vec(vec![qw, 0.0, 0.0, qz]) } + +/// Reflect orientation across the x-axis. +/// (N,7) `quat_wxyz` orientation of `N` rigid objects. +pub fn reflect_orientation_x(quat_wxyz: &ArrayView) -> Array { + let quat_wxyz = quat_wxyz.slice(s![.., -4..]); + let yaw_rad = quat_to_yaw(&quat_wxyz); + let reflected_yaw_rad = -yaw_rad; + yaw_to_quat(&reflected_yaw_rad.view()) +} + +/// Reflect orientation across the y-axis. +/// (N,7) `quat_wxyz` orientation of `N` rigid objects. +pub fn reflect_orientation_y(xyz_qwxyz: &ArrayView) -> Array { + let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); + let yaw_rad = quat_to_yaw(&quat_wxyz); + let reflected_yaw_rad = PI - yaw_rad; + yaw_to_quat(&reflected_yaw_rad.view()) +} From ffe3ad724c1be39b6fa787d6380c01d20546390b Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Sat, 6 May 2023 01:39:28 -0700 Subject: [PATCH 11/12] Improve augmentations. --- rust/src/geometry/augmentations.rs | 42 +++++++++++++++++++++++++----- rust/src/geometry/so3.rs | 26 ++++++++++++++---- 2 files changed, 56 insertions(+), 12 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 7b9dc7bc..b618943b 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -2,6 +2,8 @@ //! //! Geometric augmentations. +use itertools::Itertools; +use ndarray::{concatenate, Axis}; use polars::{ lazy::dsl::{col, GetOutput}, prelude::{DataFrame, DataType, IntoLazy}, @@ -11,7 +13,9 @@ use rand_distr::{Bernoulli, Distribution}; use crate::share::{data_frame_to_ndarray_f32, ndarray_to_series_vec}; -use super::so3::{reflect_orientation_x, reflect_orientation_y}; +use super::so3::{ + reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y, +}; /// Sample a scene reflection. /// This reflects both a point cloud and cuboids across the x-axis. @@ -40,9 +44,21 @@ pub fn sample_scene_reflection_x( .collect() .unwrap(); - let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; - let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); - let augmented_poses = reflect_orientation_x(&poses.view()); + let translation_column_names = vec!["tx_m", "ty_m", "tz_m"]; + let txyz_m = data_frame_to_ndarray_f32(cuboids.clone(), translation_column_names.clone()); + let augmentation_translation = reflect_translation_x(&txyz_m.view()); + + let orientation_column_names = vec!["qw", "qx", "qy", "qz"]; + let quat_wxyz = + data_frame_to_ndarray_f32(cuboids.clone(), orientation_column_names.clone()); + let augmented_orientation = reflect_orientation_x(&quat_wxyz.view()); + let augmented_poses = + concatenate![Axis(1), augmentation_translation, augmented_orientation]; + + let column_names = translation_column_names + .into_iter() + .chain(orientation_column_names) + .collect_vec(); let series_vec = ndarray_to_series_vec(augmented_poses, column_names); let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); (augmented_lidar, augmented_cuboids) @@ -78,9 +94,21 @@ pub fn sample_scene_reflection_y( .collect() .unwrap(); - let column_names = vec!["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"]; - let poses = data_frame_to_ndarray_f32(cuboids.clone(), column_names.clone()); - let augmented_poses = reflect_orientation_y(&poses.view()); + let translation_column_names = vec!["tx_m", "ty_m", "tz_m"]; + let txyz_m = data_frame_to_ndarray_f32(cuboids.clone(), translation_column_names.clone()); + let augmentation_translation = reflect_translation_y(&txyz_m.view()); + + let orientation_column_names = vec!["qw", "qx", "qy", "qz"]; + let quat_wxyz = + data_frame_to_ndarray_f32(cuboids.clone(), orientation_column_names.clone()); + let augmented_orientation = reflect_orientation_y(&quat_wxyz.view()); + let augmented_poses = + concatenate![Axis(1), augmentation_translation, augmented_orientation]; + + let column_names = translation_column_names + .into_iter() + .chain(orientation_column_names) + .collect_vec(); let series_vec = ndarray_to_series_vec(augmented_poses, column_names); let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); (augmented_lidar, augmented_cuboids) diff --git a/rust/src/geometry/so3.rs b/rust/src/geometry/so3.rs index 9b43db12..78d7239f 100644 --- a/rust/src/geometry/so3.rs +++ b/rust/src/geometry/so3.rs @@ -77,17 +77,33 @@ pub fn _yaw_to_quat(yaw_rad: f32) -> Array { /// Reflect orientation across the x-axis. /// (N,7) `quat_wxyz` orientation of `N` rigid objects. pub fn reflect_orientation_x(quat_wxyz: &ArrayView) -> Array { - let quat_wxyz = quat_wxyz.slice(s![.., -4..]); - let yaw_rad = quat_to_yaw(&quat_wxyz); + let yaw_rad = quat_to_yaw(quat_wxyz); let reflected_yaw_rad = -yaw_rad; yaw_to_quat(&reflected_yaw_rad.view()) } /// Reflect orientation across the y-axis. /// (N,7) `quat_wxyz` orientation of `N` rigid objects. -pub fn reflect_orientation_y(xyz_qwxyz: &ArrayView) -> Array { - let quat_wxyz = xyz_qwxyz.slice(s![.., -4..]); - let yaw_rad = quat_to_yaw(&quat_wxyz); +pub fn reflect_orientation_y(quat_wxyz: &ArrayView) -> Array { + let yaw_rad = quat_to_yaw(quat_wxyz); let reflected_yaw_rad = PI - yaw_rad; yaw_to_quat(&reflected_yaw_rad.view()) } + +/// Reflect translation across the x-axis. +pub fn reflect_translation_x(xyz_m: &ArrayView) -> Array { + let mut augmented_xyz_m = xyz_m.to_owned(); + augmented_xyz_m + .slice_mut(s![.., 1]) + .par_mapv_inplace(|y| -y); + augmented_xyz_m +} + +/// Reflect translation across the y-axis. +pub fn reflect_translation_y(xyz_m: &ArrayView) -> Array { + let mut augmented_xyz_m = xyz_m.to_owned(); + augmented_xyz_m + .slice_mut(s![.., 0]) + .par_mapv_inplace(|x| -x); + augmented_xyz_m +} From 9915d302c41e4e9e3d62d2d16c8d2434b5ec66ac Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Sat, 6 May 2023 11:11:24 -0700 Subject: [PATCH 12/12] Fix docstrings. --- rust/src/geometry/augmentations.rs | 6 +++--- rust/src/geometry/so3.rs | 4 ++-- rust/src/share.rs | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index b618943b..f4b23b18 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -11,7 +11,7 @@ use polars::{ }; use rand_distr::{Bernoulli, Distribution}; -use crate::share::{data_frame_to_ndarray_f32, ndarray_to_series_vec}; +use crate::share::{data_frame_to_ndarray_f32, ndarray_to_expr_vec}; use super::so3::{ reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y, @@ -59,7 +59,7 @@ pub fn sample_scene_reflection_x( .into_iter() .chain(orientation_column_names) .collect_vec(); - let series_vec = ndarray_to_series_vec(augmented_poses, column_names); + let series_vec = ndarray_to_expr_vec(augmented_poses, column_names); let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); (augmented_lidar, augmented_cuboids) } else { @@ -109,7 +109,7 @@ pub fn sample_scene_reflection_y( .into_iter() .chain(orientation_column_names) .collect_vec(); - let series_vec = ndarray_to_series_vec(augmented_poses, column_names); + let series_vec = ndarray_to_expr_vec(augmented_poses, column_names); let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); (augmented_lidar, augmented_cuboids) } else { diff --git a/rust/src/geometry/so3.rs b/rust/src/geometry/so3.rs index 78d7239f..b3cacc10 100644 --- a/rust/src/geometry/so3.rs +++ b/rust/src/geometry/so3.rs @@ -75,7 +75,7 @@ pub fn _yaw_to_quat(yaw_rad: f32) -> Array { } /// Reflect orientation across the x-axis. -/// (N,7) `quat_wxyz` orientation of `N` rigid objects. +/// (N,4) `quat_wxyz` orientation of `N` rigid objects. pub fn reflect_orientation_x(quat_wxyz: &ArrayView) -> Array { let yaw_rad = quat_to_yaw(quat_wxyz); let reflected_yaw_rad = -yaw_rad; @@ -83,7 +83,7 @@ pub fn reflect_orientation_x(quat_wxyz: &ArrayView) -> Array } /// Reflect orientation across the y-axis. -/// (N,7) `quat_wxyz` orientation of `N` rigid objects. +/// (N,4) `quat_wxyz` orientation of `N` rigid objects. pub fn reflect_orientation_y(quat_wxyz: &ArrayView) -> Array { let yaw_rad = quat_to_yaw(quat_wxyz); let reflected_yaw_rad = PI - yaw_rad; diff --git a/rust/src/share.rs b/rust/src/share.rs index 8416d273..03c58854 100644 --- a/rust/src/share.rs +++ b/rust/src/share.rs @@ -9,8 +9,8 @@ use polars::{ series::Series, }; -/// Convert the columns of an `ndarray::Array` into a vector of `polars::series::Series`. -pub fn ndarray_to_series_vec(arr: Array, column_names: Vec<&str>) -> Vec { +/// Convert the columns of an `ndarray::Array` into a vector of `polars` expressions. +pub fn ndarray_to_expr_vec(arr: Array, column_names: Vec<&str>) -> Vec { let num_dims = arr.shape()[1]; if num_dims != column_names.len() { panic!("Number of array columns and column names must match.");