ak_vis/color_palette/
palette.rs

1use ak_core::StructureView;
2use ak_core::geometry::AtomicNumber;
3use bevy::color::Color;
4
5#[derive(Clone, Copy, Debug, PartialEq)]
6pub struct AtomMaterial {
7    pub metallic: f32,
8    pub perceptual_roughness: f32,
9}
10
11impl AtomMaterial {
12    pub const fn new(metallic: f32, perceptual_roughness: f32) -> Self {
13        Self {
14            metallic,
15            perceptual_roughness,
16        }
17    }
18}
19
20pub const DEFAULT_ATOM_MATERIAL: AtomMaterial = AtomMaterial::new(0.0, 0.4);
21
22impl Default for AtomMaterial {
23    fn default() -> Self {
24        DEFAULT_ATOM_MATERIAL
25    }
26}
27
28#[derive(Clone, Copy, Debug, PartialEq)]
29pub struct ColorPalette {
30    pub colors: [Option<Color>; 110],
31    pub materials: [Option<AtomMaterial>; 110],
32    pub fallback: Color,
33    pub fallback_material: AtomMaterial,
34}
35
36pub trait ColorScheme {
37    fn color(&self, view: &StructureView, i: usize) -> Color;
38}
39
40impl ColorPalette {
41    pub fn get(&self, atomic_number: AtomicNumber) -> Color {
42        match self.colors[atomic_number.get() as usize] {
43            Some(color) => color,
44            None => self.fallback,
45        }
46    }
47
48    pub fn material(&self, atomic_number: AtomicNumber) -> AtomMaterial {
49        match self.materials[atomic_number.get() as usize] {
50            Some(material) => material,
51            None => self.fallback_material,
52        }
53    }
54}
55
56impl ColorScheme for ColorPalette {
57    fn color(&self, view: &StructureView, i: usize) -> Color {
58        self.get(view.numbers[i])
59    }
60}