Skip to main content

nu_command/matrix/
add.rs

1use crate::matrix::MatrixValue;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MatrixAdd;
6
7impl Command for MatrixAdd {
8    fn name(&self) -> &str {
9        "matrix add"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("matrix add")
14            .input_output_types(vec![(
15                Type::Custom("matrix".into()),
16                Type::Custom("matrix".into()),
17            )])
18            .required(
19                "other",
20                SyntaxShape::Any,
21                "The other matrix or scalar to add.",
22            )
23            .switch(
24                "broadcast",
25                "Enable broadcasting to allow compatible shapes",
26                Some('b'),
27            )
28            .category(Category::Filters)
29    }
30
31    fn description(&self) -> &str {
32        "Add a matrix or scalar to a matrix."
33    }
34
35    fn search_terms(&self) -> Vec<&str> {
36        vec!["plus", "sum"]
37    }
38
39    fn run(
40        &self,
41        engine_state: &EngineState,
42        stack: &mut Stack,
43        call: &Call,
44        input: PipelineData,
45    ) -> Result<PipelineData, ShellError> {
46        let head = call.head;
47        let other: Value = call.req(engine_state, stack, 0)?;
48        let broadcast = call.has_flag(engine_state, stack, "broadcast")?;
49        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
50
51        let result =
52            matrix.elementwise_binary(other, broadcast, head, |a, b| a + b, |a, s| a + s)?;
53
54        Ok(MatrixValue::new(result)
55            .into_value(head)
56            .into_pipeline_data())
57    }
58
59    fn examples(&self) -> Vec<Example<'static>> {
60        vec![
61            Example {
62                description: "Add a scalar to a matrix",
63                example: "matrix zeros 2 2 | matrix add 5 | matrix into-nu | to nuon",
64                result: Some(Value::test_string("[[5.0, 5.0], [5.0, 5.0]]")),
65            },
66            Example {
67                description: "Add two matrices element-wise",
68                example: "matrix identity 2 | matrix add (matrix identity 2) | matrix into-nu | to nuon",
69                result: Some(Value::test_string("[[2.0, 0.0], [0.0, 2.0]]")),
70            },
71            Example {
72                description: "Add with broadcasting a row vector",
73                example: "matrix zeros 2 3 | matrix add --broadcast ([[1.0 2.0 3.0]] | into matrix) | matrix into-nu | to nuon",
74                result: Some(Value::test_string("[[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]")),
75            },
76        ]
77    }
78}
79
80#[cfg(test)]
81mod test {
82    use super::*;
83
84    #[test]
85    fn test_examples() -> nu_test_support::Result {
86        nu_test_support::test().examples(MatrixAdd)
87    }
88}