#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
use oxiblas_core::memory::AlignedVec;
use oxiblas_core::scalar::Scalar;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TriangularKind {
Upper,
Lower,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OutOfTriangleError {
pub row: usize,
pub col: usize,
pub dim: usize,
pub kind: TriangularKind,
}
impl core::fmt::Display for OutOfTriangleError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(
f,
"index ({}, {}) is outside the stored {:?} triangle of a {}x{} packed matrix",
self.row, self.col, self.kind, self.dim, self.dim
)
}
}
#[cfg(feature = "std")]
impl std::error::Error for OutOfTriangleError {}
#[derive(Clone)]
pub struct PackedMat<T: Scalar> {
data: AlignedVec<T>,
n: usize,
kind: TriangularKind,
}
impl<T: Scalar> PackedMat<T> {
pub fn zeros(n: usize, kind: TriangularKind) -> Self
where
T: bytemuck::Zeroable,
{
let len = Self::packed_len(n);
PackedMat {
data: AlignedVec::zeros(len),
n,
kind,
}
}
pub fn filled(n: usize, kind: TriangularKind, value: T) -> Self {
let len = Self::packed_len(n);
PackedMat {
data: AlignedVec::filled(len, value),
n,
kind,
}
}
pub fn from_slice(n: usize, kind: TriangularKind, data: &[T]) -> Self {
let len = Self::packed_len(n);
assert_eq!(
data.len(),
len,
"Slice length must equal n*(n+1)/2 = {}",
len
);
PackedMat {
data: AlignedVec::from_slice(data),
n,
kind,
}
}
#[inline]
pub const fn packed_len(n: usize) -> usize {
match n.checked_add(1) {
Some(np1) => match n.checked_mul(np1) {
Some(product) => product / 2,
None => panic!(
"PackedMat: packed length overflow (n*(n+1) exceeds usize::MAX); \
requested matrix dimension is too large to allocate"
),
},
None => panic!(
"PackedMat: packed length overflow (n+1 exceeds usize::MAX); \
requested matrix dimension is too large to allocate"
),
}
}
#[inline]
pub fn dim(&self) -> usize {
self.n
}
#[inline]
pub fn kind(&self) -> TriangularKind {
self.kind
}
#[inline]
pub fn len(&self) -> usize {
self.data.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.n == 0
}
#[inline]
pub fn packed_index(&self, row: usize, col: usize) -> Option<usize> {
if row >= self.n || col >= self.n {
return None;
}
match self.kind {
TriangularKind::Upper => {
if row <= col {
Some(col * (col + 1) / 2 + row)
} else {
None
}
}
TriangularKind::Lower => {
if row >= col {
let offset = self.n * col - col * (col.saturating_sub(1)) / 2;
Some(offset + (row - col))
} else {
None
}
}
}
}
#[inline]
pub fn get(&self, row: usize, col: usize) -> Option<&T> {
self.packed_index(row, col).map(|idx| &self.data[idx])
}
#[inline]
pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
self.packed_index(row, col).map(|idx| &mut self.data[idx])
}
#[inline]
pub fn set(&mut self, row: usize, col: usize, value: T) -> Result<(), OutOfTriangleError> {
let idx = self.packed_index(row, col).ok_or(OutOfTriangleError {
row,
col,
dim: self.n,
kind: self.kind,
})?;
self.data[idx] = value;
Ok(())
}
#[inline]
pub fn as_ptr(&self) -> *const T {
self.data.as_ptr()
}
#[inline]
pub fn as_mut_ptr(&mut self) -> *mut T {
self.data.as_mut_ptr()
}
#[inline]
pub fn as_slice(&self) -> &[T] {
self.data.as_slice()
}
#[inline]
pub fn as_slice_mut(&mut self) -> &mut [T] {
self.data.as_mut_slice()
}
pub fn to_dense(&self) -> crate::Mat<T>
where
T: bytemuck::Zeroable,
{
let mut mat = crate::Mat::zeros(self.n, self.n);
for j in 0..self.n {
for i in 0..self.n {
if let Some(idx) = self.packed_index(i, j) {
mat[(i, j)] = self.data[idx];
}
}
}
mat
}
pub fn from_dense(mat: &crate::MatRef<'_, T>, kind: TriangularKind) -> Self
where
T: bytemuck::Zeroable,
{
assert_eq!(mat.nrows(), mat.ncols(), "Matrix must be square");
let n = mat.nrows();
let mut packed = Self::zeros(n, kind);
for j in 0..n {
for i in 0..n {
if let Some(idx) = packed.packed_index(i, j) {
packed.data[idx] = mat[(i, j)];
}
}
}
packed
}
#[inline]
fn diagonal_packed_index(&self, i: usize) -> usize {
match self.kind {
TriangularKind::Upper => i * (i + 1) / 2 + i,
TriangularKind::Lower => self.n * i - i * (i.saturating_sub(1)) / 2,
}
}
pub fn diagonal(&self) -> Vec<T> {
(0..self.n)
.map(|i| self.data[self.diagonal_packed_index(i)])
.collect()
}
pub fn set_diagonal(&mut self, diag: &[T]) {
assert_eq!(
diag.len(),
self.n,
"Diagonal length must match matrix dimension"
);
for (i, &val) in diag.iter().enumerate() {
let idx = self.diagonal_packed_index(i);
self.data[idx] = val;
}
}
pub fn fill(&mut self, value: T) {
for elem in self.data.as_mut_slice() {
*elem = value;
}
}
pub fn scale(&mut self, alpha: T) {
for elem in self.data.as_mut_slice() {
*elem *= alpha;
}
}
pub fn transpose(&self) -> Self
where
T: bytemuck::Zeroable,
{
let new_kind = match self.kind {
TriangularKind::Upper => TriangularKind::Lower,
TriangularKind::Lower => TriangularKind::Upper,
};
let mut result = Self::zeros(self.n, new_kind);
for j in 0..self.n {
for i in 0..self.n {
if let Some(src_idx) = self.packed_index(i, j) {
if let Some(dst_idx) = result.packed_index(j, i) {
result.data[dst_idx] = self.data[src_idx];
}
}
}
}
result
}
}
impl<T: Scalar + core::fmt::Debug> core::fmt::Debug for PackedMat<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
writeln!(f, "PackedMat {}x{} {:?} {{", self.n, self.n, self.kind)?;
for i in 0..self.n.min(8) {
write!(f, " [")?;
for j in 0..self.n.min(8) {
if j > 0 {
write!(f, ", ")?;
}
match self.get(i, j) {
Some(v) => write!(f, "{:8.4?}", v)?,
None => write!(f, " * ")?,
}
}
if self.n > 8 {
write!(f, ", ...")?;
}
writeln!(f, "]")?;
}
if self.n > 8 {
writeln!(f, " ...")?;
}
write!(f, "}}")
}
}
#[derive(Clone, Copy)]
pub struct PackedRef<'a, T: Scalar> {
ptr: *const T,
n: usize,
kind: TriangularKind,
_marker: core::marker::PhantomData<&'a T>,
}
impl<'a, T: Scalar> PackedRef<'a, T> {
#[inline]
pub unsafe fn new(ptr: *const T, n: usize, kind: TriangularKind) -> Self {
PackedRef {
ptr,
n,
kind,
_marker: core::marker::PhantomData,
}
}
#[inline]
pub fn from_slice(data: &'a [T], n: usize, kind: TriangularKind) -> Self {
let expected_len = PackedMat::<T>::packed_len(n);
assert_eq!(
data.len(),
expected_len,
"Slice length must equal n*(n+1)/2"
);
unsafe { PackedRef::new(data.as_ptr(), n, kind) }
}
#[inline]
pub fn dim(&self) -> usize {
self.n
}
#[inline]
pub fn kind(&self) -> TriangularKind {
self.kind
}
#[inline]
pub fn packed_index(&self, row: usize, col: usize) -> Option<usize> {
if row >= self.n || col >= self.n {
return None;
}
match self.kind {
TriangularKind::Upper => {
if row <= col {
Some(col * (col + 1) / 2 + row)
} else {
None
}
}
TriangularKind::Lower => {
if row >= col {
let offset = self.n * col - col * (col.saturating_sub(1)) / 2;
Some(offset + (row - col))
} else {
None
}
}
}
}
#[inline]
pub fn get(&self, row: usize, col: usize) -> Option<&T> {
self.packed_index(row, col)
.map(|idx| unsafe { &*self.ptr.add(idx) })
}
#[inline]
pub fn as_ptr(&self) -> *const T {
self.ptr
}
}
unsafe impl<'a, T: Scalar + Send> Send for PackedRef<'a, T> {}
unsafe impl<'a, T: Scalar + Sync> Sync for PackedRef<'a, T> {}
pub struct PackedMut<'a, T: Scalar> {
ptr: *mut T,
n: usize,
kind: TriangularKind,
_marker: core::marker::PhantomData<&'a mut T>,
}
impl<'a, T: Scalar> PackedMut<'a, T> {
#[inline]
pub unsafe fn new(ptr: *mut T, n: usize, kind: TriangularKind) -> Self {
PackedMut {
ptr,
n,
kind,
_marker: core::marker::PhantomData,
}
}
#[inline]
pub fn from_slice(data: &'a mut [T], n: usize, kind: TriangularKind) -> Self {
let expected_len = PackedMat::<T>::packed_len(n);
assert_eq!(
data.len(),
expected_len,
"Slice length must equal n*(n+1)/2"
);
unsafe { PackedMut::new(data.as_mut_ptr(), n, kind) }
}
#[inline]
pub fn dim(&self) -> usize {
self.n
}
#[inline]
pub fn kind(&self) -> TriangularKind {
self.kind
}
#[inline]
pub fn packed_index(&self, row: usize, col: usize) -> Option<usize> {
if row >= self.n || col >= self.n {
return None;
}
match self.kind {
TriangularKind::Upper => {
if row <= col {
Some(col * (col + 1) / 2 + row)
} else {
None
}
}
TriangularKind::Lower => {
if row >= col {
let offset = self.n * col - col * (col.saturating_sub(1)) / 2;
Some(offset + (row - col))
} else {
None
}
}
}
}
#[inline]
pub fn get(&self, row: usize, col: usize) -> Option<&T> {
self.packed_index(row, col)
.map(|idx| unsafe { &*self.ptr.add(idx) })
}
#[inline]
pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
self.packed_index(row, col)
.map(|idx| unsafe { &mut *self.ptr.add(idx) })
}
#[inline]
pub fn set(&mut self, row: usize, col: usize, value: T) -> Result<(), OutOfTriangleError> {
let idx = self.packed_index(row, col).ok_or(OutOfTriangleError {
row,
col,
dim: self.n,
kind: self.kind,
})?;
unsafe {
*self.ptr.add(idx) = value;
}
Ok(())
}
#[inline]
pub fn as_ptr(&self) -> *const T {
self.ptr
}
#[inline]
pub fn as_mut_ptr(&mut self) -> *mut T {
self.ptr
}
#[inline]
pub fn rb(&self) -> PackedRef<'_, T> {
unsafe { PackedRef::new(self.ptr, self.n, self.kind) }
}
#[inline]
pub fn rb_mut(&mut self) -> PackedMut<'_, T> {
unsafe { PackedMut::new(self.ptr, self.n, self.kind) }
}
}
unsafe impl<'a, T: Scalar + Send> Send for PackedMut<'a, T> {}
unsafe impl<'a, T: Scalar + Sync> Sync for PackedMut<'a, T> {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_packed_upper_indexing() {
let mut p: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Upper);
assert_eq!(p.packed_index(0, 0), Some(0));
assert_eq!(p.packed_index(0, 1), Some(1));
assert_eq!(p.packed_index(1, 1), Some(2));
assert_eq!(p.packed_index(0, 2), Some(3));
assert_eq!(p.packed_index(1, 2), Some(4));
assert_eq!(p.packed_index(2, 2), Some(5));
assert_eq!(p.packed_index(1, 0), None);
assert_eq!(p.packed_index(2, 0), None);
assert_eq!(p.packed_index(2, 1), None);
p.set(0, 0, 1.0).unwrap();
p.set(0, 1, 2.0).unwrap();
p.set(1, 1, 3.0).unwrap();
p.set(0, 2, 4.0).unwrap();
p.set(1, 2, 5.0).unwrap();
p.set(2, 2, 6.0).unwrap();
assert_eq!(p.get(0, 0), Some(&1.0));
assert_eq!(p.get(0, 1), Some(&2.0));
assert_eq!(p.get(1, 1), Some(&3.0));
assert_eq!(p.get(0, 2), Some(&4.0));
assert_eq!(p.get(1, 2), Some(&5.0));
assert_eq!(p.get(2, 2), Some(&6.0));
}
#[test]
fn test_packed_lower_indexing() {
let mut p: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Lower);
assert_eq!(p.packed_index(0, 0), Some(0));
assert_eq!(p.packed_index(1, 0), Some(1));
assert_eq!(p.packed_index(2, 0), Some(2));
assert_eq!(p.packed_index(1, 1), Some(3));
assert_eq!(p.packed_index(2, 1), Some(4));
assert_eq!(p.packed_index(2, 2), Some(5));
assert_eq!(p.packed_index(0, 1), None);
assert_eq!(p.packed_index(0, 2), None);
assert_eq!(p.packed_index(1, 2), None);
p.set(0, 0, 1.0).unwrap();
p.set(1, 0, 2.0).unwrap();
p.set(2, 0, 3.0).unwrap();
p.set(1, 1, 4.0).unwrap();
p.set(2, 1, 5.0).unwrap();
p.set(2, 2, 6.0).unwrap();
assert_eq!(p.get(0, 0), Some(&1.0));
assert_eq!(p.get(1, 0), Some(&2.0));
assert_eq!(p.get(2, 0), Some(&3.0));
assert_eq!(p.get(1, 1), Some(&4.0));
assert_eq!(p.get(2, 1), Some(&5.0));
assert_eq!(p.get(2, 2), Some(&6.0));
}
#[test]
fn test_packed_len() {
assert_eq!(PackedMat::<f64>::packed_len(0), 0);
assert_eq!(PackedMat::<f64>::packed_len(1), 1);
assert_eq!(PackedMat::<f64>::packed_len(2), 3);
assert_eq!(PackedMat::<f64>::packed_len(3), 6);
assert_eq!(PackedMat::<f64>::packed_len(4), 10);
assert_eq!(PackedMat::<f64>::packed_len(10), 55);
}
#[test]
fn test_packed_to_dense() {
let mut p: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Upper);
p.set(0, 0, 1.0).unwrap();
p.set(0, 1, 2.0).unwrap();
p.set(1, 1, 3.0).unwrap();
p.set(0, 2, 4.0).unwrap();
p.set(1, 2, 5.0).unwrap();
p.set(2, 2, 6.0).unwrap();
let dense = p.to_dense();
assert_eq!(dense[(0, 0)], 1.0);
assert_eq!(dense[(0, 1)], 2.0);
assert_eq!(dense[(1, 1)], 3.0);
assert_eq!(dense[(0, 2)], 4.0);
assert_eq!(dense[(1, 2)], 5.0);
assert_eq!(dense[(2, 2)], 6.0);
assert_eq!(dense[(1, 0)], 0.0);
assert_eq!(dense[(2, 0)], 0.0);
assert_eq!(dense[(2, 1)], 0.0);
}
#[test]
fn test_packed_from_dense() {
use crate::Mat;
let dense = Mat::from_rows(&[&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0], &[7.0, 8.0, 9.0]]);
let upper = PackedMat::from_dense(&dense.as_ref(), TriangularKind::Upper);
assert_eq!(upper.get(0, 0), Some(&1.0));
assert_eq!(upper.get(0, 1), Some(&2.0));
assert_eq!(upper.get(0, 2), Some(&3.0));
assert_eq!(upper.get(1, 1), Some(&5.0));
assert_eq!(upper.get(1, 2), Some(&6.0));
assert_eq!(upper.get(2, 2), Some(&9.0));
let lower = PackedMat::from_dense(&dense.as_ref(), TriangularKind::Lower);
assert_eq!(lower.get(0, 0), Some(&1.0));
assert_eq!(lower.get(1, 0), Some(&4.0));
assert_eq!(lower.get(2, 0), Some(&7.0));
assert_eq!(lower.get(1, 1), Some(&5.0));
assert_eq!(lower.get(2, 1), Some(&8.0));
assert_eq!(lower.get(2, 2), Some(&9.0));
}
#[test]
fn test_packed_diagonal() {
let mut p: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Upper);
p.set(0, 0, 1.0).unwrap();
p.set(0, 1, 10.0).unwrap();
p.set(1, 1, 2.0).unwrap();
p.set(0, 2, 20.0).unwrap();
p.set(1, 2, 30.0).unwrap();
p.set(2, 2, 3.0).unwrap();
let diag = p.diagonal();
assert_eq!(diag, vec![1.0, 2.0, 3.0]);
p.set_diagonal(&[10.0, 20.0, 30.0]);
let diag2 = p.diagonal();
assert_eq!(diag2, vec![10.0, 20.0, 30.0]);
}
#[test]
fn test_packed_transpose() {
let mut upper: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Upper);
upper.set(0, 0, 1.0).unwrap();
upper.set(0, 1, 2.0).unwrap();
upper.set(1, 1, 3.0).unwrap();
upper.set(0, 2, 4.0).unwrap();
upper.set(1, 2, 5.0).unwrap();
upper.set(2, 2, 6.0).unwrap();
let lower = upper.transpose();
assert_eq!(lower.kind(), TriangularKind::Lower);
assert_eq!(lower.get(0, 0), Some(&1.0));
assert_eq!(lower.get(1, 0), Some(&2.0)); assert_eq!(lower.get(1, 1), Some(&3.0));
assert_eq!(lower.get(2, 0), Some(&4.0)); assert_eq!(lower.get(2, 1), Some(&5.0)); assert_eq!(lower.get(2, 2), Some(&6.0));
}
#[test]
fn test_packed_ref() {
let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let pref = PackedRef::from_slice(&data, 3, TriangularKind::Upper);
assert_eq!(pref.dim(), 3);
assert_eq!(pref.get(0, 0), Some(&1.0));
assert_eq!(pref.get(0, 1), Some(&2.0));
assert_eq!(pref.get(1, 1), Some(&3.0));
assert_eq!(pref.get(0, 2), Some(&4.0));
assert_eq!(pref.get(1, 2), Some(&5.0));
assert_eq!(pref.get(2, 2), Some(&6.0));
}
#[test]
fn test_packed_mut() {
let mut data = [0.0f64; 6];
let mut pmut = PackedMut::from_slice(&mut data, 3, TriangularKind::Lower);
pmut.set(0, 0, 1.0).unwrap();
pmut.set(1, 0, 2.0).unwrap();
pmut.set(2, 0, 3.0).unwrap();
pmut.set(1, 1, 4.0).unwrap();
pmut.set(2, 1, 5.0).unwrap();
pmut.set(2, 2, 6.0).unwrap();
assert_eq!(data, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn test_packed_scale() {
let mut p: PackedMat<f64> = PackedMat::zeros(2, TriangularKind::Upper);
p.set(0, 0, 1.0).unwrap();
p.set(0, 1, 2.0).unwrap();
p.set(1, 1, 3.0).unwrap();
p.scale(2.0);
assert_eq!(p.get(0, 0), Some(&2.0));
assert_eq!(p.get(0, 1), Some(&4.0));
assert_eq!(p.get(1, 1), Some(&6.0));
}
#[test]
fn test_packed_mat_set_out_of_triangle_returns_error_not_panic() {
let mut upper: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Upper);
let err = upper
.set(1, 0, 99.0)
.expect_err("(1, 0) is below the diagonal of upper-triangular storage");
assert_eq!(
err,
OutOfTriangleError {
row: 1,
col: 0,
dim: 3,
kind: TriangularKind::Upper,
}
);
let message = err.to_string();
assert!(message.contains("(1, 0)"), "message was: {message}");
let mut lower: PackedMat<f64> = PackedMat::zeros(3, TriangularKind::Lower);
assert!(lower.set(0, 2, 1.0).is_err());
assert!(upper.set(5, 5, 1.0).is_err());
assert_eq!(upper.get(0, 0), Some(&0.0));
}
#[test]
fn test_packed_mut_set_out_of_triangle_returns_error_not_panic() {
let mut data = [0.0f64; 6];
let mut pmut = PackedMut::from_slice(&mut data, 3, TriangularKind::Lower);
let err = pmut
.set(0, 1, 42.0)
.expect_err("(0, 1) is above the diagonal of lower-triangular storage");
assert_eq!(err.row, 0);
assert_eq!(err.col, 1);
assert_eq!(err.dim, 3);
assert_eq!(err.kind, TriangularKind::Lower);
pmut.set(0, 0, 7.0).unwrap();
assert_eq!(data[0], 7.0);
}
#[test]
#[should_panic(expected = "packed length overflow")]
fn test_packed_len_overflow_panics_not_wraps() {
let _ = PackedMat::<f64>::packed_len(1usize << 32);
}
#[test]
#[should_panic(expected = "packed length overflow")]
fn test_packed_zeros_overflow_panics_not_wraps() {
let _: PackedMat<f64> = PackedMat::zeros(1usize << 32, TriangularKind::Upper);
}
#[test]
fn test_packed_len_is_still_exact_for_sane_dims() {
assert_eq!(PackedMat::<f64>::packed_len(0), 0);
assert_eq!(PackedMat::<f64>::packed_len(1), 1);
assert_eq!(PackedMat::<f64>::packed_len(3), 6);
assert_eq!(PackedMat::<f64>::packed_len(100), 5050);
}
#[test]
#[should_panic(expected = "Slice length must equal")]
fn test_packed_ref_from_slice_rejects_short_slice() {
let data = [0.0f64; 10];
let _ = PackedRef::from_slice(&data, 1000, TriangularKind::Upper);
}
#[test]
#[should_panic(expected = "Slice length must equal")]
fn test_packed_mut_from_slice_rejects_short_slice() {
let mut data = [0.0f64; 10];
let _ = PackedMut::from_slice(&mut data, 1000, TriangularKind::Upper);
}
}