use crate::error::{DatarustError, Result};
use crate::matrix::{Matrix, StrMatrix};
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Output {
pub numeric: Matrix,
pub categorical: StrMatrix,
}
impl Output {
pub fn new(numeric: Matrix, categorical: StrMatrix) -> Result<Self> {
if numeric.nrows() != categorical.nrows() {
return Err(DatarustError::ShapeMismatch {
expected: format!("{} rows", numeric.nrows()),
actual: format!("{} rows", categorical.nrows()),
});
}
Ok(Self {
numeric,
categorical,
})
}
pub fn nrows(&self) -> usize {
self.numeric.nrows()
}
pub fn ncols(&self) -> usize {
self.numeric.ncols() + self.categorical.ncols()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reports_combined_shape() {
let output = Output::new(
Matrix::new(vec![vec![1.0, 2.0], vec![3.0, 4.0]]).unwrap(),
StrMatrix::new(vec![vec!["a".into()], vec!["b".into()]]).unwrap(),
)
.unwrap();
assert_eq!(output.nrows(), 2);
assert_eq!(output.ncols(), 3);
}
#[test]
fn rejects_mismatched_row_counts() {
let error = Output::new(
Matrix::new(vec![vec![1.0], vec![2.0]]).unwrap(),
StrMatrix::new(vec![vec!["a".into()]]).unwrap(),
)
.unwrap_err();
assert!(matches!(error, DatarustError::ShapeMismatch { .. }));
}
}