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}