mdarray-linalg 0.2.0

//! Utility functions for matrix printing, shape retrieval, identity
//! generation, Kronecker product, trace, transpose operations, ...
//!
//! This module contains small user-facing utilities, plus hidden unstable
//! helpers used by backend implementation crates. It is not meant to be a
//! complete collection of linear algebra utilities at this time.

use mdarray::{Array, Dim, Layout, Shape, Slice, tensor};
use num_complex::ComplexFloat;
use num_traits::{One, Zero};

/// Displays a numeric `mdarray` in a human-readable format (NumPy-style)
pub fn pretty_print<T: ComplexFloat + std::fmt::Display, D0: Dim, D1: Dim>(mat: &Array<T, (D0, D1)>)
where
    <T as num_complex::ComplexFloat>::Real: std::fmt::Display,
{
    let shape = mat.shape();
    for i in 0..shape.dim(0) {
        for j in 0..shape.dim(1) {
            let v = mat[[i, j]];
            print!("{:>10.4} {:+.4}i  ", v.re(), v.im(),);
        }
        println!();
    }
    println!();
}

// The following backend-oriented helpers are exported for workspace backend
// crates. They are hidden from generated documentation and may be redesigned
// before the public API stabilizes.
/// Safely casts a value to `i32`
#[doc(hidden)]
pub fn into_i32<T>(x: T) -> i32
where
    T: TryInto<i32>,
    <T as TryInto<i32>>::Error: std::fmt::Debug,
{
    x.try_into().expect("dimension must fit into i32")
}

/// Make sure that matrix shapes are compatible with `C = A * B`, and
/// return the dimensions `(m, n, k)` safely cast to `i32`, where `C` is `(m
/// x n)`, and `k` is the common dimension of `A` and `B`
#[doc(hidden)]
pub fn dims3(a_shape: impl Shape, b_shape: impl Shape, c_shape: impl Shape) -> (i32, i32, i32) {
    let (m, k) = (a_shape.dim(0), a_shape.dim(1));
    let (k2, n) = (b_shape.dim(0), b_shape.dim(1));
    let (m2, n2) = (c_shape.dim(0), c_shape.dim(1));

    assert!(m == m2, "a and c must agree in number of rows");
    assert!(n == n2, "b and c must agree in number of columns");
    assert!(
        k == k2,
        "a's number of columns must be equal to b's number of rows"
    );

    (into_i32(m), into_i32(n), into_i32(k))
}

/// Make sure that matrix shapes are compatible with `A * B`, and return
/// the dimensions `(m, n)` safely cast to `i32`
#[doc(hidden)]
pub fn dims2(a_shape: impl Shape, b_shape: impl Shape) -> (i32, i32) {
    let (m, k) = (a_shape.dim(0), a_shape.dim(1));
    let (k2, n) = (b_shape.dim(0), b_shape.dim(1));

    assert!(
        k == k2,
        "a's number of columns must be equal to b's number of rows"
    );

    (into_i32(m), into_i32(n))
}

/// Transposes a matrix in-place. Dimensions stay the same, only the memory ordering changes.
/// - For square matrices: swaps elements across the main diagonal.
/// - For rectangular matrices: reshuffles data in a temporary buffer so that the
///   same `(rows, cols)` slice now represents the transposed layout.
#[doc(hidden)]
pub fn transpose_in_place<T, D0, D1, L>(c: &mut Slice<T, (D0, D1), L>)
where
    T: ComplexFloat + Default,
    D0: Dim,
    D1: Dim,
    L: Layout,
{
    let (m, n) = *c.shape();

    let m = m.size();
    let n = n.size();

    if n == m {
        for i in 0..m {
            for j in (i + 1)..n {
                c.swap(i * n + j, j * n + i);
            }
        }
    } else {
        let mut result = tensor![[T::default(); m]; n];
        for j in 0..n {
            for i in 0..m {
                result[j * m + i] = c[i * n + j];
            }
        }
        for j in 0..n {
            for i in 0..m {
                c[j * m + i] = result[j * m + i];
            }
        }
    }
}

