spatial_math/

exts.rs

// Copyright (C) 2020-2025 spatial-math authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Extension traits for spatial algebra operations.
//!
//! This module provides utility traits and methods that extend the functionality
//! of basic vector and matrix types for use in spatial algebra computations.
//! These extensions bridge the gap between general linear algebra and the
//! specific requirements of Featherstone's spatial vector algebra.

use na::Dim;

use crate::{Mat3, MatrixView, MatrixViewMut, Real, Vec3, Vec4};

/// Extension trait for 3D vectors to support cross-product matrix operations.
///
/// The cross-product matrix (also known as skew-symmetric matrix) is a fundamental
/// tool in spatial algebra for converting cross products into matrix multiplications.
/// This enables efficient computation of spatial transformations and dynamics.
pub trait Vec3Ext {
    /// Convert a 3D vector to its cross-product matrix.
    ///
    /// The cross-product matrix \[v\]× satisfies: \[v\] × u = v × u for any vector u
    ///
    /// # Mathematical Definition
    ///
    /// For vector v = [x, y, z]ᵀ, the cross-product matrix is:
    /// ```text
    /// [v]× = |  0  -z   y |
    ///       |  z   0  -x |
    ///       | -y   x   0 |
    /// ```
    fn into_cross_matrix(self) -> Mat3;
}

impl Vec3Ext for Vec3 {
    #[inline]
    fn into_cross_matrix(self) -> Mat3 {
        // Construct the skew-symmetric cross-product matrix
        // |  0  -z   y |
        // |  z   0  -x |
        // | -y   x   0 |
        Mat3::from_columns(&[
            na::vector![0.0, self.z, -self.y],
            na::vector![-self.z, 0.0, self.x],
            na::vector![self.y, -self.x, 0.0],
        ])
    }
}

/// Extension trait for unit quaternions with spatial algebra utilities.
///
/// This trait provides additional methods for creating and validating unit quaternions
/// specifically tailored for use in spatial transformations and dynamics calculations.
pub trait UnitQuatExt {
    /// Create a unit quaternion from a 4-element array without validation.
    ///
    /// # Safety
    ///
    /// This function assumes the input represents a valid unit quaternion.
    /// Use with caution and only when you're certain the input is normalized.
    ///
    /// # Arguments
    ///
    /// * `array` - [x, y, z, w] quaternion components (w is scalar part)
    fn from_array_unchecked(array: [Real; 4]) -> Self;

    /// Create a unit quaternion from a 4D vector without validation.
    ///
    /// # Safety
    ///
    /// This function assumes the input vector represents a valid unit quaternion.
    /// Use with caution and only when you're certain the input is normalized.
    /// # Arguments
    ///
    /// * `v` - 4D vector containing [x, y, z, w] quaternion components
    fn from_vec4_unchecked(v: Vec4) -> Self;

    /// Check if the array represents a valid unit quaternion.
    ///
    /// Validates that the quaternion has unit length within the specified tolerance.
    /// This is useful for input validation and debugging quaternion-related issues.
    ///
    /// # Arguments
    ///
    /// * `array` - The [x, y, z, w] array to check
    /// * `eps` - The tolerance for checking the unit length (typically 1e-6 or smaller)
    ///
    /// # Returns
    ///
    /// `true` if the quaternion is normalized within tolerance, `false` otherwise
    fn is_valid(array: [Real; 4], eps: Real) -> bool;
}

impl UnitQuatExt for na::UnitQuaternion<Real> {
    #[inline]
    fn from_array_unchecked(array: [Real; 4]) -> Self {
        na::UnitQuaternion::new_unchecked(na::Quaternion::from(array))
    }

    #[inline]
    fn from_vec4_unchecked(v: Vec4) -> Self {
        na::UnitQuaternion::new_unchecked(na::Quaternion::from(v))
    }

    #[inline]
    fn is_valid(array: [Real; 4], eps: Real) -> bool {
        let length = Vec4::from(array).norm();
        (length - 1.0).abs() < eps
    }
}

pub trait ViewMatrix {
    fn view_matrix<const R: usize, const C: usize>(
        &self,
        element_offset: usize,
    ) -> MatrixView<R, C>;
}

pub trait ViewMatrixMut {
    fn view_matrix_mut<const R: usize, const C: usize>(
        &mut self,
        element_offset: usize,
    ) -> MatrixViewMut<R, C>;
}

macro_rules! impl_view_matrix {
    ($type:ty) => {
        impl ViewMatrix for $type {
            #[inline]
            fn view_matrix<const R: usize, const C: usize>(
                &self,
                element_offset: usize,
            ) -> MatrixView<R, C> {
                MatrixView::<R, C>::from_slice(&self[element_offset..])
            }
        }
    };
}

macro_rules! impl_view_matrix_mut {
    ($type:ty) => {
        impl ViewMatrixMut for $type {
            #[inline]
            fn view_matrix_mut<const R: usize, const C: usize>(
                &mut self,
                element_offset: usize,
            ) -> MatrixViewMut<R, C> {
                MatrixViewMut::<R, C>::from_slice(&mut self[element_offset..])
            }
        }
    };
}

impl_view_matrix! {[Real]}
impl_view_matrix! {Vec<Real>}

impl_view_matrix_mut! {[Real]}
impl_view_matrix_mut! {Vec<Real>}

pub trait MatrixExt {
    /// Returns true if any element in the matrix is NaN.
    fn any_nan(&self) -> bool;
}

impl<R, C, S> MatrixExt for na::Matrix<Real, R, C, S>
where
    R: na::Dim,
    C: na::Dim,
    S: na::RawStorage<Real, R, C>,
{
    #[inline]
    fn any_nan(&self) -> bool {
        self.iter().any(
            #[inline]
            |x| x.is_nan(),
        )
    }
}

pub trait SquareMatrixExt {
    fn is_positive_definite(&self) -> bool;
}

impl<R, S> SquareMatrixExt for na::Matrix<Real, R, R, S>
where
    R: Dim,
    S: na::RawStorage<Real, R, R>,
{
    #[inline]
    fn is_positive_definite(&self) -> bool {
        for size in 1..=self.ncols() {
            let d = self.view((0, 0), (size, size)).determinant();
            if d <= 0.0 {
                return false;
            }
        }
        true
    }
}

pub trait RealSliceExt {
    /// Returns a array of the given size starting at the given offset.
    ///
    /// # Panics
    ///
    /// Panics if there are not enough elements.
    fn array<const N: usize>(&self, offset: usize) -> [Real; N];
}

impl RealSliceExt for [Real] {
    #[inline]
    fn array<const N: usize>(&self, offset: usize) -> [Real; N] {
        self[offset..offset + N]
            .try_into()
            .expect("no enough elements")
    }
}

#[cfg(test)]
mod tests {
    use super::Vec3Ext;
    use crate::exts::SquareMatrixExt;
    use crate::{SMatrix, vec3};

    #[test]
    fn test_cross_matrix() {
        let v = vec3(1.0, 2.0, 3.0);
        let m = v.into_cross_matrix();
        assert_eq!(m.column(0), vec3(0.0, 3.0, -2.0));
        assert_eq!(m.column(1), vec3(-3.0, 0.0, 1.0));
        assert_eq!(m.column(2), vec3(2.0, -1.0, 0.0));

        assert_eq!(m, v.cross_matrix());
    }

    #[test]
    fn test_positive_define() {
        let m = SMatrix::<1, 1>::new(-1.0);
        assert!(!m.is_positive_definite());

        let m = SMatrix::<2, 2>::new(1.0, 0.0, 0.0, 1.0);
        assert!(m.is_positive_definite());
    }
}