map rows and cols standalone

lib.rs

extern crate nalgebra as na;
use na::{Scalar, Matrix, MatrixMN, MatrixSlice, Matrix1, DVector, RowDVector};
use na::{U1, U3, Dynamic, Dim};
use na::allocator::Allocator;
use na::default_allocator::DefaultAllocator;
use na::storage::Storage;

trait MatrixMapExt<N, R, C, S>
where
    N: Scalar,
    R: Dim,
    C: Dim,
    S: Storage<N, R, C>,
{
    fn map_rows<N2, C2, S2, F>(&self, f: F) -> MatrixMN<N2, R, C2>
    where
        F: FnMut(MatrixSlice<N, U1, C, S::RStride, S::CStride>) -> Matrix<N2, U1, C2, S2>,
        C2: Dim,
        N2: Scalar,
        S2: Storage<N2, U1, C2>,
        DefaultAllocator: Allocator<N2, R, C2>,
    ;

    fn map_columns<N2, R2, S2, F>(&self, f: F) -> MatrixMN<N2, R2, C>
    where
        F: FnMut(MatrixSlice<N, R, U1, S::RStride, S::CStride>) -> Matrix<N2, R2, U1, S2>,
        R2: Dim,
        N2: Scalar,
        S2: Storage<N2, R2, U1>,
        DefaultAllocator: Allocator<N2, R2, C>,
    ;

    fn zip_map_columns<N2, R2, S2, N3, R3, S3, F>(
        &self,
        other: &Matrix<N2, R2, C, S2>,
        f: F,
    ) -> MatrixMN<N3, R3, C>
    where
        F: FnMut(
            MatrixSlice<N, R, U1, S::RStride, S::CStride>,
            MatrixSlice<N2, R2, U1, S2::RStride, S2::CStride>,
        ) -> Matrix<N3, R3, U1, S3>,
        R2: Dim,
        N2: Scalar,
        S2: Storage<N2, R2, C>,
        DefaultAllocator: Allocator<N2, R2, C>,
        R3: Dim,
        N3: Scalar,
        S3: Storage<N3, R3, U1>,
        DefaultAllocator: Allocator<N3, R3, C>,
    ;

    // ...and I think you get the picture with regards to zip_map_rows.
}


impl<N, R, C, S> MatrixMapExt<N, R, C, S> for Matrix<N, R, C, S>
where
    N: Scalar,
    R: Dim,
    C: Dim,
    S: Storage<N, R, C>,
{
    fn map_rows<N2, C2, S2, F>(&self, mut f: F) -> MatrixMN<N2, R, C2>
        where
            F: FnMut(MatrixSlice<N, U1, C, S::RStride, S::CStride>) -> Matrix<N2, U1, C2, S2>,
            C2: Dim,
            N2: Scalar,
            S2: Storage<N2, U1, C2>,
            DefaultAllocator: Allocator<N2, R, C2>,
    {
        // (this could probably be written to work much better for
        //  static-sized arrays, I just don't know how)
        let out_rows: Vec<_> = (0..self.nrows()).map(|i| f(self.row(i))).collect();
        MatrixMN::from_rows(&out_rows[..])
    }

    fn map_columns<N2, R2, S2, F>(&self, mut f: F) -> MatrixMN<N2, R2, C>
        where
            F: FnMut(MatrixSlice<N, R, U1, S::RStride, S::CStride>) -> Matrix<N2, R2, U1, S2>,
            R2: Dim,
            N2: Scalar,
            S2: Storage<N2, R2, U1>,
            DefaultAllocator: Allocator<N2, R2, C>,
    {
        // (this could probably be written to work much better for
        //  static-sized arrays, I just don't know how)
        let out_columns: Vec<_> = (0..self.ncols()).map(|i| f(self.column(i))).collect();
        MatrixMN::from_columns(&out_columns[..])
    }

    fn zip_map_columns<N2, R2, S2, N3, R3, S3, F>(
        &self,
        other: &Matrix<N2, R2, C, S2>,
        mut f: F,
    ) -> MatrixMN<N3, R3, C>
    where
        F: FnMut(
            MatrixSlice<N, R, U1, S::RStride, S::CStride>,
            MatrixSlice<N2, R2, U1, S2::RStride, S2::CStride>,
        ) -> Matrix<N3, R3, U1, S3>,
        R2: Dim,
        N2: Scalar,
        S2: Storage<N2, R2, C>,
        DefaultAllocator: Allocator<N2, R2, C>,
        R3: Dim,
        N3: Scalar,
        S3: Storage<N3, R3, U1>,
        DefaultAllocator: Allocator<N3, R3, C>,
    {
        assert_eq!(self.ncols(), other.ncols());

        // (this could probably be written to work much better for
        //  static-sized arrays, I just don't know how)
        let out_columns: Vec<_> =
            (0..self.ncols())
                .map(|i| f(self.column(i), other.column(i)))
                .collect();
        MatrixMN::from_columns(&out_columns[..])
    }
}

fn row_square_norms(m: &MatrixMN<f64, Dynamic, U3>) -> MatrixMN<f64, Dynamic, U1>
{
    m.map_rows(|row| Matrix1::new(row.norm_squared()))
}

fn col_square_norms(m: &MatrixMN<f64, Dynamic, U3>) -> MatrixMN<f64, U1, U3>
{
    m.map_columns(|col| Matrix1::new(col.norm_squared()))
}

fn scale_vectors<R: Dim, C: Dim>(
    vectors: &MatrixMN<f64, R, C>,
    scales: &::na::RowVectorN<f64, C>,
) -> MatrixMN<f64, R, C>
where
    DefaultAllocator: Allocator<f64, R, C> + Allocator<f64, U1, C> + Allocator<f64, R, U1>,
{
    vectors.zip_map_columns(scales, |v, scale| v * scale)
}

#[test]
fn row_col_square_norms() {
    let rows: Vec<_> = (0..6).map(|i| [i as f64, i as f64, 0.0].into()).collect();
    let matrix = MatrixMN::<f64, Dynamic, U3>::from_rows(&rows);
    assert_eq!(
        row_square_norms(&matrix),
        DVector::from_iterator(6, vec![0.0, 2.0, 8.0, 18.0, 32.0, 50.0])
    );
    assert_eq!(
        col_square_norms(&matrix),
        [55.0, 55.0, 0.0_f64].into()
    );
}

#[test]
fn scale_some_vectors() {
    let cols = ::na::Matrix3::new(
        1.0, 2.0, 3.0,
        4.0, 5.0, 6.0,
        7.0, 8.0, 9.0,
    );

    assert_eq!(
        scale_vectors(&cols, &[1.0, -1.0, 1.0].into()),
        ::na::Matrix3::new(
            1.0, -2.0, 3.0,
            4.0, -5.0, 6.0,
            7.0, -8.0, 9.0,
        )
    )
}