ak_vis/viewer/session/
faces.rs

1use smallvec::SmallVec;
2use std::collections::BTreeSet;
3
4#[derive(Clone, Debug, PartialEq)]
5pub struct Face {
6    pub atoms: SmallVec<[usize; 4]>,
7    pub color: [f32; 4],
8}
9
10#[derive(Clone, Debug, PartialEq)]
11pub struct FaceList {
12    faces: Vec<Face>,
13}
14
15#[derive(Clone, Debug, PartialEq)]
16pub struct FaceFrames {
17    frames: Vec<Option<FaceList>>,
18}
19
20impl Face {
21    pub fn new(atoms: impl IntoIterator<Item = usize>, color: [f32; 4]) -> Option<Self> {
22        let atoms: SmallVec<[usize; 4]> = atoms.into_iter().collect();
23        if atoms.len() < 3 {
24            return None;
25        }
26        let unique: BTreeSet<_> = atoms.iter().copied().collect();
27        if unique.len() != atoms.len() {
28            return None;
29        }
30        Some(Self { atoms, color })
31    }
32
33    pub fn color(self) -> bevy::color::Color {
34        bevy::color::Color::srgba(self.color[0], self.color[1], self.color[2], self.color[3])
35    }
36
37    fn canonical_atoms(&self) -> Vec<usize> {
38        canonicalize_face_atoms(&self.atoms)
39    }
40}
41
42impl FaceList {
43    pub fn new(faces: impl IntoIterator<Item = Face>) -> Self {
44        let mut seen = BTreeSet::new();
45        let mut canonical = Vec::new();
46
47        for face in faces {
48            let key = face.canonical_atoms();
49            if seen.insert(key) {
50                canonical.push(face);
51            }
52        }
53
54        Self { faces: canonical }
55    }
56
57    pub fn iter(&self) -> impl Iterator<Item = &Face> {
58        self.faces.iter()
59    }
60
61    pub fn len(&self) -> usize {
62        self.faces.len()
63    }
64
65    pub fn is_empty(&self) -> bool {
66        self.faces.is_empty()
67    }
68
69    pub fn validated_for_atom_count(&self, atom_count: usize) -> Self {
70        Self::new(
71            self.faces
72                .iter()
73                .filter(|face| face.atoms.iter().all(|&index| index < atom_count))
74                .cloned(),
75        )
76    }
77
78    pub fn merged(&self, other: &FaceList) -> Self {
79        Self::new(self.iter().cloned().chain(other.iter().cloned()))
80    }
81
82    pub fn without(&self, other: &FaceList) -> Self {
83        let removals: BTreeSet<_> = other.iter().map(Face::canonical_atoms).collect();
84        Self::new(
85            self.faces
86                .iter()
87                .filter(|face| !removals.contains(&face.canonical_atoms()))
88                .cloned(),
89        )
90    }
91}
92
93impl FaceFrames {
94    pub fn new(frame_count: usize) -> Self {
95        Self {
96            frames: vec![None; frame_count],
97        }
98    }
99
100    pub fn set(&mut self, frame_index: usize, faces: FaceList) {
101        if frame_index < self.frames.len() {
102            self.frames[frame_index] = Some(faces);
103        }
104    }
105
106    pub fn add(&mut self, frame_index: usize, faces: FaceList) {
107        if frame_index >= self.frames.len() {
108            return;
109        }
110        let merged = match self.frames[frame_index].take() {
111            Some(existing) => existing.merged(&faces),
112            None => faces,
113        };
114        self.frames[frame_index] = Some(merged);
115    }
116
117    pub fn remove(&mut self, frame_index: usize, faces: &FaceList) {
118        if frame_index >= self.frames.len() {
119            return;
120        }
121        if let Some(existing) = self.frames[frame_index].take() {
122            let updated = existing.without(faces);
123            self.frames[frame_index] = (!updated.is_empty()).then_some(updated);
124        }
125    }
126
127    pub fn clear(&mut self, frame_index: usize) {
128        if frame_index < self.frames.len() {
129            self.frames[frame_index] = None;
130        }
131    }
132
133    pub fn get(&self, frame_index: usize) -> Option<&FaceList> {
134        self.frames
135            .get(frame_index)
136            .and_then(|faces| faces.as_ref())
137    }
138
139    pub fn resize(&mut self, frame_count: usize) {
140        self.frames.resize(frame_count, None);
141    }
142}
143
144fn canonicalize_face_atoms(atoms: &[usize]) -> Vec<usize> {
145    let forward = minimum_face_rotation(atoms);
146    let mut reversed = atoms.to_vec();
147    reversed.reverse();
148    let reversed = minimum_face_rotation(&reversed);
149    if reversed < forward {
150        reversed
151    } else {
152        forward
153    }
154}
155
156fn minimum_face_rotation(atoms: &[usize]) -> Vec<usize> {
157    let mut best = atoms.to_vec();
158    for shift in 1..atoms.len() {
159        let rotated: Vec<usize> = atoms
160            .iter()
161            .cycle()
162            .skip(shift)
163            .take(atoms.len())
164            .copied()
165            .collect();
166        if rotated < best {
167            best = rotated;
168        }
169    }
170    best
171}
ak_vis/viewer/session/faces.rs

ak_vis/viewer/session/
faces.rs
