lax/
layout.rs

1//! Memory layout of matrices
2//!
3//! Different from ndarray format which consists of shape and strides,
4//! matrix format in LAPACK consists of row or column size and leading dimension.
5//!
6//! ndarray format and stride
7//! --------------------------
8//!
9//! Let us consider 3-dimensional array for explaining ndarray structure.
10//! The address of `(x,y,z)`-element in ndarray satisfies following relation:
11//!
12//! ```text
13//! shape = [Nx, Ny, Nz]
14//!     where Nx > 0, Ny > 0, Nz > 0
15//! stride = [Sx, Sy, Sz]
16//!
17//! &data[(x, y, z)] = &data[(0, 0, 0)] + Sx*x + Sy*y + Sz*z
18//!     for x < Nx, y < Ny, z < Nz
19//! ```
20//!
21//! The array is called
22//!
23//! - C-continuous if `[Sx, Sy, Sz] = [Nz*Ny, Nz, 1]`
24//! - F(Fortran)-continuous if `[Sx, Sy, Sz] = [1, Nx, Nx*Ny]`
25//!
26//! Strides of ndarray `[Sx, Sy, Sz]` take arbitrary value,
27//! e.g. it can be non-ordered `Sy > Sx > Sz`, or can be negative `Sx < 0`.
28//! If the minimum of `[Sx, Sy, Sz]` equals to `1`,
29//! the value of elements fills `data` memory region and called "continuous".
30//! Non-continuous ndarray is useful to get sub-array without copying data.
31//!
32//! Matrix layout for LAPACK
33//! -------------------------
34//!
35//! LAPACK interface focuses on the linear algebra operations for F-continuous 2-dimensional array.
36//! Under this restriction, stride becomes far simpler; we only have to consider the case `[1, S]`
37//! This `S` for a matrix `A` is called "leading dimension of the array A" in LAPACK document, and denoted by `lda`.
38//!
39
40use cauchy::Scalar;
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum MatrixLayout {
44    C { row: i32, lda: i32 },
45    F { col: i32, lda: i32 },
46}
47
48impl MatrixLayout {
49    pub fn size(&self) -> (i32, i32) {
50        match *self {
51            MatrixLayout::C { row, lda } => (row, lda),
52            MatrixLayout::F { col, lda } => (lda, col),
53        }
54    }
55
56    pub fn resized(&self, row: i32, col: i32) -> MatrixLayout {
57        match *self {
58            MatrixLayout::C { .. } => MatrixLayout::C { row, lda: col },
59            MatrixLayout::F { .. } => MatrixLayout::F { col, lda: row },
60        }
61    }
62
63    pub fn lda(&self) -> i32 {
64        std::cmp::max(
65            1,
66            match *self {
67                MatrixLayout::C { lda, .. } | MatrixLayout::F { lda, .. } => lda,
68            },
69        )
70    }
71
72    pub fn len(&self) -> i32 {
73        match *self {
74            MatrixLayout::C { row, .. } => row,
75            MatrixLayout::F { col, .. } => col,
76        }
77    }
78
79    pub fn is_empty(&self) -> bool {
80        self.len() == 0
81    }
82
83    pub fn same_order(&self, other: &MatrixLayout) -> bool {
84        match (self, other) {
85            (MatrixLayout::C { .. }, MatrixLayout::C { .. }) => true,
86            (MatrixLayout::F { .. }, MatrixLayout::F { .. }) => true,
87            _ => false,
88        }
89    }
90
91    pub fn toggle_order(&self) -> Self {
92        match *self {
93            MatrixLayout::C { row, lda } => MatrixLayout::F { lda: row, col: lda },
94            MatrixLayout::F { col, lda } => MatrixLayout::C { row: lda, lda: col },
95        }
96    }
97
98    /// Transpose without changing memory representation
99    ///
100    /// C-contigious row=2, lda=3
101    ///
102    /// ```text
103    /// [[1, 2, 3]
104    ///  [4, 5, 6]]
105    /// ```
106    ///
107    /// and F-contigious col=2, lda=3
108    ///
109    /// ```text
110    /// [[1, 4]
111    ///  [2, 5]
112    ///  [3, 6]]
113    /// ```
114    ///
115    /// have same memory representation `[1, 2, 3, 4, 5, 6]`, and this toggles them.
116    ///
117    /// ```
118    /// # use lax::layout::*;
119    /// let layout = MatrixLayout::C { row: 2, lda: 3 };
120    /// assert_eq!(layout.t(), MatrixLayout::F { col: 2, lda: 3 });
121    /// ```
122    pub fn t(&self) -> Self {
123        match *self {
124            MatrixLayout::C { row, lda } => MatrixLayout::F { col: row, lda },
125            MatrixLayout::F { col, lda } => MatrixLayout::C { row: col, lda },
126        }
127    }
128}
129
130/// In-place transpose of a square matrix by keeping F/C layout
131///
132/// Transpose for C-continuous array
133///
134/// ```rust
135/// # use lax::layout::*;
136/// let layout = MatrixLayout::C { row: 2, lda: 2 };
137/// let mut a = vec![1., 2., 3., 4.];
138/// square_transpose(layout, &mut a);
139/// assert_eq!(a, &[1., 3., 2., 4.]);
140/// ```
141///
142/// Transpose for F-continuous array
143///
144/// ```rust
145/// # use lax::layout::*;
146/// let layout = MatrixLayout::F { col: 2, lda: 2 };
147/// let mut a = vec![1., 3., 2., 4.];
148/// square_transpose(layout, &mut a);
149/// assert_eq!(a, &[1., 2., 3., 4.]);
150/// ```
151///
152/// Panics
153/// ------
154/// - If size of `a` and `layout` size mismatch
155///
156pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
157    let (m, n) = layout.size();
158    let n = n as usize;
159    let m = m as usize;
160    assert_eq!(a.len(), n * m);
161    for i in 0..m {
162        for j in (i + 1)..n {
163            let a_ij = a[i * n + j];
164            let a_ji = a[j * m + i];
165            a[i * n + j] = a_ji.conj();
166            a[j * m + i] = a_ij.conj();
167        }
168    }
169}
170
171/// Out-place transpose for general matrix
172///
173/// Inplace transpose of non-square matrices is hard.
174/// See also: https://en.wikipedia.org/wiki/In-place_matrix_transposition
175///
176/// ```rust
177/// # use lax::layout::*;
178/// let layout = MatrixLayout::C { row: 2, lda: 3 };
179/// let a = vec![1., 2., 3., 4., 5., 6.];
180/// let mut b = vec![0.0; a.len()];
181/// let l = transpose(layout, &a, &mut b);
182/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
183/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
184/// ```
185///
186/// ```rust
187/// # use lax::layout::*;
188/// let layout = MatrixLayout::F { col: 2, lda: 3 };
189/// let a = vec![1., 2., 3., 4., 5., 6.];
190/// let mut b = vec![0.0; a.len()];
191/// let l = transpose(layout, &a, &mut b);
192/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
193/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
194/// ```
195///
196/// Panics
197/// ------
198/// - If size of `a` and `layout` size mismatch
199///
200pub fn transpose<T: Scalar>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
201    let (m, n) = layout.size();
202    let transposed = layout.resized(n, m).t();
203    let m = m as usize;
204    let n = n as usize;
205    assert_eq!(from.len(), m * n);
206    assert_eq!(to.len(), m * n);
207
208    match layout {
209        MatrixLayout::C { .. } => {
210            for i in 0..m {
211                for j in 0..n {
212                    to[j * m + i] = from[i * n + j];
213                }
214            }
215        }
216        MatrixLayout::F { .. } => {
217            for i in 0..m {
218                for j in 0..n {
219                    to[i * n + j] = from[j * m + i];
220                }
221            }
222        }
223    }
224    transposed
225}