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