nac3/nac3standalone/linalg_externfns/src/lib.rs

/// Uses `nalgebra` crate to invoke `np_linalg` and `sp_linalg` functions
/// When converting between `nalgebra::Matrix` and `NDArray` following considerations are necessary
///
/// * Both `nalgebra::Matrix` and `NDArray` require their content to be stored in row-major order
/// * `NDArray` data pointer can be directly read and converted to `nalgebra::Matrix` (row and column number must be known)
/// * `nalgebra::Matrix::as_slice` returns the content of matrix in column-major order and initial data needs to be transposed before storing it in `NDArray` data pointer
mod runtime_exception;
use core::slice;
use nalgebra::DMatrix;

macro_rules! raise_exn {
    ($name:expr, $fn_name:expr, $message:expr, $param0:expr, $param1:expr, $param2:expr) => {{
        use cslice::AsCSlice;
        let name_id = $crate::runtime_exception::get_exception_id($name);
        let exn = $crate::runtime_exception::Exception {
            id: name_id,
            file: file!().as_c_slice(),
            line: line!(),
            column: column!(),
            // https://github.com/rust-lang/rfcs/pull/1719
            function: $fn_name.as_c_slice(),
            message: $message.as_c_slice(),
            param: [$param0, $param1, $param2],
        };
        #[allow(unused_unsafe)]
        unsafe {
            $crate::runtime_exception::raise(&exn)
        }
    }};
    ($name:expr, $fn_name:expr, $message:expr) => {{
        raise_exn!($name, $fn_name, $message, 0, 0, 0)
    }};
}
pub struct InputMatrix {
    pub ndims: usize,
    pub dims: *const usize,
    pub data: *mut f64,
}
impl InputMatrix {
    fn get_dims(&mut self) -> Vec<usize> {
        let dims = unsafe { slice::from_raw_parts(self.dims, self.ndims) };
        dims.to_vec()
    }
}

/// # Safety
///
/// `mat1` and `mat2` should point to a valid 1DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_dot(mat1: *mut InputMatrix, mat2: *mut InputMatrix) -> f64 {
    let mat1 = mat1.as_mut().unwrap();
    let mat2 = mat2.as_mut().unwrap();

    if !(mat1.ndims == 1 && mat2.ndims == 1) {
        let err_msg = format!(
            "expected 1D Vector Input, but received {}D and {}D input",
            mat1.ndims, mat2.ndims
        );
        raise_exn!("ValueError", "np_dot", err_msg);
    }

    let dim1 = (*mat1).get_dims();
    let dim2 = (*mat2).get_dims();

    if dim1[0] != dim2[0] {
        let err_msg = format!("shapes ({},) and ({},) not aligned", dim1[0], dim2[0]);
        raise_exn!("ValueError", "np_dot", err_msg);
    }
    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0]) };
    let data_slice2 = unsafe { slice::from_raw_parts_mut(mat2.data, dim2[0]) };

    let matrix1 = DMatrix::from_row_slice(dim1[0], 1, data_slice1);
    let matrix2 = DMatrix::from_row_slice(dim2[0], 1, data_slice2);

    matrix1.dot(&matrix2)
}

