Add vectorized rotation functions

src/lib.rs

@@ -97,7 +97,11 @@ pub fn libRustBCA(py: Python, m: &PyModule) -> PyResult<()> {
     #[cfg(feature = "parry3d")]
     m.add_function(wrap_pyfunction!(rotate_given_surface_normal_py, m)?)?;
     #[cfg(feature = "parry3d")]
+    m.add_function(wrap_pyfunction!(rotate_given_surface_normal_vec_py, m)?)?;
+    #[cfg(feature = "parry3d")]
     m.add_function(wrap_pyfunction!(rotate_back_py, m)?)?;
+    #[cfg(feature = "parry3d")]
+    m.add_function(wrap_pyfunction!(rotate_back_vec_py, m)?)?;
     Ok(())
 }
 
@@ -1410,6 +1414,8 @@ pub extern "C" fn rotate_given_surface_normal(nx: f64, ny: f64, nz: f64, ux: &mu
     *uz /= mag;
 }
 
+
+
 #[cfg(all(feature = "python", feature = "parry3d"))]
 #[pyfunction]
 /// rotate_given_surface_normal_py(nx, ny, nz, ux, uy, uz)
@@ -1433,6 +1439,46 @@ pub fn rotate_given_surface_normal_py(nx: f64, ny: f64, nz: f64, ux: f64, uy: f6
     (ux, uy, uz)
 }
 
+#[cfg(all(feature = "python", feature = "parry3d"))]
+#[pyfunction]
+/// rotate_given_surface_normal_vec_py(nx, ny, nz, ux, uy, uz)
+/// --
+///
+/// This function takes a particle direction and a normal vector and rotates from simulation to RustBCA coordinates.
+/// Args:
+///     nx (list(f64)): surface normal in global frame x-component.
+///     ny (list(f64)): surface normal in global frame y-component.
+///     nz (list(f64)): surface normal in global frame z-component.
+///     ux (list(f64)): particle direction in global frame x-component.
+///     uy (list(f64)): particle direction in global frame normal y-component.
+///     uz (list(f64)): particle direction in global frame normal z-component.
+/// Returns:
+///    direction (list(f64), list(f64), list(f64)): direction vector of particle in RustBCA coordinates.
+///    Note: non-incident particles will be returned with ux, uy, uz = (0, 0, 0)
+pub fn rotate_given_surface_normal_vec_py(nx: Vec<f64>, ny: Vec<f64>, nz: Vec<f64>, ux: Vec<f64>, uy: Vec<f64>, uz: Vec<f64>) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
+
+    let length = nx.len();
+
+    let mut ux_new = Vec::with_capacity(length);
+    let mut uy_new = Vec::with_capacity(length);
+    let mut uz_new = Vec::with_capacity(length);
+
+    (0..length).into_iter().for_each(|index| {
+
+        let mut ux_ = ux[index];
+        let mut uy_ = uy[index];
+        let mut uz_ = uz[index];
+
+        rotate_given_surface_normal(nx[index], ny[index], nz[index], &mut ux_, &mut uy_, &mut uz_);
+        ux_new.push(ux_);
+        uy_new.push(uy_);
+        uz_new.push(uz_);
+
+    });
+
+    (ux_new, uy_new, uz_new)
+}
+
 #[cfg(feature = "parry3d")]
 #[no_mangle]
 pub extern "C" fn rotate_back(nx: f64, ny: f64, nz: f64, ux: &mut f64, uy: &mut f64, uz: &mut f64) {
@@ -1488,6 +1534,43 @@ pub fn rotate_back_py(nx: f64, ny: f64, nz: f64, ux: f64, uy: f64, uz: f64) -> (
     (ux, uy, uz)
 }
 
+#[cfg(all(feature = "python", feature = "parry3d"))]
+#[pyfunction]
+/// rotate_back_vec_py(nx, ny, nz, ux, uy, uz)
+/// --
+///
+/// This function takes a RustBCA particle direction and a normal vector and rotates back from RustBCA to simulation coordinates.
+/// Args:
+///     nx (list(f64)): surface normal in global frame x-component.
+///     ny (list(f64)): surface normal in global frame y-component.
+///     nz (list(f64)): surface normal in global frame z-component.
+///     ux (list(f64)): particle direction in global frame x-component.
+///     uy (list(f64)): particle direction in global frame normal y-component.
+///     uz (list(f64)): particle direction in global frame normal z-component.
+/// Returns:
+///    direction (list(f64), list(f64), list(f64)): direction vector of particle in simulation coordinates.
+pub fn rotate_back_vec_py(nx: Vec<f64>, ny: Vec<f64>, nz: Vec<f64>, ux: Vec<f64>, uy: Vec<f64>, uz: Vec<f64>) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
+
+    let length = nx.len();
+
+    let mut ux_new = Vec::with_capacity(length);
+    let mut uy_new = Vec::with_capacity(length);
+    let mut uz_new = Vec::with_capacity(length);
+
+    (0..length).into_iter().for_each(|index| {
+
+        let mut ux_ = ux[index];
+        let mut uy_ = uy[index];
+        let mut uz_ = uz[index];
+        rotate_back(nx[index], ny[index], nz[index], &mut ux_, &mut uy_, &mut uz_);
+        ux_new.push(ux_);
+        uy_new.push(uy_);
+        uz_new.push(uz_);
+    });
+
+    (ux_new, uy_new, uz_new)
+}
+
 #[cfg(feature = "python")]
 /// A helper function to unpack a python float from a python any.
 fn unpack(python_float: &PyAny) -> f64 {