/// Conjugates a matrix in-place.
/// For complex matrices, replaces each element z with its conjugate conj(z).
/// For real matrices, this is a no-op.
#[doc(hidden)]
pub fn conjugate_in_place<T, D0, D1, L>(c: &mut Slice<T, (D0, D1), L>)
where
    T: ComplexFloat + Default,
    D0: Dim,
    D1: Dim,
    L: Layout,
{
    c.iter_mut().for_each(|elem| *elem = elem.conj());
}

/// Convert pivot indices to permutation matrix
#[doc(hidden)]
pub fn ipiv_to_perm_mat<T: ComplexFloat, D0: Dim, D1: Dim>(
    ipiv: &[i32],
    m: usize,
) -> Array<T, (D0, D1)> {
    let mut p = Array::from_elem(<(D0, D1) as Shape>::from_dims(&[m, m]), T::zero());

    for i in 0..m {
        p[[i, i]] = T::one();
    }

    // Apply row swaps according to LAPACK's ipiv convention
    for i in 0..ipiv.len() {
        let pivot_row = (ipiv[i] - 1) as usize; // LAPACK uses 1-based indexing
        if pivot_row != i {
            for j in 0..m {
                let temp = p[[i, j]];
                p[[i, j]] = p[[pivot_row, j]];
                p[[pivot_row, j]] = temp;
            }
        }
    }

    p
}

/// Given an input matrix of shape `(m × n)`, this function creates and returns
/// a new matrix of shape `(n × m)`, where each element at position `(i, j)` in the
/// original is moved to position `(j, i)` in the result.
#[doc(hidden)]
pub fn to_col_major<T, D0: Dim, D1: Dim, L>(c: &Slice<T, (D0, D1), L>) -> Array<T, (D1, D0)>
where
    T: ComplexFloat + Default + Clone,
    L: Layout,
{
    let csh = *c.shape();
    let (m, n) = (csh.dim(0), csh.dim(1));

    let shape = <(D1, D0) as Shape>::from_dims(&[n, m]);
    let mut result = Array::<T, (D1, D0)>::zeros(shape);

    for i in 0..m {
        for j in 0..n {
            result[[j, i]] = c[[i, j]];
        }
    }

    result
}

/// Computes the trace of a square matrix (sum of diagonal elements).
/// # Examples
/// ```
/// use mdarray::tensor;
/// use mdarray_linalg::utils::trace;
///
/// let a = tensor![[1., 2., 3.],
///                 [4., 5., 6.],
///                 [7., 8., 9.]];
///
/// let tr = trace(&a);
/// assert_eq!(tr, 15.0);
/// ```
pub fn trace<T, D0, D1, L>(a: &Slice<T, (D0, D1), L>) -> T
where
    T: ComplexFloat + std::ops::Add<Output = T> + Copy,
    D0: Dim,
    D1: Dim,
    L: Layout,
{
    let ash = *a.shape();
    let (m, n) = (ash.dim(0), ash.dim(1));
    assert_eq!(m, n, "trace is only defined for square matrices");

    let mut tr = T::zero();
    for i in 0..n {
        tr = tr + a[[i, i]];
    }
    tr
}

/// Creates an identity matrix of size `n x n`.
/// # Examples
/// ```
/// use mdarray::tensor;
/// use mdarray_linalg::utils::identity;
///
/// let i3 = identity::<f64, usize, usize>(3);
/// assert_eq!(i3, tensor![[1.,0.,0.],[0.,1.,0.],[0.,0.,1.]]);
/// ```
pub fn identity<T: Zero + One, D0: Dim, D1: Dim>(n: usize) -> Array<T, (D0, D1)> {
    Array::<T, (D0, D1)>::from_fn(<(D0, D1) as Shape>::from_dims(&[n, n]), |i| {
        if i[0] == i[1] { T::one() } else { T::zero() }
    })
}

/// Creates a diagonal matrix of size `n x n` with ones on a specified diagonal.
///
/// The diagonal can be shifted using `k`:
/// - `k = 0` → main diagonal (default, standard identity)
/// - `k > 0` → k-th diagonal above the main one
/// - `k < 0` → k-th diagonal below the main one
/// # Examples
/// ```
/// use mdarray::{Const, tensor};
/// use mdarray_linalg::utils::identity_k;
///
/// let i3 = identity_k::<f64, Const<3>, Const<3>>(3, 1);
/// assert_eq!(i3, tensor![[0.,1.,0.],[0.,0.,1.],[0.,0.,0.]]);
/// ```
pub fn identity_k<T: Zero + One, D0: Dim, D1: Dim>(n: usize, k: isize) -> Array<T, (D0, D1)> {
    Array::<T, (D0, D1)>::from_fn(<(D0, D1) as Shape>::from_dims(&[n, n]), |i| {
        if (i[1] as isize - i[0] as isize) == k {
            T::one()
        } else {
            T::zero()
        }
    })
}