/// # Safety
///
/// `mat1` and `mat2` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_linalg_matmul(
    mat1: *mut InputMatrix,
    mat2: *mut InputMatrix,
    out: *mut InputMatrix,
) {
    let mat1 = mat1.as_mut().unwrap();
    let mat2 = mat2.as_mut().unwrap();
    let out = out.as_mut().unwrap();

    if !(mat1.ndims == 2 && mat2.ndims == 2) {
        let err_msg = format!(
            "expected 2D Vector Input, but received {}D and {}D input",
            mat1.ndims, mat2.ndims
        );
        raise_exn!("ValueError", "np_matmul", err_msg);
    }

    let dim1 = (*mat1).get_dims();
    let dim2 = (*mat2).get_dims();

    if dim1[1] != dim2[0] {
        let err_msg = format!(
            "shapes ({},{}) and ({},{}) not aligned: {} (dim 1) != {} (dim 0)",
            dim1[0], dim1[1], dim2[0], dim2[1], dim1[1], dim2[0]
        );
        raise_exn!("ValueError", "np_matmul", err_msg);
    }

    let outdim = out.get_dims();
    let out_slice = unsafe { slice::from_raw_parts_mut(out.data, outdim[0] * outdim[1]) };
    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };
    let data_slice2 = unsafe { slice::from_raw_parts_mut(mat2.data, dim2[0] * dim2[1]) };

    let matrix1 = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let matrix2 = DMatrix::from_row_slice(dim2[0], dim2[1], data_slice2);
    let mut result = DMatrix::<f64>::zeros(outdim[0], outdim[1]);

    matrix1.mul_to(&matrix2, &mut result);
    out_slice.copy_from_slice(result.transpose().as_slice());
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_linalg_cholesky(mat1: *mut InputMatrix, out: *mut InputMatrix) {
    let mat1 = mat1.as_mut().unwrap();
    let out = out.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "np_linalg_cholesky", err_msg);
    }

    let dim1 = (*mat1).get_dims();
    if dim1[0] != dim1[1] {
        let err_msg =
            format!("last 2 dimensions of the array must be square: {0} != {1}", dim1[0], dim1[1]);
        raise_exn!("LinAlgError", "np_linalg_cholesky", err_msg);
    }

    let outdim = out.get_dims();
    let out_slice = unsafe { slice::from_raw_parts_mut(out.data, outdim[0] * outdim[1]) };
    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };

    let matrix1 = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let result = matrix1.cholesky();
    match result {
        Some(res) => {
            out_slice.copy_from_slice(res.unpack().transpose().as_slice());
        }
        None => {
            raise_exn!("LinAlgError", "np_linalg_cholesky", "Matrix is not positive definite");
        }
    };
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_linalg_qr(
    mat1: *mut InputMatrix,
    out_q: *mut InputMatrix,
    out_r: *mut InputMatrix,
) {
    let mat1 = mat1.as_mut().unwrap();
    let out_q = out_q.as_mut().unwrap();
    let out_r = out_r.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "np_linalg_cholesky", err_msg);
    }

    let dim1 = (*mat1).get_dims();
    let outq_dim = (*out_q).get_dims();
    let outr_dim = (*out_r).get_dims();

    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };
    let out_q_slice = unsafe { slice::from_raw_parts_mut(out_q.data, outq_dim[0] * outq_dim[1]) };
    let out_r_slice = unsafe { slice::from_raw_parts_mut(out_r.data, outr_dim[0] * outr_dim[1]) };

    // Refer to https://github.com/dimforge/nalgebra/issues/735
    let matrix1 = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);

    let res = matrix1.qr();
    let (q, r) = res.unpack();

    // Uses different algo need to match numpy
    out_q_slice.copy_from_slice(q.transpose().as_slice());
    out_r_slice.copy_from_slice(r.transpose().as_slice());
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_linalg_svd(
    mat1: *mut InputMatrix,
    outu: *mut InputMatrix,
    outs: *mut InputMatrix,
    outvh: *mut InputMatrix,
) {
    let mat1 = mat1.as_mut().unwrap();
    let outu = outu.as_mut().unwrap();
    let outs = outs.as_mut().unwrap();
    let outvh = outvh.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "np_linalg_svd", err_msg);
    }

    let dim1 = (*mat1).get_dims();
    let outu_dim = (*outu).get_dims();
    let outs_dim = (*outs).get_dims();
    let outvh_dim = (*outvh).get_dims();

    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };
    let out_u_slice = unsafe { slice::from_raw_parts_mut(outu.data, outu_dim[0] * outu_dim[1]) };
    let out_s_slice = unsafe { slice::from_raw_parts_mut(outs.data, outs_dim[0]) };
    let out_vh_slice =
        unsafe { slice::from_raw_parts_mut(outvh.data, outvh_dim[0] * outvh_dim[1]) };

    let matrix = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let result = matrix.svd(true, true);
    out_u_slice.copy_from_slice(result.u.unwrap().transpose().as_slice());
    out_s_slice.copy_from_slice(result.singular_values.as_slice());
    out_vh_slice.copy_from_slice(result.v_t.unwrap().transpose().as_slice());
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_linalg_inv(mat1: *mut InputMatrix, out: *mut InputMatrix) {
    let mat1 = mat1.as_mut().unwrap();
    let out = out.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "np_linalg_inv", err_msg);
    }
    let dim1 = (*mat1).get_dims();

    if dim1[0] != dim1[1] {
        let err_msg =
            format!("last 2 dimensions of the array must be square: {0} != {1}", dim1[0], dim1[1]);
        raise_exn!("LinAlgError", "np_linalg_inv", err_msg);
    }

    let outdim = out.get_dims();
    let out_slice = unsafe { slice::from_raw_parts_mut(out.data, outdim[0] * outdim[1]) };
    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };

    let matrix = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    if !matrix.is_invertible() {
        raise_exn!("LinAlgError", "np_linalg_inv", "no inverse for Singular Matrix");
    }
    let inv = matrix.try_inverse().unwrap();
    out_slice.copy_from_slice(inv.transpose().as_slice());
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn np_linalg_pinv(mat1: *mut InputMatrix, out: *mut InputMatrix) {
    let mat1 = mat1.as_mut().unwrap();
    let out = out.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "np_linalg_pinv", err_msg);
    }
    let dim1 = (*mat1).get_dims();
    let outdim = out.get_dims();
    let out_slice = unsafe { slice::from_raw_parts_mut(out.data, outdim[0] * outdim[1]) };
    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };

    let matrix = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let svd = matrix.svd(true, true);
    let inv = svd.pseudo_inverse(1e-15);

    match inv {
        Ok(m) => {
            out_slice.copy_from_slice(m.transpose().as_slice());
        }
        Err(e) => {
            raise_exn!("LinAlgError", "np_linalg_pinv", e);
        }
    }
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn sp_linalg_lu(
    mat1: *mut InputMatrix,
    out_l: *mut InputMatrix,
    out_u: *mut InputMatrix,
) {
    let mat1 = mat1.as_mut().unwrap();
    let out_l = out_l.as_mut().unwrap();
    let out_u = out_u.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "sp_linalg_lu", err_msg);
    }

    let dim1 = (*mat1).get_dims();
    let outl_dim = (*out_l).get_dims();
    let outu_dim = (*out_u).get_dims();

    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };
    let out_l_slice = unsafe { slice::from_raw_parts_mut(out_l.data, outl_dim[0] * outl_dim[1]) };
    let out_u_slice = unsafe { slice::from_raw_parts_mut(out_u.data, outu_dim[0] * outu_dim[1]) };

    let matrix = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let (_, l, u) = matrix.lu().unpack();

    out_l_slice.copy_from_slice(l.transpose().as_slice());
    out_u_slice.copy_from_slice(u.transpose().as_slice());
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn sp_linalg_schur(
    mat1: *mut InputMatrix,
    out_t: *mut InputMatrix,
    out_z: *mut InputMatrix,
) {
    let mat1 = mat1.as_mut().unwrap();
    let out_t = out_t.as_mut().unwrap();
    let out_z = out_z.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "sp_linalg_schur", err_msg);
    }

    let dim1 = (*mat1).get_dims();

    if dim1[0] != dim1[1] {
        let err_msg =
            format!("last 2 dimensions of the array must be square: {0} != {1}", dim1[0], dim1[1]);
        raise_exn!("LinAlgError", "np_linalg_schur", err_msg);
    }

    let out_t_dim = (*out_t).get_dims();
    let out_z_dim = (*out_z).get_dims();

    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };
    let out_t_slice = unsafe { slice::from_raw_parts_mut(out_t.data, out_t_dim[0] * out_t_dim[1]) };
    let out_z_slice = unsafe { slice::from_raw_parts_mut(out_z.data, out_z_dim[0] * out_z_dim[1]) };

    let matrix = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let (z, t) = matrix.schur().unpack();

    out_t_slice.copy_from_slice(t.transpose().as_slice());
    out_z_slice.copy_from_slice(z.transpose().as_slice());
}

