ecp_core/build/
assign.rs

1use ndarray::{Array1, Array2, Axis};
2
3use crate::metric::{Metric, negative_squared_distances};
4
5/// Decides which node each data vector belongs to. Finds the nearest row of
6/// `node_embeddings` for each row of `data_embeddings`, then groups the rows by
7/// node into `(offsets, data)`. Node `n`'s rows are `data[offsets[n]..offsets[n+1]]`.
8pub fn determine_node_assignments(
9    node_embeddings: &Array2<f32>,
10    data_embeddings: &Array2<f32>,
11    metric: Metric,
12    is_normalized: bool,
13) -> (Array1<u32>, Array1<u32>) {
14    let best_ids = assign_to_nearest(node_embeddings, data_embeddings, metric, is_normalized);
15    group_by_assignments(node_embeddings.nrows(), &best_ids)
16}
17
18/// Picks the nearest node for each data vector under `metric`. Returns one
19/// `node_embeddings` row index for each row of `data_embeddings`.
20fn assign_to_nearest(
21    node_embeddings: &Array2<f32>,
22    data_embeddings: &Array2<f32>,
23    metric: Metric,
24    is_normalized: bool,
25) -> Array1<u32> {
26    let scores: Array2<f32> = match metric {
27        Metric::IP => node_embeddings.dot(&data_embeddings.t()),
28        // For unit vectors, nearest-by-L2 and highest-by-dot-product agree.
29        Metric::L2 if is_normalized => node_embeddings.dot(&data_embeddings.t()),
30        Metric::L2 => negative_squared_distances(node_embeddings, data_embeddings),
31    };
32    argmax_axis0(&scores)
33}
34
35/// Argmax along axis 0: returns the highest-scoring row id for each column.
36fn argmax_axis0(matrix: &Array2<f32>) -> Array1<u32> {
37    matrix.map_axis(Axis(0), |scores| {
38        scores
39            .iter()
40            .enumerate()
41            .max_by(|(_, a), (_, b)| a.partial_cmp(b).expect("scores must not be NaN"))
42            .map(|(node, _)| node as u32)
43            .expect("scores must include at least one node")
44    })
45}
46
47/// Groups data point indices by the representative they were assigned to,
48/// stored as two flat arrays instead of a list of lists. Example:
49/// `best_ids = [1, 0, 1, 2]`, `num_reps = 3`.
50///
51/// ```text
52/// groups (the intuitive shape): [[1], [0, 2], [3]]
53/// offsets, data (what this returns): [0, 1, 3, 4], [1, 0, 2, 3]
54/// ```
55///
56/// `groups[r]` is `data[offsets[r]..offsets[r+1]]`.
57pub fn group_by_assignments(num_reps: usize, best_ids: &Array1<u32>) -> (Array1<u32>, Array1<u32>) {
58    let mut counts = vec![0u32; num_reps];
59    for &rep in best_ids.iter() {
60        counts[rep as usize] += 1;
61    }
62
63    let mut offsets = Array1::<u32>::zeros(num_reps + 1);
64    for r in 0..num_reps {
65        offsets[r + 1] = offsets[r] + counts[r];
66    }
67
68    let mut cursor: Vec<u32> = offsets.iter().take(num_reps).copied().collect();
69    let mut data = Array1::<u32>::zeros(best_ids.len());
70    for (data_point, &rep) in best_ids.iter().enumerate() {
71        let pos = cursor[rep as usize];
72        data[pos as usize] = data_point as u32;
73        cursor[rep as usize] += 1;
74    }
75
76    (offsets, data)
77}
78
79#[cfg(test)]
80#[path = "utests/assign.rs"]
81mod tests;
ecp_core/build/assign.rs

ecp_core/build/
assign.rs
