ndarray_linalg/
layout.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
//! Convert ndarray into LAPACK-compatible matrix format

use super::error::*;
use ndarray::*;

pub use lax::layout::MatrixLayout;

pub trait AllocatedArray {
    type Elem;
    fn layout(&self) -> Result<MatrixLayout>;
    fn square_layout(&self) -> Result<MatrixLayout>;
    /// Returns Ok iff the matrix is square (without computing the layout).
    fn ensure_square(&self) -> Result<()>;
    fn as_allocated(&self) -> Result<&[Self::Elem]>;
}

pub trait AllocatedArrayMut: AllocatedArray {
    fn as_allocated_mut(&mut self) -> Result<&mut [Self::Elem]>;
}

impl<A, S> AllocatedArray for ArrayBase<S, Ix2>
where
    S: Data<Elem = A>,
{
    type Elem = A;

    fn layout(&self) -> Result<MatrixLayout> {
        let shape = self.shape();
        let strides = self.strides();
        if shape[0] == strides[1] as usize {
            return Ok(MatrixLayout::F {
                col: self.ncols() as i32,
                lda: self.nrows() as i32,
            });
        }
        if shape[1] == strides[0] as usize {
            return Ok(MatrixLayout::C {
                row: self.nrows() as i32,
                lda: self.ncols() as i32,
            });
        }
        Err(LinalgError::InvalidStride {
            s0: strides[0],
            s1: strides[1],
        })
    }

    fn square_layout(&self) -> Result<MatrixLayout> {
        let l = self.layout()?;
        let (n, m) = l.size();
        if n == m {
            Ok(l)
        } else {
            Err(LinalgError::NotSquare { rows: n, cols: m })
        }
    }

    fn ensure_square(&self) -> Result<()> {
        if self.is_square() {
            Ok(())
        } else {
            Err(LinalgError::NotSquare {
                rows: self.nrows() as i32,
                cols: self.ncols() as i32,
            })
        }
    }

    fn as_allocated(&self) -> Result<&[A]> {
        self.as_slice_memory_order()
            .ok_or(LinalgError::MemoryNotCont)
    }
}

impl<A, S> AllocatedArrayMut for ArrayBase<S, Ix2>
where
    S: DataMut<Elem = A>,
{
    fn as_allocated_mut(&mut self) -> Result<&mut [A]> {
        self.as_slice_memory_order_mut()
            .ok_or(LinalgError::MemoryNotCont)
    }
}