ak_vis/render/
render_atoms.rs

1use crate::AtomVisual;
2use crate::components::DisplayAtomIdentity;
3use crate::components::FrameAtom;
4use bevy::prelude::*;
5use std::collections::HashMap;
6
7pub fn render_atoms(
8    visuals: Vec<AtomVisual>,
9    commands: &mut Commands,
10    materials: &mut ResMut<Assets<StandardMaterial>>,
11    meshes: &mut ResMut<Assets<Mesh>>,
12    quality: u32,
13) {
14    // Group atoms by (radius, color, material) to enable instancing
15    // Using bits representation for radius as HashMap key (f32 doesn't implement Hash)
16    let mut atom_groups: HashMap<(u32, [u8; 4], [u32; 2]), Vec<(DisplayAtomIdentity, Vec3)>> =
17        HashMap::new();
18
19    for atom in visuals.iter() {
20        let radius_bits = atom.radius.to_bits();
21        // Use color as bytes for HashMap key
22        let color_key = [
23            (atom.color.to_srgba().red * 255.0) as u8,
24            (atom.color.to_srgba().green * 255.0) as u8,
25            (atom.color.to_srgba().blue * 255.0) as u8,
26            (atom.color.to_srgba().alpha * 255.0) as u8,
27        ];
28
29        let material_key = [
30            atom.material.metallic.to_bits(),
31            atom.material.perceptual_roughness.to_bits(),
32        ];
33
34        let key = (radius_bits, color_key, material_key);
35        atom_groups.entry(key).or_default().push((
36            DisplayAtomIdentity {
37                atom_index: atom.atom_identity.atom_index,
38                image_offset: atom.atom_identity.image_offset,
39            },
40            Vec3::new(atom.x(), atom.y(), atom.z()),
41        ));
42    }
43
44    // Create shared meshes and materials for each unique atom type
45    for ((radius_bits, color_bytes, material_bits), positions) in atom_groups.iter() {
46        let radius = f32::from_bits(*radius_bits);
47        let color = Color::srgba(
48            color_bytes[0] as f32 / 255.0,
49            color_bytes[1] as f32 / 255.0,
50            color_bytes[2] as f32 / 255.0,
51            color_bytes[3] as f32 / 255.0,
52        );
53        let metallic = f32::from_bits(material_bits[0]);
54        let perceptual_roughness = f32::from_bits(material_bits[1]);
55
56        // Create ONE mesh for this atom type (shared by all instances)
57        let shared_mesh = meshes.add(
58            Sphere::new(radius)
59                .mesh()
60                .ico(quality)
61                .expect("Failed to create sphere mesh"),
62        );
63
64        // Create ONE material for this atom type (shared by all instances)
65        let shared_material = materials.add(StandardMaterial {
66            base_color: color,
67            metallic,
68            perceptual_roughness,
69            ..default()
70        });
71
72        // Spawn all atoms of this type - they share the same mesh handle,
73        // so Bevy will automatically instance them in a single draw call
74        for (atom_identity, position) in positions {
75            commands.spawn((
76                Mesh3d(shared_mesh.clone()),
77                MeshMaterial3d(shared_material.clone()),
78                Transform::from_translation(*position),
79                FrameAtom,
80                atom_identity.clone(),
81            ));
82        }
83    }
84}