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 super::*;
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: Copy>(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;
166            a[j * m + i] = a_ij;
167        }
168    }
169}
170
171/// Out-place transpose for general matrix
172///
173/// Examples
174/// ---------
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 (l, b) = transpose(layout, &a);
181/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
182/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
183/// ```
184///
185/// ```rust
186/// # use lax::layout::*;
187/// let layout = MatrixLayout::F { col: 2, lda: 3 };
188/// let a = vec![1., 2., 3., 4., 5., 6.];
189/// let (l, b) = transpose(layout, &a);
190/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
191/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
192/// ```
193///
194/// Panics
195/// ------
196/// - If input array size and `layout` size mismatch
197///
198pub fn transpose<T: Copy>(layout: MatrixLayout, input: &[T]) -> (MatrixLayout, Vec<T>) {
199    let (m, n) = layout.size();
200    let transposed = layout.resized(n, m).t();
201    let m = m as usize;
202    let n = n as usize;
203    assert_eq!(input.len(), m * n);
204
205    let mut out: Vec<MaybeUninit<T>> = vec_uninit(m * n);
206
207    match layout {
208        MatrixLayout::C { .. } => {
209            for i in 0..m {
210                for j in 0..n {
211                    out[j * m + i].write(input[i * n + j]);
212                }
213            }
214        }
215        MatrixLayout::F { .. } => {
216            for i in 0..m {
217                for j in 0..n {
218                    out[i * n + j].write(input[j * m + i]);
219                }
220            }
221        }
222    }
223    (transposed, unsafe { out.assume_init() })
224}
225
226/// Out-place transpose for general matrix
227///
228/// Examples
229/// ---------
230///
231/// ```rust
232/// # use lax::layout::*;
233/// let layout = MatrixLayout::C { row: 2, lda: 3 };
234/// let a = vec![1., 2., 3., 4., 5., 6.];
235/// let mut b = vec![0.0; a.len()];
236/// let l = transpose_over(layout, &a, &mut b);
237/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
238/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
239/// ```
240///
241/// ```rust
242/// # use lax::layout::*;
243/// let layout = MatrixLayout::F { col: 2, lda: 3 };
244/// let a = vec![1., 2., 3., 4., 5., 6.];
245/// let mut b = vec![0.0; a.len()];
246/// let l = transpose_over(layout, &a, &mut b);
247/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
248/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
249/// ```
250///
251/// Panics
252/// ------
253/// - If input array sizes and `layout` size mismatch
254///
255pub fn transpose_over<T: Copy>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
256    let (m, n) = layout.size();
257    let transposed = layout.resized(n, m).t();
258    let m = m as usize;
259    let n = n as usize;
260    assert_eq!(from.len(), m * n);
261    assert_eq!(to.len(), m * n);
262
263    match layout {
264        MatrixLayout::C { .. } => {
265            for i in 0..m {
266                for j in 0..n {
267                    to[j * m + i] = from[i * n + j];
268                }
269            }
270        }
271        MatrixLayout::F { .. } => {
272            for i in 0..m {
273                for j in 0..n {
274                    to[i * n + j] = from[j * m + i];
275                }
276            }
277        }
278    }
279    transposed
280}