/// # Safety
///
/// `mat1` should point to a valid 2DArray of `f64` floats in row-major order
#[no_mangle]
pub unsafe extern "C" fn sp_linalg_hessenberg(
    mat1: *mut InputMatrix,
    out_h: *mut InputMatrix,
    out_q: *mut InputMatrix,
) {
    let mat1 = mat1.as_mut().unwrap();
    let out_h = out_h.as_mut().unwrap();
    let out_q = out_q.as_mut().unwrap();

    if mat1.ndims != 2 {
        let err_msg = format!("expected 2D Vector Input, but received {}D input", mat1.ndims);
        raise_exn!("ValueError", "sp_linalg_hessenberg", err_msg);
    }

    let dim1 = (*mat1).get_dims();

    if dim1[0] != dim1[1] {
        let err_msg =
            format!("last 2 dimensions of the array must be square: {} != {}", dim1[0], dim1[1]);
        raise_exn!("LinAlgError", "sp_linalg_hessenberg", err_msg);
    }

    let out_h_dim = (*out_h).get_dims();
    let out_q_dim = (*out_q).get_dims();

    let data_slice1 = unsafe { slice::from_raw_parts_mut(mat1.data, dim1[0] * dim1[1]) };
    let out_h_slice = unsafe { slice::from_raw_parts_mut(out_h.data, out_h_dim[0] * out_h_dim[1]) };
    let out_q_slice = unsafe { slice::from_raw_parts_mut(out_q.data, out_q_dim[0] * out_q_dim[1]) };

    let matrix = DMatrix::from_row_slice(dim1[0], dim1[1], data_slice1);
    let (q, h) = matrix.hessenberg().unpack();

    out_h_slice.copy_from_slice(h.transpose().as_slice());
    out_q_slice.copy_from_slice(q.transpose().as_slice());
}