rhai_sci/matrices_and_arrays.rs
1use rhai::plugin::*;
2
3#[export_module]
4pub mod matrix_functions {
5 use crate::{
6 array_to_vec_float, if_int_convert_to_float_and_do, if_int_do_else_if_array_do, if_list_do,
7 if_matrix_convert_to_vec_array_and_do,
8 };
9 #[cfg(feature = "nalgebra")]
10 use crate::{
11 if_matrices_and_compatible_convert_to_vec_array_and_do, if_matrix_do,
12 omatrix_to_vec_dynamic, ovector_to_vec_dynamic, FOIL,
13 };
14 #[cfg(feature = "nalgebra")]
15 use nalgebralib::DMatrix;
16 use rhai::{Array, Dynamic, EvalAltResult, Map, Position, FLOAT, INT};
17 use std::collections::BTreeMap;
18
19 /// Calculates the inverse of a matrix. Fails if the matrix if not invertible, or if the
20 /// elements of the matrix aren't FLOAT or INT.
21 /// ```typescript
22 /// let x = [[ 1.0, 0.0, 2.0],
23 /// [-1.0, 5.0, 0.0],
24 /// [ 0.0, 3.0, -9.0]];
25 /// let x_inverted = inv(x);
26 /// assert_eq(x_inverted, [[0.8823529411764706, -0.11764705882352941, 0.19607843137254902],
27 /// [0.17647058823529413, 0.17647058823529413, 0.0392156862745098 ],
28 /// [0.058823529411764705, 0.058823529411764705, -0.09803921568627451]]
29 /// );
30 /// ```
31 /// ```typescript
32 /// let x = [[1, 2],
33 /// [3, 4]];
34 /// let x_inverted = inv(x);
35 /// assert_eq(x_inverted, [[-2.0, 1.0],
36 /// [1.5, -0.5]]
37 /// );
38 /// ```
39 #[cfg(feature = "nalgebra")]
40 #[rhai_fn(name = "inv", return_raw, pure)]
41 pub fn invert_matrix(matrix: &mut Array) -> Result<Array, Box<EvalAltResult>> {
42 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
43 let dm = DMatrix::from_fn(matrix_as_vec.len(), matrix_as_vec[0].len(), |i, j| {
44 if matrix_as_vec[0][0].is_float() {
45 matrix_as_vec[i][j].as_float().unwrap()
46 } else {
47 matrix_as_vec[i][j].as_int().unwrap() as FLOAT
48 }
49 });
50
51 // Try to invert
52 let dm = dm.try_inverse();
53
54 dm.map(omatrix_to_vec_dynamic).ok_or_else(|| {
55 EvalAltResult::ErrorArithmetic(
56 "Matrix cannot be inverted".to_string(),
57 Position::NONE,
58 )
59 .into()
60 })
61 })
62 }
63
64 /// Calculate the eigenvalues and eigenvectors for a matrix. Specifically, the output is an
65 /// object map with entries for real_eigenvalues, imaginary_eigenvalues, eigenvectors, and
66 /// residuals.
67 /// ```typescript
68 /// let matrix = eye(5);
69 /// let eig = eigs(matrix);
70 /// assert(sum(eig.residuals) < 0.000001);
71 /// ```
72 /// ```typescript
73 /// let matrix = [[ 0.0, 1.0],
74 /// [-2.0, -3.0]];
75 /// let eig = eigs(matrix);
76 /// assert(sum(eig.residuals) < 0.000001);
77 /// ```
78 #[cfg(feature = "nalgebra")]
79 #[rhai_fn(name = "eigs", return_raw, pure)]
80 pub fn matrix_eigs_alt(matrix: &mut Array) -> Result<Map, Box<EvalAltResult>> {
81 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
82 // Convert vec_array to omatrix
83 let dm = DMatrix::from_fn(matrix_as_vec.len(), matrix_as_vec[0].len(), |i, j| {
84 if matrix_as_vec[0][0].is_float() {
85 matrix_as_vec[i][j].as_float().unwrap()
86 } else {
87 matrix_as_vec[i][j].as_int().unwrap() as FLOAT
88 }
89 });
90
91 // Grab shape for later
92 let dms = dm.shape().1;
93
94 // Get teh eigenvalues
95 let eigenvalues = dm.complex_eigenvalues();
96
97 // Iterate through eigenvalues to get eigenvectors
98 let mut imaginary_values = vec![Dynamic::from_float(1.0); 0];
99 let mut real_values = vec![Dynamic::from_float(1.0); 0];
100 let mut residuals = vec![Dynamic::from_float(1.0); 0];
101 let mut eigenvectors = DMatrix::from_element(dms, 0, 0.0);
102 for (idx, ev) in eigenvalues.iter().enumerate() {
103 // Eigenvalue components
104 imaginary_values.push(Dynamic::from_float(ev.im));
105 real_values.push(Dynamic::from_float(ev.re));
106
107 // Get eigenvector
108 let mut A = dm.clone() - DMatrix::from_diagonal_element(dms, dms, ev.re);
109 A = A.insert_column(0, 0.0);
110 A = A.insert_row(0, 0.0);
111 A[(0, idx + 1)] = 1.0;
112 let mut b = DMatrix::from_element(dms + 1, 1, 0.0);
113 b[(0, 0)] = 1.0;
114 let eigenvector = A
115 .svd(true, true)
116 .solve(&b, 1e-10)
117 .unwrap()
118 .remove_rows(0, 1)
119 .normalize();
120
121 // Verify solution
122 residuals.push(Dynamic::from_float(
123 (dm.clone() * eigenvector.clone() - ev.re * eigenvector.clone()).amax(),
124 ));
125
126 eigenvectors.extend(eigenvector.column_iter());
127 }
128
129 let mut result = BTreeMap::new();
130 let mut vid = smartstring::SmartString::new();
131 vid.push_str("eigenvectors");
132 result.insert(
133 vid,
134 Dynamic::from_array(omatrix_to_vec_dynamic(eigenvectors)),
135 );
136 let mut did = smartstring::SmartString::new();
137 did.push_str("real_eigenvalues");
138 result.insert(did, Dynamic::from_array(real_values));
139 let mut eid = smartstring::SmartString::new();
140 eid.push_str("imaginary_eigenvalues");
141 result.insert(eid, Dynamic::from_array(imaginary_values));
142 let mut rid = smartstring::SmartString::new();
143 rid.push_str("residuals");
144 result.insert(rid, Dynamic::from_array(residuals));
145
146 Ok(result)
147 })
148 }
149
150 /// Calculates the singular value decomposition of a matrix
151 /// ```typescript
152 /// let matrix = eye(5);
153 /// let svd_results = svd(matrix);
154 /// assert_eq(svd_results, #{"s": ones([5]), "u": eye(5), "v": eye(5)});
155 /// ```
156 #[cfg(feature = "nalgebra")]
157 #[rhai_fn(name = "svd", return_raw, pure)]
158 pub fn svd_decomp(matrix: &mut Array) -> Result<Map, Box<EvalAltResult>> {
159 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
160 let dm = DMatrix::from_fn(matrix_as_vec.len(), matrix_as_vec[0].len(), |i, j| {
161 if matrix_as_vec[0][0].is::<FLOAT>() {
162 matrix_as_vec[i][j].as_float().unwrap()
163 } else {
164 matrix_as_vec[i][j].as_int().unwrap() as FLOAT
165 }
166 });
167
168 // Try ot invert
169 let svd = nalgebralib::linalg::SVD::new(dm, true, true);
170
171 let mut result = BTreeMap::new();
172 let mut uid = smartstring::SmartString::new();
173 uid.push_str("u");
174 match svd.u {
175 Some(u) => result.insert(uid, Dynamic::from_array(omatrix_to_vec_dynamic(u))),
176 None => {
177 return Err(EvalAltResult::ErrorArithmetic(
178 format!("SVD decomposition cannot be computed for this matrix."),
179 Position::NONE,
180 )
181 .into());
182 }
183 };
184
185 let mut vid = smartstring::SmartString::new();
186 vid.push_str("v");
187 match svd.v_t {
188 Some(v) => result.insert(vid, Dynamic::from_array(omatrix_to_vec_dynamic(v))),
189 None => {
190 return Err(EvalAltResult::ErrorArithmetic(
191 format!("SVD decomposition cannot be computed for this matrix."),
192 Position::NONE,
193 )
194 .into());
195 }
196 };
197
198 let mut sid = smartstring::SmartString::new();
199 sid.push_str("s");
200 result.insert(
201 sid,
202 Dynamic::from_array(ovector_to_vec_dynamic(svd.singular_values)),
203 );
204
205 Ok(result)
206 })
207 }
208
209 /// Calculates the QR decomposition of a matrix
210 /// ```typescript
211 /// let matrix = eye(5);
212 /// let qr_results = qr(matrix);
213 /// assert_eq(qr_results, #{"q": eye(5), "r": eye(5)});
214 /// ```
215 #[cfg(feature = "nalgebra")]
216 #[rhai_fn(name = "qr", return_raw, pure)]
217 pub fn qr_decomp(matrix: &mut Array) -> Result<Map, Box<EvalAltResult>> {
218 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
219 let dm = DMatrix::from_fn(matrix_as_vec.len(), matrix_as_vec[0].len(), |i, j| {
220 if matrix_as_vec[0][0].is::<FLOAT>() {
221 matrix_as_vec[i][j].as_float().unwrap()
222 } else {
223 matrix_as_vec[i][j].as_int().unwrap() as FLOAT
224 }
225 });
226
227 // Try ot invert
228 let qr = nalgebralib::linalg::QR::new(dm);
229
230 let mut result = BTreeMap::new();
231 let mut qid = smartstring::SmartString::new();
232 qid.push_str("q");
233 result.insert(qid, Dynamic::from_array(omatrix_to_vec_dynamic(qr.q())));
234
235 let mut rid = smartstring::SmartString::new();
236 rid.push_str("r");
237 result.insert(rid, Dynamic::from_array(omatrix_to_vec_dynamic(qr.r())));
238
239 Ok(result)
240 })
241 }
242
243 /// Calculates the QR decomposition of a matrix
244 /// ```typescript
245 /// let matrix = eye(5);
246 /// let h_results = hessenberg(matrix);
247 /// assert_eq(h_results, #{"h": eye(5), "q": eye(5)});
248 /// ```
249 #[cfg(feature = "nalgebra")]
250 #[rhai_fn(name = "hessenberg", return_raw, pure)]
251 pub fn hessenberg(matrix: &mut Array) -> Result<Map, Box<EvalAltResult>> {
252 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
253 let dm = DMatrix::from_fn(matrix_as_vec.len(), matrix_as_vec[0].len(), |i, j| {
254 if matrix_as_vec[0][0].is::<FLOAT>() {
255 matrix_as_vec[i][j].as_float().unwrap()
256 } else {
257 matrix_as_vec[i][j].as_int().unwrap() as FLOAT
258 }
259 });
260
261 // Try ot invert
262 let h = nalgebralib::linalg::Hessenberg::new(dm);
263
264 let mut result = BTreeMap::new();
265 let mut hid = smartstring::SmartString::new();
266 hid.push_str("h");
267 result.insert(hid, Dynamic::from_array(omatrix_to_vec_dynamic(h.h())));
268
269 let mut qid = smartstring::SmartString::new();
270 qid.push_str("q");
271 result.insert(qid, Dynamic::from_array(omatrix_to_vec_dynamic(h.q())));
272
273 Ok(result)
274 })
275 }
276
277 /// Transposes a matrix.
278 /// ```typescript
279 /// let row = [[1, 2, 3, 4]];
280 /// let column = transpose(row);
281 /// assert_eq(column, [[1],
282 /// [2],
283 /// [3],
284 /// [4]]);
285 /// ```
286 /// ```typescript
287 /// let matrix = transpose(eye(3));
288 /// assert_eq(matrix, eye(3));
289 /// ```
290 #[rhai_fn(name = "transpose", pure, return_raw)]
291 pub fn transpose(matrix: &mut Array) -> Result<Array, Box<EvalAltResult>> {
292 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
293 // Turn into Array
294 let mut out = vec![];
295 for idx in 0..matrix_as_vec[0].len() {
296 let mut new_row = vec![];
297 for jdx in 0..matrix_as_vec.len() {
298 new_row.push(matrix_as_vec[jdx][idx].clone());
299 }
300 out.push(Dynamic::from_array(new_row));
301 }
302 Ok(out)
303 })
304 }
305
306 /// Returns an array indicating the size of the matrix along each dimension, passed by reference.
307 /// ```typescript
308 /// let matrix = ones(3, 5);
309 /// assert_eq(size(matrix), [3, 5]);
310 /// ```
311 /// ```typescript
312 /// let matrix = [[[1, 2]]];
313 /// assert_eq(size(matrix), [1, 1, 2]);
314 /// ```
315 #[rhai_fn(name = "size", pure)]
316 pub fn matrix_size_by_reference(matrix: &mut Array) -> Array {
317 let mut new_matrix = matrix.clone();
318
319 let mut shape = vec![Dynamic::from_int(new_matrix.len() as INT)];
320 loop {
321 if new_matrix[0].is_array() {
322 new_matrix = new_matrix[0].clone().into_array().unwrap();
323 shape.push(Dynamic::from_int(new_matrix.len() as INT));
324 } else {
325 break;
326 }
327 }
328
329 shape
330 }
331
332 /// Return the number of dimensions in matrix, passed by reference.
333 /// ```typescript
334 /// let matrix = ones(4, 6);
335 /// let n = ndims(matrix);
336 /// assert_eq(n, 2);
337 /// ```
338 #[rhai_fn(name = "ndims")]
339 pub fn ndims_by_reference(matrix: &mut Array) -> INT {
340 matrix_size_by_reference(matrix).len() as INT
341 }
342
343 /// Returns the number of elements in a matrix, passed by reference.
344 /// ```typescript
345 /// let matrix = ones(4, 6);
346 /// let n = numel(matrix);
347 /// assert_eq(n, 24);
348 /// ```
349 /// ```typescript
350 /// let matrix = [1, [1, 2, 3]];
351 /// let n = numel(matrix);
352 /// assert_eq(n, 4);
353 /// ```
354 #[rhai_fn(name = "numel", pure)]
355 pub fn numel_by_reference(matrix: &mut Array) -> INT {
356 flatten(matrix).len() as INT
357 }
358
359 /// Returns the number of non-zero elements in a matrix, passed by reference.
360 /// ```typescript
361 /// let matrix = ones(4, 6);
362 /// let n = nnz(matrix);
363 /// assert_eq(n, 24);
364 /// ```
365 /// ```typescript
366 /// let matrix = eye(4);
367 /// let n = nnz(matrix);
368 /// assert_eq(n, 4);
369 /// ```
370 #[rhai_fn(name = "nnz", pure)]
371 pub fn nnz_by_reference(matrix: &mut Array) -> INT {
372 array_to_vec_float(&mut flatten(matrix))
373 .iter()
374 .filter(|&n| *n > 0.0)
375 .count() as INT
376 }
377
378 #[cfg(all(feature = "io"))]
379 pub mod read_write {
380 use polars::prelude::{CsvReadOptions, DataType, SerReader};
381 use rhai::{Array, Dynamic, EvalAltResult, ImmutableString, FLOAT};
382
383 /// Reads a numeric csv file from a url
384 /// ```typescript
385 /// let url = "https://raw.githubusercontent.com/plotly/datasets/master/diabetes.csv";
386 /// let x = read_matrix(url);
387 /// assert_eq(size(x), [768, 9]);
388 /// ```
389 #[rhai_fn(name = "read_matrix", return_raw)]
390 pub fn read_matrix(file_path: ImmutableString) -> Result<Array, Box<EvalAltResult>> {
391 // We will use this function later
392 fn transpose_internal<T>(v: Vec<Vec<T>>) -> Vec<Vec<T>> {
393 assert!(!v.is_empty());
394 let len = v[0].len();
395 let mut iters: Vec<_> = v.into_iter().map(|n| n.into_iter()).collect();
396 (0..len)
397 .map(|_| {
398 iters
399 .iter_mut()
400 .map(|n| n.next().unwrap())
401 .collect::<Vec<T>>()
402 })
403 .collect()
404 }
405
406 let file_path_as_str = file_path.as_str();
407
408 // Determine path is url
409 let path_is_url = url::Url::parse(file_path_as_str);
410
411 match path_is_url {
412 Err(_) => {
413 let x = CsvReadOptions::default()
414 .with_has_header(
415 csv_sniffer::Sniffer::new()
416 .sniff_path(file_path_as_str)
417 .map_err(|err| {
418 EvalAltResult::ErrorSystem(
419 format!("Cannot sniff file: {file_path_as_str}"),
420 err.into(),
421 )
422 })?
423 .dialect
424 .header
425 .has_header_row,
426 )
427 .try_into_reader_with_file_path(Some(file_path_as_str.into()))
428 .map_err(|err| {
429 EvalAltResult::ErrorSystem(
430 format!("Cannot read file as CSV: {file_path_as_str}"),
431 err.into(),
432 )
433 })?
434 .finish()
435 .map_err(|err| {
436 EvalAltResult::ErrorSystem(
437 format!("Cannot read file: {file_path_as_str}"),
438 err.into(),
439 )
440 })?
441 .drop_nulls::<String>(None)
442 .map_err(|err| {
443 EvalAltResult::ErrorSystem(
444 format!("Cannot remove null values from file: {file_path_as_str}"),
445 err.into(),
446 )
447 })?;
448
449 // Convert into vec of vec
450 let mut final_output = vec![];
451 for series in x.columns() {
452 let col: Vec<FLOAT> = series
453 .cast(&DataType::Float64)
454 .map_err(|err| {
455 EvalAltResult::ErrorArithmetic(
456 format!("Data cannot be cast to FLOAT: {err}"),
457 rhai::Position::NONE,
458 )
459 })?
460 .f64()
461 .unwrap()
462 .into_no_null_iter()
463 .map(|el| el as FLOAT)
464 .collect();
465 final_output.push(col);
466 }
467
468 final_output = transpose_internal(final_output);
469
470 let matrix_as_array = final_output
471 .into_iter()
472 .map(|x| {
473 let mut y = vec![];
474 for el in &x {
475 y.push(Dynamic::from_float(*el));
476 }
477 Dynamic::from_array(y)
478 })
479 .collect::<Array>();
480
481 Ok(matrix_as_array)
482 }
483 Ok(_) => {
484 if let Ok(_) = url::Url::parse(file_path_as_str) {
485 let file_contents =
486 minreq::get(file_path_as_str).send().map_err(|err| {
487 EvalAltResult::ErrorSystem(
488 format!("Error getting url: {file_path_as_str}"),
489 err.into(),
490 )
491 })?;
492 let temp = temp_file::with_contents(file_contents.as_bytes());
493
494 let temp_file_name: ImmutableString = temp.path().to_str().unwrap().into();
495
496 read_matrix(temp_file_name)
497 } else {
498 EvalAltResult::ErrorRuntime(
499 format!(
500 "The string {file_path_as_str} is not a valid URL or file path",
501 )
502 .into(),
503 rhai::Position::NONE,
504 )
505 .into()
506 }
507 }
508 }
509 }
510 }
511
512 /// Return a matrix of zeros. Can be called with a single integer argument (indicating the
513 /// square matrix of that size) or with an array argument (indicating the size for each dimension).
514 /// ```typescript
515 /// let matrix = zeros(3);
516 /// assert_eq(matrix, [[0.0, 0.0, 0.0],
517 /// [0.0, 0.0, 0.0],
518 /// [0.0, 0.0, 0.0]]);
519 /// ```
520 /// ```typescript
521 /// let matrix = zeros([3, 3]);
522 /// assert_eq(matrix, [[0.0, 0.0, 0.0],
523 /// [0.0, 0.0, 0.0],
524 /// [0.0, 0.0, 0.0]]);
525 /// ```
526 /// ```typescript
527 /// let matrix = zeros([3]);
528 /// assert_eq(matrix, [0.0, 0.0, 0.0]);
529 /// ```
530 /// ```typescript
531 /// let matrix = zeros([3, 3, 3]);
532 /// assert_eq(matrix, [[[0.0, 0.0, 0.0],
533 /// [0.0, 0.0, 0.0],
534 /// [0.0, 0.0, 0.0]],
535 /// [[0.0, 0.0, 0.0],
536 /// [0.0, 0.0, 0.0],
537 /// [0.0, 0.0, 0.0]],
538 /// [[0.0, 0.0, 0.0],
539 /// [0.0, 0.0, 0.0],
540 /// [0.0, 0.0, 0.0]]]);
541 /// ```
542 #[rhai_fn(name = "zeros", return_raw)]
543 pub fn zeros_single_input(n: Dynamic) -> Result<Array, Box<EvalAltResult>> {
544 if_int_do_else_if_array_do(
545 n,
546 |n| Ok(zeros_double_input(n, n)),
547 |m| {
548 if m.len() == 2 {
549 Ok(zeros_double_input(
550 m[0].as_int().unwrap(),
551 m[1].as_int().unwrap(),
552 ))
553 } else if m.len() > 2 {
554 let l = m.remove(0);
555 Ok(vec![
556 Dynamic::from_array(
557 zeros_single_input(Dynamic::from_array(m.to_vec())).unwrap()
558 );
559 l.as_int().unwrap() as usize
560 ])
561 } else {
562 Ok(vec![Dynamic::FLOAT_ZERO; m[0].as_int().unwrap() as usize])
563 }
564 },
565 )
566 }
567
568 /// Return a matrix of zeros. Arguments indicate the number of rows and columns in the matrix.
569 /// ```typescript
570 /// let matrix = zeros(3, 3);
571 /// assert_eq(matrix, [[0.0, 0.0, 0.0],
572 /// [0.0, 0.0, 0.0],
573 /// [0.0, 0.0, 0.0]]);
574 /// ```
575 #[rhai_fn(name = "zeros")]
576 pub fn zeros_double_input(nx: INT, ny: INT) -> Array {
577 let mut output = vec![];
578 for _ in 0..nx {
579 output.push(Dynamic::from_array(vec![Dynamic::FLOAT_ZERO; ny as usize]))
580 }
581 output
582 }
583
584 /// Return a matrix of ones. Can be called with a single integer argument (indicating the
585 /// square matrix of that size) or with an array argument (indicating the size for each dimension).
586 /// ```typescript
587 /// let matrix = ones(3);
588 /// assert_eq(matrix, [[1.0, 1.0, 1.0],
589 /// [1.0, 1.0, 1.0],
590 /// [1.0, 1.0, 1.0]]);
591 /// ```
592 /// ```typescript
593 /// let matrix = ones([3, 3]);
594 /// assert_eq(matrix, [[1.0, 1.0, 1.0],
595 /// [1.0, 1.0, 1.0],
596 /// [1.0, 1.0, 1.0]]);
597 /// ```
598 /// ```typescript
599 /// let matrix = ones([3]);
600 /// assert_eq(matrix, [1.0, 1.0, 1.0]);
601 /// ```
602 /// ```typescript
603 /// let matrix = ones([3, 3, 3]);
604 /// assert_eq(matrix, [[[1.0, 1.0, 1.0],
605 /// [1.0, 1.0, 1.0],
606 /// [1.0, 1.0, 1.0]],
607 /// [[1.0, 1.0, 1.0],
608 /// [1.0, 1.0, 1.0],
609 /// [1.0, 1.0, 1.0]],
610 /// [[1.0, 1.0, 1.0],
611 /// [1.0, 1.0, 1.0],
612 /// [1.0, 1.0, 1.0]]]);
613 /// ```
614 #[rhai_fn(name = "ones", return_raw)]
615 pub fn ones_single_input(n: Dynamic) -> Result<Array, Box<EvalAltResult>> {
616 crate::if_int_do_else_if_array_do(
617 n,
618 |n| Ok(ones_double_input(n, n)),
619 |m| {
620 if m.len() == 2 {
621 Ok(ones_double_input(
622 m[0].as_int().unwrap(),
623 m[1].as_int().unwrap(),
624 ))
625 } else if m.len() > 2 {
626 let l = m.remove(0);
627 Ok(vec![
628 Dynamic::from_array(
629 ones_single_input(Dynamic::from_array(m.to_vec())).unwrap()
630 );
631 l.as_int().unwrap() as usize
632 ])
633 } else {
634 Ok(vec![Dynamic::FLOAT_ONE; m[0].as_int().unwrap() as usize])
635 }
636 },
637 )
638 }
639
640 /// Return a matrix of ones. Arguments indicate the number of rows and columns in the matrix.
641 /// ```typescript
642 /// let matrix = ones(3, 3);
643 /// assert_eq(matrix, [[1.0, 1.0, 1.0],
644 /// [1.0, 1.0, 1.0],
645 /// [1.0, 1.0, 1.0]]);
646 /// ```
647 #[rhai_fn(name = "ones")]
648 pub fn ones_double_input(nx: INT, ny: INT) -> Array {
649 let mut output = vec![];
650 for _ in 0..nx {
651 output.push(Dynamic::from_array(vec![Dynamic::FLOAT_ONE; ny as usize]))
652 }
653 output
654 }
655
656 /// Returns a matrix of random values, each between zero and one. Can be called with a single integer argument (indicating the
657 /// square matrix of that size) or with an array argument (indicating the size for each dimension).
658 /// ```typescript
659 /// let matrix = rand(3);
660 /// assert_eq(size(matrix), [3, 3]);
661 /// ```
662 /// ```typescript
663 /// let matrix = rand([3, 3]);
664 /// assert_eq(size(matrix), [3, 3]);
665 /// ```
666 #[cfg(feature = "rand")]
667 #[rhai_fn(name = "rand", return_raw)]
668 pub fn rand_single_input(n: Dynamic) -> Result<Array, Box<EvalAltResult>> {
669 crate::if_int_do_else_if_array_do(
670 n,
671 |n| Ok(rand_double_input(n, n)),
672 |m| {
673 if m.len() == 2 {
674 Ok(rand_double_input(
675 m[0].as_int().unwrap(),
676 m[1].as_int().unwrap(),
677 ))
678 } else if m.len() > 2 {
679 let l = m.remove(0);
680 Ok(vec![
681 Dynamic::from_array(
682 rand_single_input(Dynamic::from_array(m.to_vec())).unwrap()
683 );
684 l.as_int().unwrap() as usize
685 ])
686 } else {
687 Ok(rand_double_input(1, m[0].as_int().unwrap())[0]
688 .clone()
689 .into_array()
690 .unwrap())
691 }
692 },
693 )
694 }
695
696 /// Return a matrix of random values, each between zero and one. Arguments indicate the number
697 /// of rows and columns in the matrix.
698 /// ```typescript
699 /// let matrix = rand(3, 3);
700 /// assert_eq(size(matrix), [3, 3]);
701 /// ```
702 #[cfg(feature = "rand")]
703 #[rhai_fn(name = "rand")]
704 pub fn rand_double_input(nx: INT, ny: INT) -> Array {
705 let mut output = vec![];
706 for _ in 0..nx {
707 let mut row = vec![];
708 for _ in 0..ny {
709 row.push(Dynamic::from_float(crate::misc_functions::rand_float()));
710 }
711 output.push(Dynamic::from_array(row))
712 }
713 output
714 }
715
716 /// Returns an identity matrix. If argument is a single number, then the output is
717 /// a square matrix. The argument can also be an array specifying the dimensions separately.
718 /// ```typescript
719 /// let matrix = eye(3);
720 /// assert_eq(matrix, [[1.0, 0.0, 0.0],
721 /// [0.0, 1.0, 0.0],
722 /// [0.0, 0.0, 1.0]]);
723 /// ```
724 /// ```typescript
725 /// let matrix = eye([3, 4]);
726 /// assert_eq(matrix, [[1.0, 0.0, 0.0, 0.0],
727 /// [0.0, 1.0, 0.0, 0.0],
728 /// [0.0, 0.0, 1.0, 0.0]]);
729 /// ```
730 #[rhai_fn(name = "eye", return_raw)]
731 pub fn eye_single_input(n: Dynamic) -> Result<Array, Box<EvalAltResult>> {
732 if_int_do_else_if_array_do(
733 n,
734 |n| Ok(eye_double_input(n, n)),
735 |m| {
736 if m.len() == 1 {
737 Ok(eye_double_input(1, m[0].as_int().unwrap())[0]
738 .clone()
739 .into_array()
740 .unwrap())
741 } else if m.len() == 2 {
742 Ok(eye_double_input(
743 m[0].as_int().unwrap(),
744 m[1].as_int().unwrap(),
745 ))
746 } else {
747 Err(EvalAltResult::ErrorMismatchDataType(
748 format!("Cannot create an identity matrix with more than 2 dimensions"),
749 format!(""),
750 Position::NONE,
751 )
752 .into())
753 }
754 },
755 )
756 }
757
758 /// Returns the identity matrix, specifying the number of rows and columns separately.
759 /// ```typescript
760 /// let matrix = eye(3, 4);
761 /// assert_eq(matrix, [[1.0, 0.0, 0.0, 0.0],
762 /// [0.0, 1.0, 0.0, 0.0],
763 /// [0.0, 0.0, 1.0, 0.0]]);
764 /// ```
765 #[rhai_fn(name = "eye")]
766 pub fn eye_double_input(nx: INT, ny: INT) -> Array {
767 let mut output = vec![];
768 for i in 0..nx {
769 let mut row = vec![];
770 for j in 0..ny {
771 if i == j {
772 row.push(Dynamic::FLOAT_ONE);
773 } else {
774 row.push(Dynamic::FLOAT_ZERO);
775 }
776 }
777 output.push(Dynamic::from_array(row))
778 }
779 output
780 }
781
782 /// Returns the contents of a multidimensional array as a 1-D array.
783 /// ```typescript
784 /// let matrix = ones(3, 5);
785 /// let flat = flatten(matrix);
786 /// assert_eq(len(flat), 15);
787 /// ```
788 /// ```typescript
789 /// let matrix = [[1.0, 2.0, 3.0], [1.0]];
790 /// let flat = flatten(matrix);
791 /// assert_eq(len(flat), 4);
792 /// ```
793 #[rhai_fn(name = "flatten", pure)]
794 pub fn flatten(matrix: &mut Array) -> Array {
795 let mut flat: Array = vec![];
796 for el in matrix {
797 if el.is_array() {
798 flat.extend(flatten(&mut el.clone().into_array().unwrap()))
799 } else {
800 flat.push(el.clone());
801 }
802 }
803 flat
804 }
805
806 /// Flip a matrix left-to-right
807 /// ```typescript
808 /// let matrix = fliplr([[1.0, 0.0],
809 /// [0.0, 2.0]]);
810 /// assert_eq(matrix, [[0.0, 1.0],
811 /// [2.0, 0.0]]);
812 /// ```
813 #[rhai_fn(name = "fliplr", return_raw)]
814 pub fn fliplr(matrix: &mut Array) -> Result<Array, Box<EvalAltResult>> {
815 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
816 let w = matrix_as_vec[0].len();
817 let h = matrix_as_vec.len();
818
819 // Turn into Array
820 let mut out = vec![];
821 for idx in 0..h {
822 let mut new_row = vec![];
823 for jdx in 0..w {
824 new_row.push(matrix_as_vec[idx][w - jdx - 1].clone());
825 }
826 out.push(Dynamic::from_array(new_row));
827 }
828 Ok(out)
829 })
830 }
831
832 /// Flip a matrix up-down
833 /// ```typescript
834 /// let matrix = flipud([[1.0, 0.0],
835 /// [0.0, 2.0]]);
836 /// assert_eq(matrix, [[0.0, 2.0],
837 /// [1.0, 0.0]]);
838 /// ```
839 #[rhai_fn(name = "flipud", return_raw)]
840 pub fn flipud(matrix: &mut Array) -> Result<Array, Box<EvalAltResult>> {
841 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
842 let w = matrix_as_vec[0].len();
843 let h = matrix_as_vec.len();
844
845 // Turn into Array
846 let mut out = vec![];
847 for idx in 0..h {
848 let mut new_row = vec![];
849 for jdx in 0..w {
850 new_row.push(matrix_as_vec[h - idx - 1][jdx].clone());
851 }
852 out.push(Dynamic::from_array(new_row));
853 }
854 Ok(out)
855 })
856 }
857
858 /// Rotate a matrix counterclockwise once
859 /// ```typescript
860 /// let matrix = rot90([[1.0, 0.0],
861 /// [0.0, 2.0]]);
862 /// assert_eq(matrix, [[0.0, 2.0],
863 /// [1.0, 0.0]]);
864 /// ```
865 #[rhai_fn(name = "rot90", return_raw)]
866 pub fn rot90_once(matrix: &mut Array) -> Result<Array, Box<EvalAltResult>> {
867 if_matrix_convert_to_vec_array_and_do(matrix, |matrix_as_vec| {
868 let w = matrix_as_vec[0].len();
869 let h = matrix_as_vec.len();
870
871 // Turn into Array
872 let mut out = vec![];
873 for idx in 0..w {
874 let mut new_row = vec![];
875 for jdx in 0..h {
876 new_row.push(matrix_as_vec[jdx][w - idx - 1].clone());
877 }
878
879 out.push(Dynamic::from_array(new_row));
880 }
881 Ok(out)
882 })
883 }
884
885 /// Rotate a matrix counterclockwise `k` times
886 /// ```typescript
887 /// let matrix = rot90([[1.0, 0.0],
888 /// [0.0, 2.0]], 2);
889 /// assert_eq(matrix, [[2.0, 0.0],
890 /// [0.0, 1.0]]);
891 /// ```
892 #[rhai_fn(name = "rot90", return_raw)]
893 pub fn rot90_ktimes(matrix: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
894 if k <= 0 {
895 return Ok(matrix.clone());
896 }
897
898 let mut result = matrix;
899 let mut result_base = Array::new();
900
901 for _ in 0..k {
902 result_base = rot90_once(result)?;
903 result = &mut result_base;
904 }
905
906 Ok(result_base)
907 }
908
909 /// Perform matrix multiplication.
910 /// ```typescript
911 /// let a = eye(3);
912 /// let b = ones(3);
913 /// let c = mtimes(a, b);
914 /// assert_eq(b, c);
915 /// ```
916 #[cfg(feature = "nalgebra")]
917 #[rhai_fn(name = "mtimes", return_raw)]
918 pub fn mtimes(matrix1: Array, matrix2: Array) -> Result<Array, Box<EvalAltResult>> {
919 if_matrices_and_compatible_convert_to_vec_array_and_do(
920 FOIL::Inside,
921 &mut matrix1.clone(),
922 &mut matrix2.clone(),
923 |matrix_as_vec1, matrix_as_vec2| {
924 let dm1 =
925 DMatrix::from_fn(matrix_as_vec1.len(), matrix_as_vec1[0].len(), |i, j| {
926 if matrix_as_vec1[0][0].is_float() {
927 matrix_as_vec1[i][j].as_float().unwrap()
928 } else {
929 matrix_as_vec1[i][j].as_int().unwrap() as FLOAT
930 }
931 });
932
933 let dm2 =
934 DMatrix::from_fn(matrix_as_vec2.len(), matrix_as_vec2[0].len(), |i, j| {
935 if matrix_as_vec2[0][0].is_float() {
936 matrix_as_vec2[i][j].as_float().unwrap()
937 } else {
938 matrix_as_vec2[i][j].as_int().unwrap() as FLOAT
939 }
940 });
941
942 // Try to multiply
943 let mat = dm1 * dm2;
944
945 // Turn into Array
946 let mut out = vec![];
947 for idx in 0..mat.shape().0 {
948 let mut new_row = vec![];
949 for jdx in 0..mat.shape().1 {
950 new_row.push(Dynamic::from_float(mat[(idx, jdx)]));
951 }
952 out.push(Dynamic::from_array(new_row));
953 }
954 Ok(out)
955 },
956 )
957 }
958
959 /// Concatenate two arrays horizontally.
960 /// ```typescript
961 /// let arr1 = eye(3);
962 /// let arr2 = eye(3);
963 /// let combined = horzcat(arr1, arr2);
964 /// assert_eq(size(combined), [3, 6]);
965 /// ```
966 #[cfg(feature = "nalgebra")]
967 #[rhai_fn(name = "horzcat", return_raw)]
968 pub fn horzcat(matrix1: Array, matrix2: Array) -> Result<Array, Box<EvalAltResult>> {
969 if_matrices_and_compatible_convert_to_vec_array_and_do(
970 FOIL::First,
971 &mut matrix1.clone(),
972 &mut matrix2.clone(),
973 |matrix_as_vec1, matrix_as_vec2| {
974 let dm1 =
975 DMatrix::from_fn(matrix_as_vec1.len(), matrix_as_vec1[0].len(), |i, j| {
976 if matrix_as_vec1[0][0].is_float() {
977 matrix_as_vec1[i][j].as_float().unwrap()
978 } else {
979 matrix_as_vec1[i][j].as_int().unwrap() as FLOAT
980 }
981 });
982
983 let dm2 =
984 DMatrix::from_fn(matrix_as_vec2.len(), matrix_as_vec2[0].len(), |i, j| {
985 if matrix_as_vec2[0][0].is_float() {
986 matrix_as_vec2[i][j].as_float().unwrap()
987 } else {
988 matrix_as_vec2[i][j].as_int().unwrap() as FLOAT
989 }
990 });
991
992 // Try to multiple
993 let w0 = dm1.shape().1;
994 let w = dm1.shape().1 + dm2.shape().1;
995 let h = dm1.shape().0;
996 let mat = DMatrix::from_fn(h, w, |i, j| {
997 if j >= w0 {
998 dm2[(i, j - w0)]
999 } else {
1000 dm1[(i, j)]
1001 }
1002 });
1003
1004 // Turn into Array
1005 let mut out = vec![];
1006 for idx in 0..h {
1007 let mut new_row = vec![];
1008 for jdx in 0..w {
1009 new_row.push(Dynamic::from_float(mat[(idx, jdx)]));
1010 }
1011 out.push(Dynamic::from_array(new_row));
1012 }
1013 Ok(out)
1014 },
1015 )
1016 }
1017
1018 /// Concatenates two array vertically.
1019 /// ```typescript
1020 /// let arr1 = eye(3);
1021 /// let arr2 = eye(3);
1022 /// let combined = vertcat(arr1, arr2);
1023 /// assert_eq(size(combined), [6, 3]);
1024 /// ```
1025 #[cfg(feature = "nalgebra")]
1026 #[rhai_fn(name = "vertcat", return_raw)]
1027 pub fn vertcat(matrix1: Array, matrix2: Array) -> Result<Array, Box<EvalAltResult>> {
1028 if_matrices_and_compatible_convert_to_vec_array_and_do(
1029 FOIL::Last,
1030 &mut matrix1.clone(),
1031 &mut matrix2.clone(),
1032 |matrix_as_vec1, matrix_as_vec2| {
1033 let dm1 =
1034 DMatrix::from_fn(matrix_as_vec1.len(), matrix_as_vec1[0].len(), |i, j| {
1035 if matrix_as_vec1[0][0].is_float() {
1036 matrix_as_vec1[i][j].as_float().unwrap()
1037 } else {
1038 matrix_as_vec1[i][j].as_int().unwrap() as FLOAT
1039 }
1040 });
1041
1042 let dm2 =
1043 DMatrix::from_fn(matrix_as_vec2.len(), matrix_as_vec2[0].len(), |i, j| {
1044 if matrix_as_vec2[0][0].is_float() {
1045 matrix_as_vec2[i][j].as_float().unwrap()
1046 } else {
1047 matrix_as_vec2[i][j].as_int().unwrap() as FLOAT
1048 }
1049 });
1050
1051 // Try to multiple
1052 let h0 = dm1.shape().0;
1053 let w = dm1.shape().1;
1054 let h = dm1.shape().0 + dm2.shape().0;
1055 let mat = DMatrix::from_fn(h, w, |i, j| {
1056 if i >= h0 {
1057 dm2[(i - h0, j)]
1058 } else {
1059 dm1[(i, j)]
1060 }
1061 });
1062
1063 // Turn into Array
1064 let mut out = vec![];
1065 for idx in 0..h {
1066 let mut new_row = vec![];
1067 for jdx in 0..w {
1068 new_row.push(Dynamic::from_float(mat[(idx, jdx)]));
1069 }
1070 out.push(Dynamic::from_array(new_row));
1071 }
1072 Ok(out)
1073 },
1074 )
1075 }
1076
1077 /// This function can be used in two distinct ways.
1078 /// 1. If the argument is an 2-D array, `diag` returns an array containing the diagonal of the array.
1079 /// 2. If the argument is a 1-D array, `diag` returns a matrix containing the argument along the
1080 /// diagonal and zeros elsewhere.
1081 /// ```typescript
1082 /// let matrix = [[1, 2, 3],
1083 /// [4, 5, 6],
1084 /// [7, 8, 9]];
1085 /// let d = diag(matrix);
1086 /// assert_eq(d, [1, 5, 9]);
1087 /// ```
1088 /// ```typescript
1089 /// let diagonal = [1.0, 2.0, 3.0];
1090 /// let matrix = diag(diagonal);
1091 /// assert_eq(matrix, [[1.0, 0.0, 0.0],
1092 /// [0.0, 2.0, 0.0],
1093 /// [0.0, 0.0, 3.0]]);
1094 /// ```
1095 #[rhai_fn(name = "diag", return_raw)]
1096 pub fn diag(matrix: Array) -> Result<Array, Box<EvalAltResult>> {
1097 if ndims_by_reference(&mut matrix.clone()) == 2 {
1098 // Turn into Vec<Array>
1099 let matrix_as_vec = matrix
1100 .into_iter()
1101 .map(|x| x.into_array().unwrap())
1102 .collect::<Vec<Array>>();
1103
1104 let mut out = vec![];
1105 for i in 0..matrix_as_vec.len() {
1106 out.push(matrix_as_vec[i][i].clone());
1107 }
1108
1109 Ok(out)
1110 } else if ndims_by_reference(&mut matrix.clone()) == 1 {
1111 let mut out = vec![];
1112 for idx in 0..matrix.len() {
1113 let mut new_row = vec![];
1114 for jdx in 0..matrix.len() {
1115 if idx == jdx {
1116 new_row.push(matrix[idx].clone());
1117 } else {
1118 if matrix[idx].is_int() {
1119 new_row.push(Dynamic::ZERO);
1120 } else {
1121 new_row.push(Dynamic::FLOAT_ZERO);
1122 }
1123 }
1124 }
1125 out.push(Dynamic::from_array(new_row));
1126 }
1127 Ok(out)
1128 } else {
1129 return Err(EvalAltResult::ErrorArithmetic(
1130 "Argument must be a 2-D matrix (to extract the diagonal) or a 1-D array (to create a matrix with that diagonal".to_string(),
1131 Position::NONE,
1132 )
1133 .into());
1134 }
1135 }
1136
1137 /// Repeats copies of a matrix
1138 /// ```typescript
1139 /// let matrix = eye(3);
1140 /// let combined = repmat(matrix, 2, 2);
1141 /// assert_eq(size(combined), [6, 6]);
1142 /// ```
1143 #[cfg(feature = "nalgebra")]
1144 #[rhai_fn(name = "repmat", return_raw)]
1145 pub fn repmat(matrix: &mut Array, nx: INT, ny: INT) -> Result<Array, Box<EvalAltResult>> {
1146 if_matrix_do(matrix, |matrix| {
1147 let mut row_matrix = matrix.clone();
1148 for _ in 1..ny {
1149 row_matrix = horzcat(row_matrix, matrix.clone())?;
1150 }
1151 let mut new_matrix = row_matrix.clone();
1152 for _ in 1..nx {
1153 new_matrix = vertcat(new_matrix, row_matrix.clone())?;
1154 }
1155 Ok(new_matrix)
1156 })
1157 }
1158
1159 /// Returns an object map containing 2-D grid coordinates based on the uni-axial coordinates
1160 /// contained in arguments x and y.
1161 /// ```typescript
1162 /// let x = [1, 2];
1163 /// let y = [3, 4];
1164 /// let g = meshgrid(x, y);
1165 /// assert_eq(g, #{"x": [[1, 2],
1166 /// [1, 2]],
1167 /// "y": [[3, 3],
1168 /// [4, 4]]});
1169 /// ```
1170 #[rhai_fn(name = "meshgrid", return_raw)]
1171 pub fn meshgrid(x: Array, y: Array) -> Result<Map, Box<EvalAltResult>> {
1172 if_list_do(&mut x.clone(), |x| {
1173 if_list_do(&mut y.clone(), |y| {
1174 let nx = x.len();
1175 let ny = y.len();
1176 let x_dyn: Array = vec![Dynamic::from_array(x.to_vec()); nx];
1177 let mut y_dyn: Array = vec![Dynamic::from_array(y.to_vec()); ny];
1178
1179 let mut result = BTreeMap::new();
1180 let mut xid = smartstring::SmartString::new();
1181 xid.push_str("x");
1182 let mut yid = smartstring::SmartString::new();
1183 yid.push_str("y");
1184 result.insert(xid, Dynamic::from_array(x_dyn));
1185 result.insert(yid, Dynamic::from_array(transpose(&mut y_dyn).unwrap()));
1186 Ok(result)
1187 })
1188 })
1189 }
1190
1191 /// Returns an array containing a number of elements linearly spaced between two bounds.
1192 /// ```typescript
1193 /// let x = linspace(1, 2, 5);
1194 /// assert_eq(x, [1.0, 1.25, 1.5, 1.75, 2.0]);
1195 /// ```
1196 #[rhai_fn(name = "linspace", return_raw)]
1197 pub fn linspace(x1: Dynamic, x2: Dynamic, n: INT) -> Result<Array, Box<EvalAltResult>> {
1198 if_int_convert_to_float_and_do(x1, |new_x1| {
1199 if_int_convert_to_float_and_do(x2.clone(), |new_x2| {
1200 let new_n = n as FLOAT;
1201
1202 let mut arr = vec![Dynamic::from_float(new_x1)];
1203 let mut counter = new_x1;
1204 let interval = (new_x2 - new_x1) / (new_n - 1.0);
1205 for _ in 0..(n - 2) {
1206 counter += interval;
1207 arr.push(Dynamic::from_float(counter));
1208 }
1209 arr.push(Dynamic::from_float(new_x2));
1210 Ok(arr)
1211 })
1212 })
1213 }
1214
1215 /// Returns an array containing a number of elements logarithmically spaced between two bounds.
1216 /// ```typescript
1217 /// let x = logspace(1, 3, 3);
1218 /// assert_eq(x, [10.0, 100.0, 1000.0]);
1219 /// ```
1220 #[rhai_fn(name = "logspace", return_raw)]
1221 pub fn logspace(a: Dynamic, b: Dynamic, n: INT) -> Result<Array, Box<EvalAltResult>> {
1222 linspace(a, b, n).map(|arr| {
1223 arr.iter()
1224 .map(|e| Dynamic::from_float((10 as FLOAT).powf(e.as_float().unwrap())))
1225 .collect::<Array>()
1226 })
1227 }
1228}