use arrow_array::types::{Float32Type, Float64Type};
use arrow_array::{FixedSizeListArray, PrimitiveArray};
use arrow_schema::Field;
use super::{
ArrowInteropError, complex64_matrix_from_owned, complex64_matrix_view,
complex64_vector_from_owned, complex64_vector_view, fixed_size_list_from_owned,
fixed_size_list_view, primitive_array_from_owned, primitive_array_view,
};
pub fn decompose_f32(
matrix: &FixedSizeListArray,
) -> Result<crate::linalg::cholesky::NdarrayCholeskyResult<f32>, ArrowInteropError> {
let matrix_view = fixed_size_list_view::<Float32Type>(matrix)?;
Ok(crate::linalg::cholesky::decompose_view(&matrix_view)?)
}
pub fn decompose_f64(
matrix: &FixedSizeListArray,
) -> Result<crate::linalg::cholesky::NdarrayCholeskyResult<f64>, ArrowInteropError> {
let matrix_view = fixed_size_list_view::<Float64Type>(matrix)?;
Ok(crate::linalg::cholesky::decompose_view(&matrix_view)?)
}
pub fn solve_f32(
matrix: &FixedSizeListArray,
rhs: &PrimitiveArray<Float32Type>,
) -> Result<PrimitiveArray<Float32Type>, ArrowInteropError> {
let matrix_view = fixed_size_list_view::<Float32Type>(matrix)?;
let rhs_view = primitive_array_view(rhs)?;
let solution = crate::linalg::cholesky::solve_view(&matrix_view, &rhs_view)?;
Ok(primitive_array_from_owned::<Float32Type>(solution))
}
pub fn solve_f64(
matrix: &FixedSizeListArray,
rhs: &PrimitiveArray<Float64Type>,
) -> Result<PrimitiveArray<Float64Type>, ArrowInteropError> {
let matrix_view = fixed_size_list_view::<Float64Type>(matrix)?;
let rhs_view = primitive_array_view(rhs)?;
let solution = crate::linalg::cholesky::solve_view(&matrix_view, &rhs_view)?;
Ok(primitive_array_from_owned::<Float64Type>(solution))
}
pub fn inverse_f32(matrix: &FixedSizeListArray) -> Result<FixedSizeListArray, ArrowInteropError> {
let matrix_view = fixed_size_list_view::<Float32Type>(matrix)?;
let output = crate::linalg::cholesky::inverse_view(&matrix_view)?;
fixed_size_list_from_owned::<Float32Type>(output)
}
pub fn inverse_f64(matrix: &FixedSizeListArray) -> Result<FixedSizeListArray, ArrowInteropError> {
let matrix_view = fixed_size_list_view::<Float64Type>(matrix)?;
let output = crate::linalg::cholesky::inverse_view(&matrix_view)?;
fixed_size_list_from_owned::<Float64Type>(output)
}
pub fn decompose_complex(
matrix: &FixedSizeListArray,
) -> Result<crate::linalg::cholesky::NdarrayComplexCholeskyResult, ArrowInteropError> {
let matrix_view = complex64_matrix_view(matrix)?;
Ok(crate::linalg::cholesky::decompose_complex_view(&matrix_view)?)
}
pub fn solve_complex(
matrix: &FixedSizeListArray,
rhs_field: &Field,
rhs: &FixedSizeListArray,
) -> Result<(Field, FixedSizeListArray), ArrowInteropError> {
let matrix_view = complex64_matrix_view(matrix)?;
let rhs_view = complex64_vector_view(rhs_field, rhs)?;
let output = crate::linalg::cholesky::solve_complex_view(&matrix_view, &rhs_view)?;
complex64_vector_from_owned("cholesky_solve_complex", output)
}
pub fn inverse_complex(
matrix: &FixedSizeListArray,
) -> Result<FixedSizeListArray, ArrowInteropError> {
let matrix_view = complex64_matrix_view(matrix)?;
let output = crate::linalg::cholesky::inverse_complex_view(&matrix_view)?;
complex64_matrix_from_owned(output)
}