ak_vis/viewer/session/
selection.rs

1use ak_core::Trajectory;
2
3use super::{ImageSelectionFrames, SelectedImageAtom, SelectionFrames};
4
5impl SelectionFrames {
6    pub fn new(traj: &Trajectory) -> Self {
7        Self {
8            frames: (0..traj.len())
9                .map(|index| vec![false; traj.view(index).positions.len()])
10                .collect(),
11            ordered: vec![Vec::new(); traj.len()],
12        }
13    }
14
15    pub fn get(&self, frame_index: usize) -> Option<&[bool]> {
16        self.frames.get(frame_index).map(Vec::as_slice)
17    }
18
19    pub fn replace(&mut self, frame_index: usize, selection: Vec<bool>) {
20        if frame_index < self.frames.len() && self.frames[frame_index].len() == selection.len() {
21            self.frames[frame_index] = selection;
22            self.ordered[frame_index] = self.frames[frame_index]
23                .iter()
24                .enumerate()
25                .filter_map(|(index, selected)| selected.then_some(index))
26                .collect();
27        }
28    }
29
30    pub fn add(&mut self, frame_index: usize, selection: &[bool]) {
31        if let Some(current) = self.frames.get_mut(frame_index) {
32            if current.len() != selection.len() {
33                return;
34            }
35            for (slot, selected) in current.iter_mut().zip(selection.iter().copied()) {
36                *slot |= selected;
37            }
38            let ordered = &mut self.ordered[frame_index];
39            for (index, selected) in selection.iter().copied().enumerate() {
40                if selected && current[index] && !ordered.contains(&index) {
41                    ordered.push(index);
42                }
43            }
44        }
45    }
46
47    pub fn remove(&mut self, frame_index: usize, selection: &[bool]) {
48        if let Some(current) = self.frames.get_mut(frame_index) {
49            if current.len() != selection.len() {
50                return;
51            }
52            for (slot, selected) in current.iter_mut().zip(selection.iter().copied()) {
53                if selected {
54                    *slot = false;
55                }
56            }
57            self.ordered[frame_index].retain(|index| current.get(*index).copied().unwrap_or(false));
58        }
59    }
60
61    pub fn clear(&mut self, frame_index: usize) {
62        if let Some(current) = self.frames.get_mut(frame_index) {
63            current.fill(false);
64            self.ordered[frame_index].clear();
65        }
66    }
67
68    pub fn append_empty_for_atom_count(&mut self, atom_count: usize) {
69        self.frames.push(vec![false; atom_count]);
70        self.ordered.push(Vec::new());
71    }
72
73    pub fn selected_main_images(&self, frame_index: usize) -> Vec<SelectedImageAtom> {
74        self.selected_indices(frame_index)
75            .into_iter()
76            .map(|atom_index| SelectedImageAtom {
77                atom_index,
78                image_offset: [0, 0, 0],
79            })
80            .collect()
81    }
82
83    pub fn selected_indices(&self, frame_index: usize) -> Vec<usize> {
84        self.ordered.get(frame_index).cloned().unwrap_or_default()
85    }
86
87    pub(crate) fn set_order(&mut self, frame_index: usize, ordered: Vec<usize>) {
88        if frame_index < self.ordered.len() {
89            self.ordered[frame_index] = ordered;
90        }
91    }
92
93    pub(crate) fn mask_from_images(atom_count: usize, atoms: &[SelectedImageAtom]) -> Vec<bool> {
94        let mut mask = vec![false; atom_count];
95        for atom in atoms {
96            if atom.atom_index < atom_count {
97                mask[atom.atom_index] = true;
98            }
99        }
100        mask
101    }
102
103    pub(crate) fn ordered_atoms_from_images(atoms: &[SelectedImageAtom]) -> Vec<usize> {
104        let mut ordered = Vec::new();
105        for atom in atoms {
106            if !ordered.contains(&atom.atom_index) {
107                ordered.push(atom.atom_index);
108            }
109        }
110        ordered
111    }
112}
113
114impl ImageSelectionFrames {
115    pub fn new(frame_count: usize) -> Self {
116        Self {
117            frames: vec![Vec::new(); frame_count],
118        }
119    }
120
121    pub fn get(&self, frame_index: usize) -> Option<&[SelectedImageAtom]> {
122        self.frames.get(frame_index).map(Vec::as_slice)
123    }
124
125    pub fn selected(&self, frame_index: usize) -> Vec<SelectedImageAtom> {
126        self.frames.get(frame_index).cloned().unwrap_or_default()
127    }
128
129    pub fn replace(&mut self, frame_index: usize, selection: Vec<SelectedImageAtom>) {
130        if frame_index < self.frames.len() {
131            self.frames[frame_index] = dedup_image_selection(selection);
132        }
133    }
134
135    pub fn add(&mut self, frame_index: usize, selection: &[SelectedImageAtom]) {
136        let Some(current) = self.frames.get_mut(frame_index) else {
137            return;
138        };
139        for atom in selection.iter().copied() {
140            if !current.contains(&atom) {
141                current.push(atom);
142            }
143        }
144    }
145
146    pub fn remove(&mut self, frame_index: usize, selection: &[SelectedImageAtom]) {
147        let Some(current) = self.frames.get_mut(frame_index) else {
148            return;
149        };
150        current.retain(|atom| !selection.contains(atom));
151    }
152
153    pub fn clear(&mut self, frame_index: usize) {
154        if frame_index < self.frames.len() {
155            self.frames[frame_index].clear();
156        }
157    }
158
159    pub fn append_empty_frame(&mut self) {
160        self.frames.push(Vec::new());
161    }
162
163    pub(super) fn replace_single(&mut self, frame_index: usize, atom: SelectedImageAtom) -> bool {
164        let Some(current) = self.frames.get_mut(frame_index) else {
165            return false;
166        };
167        if current.as_slice() == [atom] {
168            return false;
169        }
170        current.clear();
171        current.push(atom);
172        true
173    }
174
175    pub(super) fn toggle(&mut self, frame_index: usize, atom: SelectedImageAtom) -> bool {
176        let Some(current) = self.frames.get_mut(frame_index) else {
177            return false;
178        };
179        if let Some(index) = current.iter().position(|selected| *selected == atom) {
180            current.remove(index);
181        } else {
182            current.push(atom);
183        }
184        true
185    }
186}
187
188fn dedup_image_selection(selection: Vec<SelectedImageAtom>) -> Vec<SelectedImageAtom> {
189    let mut deduped = Vec::new();
190    for atom in selection {
191        if !deduped.contains(&atom) {
192            deduped.push(atom);
193        }
194    }
195    deduped
196}
ak_vis/viewer/session/selection.rs

ak_vis/viewer/session/
selection.rs