/// Computes the Kronecker product of two 2D tensors.
///
/// The Kronecker product of matrices `A (m×n)` and `B (p×q)` is defined as the
/// block matrix of size `(m*p) × (n*q)` where each element `a[i, j]` of `A`
/// multiplies the entire matrix `B`.
///
/// # Examples
/// ```
/// use mdarray::tensor;
/// use mdarray_linalg::utils::kron;
///
/// let a = tensor![[1., 2.],
///                 [3., 4.]];
///
/// let b = tensor![[0., 5.],
///                 [6., 7.]];
///
/// let k = kron(&a, &b);
///
/// assert_eq!(k, tensor![
///     [ 0.,  5.,  0., 10.],
///     [ 6.,  7., 12., 14.],
///     [ 0., 15.,  0., 20.],
///     [18., 21., 24., 28.]
/// ]);
/// ```
pub fn kron<T, D0, D1, La, Lb>(
    a: &Slice<T, (D0, D1), La>,
    b: &Slice<T, (D0, D1), Lb>,
) -> Array<T, (D0, D1)>
where
    T: ComplexFloat + std::ops::Mul<Output = T> + Copy,
    D0: Dim,
    D1: Dim,
    La: Layout,
    Lb: Layout,
{
    let ash = *a.shape();
    let (ma, na) = (ash.dim(0), ash.dim(1));

    let bsh = *b.shape();
    let (mb, nb) = (bsh.dim(0), bsh.dim(1));

    let out_shape = <(D0, D1) as Shape>::from_dims(&[ma * mb, na * nb]);

    Array::<T, (D0, D1)>::from_fn(out_shape, |idx| {
        let i = idx[0];
        let j = idx[1];

        let ai = i / mb;
        let bi = i % mb;
        let aj = j / nb;
        let bj = j % nb;

        a[[ai, aj]] * b[[bi, bj]]
    })
}

/// Converts a flat index to multidimensional coordinates.
///
/// # Examples
///
/// ```
/// use mdarray::DArray;
/// use mdarray_linalg::utils::unravel_index;
///
/// let x = DArray::<usize, 2>::from_fn([2,3], |i| i[0] + i[1]);
///
/// assert_eq!(unravel_index(&x, 0), vec![0, 0]);
/// assert_eq!(unravel_index(&x, 4), vec![1, 1]);
/// assert_eq!(unravel_index(&x, 5), vec![1, 2]);
/// ```
///
/// # Panics
///
/// Panics if `flat` is out of bounds (>= `x.len()`).
pub fn unravel_index<T, S: Shape, L: Layout>(x: &Slice<T, S, L>, mut flat: usize) -> Vec<usize> {
    let rank = x.rank();

    assert!(
        flat < x.len(),
        "flat index out of bounds: {} >= {}",
        flat,
        x.len()
    );

    let mut coords = vec![0usize; rank];

    for i in (0..rank).rev() {
        let dim = x.shape().dim(i);
        coords[i] = flat % dim;
        flat /= dim;
    }

    coords
}

/// Creates a diagonal matrix from a 1D slice, placing its elements on the main diagonal.
///
/// # Examples
/// ```
/// use mdarray::{Const, array, view};
/// use mdarray_linalg::utils::diag;
///
/// let v = view![1., 2., 3.];
/// let d = diag(&v);
/// assert_eq!(d, array![[1.,0.,0.],[0.,2.,0.],[0.,0.,3.]]);
/// ```
pub fn diag<T: Zero + One + Clone, D: Dim>(v: &Slice<T, (D,)>) -> Array<T, (D, D)> {
    let n = v.dim(0);
    Array::<T, (D, D)>::from_fn(<(D, D) as Shape>::from_dims(&[n, n]), |i| {
        if i[0] == i[1] {
            v[i[0]].clone()
        } else {
            T::zero()
        }
    })
}