herculesabqp 0.1.2

use ndarray::linalg::general_mat_vec_mul;
use ndarray::{Array1, Array2};
use sprs::prod::mul_acc_mat_vec_csr;
use sprs::{CsMat, TriMat};

#[derive(Clone, Debug)]
pub enum QuadraticMatrix {
    /// Dense quadratic matrix stored in row-major ndarray form.
    Dense(Array2<f64>),
    /// Sparse quadratic matrix stored with `sprs`.
    Sparse(CsMat<f64>),
}

impl QuadraticMatrix {
    /// Wrap a dense quadratic matrix.
    pub fn dense(matrix: Array2<f64>) -> Self {
        Self::Dense(matrix)
    }

    /// Wrap a sparse quadratic matrix stored in CSR/CSC-compatible form.
    pub fn sparse(matrix: CsMat<f64>) -> Self {
        Self::Sparse(matrix)
    }

    /// Return the matrix dimension for square quadratic forms.
    pub fn n(&self) -> usize {
        match self {
            Self::Dense(m) => m.nrows(),
            Self::Sparse(m) => m.rows(),
        }
    }

    /// Return the `(rows, cols)` shape of the matrix.
    pub fn shape(&self) -> (usize, usize) {
        match self {
            Self::Dense(m) => (m.nrows(), m.ncols()),
            Self::Sparse(m) => (m.rows(), m.cols()),
        }
    }

    /// Count structurally nonzero entries.
    pub fn nnz(&self) -> usize {
        match self {
            Self::Dense(m) => m.iter().filter(|v| v.abs() > 0.0).count(),
            Self::Sparse(m) => m.nnz(),
        }
    }

    /// Estimate density as `nnz / n^2`.
    pub fn density(&self) -> f64 {
        let n = self.n();
        self.nnz() as f64 / (n * n) as f64
    }

    /// Extract the matrix diagonal as a dense vector.
    pub fn diagonal(&self) -> Array1<f64> {
        match self {
            Self::Dense(m) => Array1::from_iter((0..m.nrows()).map(|i| m[[i, i]])),
            Self::Sparse(m) => {
                Array1::from_iter((0..m.rows()).map(|i| *m.get(i, i).unwrap_or(&0.0)))
            }
        }
    }

    /// Multiply the matrix by a dense vector and return the dense result.
    pub fn matvec(&self, x: &Array1<f64>) -> Array1<f64> {
        let mut out = Array1::<f64>::zeros(self.shape().0);
        self.matvec_into(x, &mut out);
        out
    }

    /// Multiply the matrix by a dense vector and write into `out`.
    pub fn matvec_into(&self, x: &Array1<f64>, out: &mut Array1<f64>) {
        match self {
            Self::Dense(m) => {
                debug_assert_eq!(x.len(), m.ncols());
                debug_assert_eq!(out.len(), m.nrows());
                general_mat_vec_mul(1.0, m, x, 0.0, out);
            }
            Self::Sparse(m) => {
                debug_assert_eq!(x.len(), m.cols());
                debug_assert_eq!(out.len(), m.rows());
                out.fill(0.0);
                mul_acc_mat_vec_csr(
                    m.view(),
                    x.as_slice().expect("dense vector should be contiguous"),
                    out.as_slice_mut()
                        .expect("dense vector should be contiguous"),
                );
            }
        }
    }

    /// Return the symmetric part `0.5 * (Q + Q^T)`.
    pub fn symmetrized(&self) -> Self {
        match self {
            Self::Dense(m) => Self::Dense((m + &m.t()).mapv(|x| 0.5 * x)),
            Self::Sparse(m) => {
                let sym = m + &m.transpose_view().to_owned();
                Self::Sparse(sym.map(|x| x * 0.5))
            }
        }
    }

    /// Materialize the matrix as a dense array.
    pub fn to_dense(&self) -> Array2<f64> {
        match self {
            Self::Dense(m) => m.clone(),
            Self::Sparse(m) => {
                let mut out = Array2::<f64>::zeros((m.rows(), m.cols()));
                for (i, row) in m.outer_iterator().enumerate() {
                    for (j, v) in row.iter() {
                        out[[i, j]] = *v;
                    }
                }
                out
            }
        }
    }

    /// Extract a principal submatrix while preserving dense/sparse representation.
    pub fn principal(&self, rows: &[usize], cols: &[usize]) -> Self {
        match self {
            Self::Dense(m) => {
                let mut out = Array2::<f64>::zeros((rows.len(), cols.len()));
                for (ii, &i) in rows.iter().enumerate() {
                    for (jj, &j) in cols.iter().enumerate() {
                        out[[ii, jj]] = m[[i, j]];
                    }
                }
                Self::Dense(out)
            }
            Self::Sparse(m) => {
                let missing = usize::MAX;
                let mut col_pos = vec![missing; m.cols()];
                for (jj, &j) in cols.iter().enumerate() {
                    col_pos[j] = jj;
                }
                let mut tri = TriMat::<f64>::with_capacity((rows.len(), cols.len()), cols.len());
                for (ii, &i) in rows.iter().enumerate() {
                    if let Some(row) = m.outer_view(i) {
                        for (j, v) in row.iter() {
                            let jj = col_pos[j];
                            if jj != missing {
                                tri.add_triplet(ii, jj, *v);
                            }
                        }
                    }
                }
                Self::Sparse(tri.to_csr())
            }
        }
    }

    /// Apply diagonal variable scaling `Q -> D Q D`.
    pub fn scaled(&self, scale: &[f64]) -> Self {
        match self {
            Self::Dense(m) => {
                let mut out = m.clone();
                for i in 0..m.nrows() {
                    for j in 0..m.ncols() {
                        out[[i, j]] = scale[i] * m[[i, j]] * scale[j];
                    }
                }
                Self::Dense(out)
            }
            Self::Sparse(m) => {
                let mut tri = TriMat::<f64>::with_capacity((m.rows(), m.cols()), m.nnz());
                for (i, row) in m.outer_iterator().enumerate() {
                    for (j, v) in row.iter() {
                        tri.add_triplet(i, j, scale[i] * scale[j] * v);
                    }
                }
                Self::Sparse(tri.to_csr())
            }
        }
    }
}

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

    #[test]
    fn dense_and_sparse_matvec_match() {
        let dense =
            Array2::from_shape_vec((3, 3), vec![2.0, -1.0, 0.0, -1.0, 3.0, 4.0, 0.0, 4.0, 5.0])
                .unwrap();
        let sparse = TriMat::from_triplets(
            (3, 3),
            vec![0, 0, 1, 1, 1, 2, 2],
            vec![0, 1, 0, 1, 2, 1, 2],
            vec![2.0, -1.0, -1.0, 3.0, 4.0, 4.0, 5.0],
        )
        .to_csr();
        let x = Array1::from_vec(vec![1.5, -2.0, 0.25]);

        let y_dense = QuadraticMatrix::dense(dense).matvec(&x);
        let y_sparse = QuadraticMatrix::sparse(sparse).matvec(&x);

        assert!((&y_dense - &y_sparse).mapv(|v| v * v).sum().sqrt() < 1e-12);
    }

    #[test]
    fn symmetrized_dense_matches_average_with_transpose() {
        let original = Array2::from_shape_vec((2, 2), vec![1.0, 3.0, -1.0, 2.0]).unwrap();
        let expected = (&original + &original.t()).mapv(|x| 0.5 * x);
        let sym = QuadraticMatrix::dense(original).symmetrized().to_dense();
        assert!((sym - expected).mapv(|v| v * v).sum().sqrt() < 1e-12);
    }
}