default/src/linsolve.rs

/*!
Linear equation solvers for small problems stored in Rust arrays.
*/

use crate::types::Float;
#[cfg(nightly)]
use std::mem::MaybeUninit;

/// Gaussian elimination for $AX=B$, where $A$ and $B$ are both stored in `ab`,
/// $A \in \mathbb{R}^{M \times M}$ and $X, B \in \mathbb{R}^{M \times K}$.
pub fn linsolve0<F: Float, const M: usize, const N: usize, const K: usize>(
    mut ab: [[F; N]; M],
) -> [[F; K]; M] {
    assert_eq!(M + K, N);

    let mut k = 0;

    // Convert to row-echelon form
    for h in 0..(M - 1) {
        // Find pivotable column (has some non-zero entries in rows ≥ h)
        'find_pivot: while k < N {
            let (mut î, mut v) = (h, ab[h][k].abs());
            // Find row ≥ h of maximum absolute value in this column
            for i in (h + 1)..M {
                let ṽ = ab[i][k].abs();
                if ṽ > v {
                    î = i;
                    v = ṽ;
                }
            }
            if v > F::ZERO {
                ab.swap(h, î);
                for i in (h + 1)..M {
                    let f = ab[i][k] / ab[h][k];
                    ab[i][k] = F::ZERO;
                    for j in (k + 1)..N {
                        ab[i][j] -= ab[h][j] * f;
                    }
                }
                k += 1;
                break 'find_pivot;
            }
            k += 1
        }
    }

    // Solve UAX=UB for X where UA with U presenting the transformations above an
    // upper triangular matrix.
    //
    // If the "nightly" feature is enabled, we will use an uninitialised array for a
    // little bit of efficiency.
    // This use of MaybeUninit assumes F : Copy. Otherwise undefined behaviour may occur.
    #[cfg(nightly)]
    {
        let mut x: [[MaybeUninit<F>; K]; M] = [[const { MaybeUninit::uninit() }; K]; M];
        //unsafe { std::mem::MaybeUninit::uninit().assume_init() };
        for i in (0..M).rev() {
            for 𝓁 in 0..K {
                let mut tmp = ab[i][M + 𝓁];
                for j in (i + 1)..M {
                    tmp -= ab[i][j] * unsafe { *(x[j][𝓁].assume_init_ref()) };
                }
                tmp /= ab[i][i];
                x[i][𝓁].write(tmp);
            }
        }
        unsafe {
            //core::intrinsics::assert_inhabited::<[[F; K]; M]>();
            (&x as *const _ as *const [[F; K]; M]).read()
        }
    }
    #[cfg(not(nightly))]
    {
        let mut x: [[F; K]; M] = [[F::ZERO; K]; M];
        for i in (0..M).rev() {
            for 𝓁 in 0..K {
                let mut tmp = ab[i][M + 𝓁];
                for j in (i + 1)..M {
                    tmp -= ab[i][j] * x[j][𝓁];
                }
                tmp /= ab[i][i];
                x[i][𝓁] = tmp;
            }
        }
        x
    }
}

/// Gaussian elimination for $Ax=b$, where $A$ and $b$ are both stored in `ab`,
/// $A \in \mathbb{R}^{M \times M}$ and $x, b \in \mathbb{R}^M$.
#[inline]
pub fn linsolve<F: Float, const M: usize, const N: usize>(ab: [[F; N]; M]) -> [F; M] {
    let x: [[F; 1]; M] = linsolve0(ab);
    unsafe { *((&x as *const [F; 1]) as *const [F; M]) }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn linsolve_test() {
        let ab1 = [[1.0, 2.0, 3.0], [2.0, 1.0, 6.0]];
        assert_eq!(linsolve(ab1), [3.0, 0.0]);
        let ab2 = [
            [1.0, 2.0, 0.0, 1.0],
            [4.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
        ];
        assert_eq!(linsolve(ab2), [0.0, 0.5, 0.0]);
    }
}