nu_command/matrix/
sum.rs

1use crate::matrix::MatrixValue;
2use ndarray::Axis;
3use nu_engine::command_prelude::*;
4
5#[derive(Clone)]
6pub struct MatrixSum;
7
8impl Command for MatrixSum {
9    fn name(&self) -> &str {
10        "matrix sum"
11    }
12
13    fn signature(&self) -> Signature {
14        Signature::build("matrix sum")
15            .input_output_types(vec![
16                (Type::Custom("matrix".into()), Type::Float),
17                (Type::Custom("matrix".into()), Type::Custom("matrix".into())),
18            ])
19            .named(
20                "axis",
21                SyntaxShape::Int,
22                "The axis to sum along (0-based).",
23                Some('a'),
24            )
25            .category(Category::Filters)
26    }
27
28    fn description(&self) -> &str {
29        "Sum all elements of a matrix, or sum along an axis."
30    }
31
32    fn search_terms(&self) -> Vec<&str> {
33        vec!["total", "add"]
34    }
35
36    fn run(
37        &self,
38        engine_state: &EngineState,
39        stack: &mut Stack,
40        call: &Call,
41        input: PipelineData,
42    ) -> Result<PipelineData, ShellError> {
43        let head = call.head;
44        let axis: Option<i64> = call.get_flag(engine_state, stack, "axis")?;
45        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
46
47        match axis {
48            Some(axis) => {
49                let axis = axis as usize;
50                if axis >= matrix.array.ndim() {
51                    return Err(ShellError::Generic(
52                        nu_protocol::shell_error::generic::GenericError::new(
53                            "Invalid axis",
54                            format!(
55                                "axis {} is out of bounds for a {}-dimensional array",
56                                axis,
57                                matrix.array.ndim()
58                            ),
59                            head,
60                        ),
61                    ));
62                }
63                let result = matrix.array.sum_axis(Axis(axis));
64                Ok(MatrixValue::new(result)
65                    .into_value(head)
66                    .into_pipeline_data())
67            }
68            None => {
69                let total: f64 = matrix.array.sum();
70                Ok(Value::float(total, head).into_pipeline_data())
71            }
72        }
73    }
74
75    fn examples(&self) -> Vec<Example<'static>> {
76        vec![
77            Example {
78                description: "Sum all elements of a 2x2 matrix",
79                example: "[[1 2] [3 4]] | into matrix | matrix sum",
80                result: Some(Value::test_float(10.0)),
81            },
82            Example {
83                description: "Sum along rows (axis 0)",
84                example: "[[1 2] [3 4]] | into matrix | matrix sum --axis 0 | matrix into-nu | to nuon",
85                result: Some(Value::test_string("[[4.0, 6.0]]")),
86            },
87            Example {
88                description: "Sum along columns (axis 1)",
89                example: "[[1 2] [3 4]] | into matrix | matrix sum --axis 1 | matrix into-nu | to nuon",
90                result: Some(Value::test_string("[[3.0, 7.0]]")),
91            },
92        ]
93    }
94}
95
96#[cfg(test)]
97mod test {
98    use super::*;
99
100    #[test]
101    fn test_examples() -> nu_test_support::Result {
102        nu_test_support::test().examples(MatrixSum)
103    }
104}
nu_command/matrix/sum.rs

nu_command/matrix/
sum.rs
