1use crate::matrix::MatrixValue;
2use ndarray::ArrayD;
3use nu_engine::command_prelude::*;
4
5#[derive(Clone)]
6pub struct MatrixMultiply;
7
8impl Command for MatrixMultiply {
9 fn name(&self) -> &str {
10 "matrix multiply"
11 }
12
13 fn signature(&self) -> Signature {
14 Signature::build("matrix multiply")
15 .input_output_types(vec![(
16 Type::Custom("matrix".into()),
17 Type::Custom("matrix".into()),
18 )])
19 .required(
20 "other",
21 SyntaxShape::Any,
22 "The other matrix to multiply with.",
23 )
24 .switch("swap", "Swap the left and right operands.", Some('s'))
25 .switch(
26 "multall",
27 "Multiply the input matrix with all matrices given as arguments, chaining them left to right (e.g., `$a | matrix multiply --multall $b $c` multiplies all three together).",
28 Some('a'),
29 )
30 .rest(
31 "rest",
32 SyntaxShape::Any,
33 "Additional matrices for chained multiplication (requires --multall).",
34 )
35 .category(Category::Filters)
36 }
37
38 fn description(&self) -> &str {
39 "Multiply two matrices using dot product."
40 }
41
42 fn search_terms(&self) -> Vec<&str> {
43 vec!["dot", "matmul", "product", "chain"]
44 }
45
46 fn run(
47 &self,
48 engine_state: &EngineState,
49 stack: &mut Stack,
50 call: &Call,
51 input: PipelineData,
52 ) -> Result<PipelineData, ShellError> {
53 let head = call.head;
54 let swap = call.has_flag(engine_state, stack, "swap")?;
55 let multall = call.has_flag(engine_state, stack, "multall")?;
56 let other_val: Value = call.req(engine_state, stack, 0)?;
57
58 let mut a = MatrixValue::from_value(&input.into_value(head)?)?;
59 let mut b = MatrixValue::from_value(&other_val)?;
60
61 if swap {
62 std::mem::swap(&mut a, &mut b);
63 }
64
65 let rest: Vec<Value> = call.rest(engine_state, stack, 1)?;
66
67 let mut result = multiply_arrays(a.array, b.array, head)?;
68
69 if multall {
70 for val in rest {
71 let mat = MatrixValue::from_value(&val)?;
72 result = multiply_arrays(result, mat.array, head)?;
73 }
74 }
75
76 if result.ndim() == 0 {
77 Ok(Value::float(result.first().copied().unwrap_or(0.0), head).into_pipeline_data())
78 } else {
79 Ok(MatrixValue::new(result)
80 .into_value(head)
81 .into_pipeline_data())
82 }
83 }
84
85 fn examples(&self) -> Vec<Example<'static>> {
86 vec![
87 Example {
88 description: "Multiply two 2x2 matrices",
89 example: "[[1 2] [3 4]] | into matrix | matrix multiply ([[1 0] [0 1]] | into matrix) | matrix into-nu | to nuon",
90 result: Some(Value::test_string("[[1.0, 2.0], [3.0, 4.0]]")),
91 },
92 Example {
93 description: "Multiply a matrix by its inverse gives identity",
94 example: "matrix identity 2 | matrix multiply (matrix identity 2) | matrix into-nu | to nuon",
95 result: Some(Value::test_string("[[1.0, 0.0], [0.0, 1.0]]")),
96 },
97 Example {
98 description: "Swap operands with --swap",
99 example: "[[1 2] [3 4]] | into matrix | matrix multiply --swap ([[0 1] [1 0]] | into matrix) | matrix into-nu | to nuon",
100 result: Some(Value::test_string("[[3.0, 4.0], [1.0, 2.0]]")),
101 },
102 Example {
103 description: "Chain-multiply three matrices with --multall",
104 example: "matrix identity 2 | matrix multiply --multall ([[2 0] [0 2]] | into matrix) ([[3 0] [0 3]] | into matrix) | matrix into-nu | to nuon",
105 result: Some(Value::test_string("[[6.0, 0.0], [0.0, 6.0]]")),
106 },
107 ]
108 }
109}
110
111fn multiply_arrays(a: ArrayD<f64>, b: ArrayD<f64>, head: Span) -> Result<ArrayD<f64>, ShellError> {
112 match (a.ndim(), b.ndim()) {
113 (1, 1) => {
114 if a.len() != b.len() {
115 return Err(ShellError::Generic(
116 nu_protocol::shell_error::generic::GenericError::new(
117 "Shape mismatch",
118 format!(
119 "vectors must have the same length: {} vs {}",
120 a.len(),
121 b.len()
122 ),
123 head,
124 ),
125 ));
126 }
127 let dot: f64 = ndarray::Zip::from(&a)
128 .and(&b)
129 .fold(0.0, |acc, &x, &y| acc + x * y);
130 Ok(ArrayD::from_shape_vec(vec![], vec![dot]).map_err(|e| {
131 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
132 "Shape error",
133 e.to_string(),
134 head,
135 ))
136 })?)
137 }
138 (2, 1) => {
139 let a_view = a
140 .view()
141 .into_dimensionality::<ndarray::Ix2>()
142 .map_err(|e| {
143 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
144 "Dimension error",
145 e.to_string(),
146 head,
147 ))
148 })?;
149 let b_view = b
150 .view()
151 .into_dimensionality::<ndarray::Ix1>()
152 .map_err(|_| {
153 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
154 "Dimension error",
155 "expected a 1D vector",
156 head,
157 ))
158 })?;
159 if a_view.shape()[1] != b_view.shape()[0] {
160 return Err(ShellError::Generic(
161 nu_protocol::shell_error::generic::GenericError::new(
162 "Shape mismatch",
163 format!(
164 "inner dimensions do not match: ({} x {}) dot {}",
165 a_view.shape()[0],
166 a_view.shape()[1],
167 b_view.shape()[0],
168 ),
169 head,
170 ),
171 ));
172 }
173 let result = a_view.dot(&b_view);
174 Ok(result.into_dyn())
175 }
176 (1, 2) => {
177 let a_view = a
178 .view()
179 .into_dimensionality::<ndarray::Ix1>()
180 .map_err(|_| {
181 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
182 "Dimension error",
183 "expected a 1D vector",
184 head,
185 ))
186 })?;
187 let b_view = b
188 .view()
189 .into_dimensionality::<ndarray::Ix2>()
190 .map_err(|e| {
191 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
192 "Dimension error",
193 e.to_string(),
194 head,
195 ))
196 })?;
197 if a_view.shape()[0] != b_view.shape()[0] {
198 return Err(ShellError::Generic(
199 nu_protocol::shell_error::generic::GenericError::new(
200 "Shape mismatch",
201 format!(
202 "inner dimensions do not match: {} dot {}",
203 a_view.shape()[0],
204 b_view.shape()[0],
205 ),
206 head,
207 ),
208 ));
209 }
210 let result = a_view.dot(&b_view);
211 Ok(result.into_dyn())
212 }
213 (2, 2) => {
214 let a_view = a
215 .view()
216 .into_dimensionality::<ndarray::Ix2>()
217 .map_err(|e| {
218 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
219 "Dimension error",
220 e.to_string(),
221 head,
222 ))
223 })?;
224 let b_view = b
225 .view()
226 .into_dimensionality::<ndarray::Ix2>()
227 .map_err(|e| {
228 ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
229 "Dimension error",
230 e.to_string(),
231 head,
232 ))
233 })?;
234 if a_view.shape()[1] != b_view.shape()[0] {
235 return Err(ShellError::Generic(
236 nu_protocol::shell_error::generic::GenericError::new(
237 "Shape mismatch",
238 format!(
239 "inner dimensions do not match: ({} x {}) dot ({} x {})",
240 a_view.shape()[0],
241 a_view.shape()[1],
242 b_view.shape()[0],
243 b_view.shape()[1],
244 ),
245 head,
246 ),
247 ));
248 }
249 let result = a_view.dot(&b_view);
250 Ok(result.into_dyn())
251 }
252 _ => Err(ShellError::Generic(
253 nu_protocol::shell_error::generic::GenericError::new(
254 "Unsupported dimensions",
255 format!(
256 "matrix multiply only supports 1D and 2D arrays, got shapes {:?} and {:?}",
257 a.shape(),
258 b.shape()
259 ),
260 head,
261 ),
262 )),
263 }
264}
265
266#[cfg(test)]
267mod test {
268 use super::*;
269
270 #[test]
271 fn test_examples() -> nu_test_support::Result {
272 nu_test_support::test().examples(MatrixMultiply)
273 }
274
275 #[test]
276 fn test_incompatible_matrix_dimensions_error() {
277 let head = Span::test_data();
278 let a = ArrayD::from_shape_vec(vec![2, 3], vec![1.0; 6]).unwrap();
280 let b = ArrayD::from_shape_vec(vec![4, 2], vec![1.0; 8]).unwrap();
281 let result = multiply_arrays(a, b, head);
282 assert!(result.is_err());
283 let err = result.unwrap_err();
284 let msg = format!("{err}");
285 assert!(
286 msg.contains("Shape mismatch"),
287 "expected shape mismatch error, got: {msg}"
288 );
289 }
290
291 #[test]
292 fn test_incompatible_matrix_vector_dimensions_error() {
293 let head = Span::test_data();
294 let a = ArrayD::from_shape_vec(vec![2, 3], vec![1.0; 6]).unwrap();
296 let b = ArrayD::from_shape_vec(vec![4], vec![1.0; 4]).unwrap();
297 let result = multiply_arrays(a, b, head);
298 assert!(result.is_err());
299 let err = result.unwrap_err();
300 let msg = format!("{err}");
301 assert!(
302 msg.contains("Shape mismatch"),
303 "expected shape mismatch error, got: {msg}"
304 );
305 }
306
307 #[test]
308 fn test_incompatible_vector_matrix_dimensions_error() {
309 let head = Span::test_data();
310 let a = ArrayD::from_shape_vec(vec![3], vec![1.0; 3]).unwrap();
312 let b = ArrayD::from_shape_vec(vec![4, 2], vec![1.0; 8]).unwrap();
313 let result = multiply_arrays(a, b, head);
314 assert!(result.is_err());
315 let err = result.unwrap_err();
316 let msg = format!("{err}");
317 assert!(
318 msg.contains("Shape mismatch"),
319 "expected shape mismatch error, got: {msg}"
320 );
321 }
322}