ferritin_bevy/
colors.rs

1//! Color schemes for protein structure rendering.
2
3use crate::selection::{AtomMask, evaluate_selector};
4use bevy::color::Srgba;
5use bevy::math::Vec4;
6use bevy::prelude::{Color, Mesh};
7use ferritin_core::AtomCollection;
8use ferritin_molviewspec::molviewspec::nodes::{
9    ColorNamesT, ColorT, ColorThemeT, ComponentSelector,
10};
11
12/// Represents different color schemes for rendering atoms.
13#[derive(Clone)]
14pub enum ColorScheme {
15    /// A solid, single color for all atoms.
16    Solid(Color),
17    /// Colors atoms based on their element type (CPK scheme).
18    ByAtomType,
19    /// Color from a MolViewSpec `ColorT` (named or hex).
20    MvsColor(ColorT),
21}
22
23impl ColorScheme {
24    pub fn get_color(&self, element: &ferritin_core::info::elements::Element) -> Color {
25        match self {
26            ColorScheme::Solid(color) => *color,
27            ColorScheme::ByAtomType => Color::Srgba(element_color(element)),
28            ColorScheme::MvsColor(c) => color_t_to_bevy(c),
29        }
30    }
31}
32
33/// Discriminates whether a vertex_map contains atom indices or residue iteration indices.
34pub enum VertexMapKind {
35    /// Each entry is a 0-based atom index into the `AtomCollection`.
36    AtomMap,
37    /// Each entry is a 0-based amino-acid residue iteration index.
38    /// Coloring is approximate: a residue is colored if any of its atoms passes the selector.
39    ResidueMap,
40}
41
42/// A rotating palette of visually distinct chain colors.
43const CHAIN_PALETTE: &[Srgba] = &[
44    Srgba::new(0.290, 0.565, 0.886, 1.0), // blue
45    Srgba::new(0.886, 0.565, 0.290, 1.0), // orange
46    Srgba::new(0.290, 0.765, 0.404, 1.0), // green
47    Srgba::new(0.765, 0.290, 0.404, 1.0), // rose
48    Srgba::new(0.686, 0.290, 0.886, 1.0), // purple
49    Srgba::new(0.290, 0.765, 0.765, 1.0), // teal
50    Srgba::new(0.886, 0.886, 0.290, 1.0), // gold
51    Srgba::new(0.886, 0.290, 0.290, 1.0), // red
52];
53
54/// Compute per-vertex initial colors from a `ColorThemeT`.
55///
56/// Returns a flat `Vec<Vec4>` parallel to `vertex_map`; callers then apply
57/// inline Color-node overrides on top of this base.
58pub fn initial_colors_from_theme(
59    theme: &ColorThemeT,
60    model: &AtomCollection,
61    vertex_map: &[usize],
62    map_kind: &VertexMapKind,
63) -> Vec<Vec4> {
64    let n_verts = vertex_map.len();
65    let white = Vec4::new(1.0, 1.0, 1.0, 1.0);
66
67    match theme {
68        ColorThemeT::Uniform => vec![white; n_verts],
69
70        ColorThemeT::ElementSymbol => vertex_map
71            .iter()
72            .map(|&atom_idx| {
73                let elem = model.get_element(atom_idx);
74                let c = element_color(elem);
75                Vec4::new(c.red, c.green, c.blue, c.alpha)
76            })
77            .collect(),
78
79        ColorThemeT::ChainId => {
80            // Assign each unique chain an index in encounter order; wrap palette.
81            let mut chain_order: Vec<String> = Vec::new();
82            let mut colors: Vec<Vec4> = Vec::with_capacity(n_verts);
83            for &atom_idx in vertex_map {
84                let chain = match map_kind {
85                    VertexMapKind::AtomMap => model.get_chain_id(atom_idx).to_string(),
86                    VertexMapKind::ResidueMap => {
87                        let residue_starts = model.get_residue_start_indices();
88                        let a_idx = residue_starts
89                            .and_then(|rs| rs.get(atom_idx))
90                            .copied()
91                            .unwrap_or(atom_idx);
92                        model.get_chain_id(a_idx).to_string()
93                    }
94                };
95                let idx = if let Some(pos) = chain_order.iter().position(|c| *c == chain) {
96                    pos
97                } else {
98                    chain_order.push(chain);
99                    chain_order.len() - 1
100                };
101                let c = CHAIN_PALETTE[idx % CHAIN_PALETTE.len()];
102                colors.push(Vec4::new(c.red, c.green, c.blue, c.alpha));
103            }
104            colors
105        }
106
107        ColorThemeT::SecondaryStructure => {
108            // CA-distance heuristic: collect per-residue CA positions, detect helix/strand/coil,
109            // then map each vertex back to its residue's color.
110            let residue_ss = detect_secondary_structure_colors(model);
111            vertex_map
112                .iter()
113                .map(|&idx| {
114                    let res_idx = match map_kind {
115                        VertexMapKind::AtomMap => {
116                            // Find which residue owns this atom.
117                            model
118                                .get_residue_start_indices()
119                                .and_then(|rs| rs.partition_point(|&s| s <= idx).checked_sub(1))
120                                .unwrap_or(0)
121                        }
122                        VertexMapKind::ResidueMap => idx,
123                    };
124                    residue_ss.get(res_idx).copied().unwrap_or(white)
125                })
126                .collect()
127        }
128
129        ColorThemeT::Uncertainty => match model.get_b_factors() {
130            Some(values) => {
131                let (min, max) = min_max(values);
132                vertex_map
133                    .iter()
134                    .map(|&idx| {
135                        let atom_idx = resolve_atom_index(model, idx, map_kind);
136                        let v = values[atom_idx];
137                        let t = if max > min {
138                            (v - min) / (max - min)
139                        } else {
140                            0.5
141                        };
142                        uncertainty_gradient(t)
143                    })
144                    .collect()
145            }
146            None => {
147                bevy::log::warn!(
148                    "Uncertainty color theme requested but no per-atom b_factor/pLDDT \
149                     data is available; falling back to uniform white"
150                );
151                vec![white; n_verts]
152            }
153        },
154
155        ColorThemeT::PlddtConfidence => match model.get_b_factors() {
156            Some(values) => vertex_map
157                .iter()
158                .map(|&idx| {
159                    let atom_idx = resolve_atom_index(model, idx, map_kind);
160                    plddt_band_color(values[atom_idx])
161                })
162                .collect(),
163            None => {
164                bevy::log::warn!(
165                    "PlddtConfidence color theme requested but no per-atom confidence \
166                     data is available; falling back to uniform white"
167                );
168                vec![white; n_verts]
169            }
170        },
171
172        ColorThemeT::Occupancy => match model.get_occupancies() {
173            Some(values) => vertex_map
174                .iter()
175                .map(|&idx| {
176                    let atom_idx = resolve_atom_index(model, idx, map_kind);
177                    occupancy_color(values[atom_idx])
178                })
179                .collect(),
180            None => {
181                bevy::log::warn!(
182                    "Occupancy color theme requested but no per-atom occupancy data is \
183                     available; falling back to uniform white"
184                );
185                vec![white; n_verts]
186            }
187        },
188    }
189}
190
191/// Resolve a `vertex_map` entry to an atom index, per `map_kind`.
192///
193/// For `ResidueMap`, `idx` is an amino-acid iteration index; map it to that
194/// residue's first atom, since per-residue scalar fields (B-factor/pLDDT/occupancy)
195/// are approximated by a single representative atom.
196fn resolve_atom_index(model: &AtomCollection, idx: usize, map_kind: &VertexMapKind) -> usize {
197    match map_kind {
198        VertexMapKind::AtomMap => idx,
199        VertexMapKind::ResidueMap => model
200            .get_residue_start_indices()
201            .and_then(|rs| rs.get(idx))
202            .copied()
203            .unwrap_or(idx),
204    }
205}
206
207fn min_max(values: &[f32]) -> (f32, f32) {
208    let min = values.iter().copied().fold(f32::INFINITY, f32::min);
209    let max = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
210    (min, max)
211}
212
213/// Blue (low) -> white (mid) -> red (high) gradient over a normalized `t` in `[0, 1]`.
214fn uncertainty_gradient(t: f32) -> Vec4 {
215    let t = t.clamp(0.0, 1.0);
216    if t < 0.5 {
217        let u = t / 0.5;
218        Vec4::new(u, u, 1.0, 1.0)
219    } else {
220        let u = (t - 0.5) / 0.5;
221        Vec4::new(1.0, 1.0 - u, 1.0 - u, 1.0)
222    }
223}
224
225/// AlphaFold-style pLDDT confidence banding (assumes a 0-100 scale).
226fn plddt_band_color(v: f32) -> Vec4 {
227    if v > 90.0 {
228        Vec4::new(0.0, 0.325, 0.851, 1.0) // very high: blue
229    } else if v > 70.0 {
230        Vec4::new(0.325, 0.788, 0.855, 1.0) // confident: teal
231    } else if v > 50.0 {
232        Vec4::new(1.0, 0.859, 0.0, 1.0) // low: yellow
233    } else {
234        Vec4::new(1.0, 0.494, 0.0, 1.0) // very low: orange
235    }
236}
237
238/// Full occupancy (1.0) -> white; partial/absent occupancy -> red, to make
239/// alternate-conformer atoms stand out.
240fn occupancy_color(v: f32) -> Vec4 {
241    let v = v.clamp(0.0, 1.0);
242    Vec4::new(1.0, v, v, 1.0)
243}
244
245/// Simple CA-distance secondary-structure heuristic (same logic as `Structure::detect_secondary_structure`).
246/// Returns one `Vec4` color per amino-acid residue in iteration order.
247fn detect_secondary_structure_colors(model: &AtomCollection) -> Vec<Vec4> {
248    // Canonical SS colors matching the cartoon renderer.
249    let helix_color = Vec4::new(1.0, 0.6, 0.6, 1.0); // salmon
250    let sheet_color = Vec4::new(0.6, 0.6, 1.0, 1.0); // periwinkle
251    let loop_color = Vec4::new(0.6, 0.9, 0.6, 1.0); // light-green
252
253    // Collect CA positions for each amino-acid residue.
254    let ca_positions: Vec<Option<[f32; 3]>> = model
255        .iter_residues_aminoacid()
256        .map(|res| res.find_atom_by_name("CA").map(|a| *a.coords()))
257        .collect();
258
259    let n = ca_positions.len();
260    let mut ss_colors = vec![loop_color; n];
261
262    for i in 1..n.saturating_sub(1) {
263        let (Some(prev), Some(curr), Some(next)) =
264            (ca_positions[i - 1], ca_positions[i], ca_positions[i + 1])
265        else {
266            continue;
267        };
268        let d1 = dist(prev, curr);
269        let d2 = dist(curr, next);
270        if d1 < 4.2 && d2 < 4.2 {
271            ss_colors[i] = helix_color;
272        } else if d1 < 4.5 && d2 < 4.5 {
273            ss_colors[i] = sheet_color;
274        }
275    }
276    ss_colors
277}
278
279fn dist(a: [f32; 3], b: [f32; 3]) -> f32 {
280    let dx = a[0] - b[0];
281    let dy = a[1] - b[1];
282    let dz = a[2] - b[2];
283    (dx * dx + dy * dy + dz * dz).sqrt()
284}
285
286/// Paint vertex colors on `mesh` by applying MVS color nodes in order.
287///
288/// `initial` is the per-vertex base color (e.g. from `initial_colors_from_theme`).
289/// Each `(ComponentSelector, ColorT)` pair selects atoms and overrides their color.
290/// Later nodes override earlier nodes on overlapping vertices.
291///
292/// `vertex_map[vertex_idx]` holds the atom index (for `VertexMapKind::AtomMap`) or the
293/// amino-acid residue iteration index (for `VertexMapKind::ResidueMap`) that produced
294/// that vertex.
295pub fn apply_mvs_colors(
296    mesh: &mut Mesh,
297    color_nodes: &[(ComponentSelector, ColorT)],
298    model: &AtomCollection,
299    vertex_map: &[usize],
300    map_kind: VertexMapKind,
301) {
302    let n_verts = vertex_map.len();
303    if n_verts == 0 {
304        return;
305    }
306
307    // Default base: white for all vertices (unchanged from existing behaviour when
308    // no theme is specified and no Color nodes are present).
309    let mut vertex_colors: Vec<Vec4> = vec![Vec4::new(1.0, 1.0, 1.0, 1.0); n_verts];
310
311    apply_color_overrides(
312        &mut vertex_colors,
313        color_nodes,
314        model,
315        vertex_map,
316        &map_kind,
317    );
318
319    mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, vertex_colors);
320}
321
322/// Apply theme + Color-node overrides in one call.
323///
324/// Equivalent to `initial_colors_from_theme` followed by `apply_color_overrides`,
325/// then writing to the mesh.
326pub fn apply_mvs_colors_with_theme(
327    mesh: &mut Mesh,
328    theme: &ColorThemeT,
329    color_nodes: &[(ComponentSelector, ColorT)],
330    model: &AtomCollection,
331    vertex_map: &[usize],
332    map_kind: VertexMapKind,
333) {
334    let n_verts = vertex_map.len();
335    if n_verts == 0 {
336        return;
337    }
338    let mut vertex_colors = initial_colors_from_theme(theme, model, vertex_map, &map_kind);
339    apply_color_overrides(
340        &mut vertex_colors,
341        color_nodes,
342        model,
343        vertex_map,
344        &map_kind,
345    );
346    mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, vertex_colors);
347}
348
349/// Overlay `color_nodes` on top of an existing `vertex_colors` array in-place.
350fn apply_color_overrides(
351    vertex_colors: &mut [Vec4],
352    color_nodes: &[(ComponentSelector, ColorT)],
353    model: &AtomCollection,
354    vertex_map: &[usize],
355    map_kind: &VertexMapKind,
356) {
357    for (selector, color_t) in color_nodes {
358        let atom_mask = evaluate_selector(selector, model);
359        let bevy_color = color_t_to_bevy(color_t).to_srgba();
360        let cv = Vec4::new(
361            bevy_color.red,
362            bevy_color.green,
363            bevy_color.blue,
364            bevy_color.alpha,
365        );
366
367        match map_kind {
368            VertexMapKind::AtomMap => {
369                for (v, &atom_idx) in vertex_map.iter().enumerate() {
370                    if atom_mask.0.get(atom_idx).copied().unwrap_or(false) {
371                        vertex_colors[v] = cv;
372                    }
373                }
374            }
375            VertexMapKind::ResidueMap => {
376                // Build a per-(residue-iter-idx) boolean: true if any atom in that residue passes.
377                let residue_passes = residue_mask_from_atom_mask(model, &atom_mask);
378                for (v, &res_iter_idx) in vertex_map.iter().enumerate() {
379                    if residue_passes.get(res_iter_idx).copied().unwrap_or(false) {
380                        vertex_colors[v] = cv;
381                    }
382                }
383            }
384        }
385    }
386}
387
388/// Build a per-residue boolean mask (indexed by ALL-residue iteration order)
389/// from a per-atom mask: residue passes if ANY of its atoms passes.
390fn residue_mask_from_atom_mask(model: &AtomCollection, atom_mask: &AtomMask) -> Vec<bool> {
391    model
392        .iter_residues()
393        .map(|residue| {
394            residue
395                .atom_range()
396                .any(|i| atom_mask.0.get(i).copied().unwrap_or(false))
397        })
398        .collect()
399}
400
401/// Convert a MolViewSpec `ColorT` to a Bevy `Color`.
402pub fn color_t_to_bevy(color: &ColorT) -> Color {
403    match color {
404        ColorT::Hex(s) => parse_hex_color(s),
405        ColorT::Named(name) => named_color_to_bevy(name),
406    }
407}
408
409/// Parse `#rrggbb` (case-insensitive, leading `#` optional).
410/// Returns `Color::WHITE` with a `warn!` on any parse error.
411pub fn parse_hex_color(s: &str) -> Color {
412    let s = s.trim_start_matches('#');
413    if s.len() == 6
414        && let (Ok(r), Ok(g), Ok(b)) = (
415            u8::from_str_radix(&s[0..2], 16),
416            u8::from_str_radix(&s[2..4], 16),
417            u8::from_str_radix(&s[4..6], 16),
418        )
419    {
420        return Color::srgb(r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0);
421    }
422    bevy::log::warn!("Invalid hex color '{}'; falling back to white", s);
423    Color::WHITE
424}
425
426fn named_color_to_bevy(name: &ColorNamesT) -> Color {
427    let (r, g, b) = css_named_color(name);
428    Color::srgb(r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0)
429}
430
431/// CSS named color RGB values (W3C spec, <https://www.w3.org/TR/css-color-3/>).
432#[rustfmt::skip]
433pub fn css_named_color(name: &ColorNamesT) -> (u8, u8, u8) {
434    match name {
435        ColorNamesT::Aliceblue           => (240, 248, 255),
436        ColorNamesT::Antiquewhite        => (250, 235, 215),
437        ColorNamesT::Aqua                => (0,   255, 255),
438        ColorNamesT::Aquamarine          => (127, 255, 212),
439        ColorNamesT::Azure               => (240, 255, 255),
440        ColorNamesT::Beige               => (245, 245, 220),
441        ColorNamesT::Bisque              => (255, 228, 196),
442        ColorNamesT::Black               => (0,   0,   0  ),
443        ColorNamesT::Blanchedalmond      => (255, 235, 205),
444        ColorNamesT::Blue                => (0,   0,   255),
445        ColorNamesT::Blueviolet          => (138, 43,  226),
446        ColorNamesT::Brown               => (165, 42,  42 ),
447        ColorNamesT::Burlywood           => (222, 184, 135),
448        ColorNamesT::Cadetblue           => (95,  158, 160),
449        ColorNamesT::Chartreuse          => (127, 255, 0  ),
450        ColorNamesT::Chocolate           => (210, 105, 30 ),
451        ColorNamesT::Coral               => (255, 127, 80 ),
452        ColorNamesT::Cornflowerblue      => (100, 149, 237),
453        ColorNamesT::Cornsilk            => (255, 248, 220),
454        ColorNamesT::Crimson             => (220, 20,  60 ),
455        ColorNamesT::Cyan                => (0,   255, 255),
456        ColorNamesT::Darkblue            => (0,   0,   139),
457        ColorNamesT::Darkcyan            => (0,   139, 139),
458        ColorNamesT::Darkgoldenrod       => (184, 134, 11 ),
459        ColorNamesT::Darkgray            => (169, 169, 169),
460        ColorNamesT::Darkgreen           => (0,   100, 0  ),
461        ColorNamesT::Darkgrey            => (169, 169, 169),
462        ColorNamesT::Darkkhaki           => (189, 183, 107),
463        ColorNamesT::Darkmagenta         => (139, 0,   139),
464        ColorNamesT::Darkolivegreen      => (85,  107, 47 ),
465        ColorNamesT::Darkorange          => (255, 140, 0  ),
466        ColorNamesT::Darkorchid          => (153, 50,  204),
467        ColorNamesT::Darkred             => (139, 0,   0  ),
468        ColorNamesT::Darksalmon          => (233, 150, 122),
469        ColorNamesT::Darkseagreen        => (143, 188, 143),
470        ColorNamesT::Darkslateblue       => (72,  61,  139),
471        ColorNamesT::Darkslategray       => (47,  79,  79 ),
472        ColorNamesT::Darkslategrey       => (47,  79,  79 ),
473        ColorNamesT::Darkturquoise       => (0,   206, 209),
474        ColorNamesT::Darkviolet          => (148, 0,   211),
475        ColorNamesT::Deeppink            => (255, 20,  147),
476        ColorNamesT::Deepskyblue         => (0,   191, 255),
477        ColorNamesT::Dimgray             => (105, 105, 105),
478        ColorNamesT::Dimgrey             => (105, 105, 105),
479        ColorNamesT::Dodgerblue          => (30,  144, 255),
480        ColorNamesT::Firebrick           => (178, 34,  34 ),
481        ColorNamesT::Floralwhite         => (255, 250, 240),
482        ColorNamesT::Forestgreen         => (34,  139, 34 ),
483        ColorNamesT::Fuchsia             => (255, 0,   255),
484        ColorNamesT::Gainsboro           => (220, 220, 220),
485        ColorNamesT::Ghostwhite          => (248, 248, 255),
486        ColorNamesT::Gold                => (255, 215, 0  ),
487        ColorNamesT::Goldenrod           => (218, 165, 32 ),
488        ColorNamesT::Gray                => (128, 128, 128),
489        ColorNamesT::Green               => (0,   128, 0  ),
490        ColorNamesT::Greenyellow         => (173, 255, 47 ),
491        ColorNamesT::Grey                => (128, 128, 128),
492        ColorNamesT::Honeydew            => (240, 255, 240),
493        ColorNamesT::Hotpink             => (255, 105, 180),
494        ColorNamesT::Indianred           => (205, 92,  92 ),
495        ColorNamesT::Indigo              => (75,  0,   130),
496        ColorNamesT::Ivory               => (255, 255, 240),
497        ColorNamesT::Khaki               => (240, 230, 140),
498        ColorNamesT::Lavender            => (230, 230, 250),
499        ColorNamesT::Lavenderblush       => (255, 240, 245),
500        ColorNamesT::Lawngreen           => (124, 252, 0  ),
501        ColorNamesT::Lemonchiffon        => (255, 250, 205),
502        ColorNamesT::Lightblue           => (173, 216, 230),
503        ColorNamesT::Lightcoral          => (240, 128, 128),
504        ColorNamesT::Lightcyan           => (224, 255, 255),
505        ColorNamesT::Lightgoldenrodyellow=> (250, 250, 210),
506        ColorNamesT::Lightgray           => (211, 211, 211),
507        ColorNamesT::Lightgreen          => (144, 238, 144),
508        ColorNamesT::Lightgrey           => (211, 211, 211),
509        ColorNamesT::Lightpink           => (255, 182, 193),
510        ColorNamesT::Lightsalmon         => (255, 160, 122),
511        ColorNamesT::Lightseagreen       => (32,  178, 170),
512        ColorNamesT::Lightskyblue        => (135, 206, 250),
513        ColorNamesT::Lightslategray      => (119, 136, 153),
514        ColorNamesT::Lightslategrey      => (119, 136, 153),
515        ColorNamesT::Lightsteelblue      => (176, 196, 222),
516        ColorNamesT::Lightyellow         => (255, 255, 224),
517        ColorNamesT::Lime                => (0,   255, 0  ),
518        ColorNamesT::Limegreen           => (50,  205, 50 ),
519        ColorNamesT::Linen               => (250, 240, 230),
520        ColorNamesT::Magenta             => (255, 0,   255),
521        ColorNamesT::Maroon              => (128, 0,   0  ),
522        ColorNamesT::Mediumaquamarine    => (102, 205, 170),
523        ColorNamesT::Mediumblue          => (0,   0,   205),
524        ColorNamesT::Mediumorchid        => (186, 85,  211),
525        ColorNamesT::Mediumpurple        => (147, 112, 219),
526        ColorNamesT::Mediumseagreen      => (60,  179, 113),
527        ColorNamesT::Mediumslateblue     => (123, 104, 238),
528        ColorNamesT::Mediumspringgreen   => (0,   250, 154),
529        ColorNamesT::Mediumturquoise     => (72,  209, 204),
530        ColorNamesT::Mediumvioletred     => (199, 21,  133),
531        ColorNamesT::Midnightblue        => (25,  25,  112),
532        ColorNamesT::Mintcream           => (245, 255, 250),
533        ColorNamesT::Mistyrose           => (255, 228, 225),
534        ColorNamesT::Moccasin            => (255, 228, 181),
535        ColorNamesT::Navajowhite         => (255, 222, 173),
536        ColorNamesT::Navy                => (0,   0,   128),
537        ColorNamesT::Oldlace             => (253, 245, 230),
538        ColorNamesT::Olive               => (128, 128, 0  ),
539        ColorNamesT::Olivedrab           => (107, 142, 35 ),
540        ColorNamesT::Orange              => (255, 165, 0  ),
541        ColorNamesT::Orangered           => (255, 69,  0  ),
542        ColorNamesT::Orchid              => (218, 112, 214),
543        ColorNamesT::Palegoldenrod       => (238, 232, 170),
544        ColorNamesT::Palegreen           => (152, 251, 152),
545        ColorNamesT::Paleturquoise       => (175, 238, 238),
546        ColorNamesT::Palevioletred       => (219, 112, 147),
547        ColorNamesT::Papayawhip          => (255, 239, 213),
548        ColorNamesT::Peachpuff           => (255, 218, 185),
549        ColorNamesT::Peru                => (205, 133, 63 ),
550        ColorNamesT::Pink                => (255, 192, 203),
551        ColorNamesT::Plum                => (221, 160, 221),
552        ColorNamesT::Powderblue          => (176, 224, 230),
553        ColorNamesT::Purple              => (128, 0,   128),
554        ColorNamesT::Red                 => (255, 0,   0  ),
555        ColorNamesT::Rosybrown           => (188, 143, 143),
556        ColorNamesT::Royalblue           => (65,  105, 225),
557        ColorNamesT::Saddlebrown         => (139, 69,  19 ),
558        ColorNamesT::Salmon              => (250, 128, 114),
559        ColorNamesT::Sandybrown          => (244, 164, 96 ),
560        ColorNamesT::Seagreen            => (46,  139, 87 ),
561        ColorNamesT::Seashell            => (255, 245, 238),
562        ColorNamesT::Sienna              => (160, 82,  45 ),
563        ColorNamesT::Silver              => (192, 192, 192),
564        ColorNamesT::Skyblue             => (135, 206, 235),
565        ColorNamesT::Slateblue           => (106, 90,  205),
566        ColorNamesT::Slategray           => (112, 128, 144),
567        ColorNamesT::Slategrey           => (112, 128, 144),
568        ColorNamesT::Snow                => (255, 250, 250),
569        ColorNamesT::Springgreen         => (0,   255, 127),
570        ColorNamesT::Steelblue           => (70,  130, 180),
571        ColorNamesT::Tan                 => (210, 180, 140),
572        ColorNamesT::Teal                => (0,   128, 128),
573        ColorNamesT::Thistle             => (216, 191, 216),
574        ColorNamesT::Tomato              => (255, 99,  71 ),
575        ColorNamesT::Turquoise           => (64,  224, 208),
576        ColorNamesT::Violet              => (238, 130, 238),
577        ColorNamesT::Wheat               => (245, 222, 179),
578        ColorNamesT::White               => (255, 255, 255),
579        ColorNamesT::Whitesmoke          => (245, 245, 245),
580        ColorNamesT::Yellow              => (255, 255, 0  ),
581        ColorNamesT::Yellowgreen         => (154, 205, 50 ),
582    }
583}
584
585/// CPK element colors.
586#[rustfmt::skip]
587pub fn element_color(element: &ferritin_core::info::elements::Element) -> Srgba {
588    use ferritin_core::info::elements::Element;
589    match element {
590        Element::H  => Srgba::rgb(1.0, 1.0, 1.0),
591        Element::C  => Srgba::rgb(0.5, 0.5, 0.5),
592        Element::N  => Srgba::rgb(0.0, 0.0, 1.0),
593        Element::O  => Srgba::rgb(1.0, 0.0, 0.0),
594        Element::P  => Srgba::rgb(1.0, 0.5, 0.0),
595        Element::S  => Srgba::rgb(1.0, 1.0, 0.0),
596        Element::Cl => Srgba::rgb(0.0, 1.0, 0.0),
597        Element::Fe => Srgba::rgb(0.6, 0.0, 0.0),
598        Element::Ca => Srgba::rgb(0.5, 0.5, 0.5),
599        Element::Mg => Srgba::rgb(0.5, 1.0, 0.0),
600        Element::Na => Srgba::rgb(0.0, 0.0, 1.0),
601        Element::K  => Srgba::rgb(0.8, 0.6, 1.0),
602        Element::Zn => Srgba::rgb(0.6, 0.6, 0.6),
603        Element::Cu => Srgba::rgb(0.8, 0.4, 0.0),
604        Element::F  => Srgba::rgb(0.7, 1.0, 1.0),
605        Element::Br => Srgba::rgb(0.6, 0.1, 0.1),
606        Element::I  => Srgba::rgb(0.4, 0.0, 0.7),
607        Element::B  => Srgba::rgb(1.0, 0.7, 0.7),
608        Element::Se => Srgba::rgb(1.0, 0.5, 0.0),
609        _           => Srgba::rgb(0.5, 0.5, 0.5),
610    }
611}
612
613#[cfg(test)]
614mod tests {
615    use super::*;
616    use bevy::asset::RenderAssetUsages;
617    use bevy::prelude::Mesh;
618    use bevy::render::mesh::{PrimitiveTopology, VertexAttributeValues};
619    use ferritin_core::AtomCollection;
620    use ferritin_core::info::elements::Element;
621    use ferritin_molviewspec::molviewspec::nodes::{
622        ColorNamesT, ColorT, ComponentExpression, ComponentSelector, ComponentSelectorT,
623    };
624
625    /// 9-atom mixed collection for selector/color tests.
626    /// Atoms: ALA(3, A/1), GLY(2, A/2), HOH(1, B/3 hetero), ZN(1, B/4 hetero), LIG(2, B/5 hetero)
627    fn make_test_collection() -> AtomCollection {
628        let coords = vec![[0.0f32; 3]; 9];
629        let res_ids: Vec<i32> = vec![1, 1, 1, 2, 2, 3, 4, 5, 5];
630        let res_names: Vec<String> = ["ALA", "ALA", "ALA", "GLY", "GLY", "HOH", "ZN", "LIG", "LIG"]
631            .iter()
632            .map(|s| s.to_string())
633            .collect();
634        let is_hetero = vec![false, false, false, false, false, true, true, true, true];
635        let elements = vec![
636            Element::N,
637            Element::C,
638            Element::C,
639            Element::N,
640            Element::C,
641            Element::O,
642            Element::Zn,
643            Element::C,
644            Element::O,
645        ];
646        let atom_names: Vec<String> = ["N", "CA", "C", "N", "CA", "O", "ZN", "C1", "O1"]
647            .iter()
648            .map(|s| s.to_string())
649            .collect();
650        let chain_ids: Vec<String> = ["A", "A", "A", "A", "A", "B", "B", "B", "B"]
651            .iter()
652            .map(|s| s.to_string())
653            .collect();
654        AtomCollection::new(
655            9, coords, res_ids, res_names, is_hetero, elements, atom_names, chain_ids, None,
656        )
657    }
658
659    /// Build a minimal Mesh with `n_verts` vertices and a vertex_map [0, 0, 1, 1, ..].
660    fn make_dummy_mesh_and_map(n_verts: usize, n_atoms: usize) -> (Mesh, Vec<usize>) {
661        let mesh = Mesh::new(PrimitiveTopology::TriangleList, RenderAssetUsages::all());
662        // 2 vertices per atom, cycling through atom indices
663        let map = (0..n_verts).map(|v| (v / 2).min(n_atoms - 1)).collect();
664        (mesh, map)
665    }
666
667    fn get_colors(mesh: &Mesh) -> Option<Vec<[f32; 4]>> {
668        match mesh.attribute(Mesh::ATTRIBUTE_COLOR)? {
669            VertexAttributeValues::Float32x4(c) => Some(c.clone()),
670            _ => None,
671        }
672    }
673
674    // T2-41: Single color node (All, hex) → all vertices get that color
675    #[test]
676    fn test_apply_single_color_all() {
677        let ac = make_test_collection();
678        let (mut mesh, map) = make_dummy_mesh_and_map(18, 9);
679        let color_nodes = vec![(
680            ComponentSelector::Selector(ComponentSelectorT::All),
681            ColorT::Hex("#4b7fcc".to_string()),
682        )];
683        apply_mvs_colors(&mut mesh, &color_nodes, &ac, &map, VertexMapKind::AtomMap);
684        let colors = get_colors(&mesh).expect("COLOR attribute must be set");
685        assert_eq!(colors.len(), 18);
686        let expected_r = 0x4b as f32 / 255.0;
687        let expected_g = 0x7f as f32 / 255.0;
688        let expected_b = 0xcc as f32 / 255.0;
689        for c in &colors {
690            assert!((c[0] - expected_r).abs() < 1e-4, "red mismatch");
691            assert!((c[1] - expected_g).abs() < 1e-4, "green mismatch");
692            assert!((c[2] - expected_b).abs() < 1e-4, "blue mismatch");
693        }
694    }
695
696    // T2-42: Two color nodes; later (expression) overrides earlier (All)
697    #[test]
698    fn test_apply_two_colors_base_then_override() {
699        let ac = make_test_collection();
700        // 2 verts per atom, 9 atoms = 18 verts; map: v/2 → atom_idx
701        let (mut mesh, map) = make_dummy_mesh_and_map(18, 9);
702        // ALA residue = atoms 0,1,2 → verts 0..5
703        let color_nodes = vec![
704            (
705                ComponentSelector::Selector(ComponentSelectorT::All),
706                ColorT::Named(ColorNamesT::Blue),
707            ),
708            (
709                ComponentSelector::Expression(ComponentExpression {
710                    auth_seq_id: Some(1), // ALA residue (atoms 0,1,2)
711                    ..Default::default()
712                }),
713                ColorT::Named(ColorNamesT::Red),
714            ),
715        ];
716        apply_mvs_colors(&mut mesh, &color_nodes, &ac, &map, VertexMapKind::AtomMap);
717        let colors = get_colors(&mesh).unwrap();
718        // First 6 verts (atoms 0,1,2 = ALA/res 1) should be red
719        for c in &colors[0..6] {
720            assert!(c[0] > 0.99, "expected red for ALA res 1 verts");
721            assert!(c[2] < 0.01, "expected no blue for ALA res 1 verts");
722        }
723        // Remaining verts should be blue
724        for c in &colors[6..] {
725            assert!(c[2] > 0.99, "expected blue for non-res-1 verts");
726            assert!(c[0] < 0.01, "expected no red for non-res-1 verts");
727        }
728    }
729
730    // T2-43: Order matters — (expression, red) then (All, blue) → all blue
731    #[test]
732    fn test_apply_colors_order_matters_reversed() {
733        let ac = make_test_collection();
734        let (mut mesh, map) = make_dummy_mesh_and_map(18, 9);
735        let color_nodes = vec![
736            (
737                ComponentSelector::Expression(ComponentExpression {
738                    auth_seq_id: Some(1),
739                    ..Default::default()
740                }),
741                ColorT::Named(ColorNamesT::Red),
742            ),
743            (
744                ComponentSelector::Selector(ComponentSelectorT::All),
745                ColorT::Named(ColorNamesT::Blue),
746            ),
747        ];
748        apply_mvs_colors(&mut mesh, &color_nodes, &ac, &map, VertexMapKind::AtomMap);
749        let colors = get_colors(&mesh).unwrap();
750        for (i, c) in colors.iter().enumerate() {
751            assert!(
752                c[2] > 0.99,
753                "vert {i}: expected all-blue after All override"
754            );
755            assert!(c[0] < 0.01, "vert {i}: expected no red");
756        }
757    }
758
759    // T2-44: Middle non-matching selector is a no-op; last All wins
760    #[test]
761    fn test_apply_colors_middle_no_match() {
762        let ac = make_test_collection();
763        let (mut mesh, map) = make_dummy_mesh_and_map(18, 9);
764        let color_nodes = vec![
765            (
766                ComponentSelector::Selector(ComponentSelectorT::All),
767                ColorT::Named(ColorNamesT::Blue),
768            ),
769            (
770                ComponentSelector::Expression(ComponentExpression {
771                    auth_asym_id: Some("Z".to_string()), // nonexistent chain
772                    ..Default::default()
773                }),
774                ColorT::Named(ColorNamesT::Red),
775            ),
776            (
777                ComponentSelector::Selector(ComponentSelectorT::All),
778                ColorT::Named(ColorNamesT::Green),
779            ),
780        ];
781        apply_mvs_colors(&mut mesh, &color_nodes, &ac, &map, VertexMapKind::AtomMap);
782        let colors = get_colors(&mesh).unwrap();
783        for (i, c) in colors.iter().enumerate() {
784            // Green = (0, 128, 0) / 255 ≈ (0, 0.502, 0)
785            assert!(c[0] < 0.01, "vert {i}: expected no red");
786            assert!(c[1] > 0.49, "vert {i}: expected green channel");
787            assert!(c[2] < 0.01, "vert {i}: expected no blue");
788        }
789    }
790
791    // T2-45: No color nodes → white fallback on all vertices
792    #[test]
793    fn test_apply_colors_no_color_nodes() {
794        let ac = make_test_collection();
795        let (mut mesh, map) = make_dummy_mesh_and_map(6, 9);
796        apply_mvs_colors(&mut mesh, &[], &ac, &map, VertexMapKind::AtomMap);
797        let colors = get_colors(&mesh).unwrap();
798        for (i, c) in colors.iter().enumerate() {
799            assert!((c[0] - 1.0).abs() < 1e-4, "vert {i}: expected white (red)");
800            assert!(
801                (c[1] - 1.0).abs() < 1e-4,
802                "vert {i}: expected white (green)"
803            );
804            assert!((c[2] - 1.0).abs() < 1e-4, "vert {i}: expected white (blue)");
805        }
806    }
807
808    // T2-46: apply_mvs_colors never changes vertex count
809    #[test]
810    fn test_apply_colors_preserves_vertex_count() {
811        let ac = make_test_collection();
812        let (mut mesh, map) = make_dummy_mesh_and_map(18, 9);
813        let before = map.len();
814        let color_nodes = vec![(
815            ComponentSelector::Selector(ComponentSelectorT::All),
816            ColorT::Named(ColorNamesT::Blue),
817        )];
818        apply_mvs_colors(&mut mesh, &color_nodes, &ac, &map, VertexMapKind::AtomMap);
819        let colors = get_colors(&mesh).unwrap();
820        assert_eq!(colors.len(), before, "vertex count must not change");
821    }
822
823    // T2-48: hex lowercase
824    #[test]
825    fn test_parse_hex_lowercase() {
826        let color = parse_hex_color("#1b9e77");
827        let srgb = color.to_srgba();
828        assert!((srgb.red - 0x1b as f32 / 255.0).abs() < 1e-4);
829        assert!((srgb.green - 0x9e as f32 / 255.0).abs() < 1e-4);
830        assert!((srgb.blue - 0x77 as f32 / 255.0).abs() < 1e-4);
831    }
832
833    // T2-42: hex uppercase
834    #[test]
835    fn test_parse_hex_uppercase() {
836        let c1 = parse_hex_color("#FF0000");
837        let c2 = parse_hex_color("#ff0000");
838        let s1 = c1.to_srgba();
839        let s2 = c2.to_srgba();
840        assert!((s1.red - s2.red).abs() < 1e-6);
841    }
842
843    // T2-43: hex without leading #
844    #[test]
845    fn test_parse_hex_no_hash() {
846        let color = parse_hex_color("0000ff");
847        let srgb = color.to_srgba();
848        assert!(srgb.blue > 0.99);
849        assert!(srgb.red < 0.01);
850    }
851
852    // T2-44: invalid hex → white
853    #[test]
854    fn test_parse_hex_invalid_white_fallback() {
855        let color = parse_hex_color("not-a-color");
856        let srgb = color.to_srgba();
857        assert!((srgb.red - 1.0).abs() < 1e-4);
858        assert!((srgb.green - 1.0).abs() < 1e-4);
859        assert!((srgb.blue - 1.0).abs() < 1e-4);
860    }
861
862    // T2-45: named blue
863    #[test]
864    fn test_named_blue() {
865        let color = color_t_to_bevy(&ColorT::Named(ColorNamesT::Blue));
866        let srgb = color.to_srgba();
867        assert!(srgb.blue > 0.99);
868        assert!(srgb.red < 0.01);
869        assert!(srgb.green < 0.01);
870    }
871
872    // T2-46: named red
873    #[test]
874    fn test_named_red() {
875        let color = color_t_to_bevy(&ColorT::Named(ColorNamesT::Red));
876        let srgb = color.to_srgba();
877        assert!(srgb.red > 0.99);
878        assert!(srgb.blue < 0.01);
879    }
880
881    // T2-47: named black
882    #[test]
883    fn test_named_black() {
884        let color = color_t_to_bevy(&ColorT::Named(ColorNamesT::Black));
885        let srgb = color.to_srgba();
886        assert!(srgb.red < 0.01);
887        assert!(srgb.green < 0.01);
888        assert!(srgb.blue < 0.01);
889    }
890
891    // T2-48: named white
892    #[test]
893    fn test_named_white() {
894        let color = color_t_to_bevy(&ColorT::Named(ColorNamesT::White));
895        let srgb = color.to_srgba();
896        assert!((srgb.red - 1.0).abs() < 1e-4);
897        assert!((srgb.green - 1.0).abs() < 1e-4);
898        assert!((srgb.blue - 1.0).abs() < 1e-4);
899    }
900
901    // T2-49: ColorScheme::MvsColor round-trip
902    #[test]
903    fn test_color_scheme_mvscolor() {
904        let scheme = ColorScheme::MvsColor(ColorT::Hex("#ff0000".to_string()));
905        let color = scheme.get_color(&ferritin_core::info::elements::Element::C);
906        let srgb = color.to_srgba();
907        assert!(srgb.red > 0.99);
908        assert!(srgb.blue < 0.01);
909    }
910
911    // T2-50: ElementSymbol theme gives CPK colors — carbon gray, oxygen red, nitrogen blue
912    #[test]
913    fn test_element_symbol_theme_cpk() {
914        let ac = make_test_collection();
915        // vertex_map: one vertex per atom, AtomMap
916        let vertex_map: Vec<usize> = (0..9).collect();
917        let colors = initial_colors_from_theme(
918            &ColorThemeT::ElementSymbol,
919            &ac,
920            &vertex_map,
921            &VertexMapKind::AtomMap,
922        );
923        assert_eq!(colors.len(), 9);
924        // ALA atoms 0-2 are all N,C,C — N is blue-ish (high blue), C is gray
925        // Just verify colors differ from white (1,1,1) meaning the theme had an effect.
926        let has_non_white = colors.iter().any(|c| {
927            (c[0] - 1.0).abs() > 0.05 || (c[1] - 1.0).abs() > 0.05 || (c[2] - 1.0).abs() > 0.05
928        });
929        assert!(
930            has_non_white,
931            "ElementSymbol theme should produce non-white colors"
932        );
933    }
934
935    // T2-51: ChainId theme gives distinct colors per chain
936    #[test]
937    fn test_chain_id_theme_distinct_per_chain() {
938        let ac = make_test_collection();
939        // atoms 0-4 are chain A, atoms 5-8 are chain B
940        let vertex_map: Vec<usize> = (0..9).collect();
941        let colors = initial_colors_from_theme(
942            &ColorThemeT::ChainId,
943            &ac,
944            &vertex_map,
945            &VertexMapKind::AtomMap,
946        );
947        // All chain-A atoms should share the same color
948        assert_eq!(colors[0], colors[1], "same chain → same color");
949        assert_eq!(colors[0], colors[4], "same chain → same color");
950        // Chain A and chain B should differ
951        assert_ne!(colors[0], colors[5], "different chains → different colors");
952    }
953
954    // T2-52: apply_mvs_colors_with_theme — theme base + color node override
955    #[test]
956    fn test_theme_with_color_node_override() {
957        let ac = make_test_collection();
958        let (mut mesh, map) = make_dummy_mesh_and_map(18, 9);
959        // Override ALL with red on top of ElementSymbol theme
960        let color_nodes = vec![(
961            ComponentSelector::Selector(ComponentSelectorT::All),
962            ColorT::Hex("#ff0000".to_string()),
963        )];
964        apply_mvs_colors_with_theme(
965            &mut mesh,
966            &ColorThemeT::ElementSymbol,
967            &color_nodes,
968            &ac,
969            &map,
970            VertexMapKind::AtomMap,
971        );
972        let colors = get_colors(&mesh).expect("COLOR attribute must be set");
973        // All vertices overridden to red
974        for c in &colors {
975            assert!(c[0] > 0.99, "expected red");
976            assert!(c[1] < 0.01, "expected no green");
977            assert!(c[2] < 0.01, "expected no blue");
978        }
979    }
980
981    // T2-53: Uniform theme → all white
982    #[test]
983    fn test_uniform_theme_is_white() {
984        let ac = make_test_collection();
985        let vertex_map: Vec<usize> = (0..9).collect();
986        let colors = initial_colors_from_theme(
987            &ColorThemeT::Uniform,
988            &ac,
989            &vertex_map,
990            &VertexMapKind::AtomMap,
991        );
992        for c in colors {
993            assert!((c[0] - 1.0).abs() < 1e-6);
994            assert!((c[1] - 1.0).abs() < 1e-6);
995            assert!((c[2] - 1.0).abs() < 1e-6);
996        }
997    }
998
999    // ferritin-ujk: data-driven themes fall back to uniform white while AtomCollection
1000    // has no per-atom b_factor/occupancy data (tracked separately from the Model layer).
1001    #[test]
1002    fn test_data_driven_themes_fall_back_to_uniform() {
1003        let ac = make_test_collection();
1004        let vertex_map: Vec<usize> = (0..9).collect();
1005        for theme in [
1006            ColorThemeT::Uncertainty,
1007            ColorThemeT::PlddtConfidence,
1008            ColorThemeT::Occupancy,
1009        ] {
1010            let colors =
1011                initial_colors_from_theme(&theme, &ac, &vertex_map, &VertexMapKind::AtomMap);
1012            assert_eq!(colors.len(), 9);
1013            for c in colors {
1014                assert!(
1015                    (c[0] - 1.0).abs() < 1e-6,
1016                    "{:?} should fall back to white",
1017                    theme
1018                );
1019                assert!(
1020                    (c[1] - 1.0).abs() < 1e-6,
1021                    "{:?} should fall back to white",
1022                    theme
1023                );
1024                assert!(
1025                    (c[2] - 1.0).abs() < 1e-6,
1026                    "{:?} should fall back to white",
1027                    theme
1028                );
1029            }
1030        }
1031    }
1032
1033    // ferritin-clv: once b_factor is attached, Uncertainty maps the observed
1034    // range to a blue-white-red gradient instead of falling back to white.
1035    #[test]
1036    fn test_uncertainty_theme_uses_real_gradient() {
1037        let ac = make_test_collection()
1038            .with_b_factor(vec![0.0, 12.5, 25.0, 37.5, 50.0, 62.5, 75.0, 87.5, 100.0]);
1039        let vertex_map: Vec<usize> = (0..9).collect();
1040        let colors = initial_colors_from_theme(
1041            &ColorThemeT::Uncertainty,
1042            &ac,
1043            &vertex_map,
1044            &VertexMapKind::AtomMap,
1045        );
1046        // Lowest value -> blue-heavy, highest value -> red-heavy, not uniform white.
1047        let low = colors[0];
1048        let high = colors[8];
1049        assert!(
1050            low[2] > low[0],
1051            "lowest b_factor should read blue-leaning: {:?}",
1052            low
1053        );
1054        assert!(
1055            high[0] > high[2],
1056            "highest b_factor should read red-leaning: {:?}",
1057            high
1058        );
1059        assert_ne!(low, high);
1060    }
1061
1062    // ferritin-clv: PlddtConfidence buckets values into AlphaFold-style bands.
1063    #[test]
1064    fn test_plddt_confidence_theme_bands_values() {
1065        let ac = make_test_collection()
1066            .with_b_factor(vec![95.0, 95.0, 95.0, 80.0, 80.0, 60.0, 60.0, 30.0, 30.0]);
1067        let vertex_map: Vec<usize> = (0..9).collect();
1068        let colors = initial_colors_from_theme(
1069            &ColorThemeT::PlddtConfidence,
1070            &ac,
1071            &vertex_map,
1072            &VertexMapKind::AtomMap,
1073        );
1074        // Very-high (>90) and very-low (<50) bands must render as distinct colors.
1075        assert_ne!(colors[0], colors[7]);
1076        // Same band -> identical color.
1077        assert_eq!(colors[0], colors[1]);
1078        assert_eq!(colors[7], colors[8]);
1079    }
1080
1081    // ferritin-clv: Occupancy maps partial occupancy toward red, full occupancy toward white.
1082    #[test]
1083    fn test_occupancy_theme_uses_real_gradient() {
1084        let ac = make_test_collection()
1085            .with_occupancy(vec![1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.5, 0.2, 0.2]);
1086        let vertex_map: Vec<usize> = (0..9).collect();
1087        let colors = initial_colors_from_theme(
1088            &ColorThemeT::Occupancy,
1089            &ac,
1090            &vertex_map,
1091            &VertexMapKind::AtomMap,
1092        );
1093        assert_eq!(
1094            colors[0],
1095            Vec4::new(1.0, 1.0, 1.0, 1.0),
1096            "full occupancy is white"
1097        );
1098        assert!(
1099            colors[7][1] < colors[0][1],
1100            "partial occupancy should read less white/more red"
1101        );
1102    }
1103}
ferritin_bevy/colors.rs

ferritin_bevy/
colors.rs
