File: zip_many.rs

package info (click to toggle)
rust-ndarray 0.16.1-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 2,028 kB
  • sloc: sh: 30; makefile: 2
file content (49 lines) | stat: -rw-r--r-- 1,366 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
#![allow(
    clippy::many_single_char_names, clippy::deref_addrof, clippy::unreadable_literal, clippy::many_single_char_names
)]

use ndarray::prelude::*;
use ndarray::Zip;

fn main()
{
    let n = 6;

    let mut a = Array::<f32, _>::zeros((n, n));
    let mut b = Array::<f32, _>::zeros((n, n));
    for ((i, j), elt) in b.indexed_iter_mut() {
        *elt = 1. / (1. + (i + 2 * j) as f32);
    }
    let c = Array::<f32, _>::from_elem((n, n + 1), 1.7);
    let c = c.slice(s![.., ..-1]);

    // Using Zip for arithmetic ops across a, b, c
    Zip::from(&mut a)
        .and(&b)
        .and(&c)
        .for_each(|a, &b, &c| *a = b + c);
    assert_eq!(a, &b + &c);

    // and this is how to do the *same thing* with azip!()
    azip!((a in &mut a, &b in &b, &c in c) *a = b + c);

    println!("{:8.4}", a);

    // sum of each row
    let mut sums = Array::zeros(a.nrows());
    Zip::from(a.rows())
        .and(&mut sums)
        .for_each(|row, sum| *sum = row.sum());
    // show sums as a column matrix
    println!("{:8.4}", sums.insert_axis(Axis(1)));

    // sum of each 2x2 chunk
    let chunk_sz = (2, 2);
    let nchunks = (n / chunk_sz.0, n / chunk_sz.1);
    let mut sums = Array::zeros(nchunks);

    Zip::from(a.exact_chunks(chunk_sz))
        .and(&mut sums)
        .for_each(|chunk, sum| *sum = chunk.sum());
    println!("{:8.4}", sums);
}