Initialize Ball and Sticks (#7)

* update bevy example for atom collection

* begin port to atom colleciton

* get spheres with color by element

* pass tests

* simplify running of test

* remove unused code

* get code for bond cylinders in place.

* add bond functionality

ferritin-bevy/Cargo.toml

@@ -7,9 +7,9 @@ license.workspace = true
 description.workspace = true
 
 [dependencies]
+bevy = "0.14"
+bon = "2.3.0"
 serde.workspace = true
 ferritin-core = { path = "../ferritin-core" }
+itertools.workspace = true
 pdbtbx.workspace = true
-bevy = "0.14.2"                               # Specific to bevy crate
-bon = "2.3.0"                                 # Specific to bevy crate
-glam = "0.29.0"                               # Specific to bevy crate

ferritin-bevy/Readme.md

@@ -4,7 +4,8 @@
 ## Examples
 
 ```sh
-cargo run --example bevy_protein_mesh_custom
+cargo run --example basic_spheres
+cargo run --example basic_ballandstick
 ```
 
 ![](docs/images/protein_01.png)

ferritin-bevy/examples/bevy_protein_mesh_custom.rs → ferritin-bevy/examples/basic_spheres.rs

@@ -41,13 +41,6 @@ fn main() {
             }),
         ))
         .add_systems(Startup, setup)
-        // .add_systems(
-        //     Update,
-        //     (
-        //         update_protein_meshes,
-        //         focus_camera_on_proteins,
-        //     ),
-        // )
         .run();
 }
 
@@ -129,59 +122,3 @@ fn setup(mut commands: Commands) {
         ..default()
     });
 }
-
-// fn update_protein_meshes(
-//     mut commands: Commands,
-//     mut meshes: ResMut<Assets<Mesh>>,
-//     mut materials: ResMut<Assets<StandardMaterial>>,
-//     query: Query<(Entity, &Structure), (With<Structure>, Without<Handle<Mesh>>)>,
-// ) {
-//     println!("I'm in the update_protein_mesh function!");
-//     println!("Total entities with Structure: {}", query.iter().count());
-//     for (entity, protein) in query.iter() {
-//         println!("Working on {:?}", entity);
-//         let mesh = protein.to_mesh();
-//         let mesh_handle = meshes.add(mesh);
-//         let material = materials.add(StandardMaterial {
-//             base_color: Color::srgb(0.8, 0.7, 0.6),
-//             metallic: 0.1,
-//             perceptual_roughness: 0.5,
-//             ..default()
-//         });
-
-//         commands.entity(entity).insert(PbrBundle {
-//             mesh: mesh_handle,
-//             material,
-//             ..default()
-//         });
-//     }
-// }
-
-// fn focus_camera_on_proteins(
-//     mut camera_query: Query<&mut Transform, (With<Camera>, Without<Structure>)>,
-//     protein_query: Query<&Transform, With<Structure>>,
-// ) {
-//     if let Ok(mut camera_transform) = camera_query.get_single_mut() {
-//         if let Some(center) = calculate_center_of_proteins(&protein_query) {
-//             let camera_position = center + Vec3::new(0.0, 50.0, 100.0);
-//             camera_transform.translation = camera_position;
-//             camera_transform.look_at(center, Vec3::Y);
-//         }
-//     }
-// }
-
-// fn calculate_center_of_proteins(
-//     protein_query: &Query<&Transform, With<Structure>>,
-// ) -> Option<Vec3> {
-//     let mut total = Vec3::ZERO;
-//     let mut count = 0;
-//     for transform in protein_query.iter() {
-//         total += transform.translation;
-//         count += 1;
-//     }
-//     if count > 0 {
-//         Some(total / count as f32)
-//     } else {
-//         None
-//     }
-// }

ferritin-bevy/src/colors.rs

@@ -2,12 +2,11 @@
 //!
 //! This module defines the color mapping used for rendering.
 use bevy::prelude::Color;
