nu_command/matrix/

value.rs

use ndarray::ArrayD;
use nu_protocol::{
    CellPathMutation, CustomValue, ShellError, Span, Type, Value,
    ast::{Comparison, Math, Operator, PathMember},
    casing::Casing,
};
use serde::{Deserialize, Serialize};
use std::any::Any;
use std::cmp::Ordering;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MatrixValue {
    pub array: ArrayD<f64>,
}

#[typetag::serde]
impl CustomValue for MatrixValue {
    fn clone_value(&self, span: Span) -> Value {
        Value::custom(Box::new(self.clone()), span)
    }

    fn type_name(&self) -> String {
        "matrix".to_string()
    }

    fn to_base_value(&self, span: Span) -> Result<Value, ShellError> {
        Ok(ndarray_to_value(&self.array, span))
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_mut_any(&mut self) -> &mut dyn Any {
        self
    }

    fn partial_cmp(&self, other: &Value) -> Option<Ordering> {
        match other {
            Value::Custom { val, .. } if val.type_name() == self.type_name() => {
                let other_matrix = val.as_any().downcast_ref::<MatrixValue>()?;
                if self.array.shape() != other_matrix.array.shape() {
                    return None;
                }
                if ndarray::Zip::from(&self.array)
                    .and(&other_matrix.array)
                    .all(|a, b| a == b)
                {
                    Some(Ordering::Equal)
                } else if ndarray::Zip::from(&self.array)
                    .and(&other_matrix.array)
                    .all(|a, b| *a <= *b)
                {
                    Some(Ordering::Less)
                } else if ndarray::Zip::from(&self.array)
                    .and(&other_matrix.array)
                    .all(|a, b| *a >= *b)
                {
                    Some(Ordering::Greater)
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    fn follow_path_string(
        &self,
        self_span: Span,
        column_name: String,
        path_span: Span,
        _optional: bool,
        casing: Casing,
    ) -> Result<Value, ShellError> {
        let col = match casing {
            Casing::Sensitive => column_name,
            Casing::Insensitive => column_name.to_lowercase(),
        };

        match col.as_str() {
            "shape" => Ok(Value::list(
                self.array
                    .shape()
                    .iter()
                    .map(|d| Value::int(*d as i64, path_span))
                    .collect(),
                path_span,
            )),
            "ndim" => Ok(Value::int(self.array.ndim() as i64, path_span)),
            "size" => Ok(Value::int(self.array.len() as i64, path_span)),
            _ => Err(ShellError::CantFindColumn {
                col_name: col,
                span: Some(path_span),
                src_span: self_span,
            }),
        }
    }

    fn follow_path_int(
        &self,
        self_span: Span,
        index: usize,
        path_span: Span,
        _optional: bool,
    ) -> Result<Value, ShellError> {
        if self.array.ndim() == 0 {
            return Err(ShellError::IncompatiblePathAccess {
                type_name: self.type_name(),
                span: path_span,
            });
        }
        if index >= self.array.shape()[0] {
            return Err(ShellError::AccessBeyondEnd {
                max_idx: self.array.shape()[0] - 1,
                span: self_span,
            });
        }
        let subview = self.array.index_axis(ndarray::Axis(0), index);
        if subview.ndim() == 0 {
            Ok(Value::float(*subview.first().unwrap_or(&0.0), path_span))
        } else {
            Ok(ndarray_to_value(&subview.to_owned(), path_span))
        }
    }

    fn update_data_at_cell_path(
        &self,
        cell_path: &[PathMember],
        new_val: Value,
        action: &CellPathMutation,
        head: Span,
    ) -> Result<Value, ShellError> {
        let mut base = self.to_base_value(head)?;
        base.mutate_data_at_cell_path(cell_path, new_val, action)?;
        match base {
            Value::List { vals, .. } => {
                MatrixValue::from_list_of_lists(&vals, head).map(|m| m.into_value(head))
            }
            other => Ok(other),
        }
    }

    fn operation(
        &self,
        lhs_span: Span,
        operator: Operator,
        op: Span,
        right: &Value,
    ) -> Result<Value, ShellError> {
        match operator {
            Operator::Math(Math::Add) => {
                matrix_math_op(self, right, operator, op, lhs_span, |a, b| a + b)
            }
            Operator::Math(Math::Subtract) => {
                matrix_math_op(self, right, operator, op, lhs_span, |a, b| a - b)
            }
            Operator::Math(Math::Multiply) => {
                matrix_math_op(self, right, operator, op, lhs_span, |a, b| a * b)
            }
            Operator::Math(Math::Divide) => {
                matrix_math_op(self, right, operator, op, lhs_span, |a, b| a / b)
            }
            Operator::Comparison(comparison @ Comparison::Equal)
            | Operator::Comparison(comparison @ Comparison::NotEqual)
            | Operator::Comparison(comparison @ Comparison::LessThan)
            | Operator::Comparison(comparison @ Comparison::GreaterThan)
            | Operator::Comparison(comparison @ Comparison::LessThanOrEqual)
            | Operator::Comparison(comparison @ Comparison::GreaterThanOrEqual) => {
                compare_matrix(self, right, op, lhs_span, comparison)
            }
            _ => Err(ShellError::OperatorUnsupportedType {
                op: operator,
                unsupported: Type::Custom(self.type_name().into()),
                op_span: op,
                unsupported_span: lhs_span,
                help: None,
            }),
        }
    }
}

impl MatrixValue {
    pub fn new(array: ArrayD<f64>) -> Self {
        Self { array }
    }

    pub fn into_value(self, span: Span) -> Value {
        Value::custom(Box::new(self), span)
    }

    pub fn from_value(value: &Value) -> Result<Self, ShellError> {
        let span = value.span();
        match value {
            Value::Custom { val, .. } => {
                val.as_any().downcast_ref::<Self>().cloned().ok_or_else(|| {
                    ShellError::CantConvert {
                        to_type: "matrix".into(),
                        from_type: val.type_name(),
                        span,
                        help: Some("expected a matrix value".into()),
                    }
                })
            }
            x => Err(ShellError::CantConvert {
                to_type: "matrix".into(),
                from_type: x.get_type().to_string(),
                span,
                help: None,
            }),
        }
    }

    pub fn from_shape_vec(
        shape: Vec<usize>,
        data: Vec<f64>,
        span: Span,
    ) -> Result<Self, ShellError> {
        ArrayD::from_shape_vec(shape, data)
            .map(Self::new)
            .map_err(|e| {
                ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
                    "Matrix shape error",
                    e.to_string(),
                    span,
                ))
            })
    }

    pub fn from_list_of_lists(values: &[Value], span: Span) -> Result<Self, ShellError> {
        let mut rows: Vec<Vec<f64>> = Vec::new();
        let mut ncols: Option<usize> = None;

        for (i, value) in values.iter().enumerate() {
            match value {
                Value::List { vals, .. } => {
                    let row: Result<Vec<f64>, ShellError> =
                        vals.iter().map(|v| value_to_f64(v, span)).collect();
                    let row = row?;
                    if let Some(expected) = ncols {
                        if row.len() != expected {
                            return Err(ShellError::Generic(
                                nu_protocol::shell_error::generic::GenericError::new(
                                    "Inconsistent row lengths",
                                    format!(
                                        "row {} has {} elements, expected {}",
                                        i,
                                        row.len(),
                                        expected
                                    ),
                                    span,
                                ),
                            ));
                        }
                    } else {
                        ncols = Some(row.len());
                    }
                    rows.push(row);
                }
                _ => {
                    return Err(ShellError::Generic(
                        nu_protocol::shell_error::generic::GenericError::new(
                            "Invalid matrix input",
                            format!("row {} is not a list", i),
                            span,
                        ),
                    ));
                }
            }
        }

        let nrows = rows.len();
        let ncols = ncols.unwrap_or(0);
        let flat: Vec<f64> = rows.into_iter().flatten().collect();

        Self::from_shape_vec(vec![nrows, ncols], flat, span)
    }

    pub fn from_list_of_records(values: &[Value], span: Span) -> Result<Self, ShellError> {
        if values.is_empty() {
            let array = ArrayD::from_shape_vec(vec![0, 0], vec![]).map_err(|e| {
                ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
                    "Matrix shape error",
                    e.to_string(),
                    span,
                ))
            })?;
            return Ok(Self::new(array));
        }

        let first_record = match &values[0] {
            Value::Record { val, .. } => val,
            _ => {
                return Err(ShellError::Generic(
                    nu_protocol::shell_error::generic::GenericError::new(
                        "Invalid matrix input",
                        "expected a list of records",
                        span,
                    ),
                ));
            }
        };

        let cols: Vec<String> = first_record.columns().cloned().collect();
        let ncols = cols.len();
        let nrows = values.len();

        let mut data = Vec::with_capacity(nrows * ncols);

        for (i, value) in values.iter().enumerate() {
            match value {
                Value::Record { val, .. } => {
                    for col in &cols {
                        let element = val.get(col).ok_or_else(|| {
                            ShellError::Generic(
                                nu_protocol::shell_error::generic::GenericError::new(
                                    "Missing column",
                                    format!("row {} is missing column '{}'", i, col),
                                    span,
                                ),
                            )
                        })?;
                        data.push(value_to_f64(element, span)?);
                    }
                }
                _ => {
                    return Err(ShellError::Generic(
                        nu_protocol::shell_error::generic::GenericError::new(
                            "Invalid matrix input",
                            format!("row {} is not a record", i),
                            span,
                        ),
                    ));
                }
            }
        }

        Self::from_shape_vec(vec![nrows, ncols], data, span)
    }

    /// Apply an element-wise binary operation with another matrix or scalar.
    ///
    /// When `broadcast` is true, `other` may be broadcast to this matrix's shape.
    pub fn elementwise_binary<F, G>(
        self,
        other: Value,
        broadcast: bool,
        head: Span,
        f_matrix: F,
        f_scalar: G,
    ) -> Result<ArrayD<f64>, ShellError>
    where
        F: FnOnce(ArrayD<f64>, ArrayD<f64>) -> ArrayD<f64>,
        G: FnOnce(ArrayD<f64>, f64) -> ArrayD<f64>,
    {
        match other {
            Value::Int { val, .. } => Ok(f_scalar(self.array, val as f64)),
            Value::Float { val, .. } => Ok(f_scalar(self.array, val)),
            Value::Custom { .. } => {
                let other_matrix = MatrixValue::from_value(&other)?;
                if broadcast {
                    let target_shape = self.array.shape().to_vec();
                    let other_view = other_matrix
                        .array
                        .broadcast(target_shape.as_slice())
                        .ok_or_else(|| {
                            ShellError::Generic(
                                nu_protocol::shell_error::generic::GenericError::new(
                                    "Broadcast error",
                                    "shapes are not compatible for broadcasting",
                                    head,
                                ),
                            )
                        })?;
                    Ok(f_matrix(self.array, other_view.to_owned().into_dyn()))
                } else if self.array.shape() == other_matrix.array.shape() {
                    Ok(f_matrix(self.array, other_matrix.array))
                } else {
                    Err(ShellError::Generic(
                        nu_protocol::shell_error::generic::GenericError::new(
                            "Shape mismatch",
                            format!(
                                "shapes do not match: {:?} vs {:?}. Use --broadcast to enable broadcasting.",
                                self.array.shape(),
                                other_matrix.array.shape()
                            ),
                            head,
                        ),
                    ))
                }
            }
            _ => Err(ShellError::Generic(
                nu_protocol::shell_error::generic::GenericError::new(
                    "Invalid argument",
                    "expected a matrix, int, or float",
                    head,
                ),
            )),
        }
    }

    pub fn test_value(rows: &[&[f64]]) -> Value {
        let nrows = rows.len();
        let ncols = if nrows > 0 { rows[0].len() } else { 0 };
        let flat: Vec<f64> = rows.iter().flat_map(|r| r.iter()).copied().collect();
        let array = ArrayD::from_shape_vec(vec![nrows, ncols], flat)
            .expect("test value shape must be valid");
        Value::test_custom_value(Box::new(Self { array }))
    }
}

/// Convert a numeric Value to f64.
pub(crate) fn value_to_f64(value: &Value, span: Span) -> Result<f64, ShellError> {
    match value {
        Value::Int { val, .. } => Ok(*val as f64),
        Value::Float { val, .. } => Ok(*val),
        Value::String { val, .. } => val.parse::<f64>().map_err(|_| ShellError::CantConvert {
            to_type: "float".into(),
            from_type: "string".into(),
            span,
            help: None,
        }),
        _ => Err(ShellError::CantConvert {
            to_type: "float".into(),
            from_type: value.get_type().to_string(),
            span,
            help: None,
        }),
    }
}

/// Convert a list of numeric Values to `Vec<f64>`.
pub(crate) fn values_to_f64s(vals: &[Value], span: Span) -> Result<Vec<f64>, ShellError> {
    vals.iter().map(|v| value_to_f64(v, span)).collect()
}

/// Parse positive dimensions from i64 values, rejecting zero and negatives.
pub(crate) fn positive_dim(dim: i64, span: Span) -> Result<usize, ShellError> {
    if dim > 0 {
        Ok(dim as usize)
    } else {
        Err(ShellError::Generic(
            nu_protocol::shell_error::generic::GenericError::new(
                "Invalid dimensions",
                "dimensions must be positive integers",
                span,
            ),
        ))
    }
}

fn ndarray_to_value(array: &ArrayD<f64>, span: Span) -> Value {
    if array.ndim() == 0 {
        return Value::float(array.first().copied().unwrap_or(0.0), span);
    }

    if array.ndim() == 1 {
        let list: Vec<Value> = array.iter().map(|v| Value::float(*v, span)).collect();
        return Value::list(list, span);
    }

    if array.ndim() == 2 {
        let rows: Vec<Value> = array
            .axis_iter(ndarray::Axis(0))
            .map(|row| {
                let vals: Vec<Value> = row.iter().map(|v| Value::float(*v, span)).collect();
                Value::list(vals, span)
            })
            .collect();
        return Value::list(rows, span);
    }

    let sub_results: Vec<Value> = array
        .axis_iter(ndarray::Axis(0))
        .map(|sub| ndarray_to_value(&sub.to_owned(), span))
        .collect();
    Value::list(sub_results, span)
}

fn matrix_math_op<F>(
    left: &MatrixValue,
    right: &Value,
    operator: Operator,
    op_span: Span,
    lhs_span: Span,
    f: F,
) -> Result<Value, ShellError>
where
    F: Fn(f64, f64) -> f64,
{
    match right {
        Value::Int { val, .. } => {
            let result = left.array.map(|v| f(*v, *val as f64));
            Ok(MatrixValue::new(result).into_value(op_span))
        }
        Value::Float { val, .. } => {
            let result = left.array.map(|v| f(*v, *val));
            Ok(MatrixValue::new(result).into_value(op_span))
        }
        Value::Custom { val, .. } => {
            let other = val.as_any().downcast_ref::<MatrixValue>().ok_or_else(|| {
                ShellError::OperatorIncompatibleTypes {
                    op: operator,
                    lhs: Type::Custom("matrix".into()),
                    rhs: Type::Custom(val.type_name().into()),
                    op_span,
                    lhs_span,
                    rhs_span: right.span(),
                    help: None,
                }
            })?;
            if left.array.shape() != other.array.shape() {
                return Err(ShellError::OperatorIncompatibleTypes {
                    op: operator,
                    lhs: Type::Custom("matrix".into()),
                    rhs: Type::Custom("matrix".into()),
                    op_span,
                    lhs_span,
                    rhs_span: right.span(),
                    help: Some("shapes do not match"),
                });
            }
            let shape: Vec<usize> = left.array.shape().to_vec();
            let mut result = ArrayD::zeros(shape);
            ndarray::Zip::from(&mut result)
                .and(&left.array)
                .and(&other.array)
                .for_each(|r, &a, &b| *r = f(a, b));
            Ok(MatrixValue::new(result).into_value(op_span))
        }
        _ => Err(ShellError::OperatorIncompatibleTypes {
            op: operator,
            lhs: Type::Custom("matrix".into()),
            rhs: right.get_type(),
            op_span,
            lhs_span,
            rhs_span: right.span(),
            help: Some("expected a matrix or scalar"),
        }),
    }
}

fn compare_matrix(
    left: &MatrixValue,
    right: &Value,
    op_span: Span,
    lhs_span: Span,
    comparison: Comparison,
) -> Result<Value, ShellError> {
    let op = Operator::Comparison(comparison);

    match right {
        Value::Custom { val, .. } if val.type_name() == "matrix" => {
            let other = val.as_any().downcast_ref::<MatrixValue>().ok_or_else(|| {
                ShellError::OperatorIncompatibleTypes {
                    op,
                    lhs: Type::Custom("matrix".into()),
                    rhs: Type::Custom(val.type_name().into()),
                    op_span,
                    lhs_span,
                    rhs_span: right.span(),
                    help: None,
                }
            })?;

            if left.array.shape() != other.array.shape() {
                return shape_mismatch_comparison(comparison, op, op_span, lhs_span, right.span());
            }

            let all_match = match comparison {
                Comparison::Equal => ndarray::Zip::from(&left.array)
                    .and(&other.array)
                    .all(|a, b| a == b),
                Comparison::NotEqual => !ndarray::Zip::from(&left.array)
                    .and(&other.array)
                    .all(|a, b| a == b),
                Comparison::LessThan => ndarray::Zip::from(&left.array)
                    .and(&other.array)
                    .all(|a, b| a < b),
                Comparison::GreaterThan => ndarray::Zip::from(&left.array)
                    .and(&other.array)
                    .all(|a, b| a > b),
                Comparison::LessThanOrEqual => ndarray::Zip::from(&left.array)
                    .and(&other.array)
                    .all(|a, b| a <= b),
                Comparison::GreaterThanOrEqual => ndarray::Zip::from(&left.array)
                    .and(&other.array)
                    .all(|a, b| a >= b),
                _ => {
                    return Err(ShellError::OperatorUnsupportedType {
                        op,
                        unsupported: Type::Custom("matrix".into()),
                        op_span,
                        unsupported_span: lhs_span,
                        help: None,
                    });
                }
            };

            Ok(Value::bool(all_match, op_span))
        }
        Value::Int { val, .. } => {
            let s = *val as f64;
            Ok(Value::bool(
                compare_all_to_scalar(&left.array, s, comparison, op, op_span, lhs_span)?,
                op_span,
            ))
        }
        Value::Float { val, .. } => Ok(Value::bool(
            compare_all_to_scalar(&left.array, *val, comparison, op, op_span, lhs_span)?,
            op_span,
        )),
        _ => Err(ShellError::OperatorIncompatibleTypes {
            op,
            lhs: Type::Custom("matrix".into()),
            rhs: right.get_type(),
            op_span,
            lhs_span,
            rhs_span: right.span(),
            help: Some("expected a matrix or numeric scalar"),
        }),
    }
}

/// Equality on mismatched shapes is false / not-equal is true; ordering requires matching shapes.
fn shape_mismatch_comparison(
    comparison: Comparison,
    op: Operator,
    op_span: Span,
    lhs_span: Span,
    rhs_span: Span,
) -> Result<Value, ShellError> {
    match comparison {
        Comparison::Equal => Ok(Value::bool(false, op_span)),
        Comparison::NotEqual => Ok(Value::bool(true, op_span)),
        Comparison::LessThan
        | Comparison::GreaterThan
        | Comparison::LessThanOrEqual
        | Comparison::GreaterThanOrEqual => Err(ShellError::OperatorIncompatibleTypes {
            op,
            lhs: Type::Custom("matrix".into()),
            rhs: Type::Custom("matrix".into()),
            op_span,
            lhs_span,
            rhs_span,
            help: Some("cannot compare matrices with different shapes"),
        }),
        _ => Err(ShellError::OperatorUnsupportedType {
            op,
            unsupported: Type::Custom("matrix".into()),
            op_span,
            unsupported_span: lhs_span,
            help: None,
        }),
    }
}

fn compare_all_to_scalar(
    array: &ArrayD<f64>,
    scalar: f64,
    comparison: Comparison,
    op: Operator,
    op_span: Span,
    lhs_span: Span,
) -> Result<bool, ShellError> {
    match comparison {
        Comparison::Equal => Ok(array.iter().all(|v| (*v - scalar).abs() < f64::EPSILON)),
        Comparison::NotEqual => Ok(array.iter().any(|v| (*v - scalar).abs() >= f64::EPSILON)),
        Comparison::LessThan => Ok(array.iter().all(|&v| v < scalar)),
        Comparison::GreaterThan => Ok(array.iter().all(|&v| v > scalar)),
        Comparison::LessThanOrEqual => Ok(array.iter().all(|&v| v <= scalar)),
        Comparison::GreaterThanOrEqual => Ok(array.iter().all(|&v| v >= scalar)),
        _ => Err(ShellError::OperatorUnsupportedType {
            op,
            unsupported: Type::Custom("matrix".into()),
            op_span,
            unsupported_span: lhs_span,
            help: None,
        }),
    }
}