nu_command/matrix/
set_col.rs

1use crate::matrix::MatrixValue;
2use crate::matrix::value::values_to_f64s;
3use nu_engine::command_prelude::*;
4
5#[derive(Clone)]
6pub struct MatrixSetCol;
7
8impl Command for MatrixSetCol {
9    fn name(&self) -> &str {
10        "matrix set-col"
11    }
12
13    fn signature(&self) -> Signature {
14        Signature::build("matrix set-col")
15            .input_output_types(vec![(
16                Type::Custom("matrix".into()),
17                Type::Custom("matrix".into()),
18            )])
19            .required(
20                "index",
21                SyntaxShape::Int,
22                "The column index to replace (0-based).",
23            )
24            .required(
25                "replacement",
26                SyntaxShape::List(Box::new(SyntaxShape::Number)),
27                "The new column values as a list of numbers.",
28            )
29            .category(Category::Filters)
30    }
31
32    fn description(&self) -> &str {
33        "Replace a column in a 2D matrix."
34    }
35
36    fn search_terms(&self) -> Vec<&str> {
37        vec!["column", "replace"]
38    }
39
40    fn run(
41        &self,
42        engine_state: &EngineState,
43        stack: &mut Stack,
44        call: &Call,
45        input: PipelineData,
46    ) -> Result<PipelineData, ShellError> {
47        let head = call.head;
48        let index: usize = call.req::<i64>(engine_state, stack, 0)? as usize;
49        let replacement: Value = call.req(engine_state, stack, 1)?;
50        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
51
52        if matrix.array.ndim() != 2 {
53            return Err(ShellError::Generic(
54                nu_protocol::shell_error::generic::GenericError::new(
55                    "Invalid dimensions",
56                    format!(
57                        "set-col requires a 2D matrix, got {} dimensions",
58                        matrix.array.ndim()
59                    ),
60                    head,
61                ),
62            ));
63        }
64
65        let replacement_vals = match replacement {
66            Value::List { vals, .. } => vals,
67            _ => {
68                return Err(ShellError::Generic(
69                    nu_protocol::shell_error::generic::GenericError::new(
70                        "Invalid replacement",
71                        "expected a list of numbers",
72                        head,
73                    ),
74                ));
75            }
76        };
77
78        let nrows = matrix.array.shape()[0];
79        let ncols = matrix.array.shape()[1];
80
81        if replacement_vals.len() != nrows {
82            return Err(ShellError::Generic(
83                nu_protocol::shell_error::generic::GenericError::new(
84                    "Size mismatch",
85                    format!(
86                        "replacement has {} elements, but column has {} rows",
87                        replacement_vals.len(),
88                        nrows
89                    ),
90                    head,
91                ),
92            ));
93        }
94
95        if index >= ncols {
96            return Err(ShellError::Generic(
97                nu_protocol::shell_error::generic::GenericError::new(
98                    "Index out of bounds",
99                    format!(
100                        "column index {} is out of bounds, matrix has {} columns",
101                        index, ncols
102                    ),
103                    head,
104                ),
105            ));
106        }
107
108        let floats = values_to_f64s(&replacement_vals, head)?;
109
110        let mut new_array = matrix.array;
111        for (i, &val) in floats.iter().enumerate() {
112            new_array[[i, index]] = val;
113        }
114
115        Ok(MatrixValue::new(new_array)
116            .into_value(head)
117            .into_pipeline_data())
118    }
119
120    fn examples(&self) -> Vec<Example<'static>> {
121        vec![
122            Example {
123                description: "Replace the first column of a 2x3 matrix",
124                example: "matrix zeros 2 3 | matrix set-col 0 [1.0 2.0] | matrix into-nu | to nuon",
125                result: Some(Value::test_string("[[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]")),
126            },
127            Example {
128                description: "Replace the second column of a 2x2 matrix",
129                example: "matrix zeros 2 2 | matrix set-col 1 [4.0 5.0] | matrix into-nu | to nuon",
130                result: Some(Value::test_string("[[0.0, 4.0], [0.0, 5.0]]")),
131            },
132        ]
133    }
134}
135
136#[cfg(test)]
137mod test {
138    use super::*;
139
140    #[test]
141    fn test_examples() -> nu_test_support::Result {
142        nu_test_support::test().examples(MatrixSetCol)
143    }
144}
nu_command/matrix/set_col.rs

nu_command/matrix/
set_col.rs
