brain2qwerty 0.0.1

Brain2Qwerty V1/V2 MEG neural decoding inference in Rust (parity-tested vs Python)
Documentation 
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

//! Fourier positional embedding (neuraltrain FourierEmbModel).

use std::f64::consts::PI;

use crate::tensor::Tensor;

pub struct FourierEmb {
    pub n_freqs: usize,
    pub n_dims: usize,
    pub margin: f32,
    pub pos: Vec<f32>,
    pub total_dim: usize,
}

impl FourierEmb {
    pub fn new(n_freqs: usize, n_dims: usize, margin: f32) -> Self {
        let width = 1.0 + 2.0 * margin as f64;
        let pos: Vec<f32> = (0..n_freqs)
            .map(|f| (2.0 * PI * f as f64 / width) as f32)
            .collect();
        let total_dim = n_freqs.pow(n_dims as u32) * 2;
        Self {
            n_freqs,
            n_dims,
            margin,
            pos,
            total_dim,
        }
    }

    pub fn forward(&self, positions: &Tensor) -> Tensor {
        assert_eq!(positions.ndim(), 3);
        let (b, c, d) = (positions.shape[0], positions.shape[1], positions.shape[2]);
        assert_eq!(d, self.n_dims);
        let grid_size = self.n_freqs.pow(self.n_dims as u32);
        let mut emb = vec![0.0f32; b * c * self.total_dim];

        for bi in 0..b {
            for ci in 0..c {
                let pos_base = (bi * c + ci) * d;
                let mut coords = vec![0.0f32; self.n_dims];
                for di in 0..self.n_dims {
                    coords[di] = positions.data[pos_base + di] + self.margin;
                }
                let mut loc_grid = vec![0.0f32; grid_size];
                self.fill_outer_sum(&coords, &mut loc_grid);
                let out_base = (bi * c + ci) * self.total_dim;
                for g in 0..grid_size {
                    let v = loc_grid[g];
                    emb[out_base + g] = v.cos();
                    emb[out_base + grid_size + g] = v.sin();
                }
            }
        }
        Tensor::from_vec(emb, vec![b, c, self.total_dim])
    }

    fn fill_outer_sum(&self, coords: &[f32], out: &mut [f32]) {
        match self.n_dims {
            1 => {
                for (f, p) in self.pos.iter().enumerate() {
                    out[f] = coords[0] * p;
                }
            }
            2 => {
                let n = self.n_freqs;
                for f0 in 0..n {
                    for f1 in 0..n {
                        out[f0 * n + f1] = coords[0] * self.pos[f0] + coords[1] * self.pos[f1];
                    }
                }
            }
            _ => {
                self.fill_outer_sum_nd(coords, out, 0, 0, 0.0);
            }
        }
    }

    fn fill_outer_sum_nd(&self, coords: &[f32], out: &mut [f32], dim: usize, idx: usize, acc: f32) {
        if dim == self.n_dims {
            out[idx] = acc;
            return;
        }
        let stride = self.n_freqs.pow((self.n_dims - dim - 1) as u32);
        for f in 0..self.n_freqs {
            let new_acc = acc + coords[dim] * self.pos[f];
            self.fill_outer_sum_nd(coords, out, dim + 1, idx + f * stride, new_acc);
        }
    }
}