-use pdbtbx::Atom;
 
 /// Represents different color schemes for rendering atoms.
 #[derive(Clone)]
 pub enum ColorScheme {
-    /// A solid, sinel color for all atoms.
+    /// A solid, single color for all atoms.
     Solid(Color),
     /// Colors atoms based on their element type.
     ByAtomType,
@@ -26,11 +25,11 @@ pub enum ColorScheme {
 // ColorScheme::ByResidueType(func) => func(residue),
 // ColorScheme::Custom(func) => func(atom, residue, chain),
 impl ColorScheme {
-    pub fn get_color(&self, atom: &Atom) -> Color {
+    pub fn get_color(&self, atom: &str) -> Color {
         match &self {
             ColorScheme::Solid(color) => *color,
             ColorScheme::ByAtomType => {
-                match atom.element().expect("expect atom").symbol() {
+                match atom {
                     "C" => Color::srgb(0.5, 0.5, 0.5), // Carbon: Gray
                     "N" => Color::srgb(0.0, 0.0, 1.0), // Nitrogen: Blue
                     "O" => Color::srgb(1.0, 0.0, 0.0), // Oxygen: Red
@@ -45,38 +44,12 @@ impl ColorScheme {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use pdbtbx::Atom;
-
     #[test]
     fn test_get_color() {
         let by_atom_scheme = ColorScheme::ByAtomType;
-        let create_atom =
-            |element: &str| Atom::new(true, 1, "", 0.0, 0.0, 0.0, 0.0, 0.0, element, 1).unwrap();
-        let carbon_atom = create_atom("C");
-        let nitrogen_atom = create_atom("N");
-        let oxygen_atom = create_atom("O");
-        let sulfur_atom = create_atom("S");
-        // let other_atom = create_atom("X");
-        // Test ByAtomType color scheme
-        assert_eq!(
-            by_atom_scheme.get_color(&carbon_atom),
-            Color::srgb(0.5, 0.5, 0.5)
-        );
-        assert_eq!(
-            by_atom_scheme.get_color(&nitrogen_atom),
-            Color::srgb(0.0, 0.0, 1.0)
-        );
-        assert_eq!(
-            by_atom_scheme.get_color(&oxygen_atom),
-            Color::srgb(1.0, 0.0, 0.0)
-        );
-        assert_eq!(
-            by_atom_scheme.get_color(&sulfur_atom),
-            Color::srgb(1.0, 1.0, 0.0)
-        );
-        // assert_eq!(
-        //     by_atom_scheme.get_color(&other_atom),
-        //     Color::srgb(1.0, 1.0, 1.0)
-        // );
+        assert_eq!(by_atom_scheme.get_color("C"), Color::srgb(0.5, 0.5, 0.5));
+        assert_eq!(by_atom_scheme.get_color("N"), Color::srgb(0.0, 0.0, 1.0));
+        assert_eq!(by_atom_scheme.get_color("O"), Color::srgb(1.0, 0.0, 0.0));
+        assert_eq!(by_atom_scheme.get_color("S"), Color::srgb(1.0, 1.0, 0.0));
     }
 }

ferritin-bevy/src/plugin.rs

@@ -4,6 +4,7 @@
 use super::{ColorScheme, RenderOptions, Structure};
 use bevy::prelude::*;
 // use pdbtbx::StrictnessLevel;
+use ferritin_core::AtomCollection;
 use std::path::Path;
 use std::path::PathBuf;
 
@@ -78,7 +79,7 @@ fn load_initial_proteins(
             // StrictnessLevel::Medium,
         ) {
             let structure = Structure::builder()
-                .pdb(pdb)
+                .pdb(AtomCollection::from(&pdb))
                 .rendertype(settings.render_type.clone())
                 .color_scheme(settings.color_scheme.clone())
                 .material(settings.material.clone())

ferritin-bevy/src/structure.rs

@@ -7,12 +7,14 @@
 // use bevy::prelude::*;
 use super::ColorScheme;
 use bevy::asset::Assets;
+use bevy::math::Vec4;
 use bevy::prelude::{
-    default, Color, ColorToComponents, Component, Mesh, MeshBuilder, Meshable, PbrBundle, Sphere,
-    StandardMaterial, Vec3,
+    default, Color, Component, Cylinder, Mesh, MeshBuilder, Meshable, PbrBundle, Quat, Sphere,
+    StandardMaterial, Transform, Vec3,
 };
 use bon::Builder;
-use pdbtbx::PDB;
+use ferritin_core::{AtomCollection, Bond};
+use pdbtbx::Element;
 
 /// Enum representing various rendering options.
 ///
@@ -31,7 +33,7 @@ pub enum RenderOptions {
 /// Define Everything Needed to render
 #[derive(Builder, Component)]
 pub struct Structure {
-    pdb: PDB,
+    pdb: AtomCollection,
     #[builder(default = RenderOptions::Solid)]
     rendertype: RenderOptions,
     #[builder(default = ColorScheme::Solid(Color::WHITE))]
@@ -71,32 +73,97 @@ impl Structure {
         todo!()
     }
     fn render_ballandstick(&self) -> Mesh {
-        todo!()
+        let radius = 0.5;
+        let mut combined_mesh = self
+            .pdb
+            .iter_coords_and_elements()
+            .map(|(coord, element_str)| {
+                let center = Vec3::new(coord[0], coord[1], coord[2]);
+                let mut sphere_mesh = Sphere::new(radius).mesh().build();
+                let vertex_count = sphere_mesh.count_vertices();
+                // let element = Element::from_symbol(element_str).expect("Element not recognized");
+                let color = self.color_scheme.get_color(element_str).to_srgba();
+                let color_array =
+                    vec![Vec4::new(color.red, color.green, color.blue, color.alpha); vertex_count];
+                sphere_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, color_array);
+                sphere_mesh = sphere_mesh.translated_by(center);
+                sphere_mesh.compute_smooth_normals();
+                sphere_mesh
+            })
+            .reduce(|mut acc, mesh| {
+                acc.merge(&mesh);
+                acc
+            })
+            .unwrap();
+
+        // Add bond cylinders using iterators
+        if let Some(bonds) = self.pdb.bonds() {
+            let coords = self.pdb.coords();
+            bonds
+                .iter()
+                .filter_map(|bond| {
+                    let (atom1, atom2) = bond.get_atom_indices();
+                    let pos1 = Vec3::from_array(*coords.get(atom1 as usize)?);
+                    let pos2 = Vec3::from_array(*coords.get(atom2 as usize)?);
+
+                    // Calculate cylinder properties
+                    let center = (pos1 + pos2) / 2.0;
+                    let direction = pos2 - pos1;
+                    let height = direction.length();
+                    let rotation = Quat::from_rotation_arc(Vec3::Y, direction.normalize());
+
+                    // Create and transform cylinder mesh
+                    let mut cylinder_mesh = Cylinder {
+                        radius: 0.5,
+                        half_height: height / 2.0, // Note: we divide height by 2 since it expects half_height
+                    }
+                    .mesh()
+                    .build();
+
+                    // Apply transformation
+                    cylinder_mesh = cylinder_mesh.transformed_by(Transform {
+                        translation: center,
+                        rotation,
+                        ..default()
+                    });
+
+                    // Add colors
+                    let cylinder_vertex_count = cylinder_mesh.count_vertices();
+                    let cylinder_colors =
+                        vec![Vec4::new(0.5, 0.5, 0.5, 0.5); cylinder_vertex_count];
+                    cylinder_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, cylinder_colors);
+                    Some(cylinder_mesh)
+                })
+                .for_each(|cylinder_mesh| {
+                    combined_mesh.merge(&cylinder_mesh);
+                });
+        } else {
+            println!("No-Bonds found!!")
+        }
+
+        combined_mesh
     }
     /// Internal fn for rendering spheres.
     fn render_spheres(&self) -> Mesh {
-        let mut meshes = Vec::new();
-        for atom in self.pdb.atoms() {
-            let (x, y, z) = atom.pos();
-            let center = Vec3::new(x as f32, y as f32, z as f32);
-            let radius = atom
-                .element()
-                .expect("Atom Element not Defined")
-                .atomic_radius()
-                .van_der_waals
-                .expect("Van der waals not defined") as f32;
-            let color = self.color_scheme.get_color(atom).to_srgba();
-            let mut sphere_mesh = Sphere::new(radius).mesh().build();
-            let vertex_count = sphere_mesh.count_vertices();
-            let color_array = vec![color.to_vec4(); vertex_count];
-            sphere_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, color_array);
-            sphere_mesh = sphere_mesh.translated_by(center);
-            sphere_mesh.compute_smooth_normals();
-            meshes.push(sphere_mesh);
-        }
-        // combine all the meshes together
-        meshes
-            .into_iter()
+        self.pdb
+            .iter_coords_and_elements()
+            .map(|(coord, element_str)| {
+                let center = Vec3::new(coord[0], coord[1], coord[2]);
+                let element = Element::from_symbol(element_str).expect("Element not recognized");
+                let radius = element
+                    .atomic_radius()
+                    .van_der_waals
+                    .expect("Van der waals not defined") as f32;
+                let mut sphere_mesh = Sphere::new(radius).mesh().build();
+                let vertex_count = sphere_mesh.count_vertices();
+                let color = self.color_scheme.get_color(element_str).to_srgba();
+                let color_array =
+                    vec![Vec4::new(color.red, color.green, color.blue, color.alpha); vertex_count];
+                sphere_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, color_array);
+                sphere_mesh = sphere_mesh.translated_by(center);
+                sphere_mesh.compute_smooth_normals();
+                sphere_mesh
+            })
             .reduce(|mut acc, mesh| {
                 acc.merge(&mesh);
                 acc
@@ -111,8 +178,8 @@ mod tests {
     #[test]
     fn test_pdb_to_mesh() {
         let (pdb, _errors) = pdbtbx::open("examples/1fap.cif").unwrap();
-        let structure = Structure::builder().pdb(pdb).build();
-        assert_eq!(structure.pdb.atom_count(), 2154);
+        let structure = Structure::builder().pdb(AtomCollection::from(&pdb)).build();
+        assert_eq!(structure.pdb.size(), 2154);
         let mesh = structure.to_mesh();
         assert_eq!(mesh.count_vertices(), 779748);
     }

ferritin-core/src/core/atomcollection.rs

@@ -27,6 +27,18 @@ pub struct AtomCollection {
 }
 
 impl AtomCollection {
+    pub fn size(&self) -> usize {
+        self.size
+    }
+    pub fn bonds(&self) -> Option<&Vec<Bond>> {
+        self.bonds.as_ref()
+    }
+    pub fn coords(&self) -> &Vec<[f32; 3]> {
+        self.coords.as_ref()
+    }
+    pub fn iter_coords_and_elements(&self) -> impl Iterator<Item = (&[f32; 3], &String)> {
+        izip!(&self.coords, &self.elements)
+    }
     pub fn calculate_displacement(&self) {
         // Measure the displacement vector, i.e. the vector difference, from
         // one array of atom coordinates to another array of coordinates.
@@ -370,6 +382,12 @@ pub struct Bond {
     // has_setting
 }
 
+impl Bond {
+    pub fn get_atom_indices(&self) -> (i32, i32) {
+        (self.atom1, self.atom2)
+    }
+}
+
 #[repr(u8)]
 #[derive(Debug, PartialEq)]
 /// BondOrder:

ferritin-core/src/core/mod.rs

@@ -1,2 +1,4 @@
 mod atomcollection;
 mod constants;
+
+pub use atomcollection::{AtomCollection, Bond};

ferritin-core/src/lib.rs

@@ -1 +1,2 @@
 mod core;
+pub use core::{AtomCollection, Bond};