second attempt at tree growth - better

src/main.rs

@@ -1,4 +1,4 @@
-use std::f32::consts::PI;
+use std::f32::consts::{PI, TAU};
 use std::time::{Duration, SystemTime};
 
 use bevy::asset::RenderAssetUsages;
@@ -11,8 +11,11 @@ use bevy::time::common_conditions::on_timer;
 mod camera;
 mod debugging;
 mod tree;
+mod tree2;
 mod tree_generation;
+mod tree_generation2;
 mod tree_rendering;
+mod tree_rendering2;
 
 #[derive(States, Debug, Clone, PartialEq, Eq, Hash)]
 enum RunningGrowthState {
@@ -59,8 +62,46 @@ fn update_tree_grow_state(
     }
 }
 
+fn get_tree_params() -> tree_generation2::TreeParams {
+    // tree_generation2::TreeParams{
+    //     apical_dominance: 2.,
+    //     bud_spacing_distance: 0.3,
+    //     transport_cost_factor: 20000.,
+    //     length_to_area_growth_factor: 0.0005,
+    //     random_growth_factor: 0.05,
+    //     gravitropism_strength: 0.005,
+    //     straight_growth_factor: 0.01,
+    //     buds_per_node: 3,
+    //     lateral_gsa: PI * (180. - 60.)/180.,
+    //     auxin_degrade_percent: 0.3,
+    //     canalization_threshold: 0.6,
+    //     bud_auxin_ramp: 0.01,
+    //     initial_branch_radius: 0.001,
+    //     phyllotactic_angle: 2./5. * TAU,
+    //     phyllotactic_angle_jitter: 0.05 * TAU,
+    // }
+
+    tree_generation2::TreeParams{
+        apical_dominance: 1.5,
+        bud_spacing_distance: 0.2,
+        transport_cost_factor: 20000.,
+        length_to_area_growth_factor: 0.0005,
+        random_growth_factor: 0.05,
+        gravitropism_strength: 0.002,
+        straight_growth_factor: 0.005,
+        buds_per_node: 1,
+        lateral_gsa: PI * (180. - 60.)/180.,
+        auxin_degrade_percent: 0.2,
+        canalization_threshold: 0.5,
+        bud_auxin_ramp: 0.01,
+        initial_branch_radius: 0.001,
+        phyllotactic_angle: 2./5. * TAU,
+        phyllotactic_angle_jitter: 0.05 * TAU,
+    }
+}
+
 fn update_tree(
-    mut query: Query<(Entity, &mut tree::Tree, &mut Mesh3d)>,
+    mut query: Query<(Entity, &mut tree2::Tree, &mut Mesh3d)>,
     mut commands: Commands,
     mut meshes: ResMut<Assets<Mesh>>,
 ) {
@@ -79,10 +120,10 @@ fn update_tree(
 
     info!("here");
     for (entity, mut tree, mut mesh) in query {
-        tree_generation::grow_tree(tree.as_mut());
+        tree_generation2::grow_tree(tree.as_mut(), &get_tree_params());
 
         let t1 = SystemTime::now();
-        let tree_mesh = tree_rendering::generate_tree(&tree);
+        let tree_mesh = tree_rendering2::generate_tree(&tree);
         let t2 = SystemTime::now();
         meshes.remove(&mesh.0);
         mesh.0 = meshes.add(tree_mesh);
@@ -142,8 +183,8 @@ fn setup(
         },
     ));
 
-    let tree = tree_generation::initial_tree();
+    let tree = tree_generation2::initial_tree();
-    let tree_mesh = tree_rendering::generate_tree(&tree);
+    let tree_mesh = tree_rendering2::generate_tree(&tree);
 
     // let normal_material = materials.add(Color::srgb_u8(255, 0, 0));
     // draw_normals(
@@ -157,7 +198,8 @@ fn setup(
     commands.spawn((
         tree,
         Mesh3d(meshes.add(tree_mesh)),
-        MeshMaterial3d(materials.add(Color::srgb_u8(124, 144, 255))),
+        MeshMaterial3d(materials.add(Color::srgb_u8(255, 255, 255))),
+        // MeshMaterial3d(materials.add(Color::srgb_u8(124, 144, 255))),
         // MeshMaterial3d(materials.add(StandardMaterial {
         //     base_color: Color::srgba_u8(124, 144, 255, 255),
         //     alpha_mode: AlphaMode::Blend,
@@ -200,103 +242,103 @@ fn draw_normals(
     }
 }
 
-fn initial_tree() -> tree::Tree {
+// fn initial_tree() -> tree::Tree {
-    let mut tree = tree::Tree::default();
+//     let mut tree = tree::Tree::default();
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        0,
+//         0,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 0,
+//             id: 0,
-            kind: tree::TreeNodeKind::Extension,
+//             kind: tree::TreeNodeKind::Extension,
-            parent_id: 0,
+//             parent_id: 0,
-            radius: 0.5,
+//             radius: 0.5,
-            offset: 2.,
+//             offset: 2.,
-            angle: 0.,
+//             angle: 0.,
-            rotation: 0.,
+//             rotation: 0.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        1,
+//         1,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 1,
+//             id: 1,
-            kind: tree::TreeNodeKind::Extension,
+//             kind: tree::TreeNodeKind::Extension,
-            parent_id: 0,
+//             parent_id: 0,
-            radius: 0.3,
+//             radius: 0.3,
-            offset: 2.,
+//             offset: 2.,
-            angle: 0.,
+//             angle: 0.,
-            rotation: 0.,
+//             rotation: 0.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        2,
+//         2,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 2,
+//             id: 2,
-            kind: tree::TreeNodeKind::Extension,
+//             kind: tree::TreeNodeKind::Extension,
-            parent_id: 1,
+//             parent_id: 1,
-            radius: 0.25,
+//             radius: 0.25,
-            offset: 2.,
+//             offset: 2.,
-            angle: -0.4,
+//             angle: -0.4,
-            rotation: PI / 2.,
+//             rotation: PI / 2.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        3,
+//         3,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 3,
+//             id: 3,
-            kind: tree::TreeNodeKind::EndCut,
+//             kind: tree::TreeNodeKind::EndCut,
-            parent_id: 2,
+//             parent_id: 2,
-            radius: 0.2,
+//             radius: 0.2,
-            offset: 0.,
+//             offset: 0.,
-            angle: 0.4,
+//             angle: 0.4,
-            rotation: PI / 2.,
+//             rotation: PI / 2.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        4,
+//         4,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 4,
+//             id: 4,
-            kind: tree::TreeNodeKind::StartBranch,
+//             kind: tree::TreeNodeKind::StartBranch,
-            parent_id: 1,
+//             parent_id: 1,
-            radius: 0.2,
+//             radius: 0.2,
-            offset: 2.,
+//             offset: 2.,
-            angle: -0.6,
+//             angle: -0.6,
-            rotation: -PI / 2.,
+//             rotation: -PI / 2.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        5,
+//         5,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 5,
+//             id: 5,
-            kind: tree::TreeNodeKind::Extension,
+//             kind: tree::TreeNodeKind::Extension,
-            parent_id: 4,
+//             parent_id: 4,
-            radius: 0.1,
+//             radius: 0.1,
-            offset: 1.,
+//             offset: 1.,
-            angle: 0.2,
+//             angle: 0.2,
-            rotation: -PI / 2.,
+//             rotation: -PI / 2.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        6,
+//         6,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 6,
+//             id: 6,
-            kind: tree::TreeNodeKind::EndMeristem,
+//             kind: tree::TreeNodeKind::EndMeristem,
-            parent_id: 5,
+//             parent_id: 5,
-            radius: 0.05,
+//             radius: 0.05,
-            offset: 0.03,
+//             offset: 0.03,
-            angle: 0.2,
+//             angle: 0.2,
-            rotation: -PI / 2.,
+//             rotation: -PI / 2.,
-        },
+//         },
-    );
+//     );
-    tree.nodes.insert(
+//     tree.nodes.insert(
-        7,
+//         7,
-        tree::TreeNode {
+//         tree::TreeNode {
-            id: 7,
+//             id: 7,
-            kind: tree::TreeNodeKind::EndBranch,
+//             kind: tree::TreeNodeKind::EndBranch,
-            parent_id: 4,
+//             parent_id: 4,
-            radius: 0.04,
+//             radius: 0.04,
-            offset: 0.5,
+//             offset: 0.5,
-            angle: -0.8,
+//             angle: -0.8,
-            rotation: PI / 4.,
+//             rotation: PI / 4.,
-        },
+//         },
-    );
+//     );
-    tree
+//     tree
-}
+// }

src/tree.rs

@@ -20,11 +20,13 @@ pub struct Tree {
 
 #[derive(Debug, Copy, Clone)]
 pub enum TreeNodeKind {
-    StartBranch,
+    Root,
+    Branch,
     Extension,
-    EndMeristem,
+    BranchingExtension,
-    EndBranch,
+    Meristem,
-    EndCut,
+    ApicalMeristem,
+    Cut,
 }
 
 #[derive(Debug, Clone)]
@@ -36,8 +38,17 @@ pub struct TreeNode {
     pub offset: f32,
     pub angle: f32,
     pub rotation: f32,
+    pub age: f32,
+    pub auxin_level: f32,
 }
 
+struct AccumulatedTreeNode {
+    transform: Transform,
+    transport_cost: f32,
+    auxin_export_score: f32,
+}
+
+
 impl Default for Tree {
     fn default() -> Self {
         Tree {
@@ -49,10 +60,12 @@ impl Default for Tree {
 }
 
 impl Tree {
-    pub fn add_node(&mut self, mut tree_node: TreeNode) {
+    pub fn add_node(&mut self, mut tree_node: TreeNode) -> u32 {
-        tree_node.id = self.next_node_id;
+        let tree_node_id = self.next_node_id;
+        tree_node.id = tree_node_id;
         self.nodes.insert(self.next_node_id, tree_node);
         self.next_node_id += 1;
+        tree_node_id
     }
 
     pub fn get_node(&self, id: u32) -> Option<&TreeNode> {
@@ -70,9 +83,13 @@ impl Tree {
             if curr_node.id == curr_node.parent_id {
                 break;
             }
-            curr_node = &self.nodes[&curr_node.parent_id];
-            // info!("transform: {:?}", curr_node);
             transform = curr_node.get_transform() * transform;
+            let next_node = &self.nodes[&curr_node.parent_id];
+            // info!("transform: {:?}", curr_node);
+            // if !curr_node.is_branch_begin() {
+            // }
+
+            curr_node = next_node;
         }
         transform
     }
@@ -88,6 +105,8 @@ impl TreeNode {
             offset: 0.,
             angle: 0.,
             rotation: 0.,
+            age: 0.,
+            auxin_level: 0.,
         }
     }
 
@@ -119,21 +138,21 @@ impl TreeNode {
     pub fn is_end(&self) -> bool {
         matches!(
             self.kind,
-            TreeNodeKind::EndMeristem | TreeNodeKind::EndCut | TreeNodeKind::EndBranch
+            TreeNodeKind::Meristem | TreeNodeKind::ApicalMeristem | TreeNodeKind::Cut
         )
     }
 
     pub fn is_meristem(&self) -> bool {
         matches!(
             self.kind,
-            TreeNodeKind::EndMeristem | TreeNodeKind::EndBranch
+            TreeNodeKind::Meristem | TreeNodeKind::ApicalMeristem
         )
     }
 
     pub fn is_extension(&self) -> bool {
         matches!(
             self.kind,
-            TreeNodeKind::Extension | TreeNodeKind::StartBranch
+            TreeNodeKind::Extension | TreeNodeKind::BranchingExtension
         )
     }
 
@@ -141,7 +160,7 @@ impl TreeNode {
         self.id == self.parent_id
             || matches!(
                 self.kind,
-                TreeNodeKind::StartBranch | TreeNodeKind::EndBranch
+                TreeNodeKind::Branch | TreeNodeKind::Root
             )
     }
 }

src/tree2.rs

@@ -0,0 +1,319 @@
+use std::collections::HashMap;
+
+use bevy::prelude::*;
+
+#[derive(Debug, Clone)]
+pub struct Base {
+    pub radius: f32,
+}
+
+#[derive(Debug, Clone)]
+pub struct Bud {
+    pub is_active: bool,
+    pub age: f32,
+    pub azimuthal_angle: f32,
+    pub auxin_production_level: f32,
+    // pub export_score: f32,
+}
+
+#[derive(Debug, Clone)]
+pub enum MeristemKind {
+    Apical,
+    Lateral,
+    Root,
+}
+
+#[derive(Debug, Clone)]
+pub struct Meristem {
+    pub kind: MeristemKind,
+    pub length: f32,
+    pub inclination_angle: f32,
+    pub azimuthal_angle: f32,
+    pub gravitropic_set_angle: f32, /* relative to gravity */
+    pub auxin_production_level: f32,
+    pub phyllotactic_angle: f32,
+}
+
+#[derive(Debug, Clone)]
+pub struct Stem {
+    pub radius: f32,
+    pub length: f32,
+    pub inclination_angle: f32,
+    pub azimuthal_angle: f32,
+    pub age: f32,
+    pub auxin_level: f32,
+}
+
+#[derive(Debug, Clone)]
+pub struct CutStem {
+    pub radius: f32,
+    pub length: f32,
+    pub inclination_angle: f32,
+    pub azimuthal_angle: f32,
+    pub cut_inclination_angle: f32,
+    pub cut_azimuthal_angle: f32,
+}
+
+#[derive(Debug, Clone)]
+pub struct Branch {
+    pub radius: f32,
+    pub inclination_angle: f32,
+    pub azimuthal_angle: f32,
+    pub auxin_level: f32,
+}
+
+#[derive(Debug, Clone)]
+pub enum TreeNodeData {
+    Base(Base),
+    Bud(Bud),
+    Meristem(Meristem),
+    Stem(Stem),
+    CutStem(CutStem),
+    Branch(Branch),
+}
+
+#[derive(Debug, Clone)]
+pub struct TreeNode {
+    // id: u32,
+    pub parent_id: u32,
+    pub data: TreeNodeData,
+}
+
+// struct AccumulatedTreeNode {
+//     transform: Transform,
+//     transport_cost: f32,
+//     auxin_export_score: f32,
+// }
+
+#[derive(Component)]
+pub struct Tree {
+    pub age: f32,
+    pub nodes: HashMap<u32, TreeNode>,
+    next_node_id: u32,
+}
+
+impl Default for Tree {
+    fn default() -> Self {
+        Tree {
+            age: 0.,
+            nodes: HashMap::new(),
+            next_node_id: 0,
+        }
+    }
+}
+
+impl Tree {
+    pub fn add_base(&mut self, radius: f32) {
+        assert!(self.next_node_id == 0);
+        self.nodes.insert(
+            0,
+            TreeNode {
+                parent_id: 0,
+                data: TreeNodeData::Base(Base { radius }),
+            },
+        );
+        self.next_node_id += 1;
+    }
+
+    pub fn add_node<T: Into<TreeNodeData>>(&mut self, parent_id: u32, data: T) -> u32 {
+        let tree_node_id = self.next_node_id;
+        self.nodes.insert(
+            self.next_node_id,
+            TreeNode {
+                parent_id,
+                data: data.into(),
+            },
+        );
+        self.next_node_id += 1;
+        tree_node_id
+    }
+
+    pub fn get_node(&self, id: u32) -> Option<&TreeNode> {
+        self.nodes.get(&id)
+    }
+
+    pub fn get_node_mut(&mut self, id: u32) -> Option<&mut TreeNode> {
+        self.nodes.get_mut(&id)
+    }
+
+    pub fn get_parent_node(&self, id: u32)  -> Option<&TreeNode> {
+        let node = self.get_node(id)?;
+        self.get_node(node.parent_id)
+    }
+
+    // Does not include current node transform
+    pub fn get_accumulated_transform(&self, id: u32) -> Transform {
+        let mut transform = Transform::IDENTITY;
+        let mut curr_node = &self.nodes[&id];
+        loop {
+            if curr_node.is_base() {
+                break;
+            }
+            curr_node = &self.nodes[&curr_node.parent_id];
+            transform = curr_node.get_transform() * transform;
+        }
+        transform
+    }
+
+    pub fn get_accumulated_transport_cost(&self, id: u32) -> f32 {
+        let mut transport_cost = 0.;
+        let mut curr_node = &self.nodes[&id];
+        loop {
+            if curr_node.is_base() {
+                break;
+            }
+            let parent_node = &self.nodes[&curr_node.parent_id];
+            // info!("transport {}", self.get_transport_cost(curr_node, parent_node));
+            transport_cost += self.get_transport_cost(curr_node, parent_node);
+
+            curr_node = parent_node;
+        }
+        transport_cost
+    }
+
+    pub fn get_transport_cost(&self, curr_node: &TreeNode, parent_node: &TreeNode) -> f32 {        
+        if let TreeNodeData::Branch(branch) = &curr_node.data {
+            if let TreeNodeData::Stem(stem) = &parent_node.data {
+                return stem.radius.powi(2) / branch.radius.powi(2)
+            }
+        }
+
+        let avg_radius = (parent_node.get_radius() + curr_node.get_radius()) / 2.;
+        let length = curr_node.get_length();
+        length / avg_radius.powi(3) // TODO: should this be 4?
+    }
+}
+
+impl TreeNode {
+    pub fn is_base(&self) -> bool {
+        return matches!(self.data, TreeNodeData::Base(_))
+    }
+
+    pub fn get_transform(&self) -> Transform {
+        let (azimuthal_angle, inclination_angle, length) = match &self.data {
+            TreeNodeData::Base(_) => (0., 0., 0.),
+            TreeNodeData::Bud(_) => (0., 0., 0.),
+            TreeNodeData::Meristem(meristem) => (
+                meristem.azimuthal_angle,
+                meristem.inclination_angle,
+                meristem.length,
+            ),
+            TreeNodeData::Stem(stem) => (stem.azimuthal_angle, stem.inclination_angle, stem.length),
+            TreeNodeData::CutStem(cut_stem) => (
+                cut_stem.azimuthal_angle,
+                cut_stem.inclination_angle,
+                cut_stem.length,
+            ),
+            TreeNodeData::Branch(branch) => (branch.azimuthal_angle, branch.inclination_angle, 0.),
+        };
+        let mut transform = Transform::IDENTITY;
+        transform.rotate_local_y(azimuthal_angle);
+        transform.rotate_local_x(inclination_angle);
+        transform.rotate_local_y(-azimuthal_angle);
+        transform.translation += transform.up() * length;
+        transform
+    }
+
+    pub fn get_radius(&self) -> f32 {
+        match &self.data {
+            TreeNodeData::Base(base) => base.radius,
+            TreeNodeData::Bud(_) => 0.,
+            TreeNodeData::Meristem(_) => 0.,
+            TreeNodeData::Stem(stem) => stem.radius,
+            TreeNodeData::CutStem(cut_stem) => cut_stem.radius,
+            TreeNodeData::Branch(branch) => branch.radius,
+        }
+    }
+
+    pub fn get_length(&self) -> f32 {
+        match &self.data {
+            TreeNodeData::Base(_) => 0.,
+            TreeNodeData::Bud(_) => 0.,
+            TreeNodeData::Meristem(meristem) => meristem.length,
+            TreeNodeData::Stem(stem) => stem.length,
+            TreeNodeData::CutStem(cut_stem) => cut_stem.length,
+            TreeNodeData::Branch(_) => 0.,
+        }
+    }
+
+    pub fn get_auxin_level(&self) -> f32 {
+        match &self.data {
+            TreeNodeData::Stem(stem) => stem.auxin_level,
+            TreeNodeData::Branch(branch) => branch.auxin_level,
+            TreeNodeData::Meristem(meristem) => meristem.auxin_production_level,
+            TreeNodeData::Bud(bud) => bud.auxin_production_level,
+            _ => 1.,
+        }
+    }
+
+    pub fn is_meristem(&self) -> bool {
+        matches!(self.data, TreeNodeData::Meristem(_))
+    }
+
+    pub fn is_bud(&self) -> bool {
+        matches!(self.data, TreeNodeData::Bud(_))
+    }
+
+    pub fn is_branch_end(&self) -> bool {
+        matches!(
+            self.data,
+            TreeNodeData::Meristem(_) | TreeNodeData::CutStem(_)
+        )
+    }
+
+    pub fn is_branch_begin(&self) -> bool {
+        self.is_base() || matches!(self.data, TreeNodeData::Branch(_))
+    }
+
+    pub fn add_area(&mut self, area: f32) {
+        let radius = match &mut self.data {
+            TreeNodeData::Base(base) => &mut base.radius,
+            TreeNodeData::Bud(_) => return,
+            TreeNodeData::Meristem(_) => return,
+            TreeNodeData::Stem(stem) => &mut stem.radius,
+            TreeNodeData::CutStem(cut_stem) => &mut cut_stem.radius,
+            TreeNodeData::Branch(branch) => &mut branch.radius,
+        };
+
+        *radius = (radius.powi(2) + area).sqrt();
+    }
+
+    // pub fn new_base(radius: f32) -> Self {
+    //     TreeNode {
+    //         parent_id: 0,
+    //         data: TreeNodeData::Base(Base { radius }),
+    //     }
+    // }
+
+    // pub fn new_meristem(radius: f32) -> Self {
+    //     TreeNode {
+    //         parent_id: 0,
+    //         data: TreeNodeData::Base(Base { radius }),
+    //     }
+    // }
+}
+
+impl From<Bud> for TreeNodeData {
+    fn from(bud: Bud) -> Self {
+        TreeNodeData::Bud(bud)
+    }
+}
+
+impl From<Meristem> for TreeNodeData {
+    fn from(meristem: Meristem) -> Self {
+        TreeNodeData::Meristem(meristem)
+    }
+}
+
+impl From<Stem> for TreeNodeData {
+    fn from(stem: Stem) -> Self {
+        TreeNodeData::Stem(stem)
+    }
+}
+
+impl From<Branch> for TreeNodeData {
+    fn from(branch: Branch) -> Self {
+        TreeNodeData::Branch(branch)
+    }
+}
+

src/tree_generation.rs

@@ -1,26 +1,66 @@
-use std::f32::consts::TAU;
+use std::f32::consts::{PI, TAU};
 
 use crate::tree::{self, Tree, TreeNode, TreeNodeKind};
 use bevy::prelude::*;
 use rand::{distr::StandardUniform, Rng};
 
+
+
+
+
+pub struct TreeParams {
+    apical_dominance: f32, /* 0 to 1 */
+    gravitropism_strength: f32, /* 0 to 1 */
+    apical_dominance_reduction_with_age: f32, /* 0 to 1 */
+    bud_spacing_distance: f32,
+    canalization_threshold: f32, /* 0 to 1 */
+    buds_per_node: u32,
+}
+
+
+
+
+
+
+
 pub fn initial_tree() -> tree::Tree {
     let mut tree = tree::Tree::default();
     tree.add_node(tree::TreeNode {
         id: 0,
-        kind: tree::TreeNodeKind::EndMeristem,
+        kind: tree::TreeNodeKind::Root,
         parent_id: 0,
         radius: 0.05,
+        offset: 0.,
+        angle: 0.,
+        rotation: 0.,
+        age: 0.,
+        auxin_level: 1.,
+    });
+    tree.add_node(tree::TreeNode {
+        id: 1,
+        kind: tree::TreeNodeKind::ApicalMeristem,
+        parent_id: 0,
+        radius: 0.03,
         offset: 0.5,
         angle: 0.,
         rotation: 0.,
+        age: 0.,
+        auxin_level: 1.,
     });
     tree
 }
 
 const MAX_MERISTEM_LENGTH: f32 = 1.;
+const AUXIN_DROP: f32 = 0.1;
+const AUXIN_MERISTEM: f32 = 1.;
+const AUXIN_DROP_PER_NODE_FACTOR: f32 = 2.;
 
 pub fn grow_tree(tree: &mut tree::Tree) {
+    for (_, node) in tree.nodes.iter_mut() {
+        node.age += 1.;
+        node.auxin_level = (node.auxin_level - AUXIN_DROP).max(0.);
+    }
+
     let meristems_ids: Vec<u32> = tree
         .nodes
         .iter()
@@ -30,6 +70,17 @@ pub fn grow_tree(tree: &mut tree::Tree) {
     for meristem_id in meristems_ids {
         grow_meristem(meristem_id, tree);
     }
+
+    info!("nodes: {:?}", tree.nodes);
+    for node_id in tree
+        .nodes
+        .iter()
+        .filter(|(_, node)| node.is_extension() && node.auxin_level < 0.05 )
+        .map(|(i, node)| *i)
+        .collect::<Vec<u32>>()
+    {
+        add_branch(node_id, tree);
+    }
 }
 
 pub fn grow_meristem(node_id: u32, tree: &mut Tree) {
@@ -37,7 +88,7 @@ pub fn grow_meristem(node_id: u32, tree: &mut Tree) {
 
     let accumulated_transform = tree.get_accumulated_transform(node_id);
 
-    let grow_factor = 0.001;
+    let grow_factor = 0.0005;
     let node = tree.get_node_mut(node_id).unwrap();
     node.offset += 0.1;
 
@@ -46,13 +97,16 @@ pub fn grow_meristem(node_id: u32, tree: &mut Tree) {
 
     let accumulated_rotation_inverse = accumulated_transform.rotation.inverse();
     let up_vector = accumulated_rotation_inverse * Vec3::Y;
-    let towards_up_vector = up_vector - original_vector;
+    info!("up: {}", up_vector);
+    // let towards_up_vector = up_vector - original_vector;
+    // let gravity_vector = -up_vector * 0.01;
 
     let random_angle = rng.random::<f32>() * TAU;
     let random_vector = Vec3::new(random_angle.cos(), 0.0, random_angle.sin());
 
     let random_magnitude = rng.sample::<f32, StandardUniform>(StandardUniform) - 0.5;
-    let new_vector = (original_vector + towards_up_vector*0.05 + random_vector*random_magnitude*0.1).normalize();
+    let new_vector =
+        (original_vector + up_vector * 0.005 + random_vector * random_magnitude * 0.05).normalize();
 
     let r = new_vector.length();
     let inclination = (new_vector.y / r).acos();
@@ -65,28 +119,99 @@ pub fn grow_meristem(node_id: u32, tree: &mut Tree) {
         let new_meristem_length = node.offset - MAX_MERISTEM_LENGTH;
         let new_meristem_radius = node.radius * 0.5;
         node.offset -= new_meristem_length;
+        let existing_node_kind = node.kind;
         node.kind = TreeNodeKind::Extension;
         let parent_id = node.id;
-        let angle = node.angle;
+        // let angle = node.angle;
-        let rotation = node.rotation;
+        // let rotation = node.rotation;
         tree.add_node(TreeNode {
             id: 0,
-            kind: TreeNodeKind::EndMeristem,
+            kind: existing_node_kind,
             parent_id: parent_id,
             radius: new_meristem_radius,
             offset: new_meristem_length,
-            angle: angle,
+            angle: 0.,
-            rotation: rotation,
+            rotation: 0.,
+            age: 0.,
+            auxin_level: 1.,
         });
+        // add_branch(node_id, tree);
     }
 
+    // let mut auxin_level = 1.;
+    let mut auxin_to_add = 1.;
     let mut growth_node_id = node_id;
     loop {
         let node = tree.get_node_mut(growth_node_id).unwrap();
+
+        let auxin_to_add_node = auxin_to_add / 2.;
+        info!("auxin: {} {}", node.auxin_level, auxin_to_add_node);
+        auxin_to_add /= 2.;
+        // if auxin_to_add_node.is_sign_positive() {
+        node.auxin_level = (node.auxin_level + auxin_to_add_node).min(1.);
+        // }
+
+        // node.auxin_level = (node.auxin_level + auxin_level*0.5).max(1.);
         node.radius = (node.radius.powi(2) + grow_factor).sqrt();
         if node.is_root() {
             break;
         }
         growth_node_id = node.parent_id;
+        // auxin_level = node.auxin_level;
+        // auxin_level /= AUXIN_DROP_PER_NODE_FACTOR;
     }
 }
+
+pub fn add_branch(node_id: u32, tree: &mut Tree) {
+    // let random_threshold = rand::rng().random::<f32>();
+    // if random_threshold < 0.8 {
+    //     return;
+    // }
+
+    let random_rotation = rand::rng().random::<f32>() * TAU;
+
+    let branch1_id = tree.add_node(TreeNode {
+        id: 0,
+        kind: TreeNodeKind::Branch,
+        parent_id: node_id,
+        radius: 0.05,
+        offset: 0.05,
+        angle: PI / 2. * 0.8,
+        rotation: random_rotation,
+        age: 0.,
+        auxin_level: 0.,
+    });
+    tree.add_node(TreeNode {
+        id: 0,
+        kind: TreeNodeKind::Meristem,
+        parent_id: branch1_id,
+        radius: 0.035,
+        offset: 0.05,
+        angle: 0.,
+        rotation: 0.,
+        age: 0.,
+        auxin_level: 0.,
+    });
+    let branch2_id = tree.add_node(TreeNode {
+        id: 0,
+        kind: TreeNodeKind::Branch,
+        parent_id: node_id,
+        radius: 0.05,
+        offset: 0.05,
+        angle: PI / 2. * 0.8,
+        rotation: random_rotation - PI,
+        age: 0.,
+        auxin_level: 0.,
+    });
+    tree.add_node(TreeNode {
+        id: 0,
+        kind: TreeNodeKind::Meristem,
+        parent_id: branch2_id,
+        radius: 0.035,
+        offset: 0.05,
+        angle: 0.,
+        rotation: 0.,
+        age: 0.,
+        auxin_level: 0.,
+    });
+}

src/tree_generation2.rs

@@ -0,0 +1,401 @@
+use bevy::prelude::*;
+use rand::{distr::StandardUniform, Rng};
+use std::{f32::consts::{PI, TAU}, time::SystemTime};
+
+use crate::tree2::{Branch, Bud, Meristem, MeristemKind, Stem, Tree, TreeNode, TreeNodeData};
+
+pub struct TreeParams {
+    pub apical_dominance: f32,
+    pub gravitropism_strength: f32,
+    // apical_dominance_reduction_with_age: f32, /* 0 to 1 */
+    pub bud_spacing_distance: f32,
+    pub canalization_threshold: f32,
+    pub bud_auxin_ramp: f32,
+    pub buds_per_node: u32,
+    pub lateral_gsa: f32,
+    pub transport_cost_factor: f32,
+    pub length_to_area_growth_factor: f32,
+    pub random_growth_factor: f32,
+    pub straight_growth_factor: f32,
+    pub auxin_degrade_percent: f32, /* 0 to 1 */
+    pub initial_branch_radius: f32,
+    pub phyllotactic_angle: f32,
+    pub phyllotactic_angle_jitter: f32,
+}
+
+pub fn initial_tree() -> Tree {
+    let mut tree = Tree::default();
+    tree.add_base(0.001);
+    tree.add_node(
+        0,
+        Meristem {
+            kind: MeristemKind::Apical,
+            length: 0.00001,
+            inclination_angle: 0.,
+            azimuthal_angle: 0.,
+            gravitropic_set_angle: PI,
+            auxin_production_level: 1.,
+            phyllotactic_angle: 0.,
+        },
+    );
+
+    // tree.add_node(
+    //     0,
+    //     Stem {
+    //         radius: 0.03,
+    //         length: 1.,
+    //         inclination_angle: 0.,
+    //         azimuthal_angle: 0.,
+    //         age: 0.,
+    //         auxin_level: 0.,
+    //     },
+    // );
+    // tree.add_node(
+    //     1,
+    //     Meristem {
+    //         kind: MeristemKind::Apical,
+    //         length: 0.5,
+    //         inclination_angle: 0.,
+    //         azimuthal_angle: 0.,
+    //         gravitropic_set_angle: PI,
+    //         auxin_production_level: 1.,
+    //     },
+    // );
+    // tree.add_node(
+    //     1,
+    //     Branch {
+    //         radius: 0.03,
+    //         inclination_angle: PI / 4.,
+    //         azimuthal_angle: 0.,
+    //         auxin_level: 0.,
+    //     },
+    // );
+    // tree.add_node(
+    //     3,
+    //     Stem {
+    //         radius: 0.03,
+    //         length: 1.,
+    //         inclination_angle: 0.,
+    //         azimuthal_angle: 0.,
+    //         age: 0.,
+    //         auxin_level: 0.,
+    //     },
+    // );
+    // tree.add_node(
+    //     4,
+    //     Meristem {
+    //         kind: MeristemKind::Lateral,
+    //         length: 0.5,
+    //         inclination_angle: 0.,
+    //         azimuthal_angle: 0.,
+    //         gravitropic_set_angle: PI,
+    //         auxin_production_level: 1.,
+    //     },
+    // );
+    // tree.add_node(
+    //     4,
+    //     Branch {
+    //         radius: 0.03,
+    //         inclination_angle: PI / 4.,
+    //         azimuthal_angle: 0.,
+    //         auxin_level: 0.,
+    //     },
+    // );
+    // tree.add_node(
+    //     6,
+    //     Meristem {
+    //         kind: MeristemKind::Lateral,
+    //         length: 0.5,
+    //         inclination_angle: 0.,
+    //         azimuthal_angle: 0.,
+    //         gravitropic_set_angle: PI,
+    //         auxin_production_level: 1.,
+    //     },
+    // );
+    tree
+}
+
+pub fn grow_tree(tree: &mut Tree, params: &TreeParams) {
+    let t1 = SystemTime::now();
+    for (_, node) in tree.nodes.iter_mut() {
+        age_node(node, 1.);
+        degrade_auxin(node, params.auxin_degrade_percent);
+    }
+    let t2 = SystemTime::now();
+    let meristems_ids: Vec<u32> = tree
+        .nodes
+        .iter()
+        .filter(|(_, n)| n.is_meristem())
+        .map(|(i, _)| *i)
+        .collect();
+    for meristem_id in meristems_ids {
+        grow_meristem(meristem_id, tree, params);
+    }
+    let t3 = SystemTime::now();
+    let bud_ids: Vec<u32> = tree
+        .nodes
+        .iter()
+        .filter(|(_, n)| n.is_bud())
+        .map(|(i, _)| *i)
+        .collect();
+    for bud_id in bud_ids {
+        update_bud(bud_id, tree, params);
+    }
+    let t4 = SystemTime::now();
+
+    let iterate_duration = t2.duration_since(t1).unwrap();
+    let growth_duration = t3.duration_since(t2).unwrap();
+    let bud_check_duration = t4.duration_since(t3).unwrap();
+    info!("dur1: {:?} {:?} {:?}", iterate_duration, growth_duration, bud_check_duration);
+
+    // info!("{:?}", tree.nodes);
+
+    // for node_id in tree
+    //     .nodes
+    //     .iter()
+    //     .filter(|(_, node)| node.is_extension() && node.auxin_level < 0.05 )
+    //     .map(|(i, node)| *i)
+    //     .collect::<Vec<u32>>()
+    // {
+    //     add_branch(node_id, tree);
+    // }
+}
+
+fn age_node(node: &mut TreeNode, age_increment: f32) {
+    match &mut node.data {
+        crate::tree2::TreeNodeData::Stem(stem) => stem.age += age_increment,
+        _ => {}
+    }
+}
+
+fn degrade_auxin(node: &mut TreeNode, degrade_percent: f32) {
+    if let TreeNodeData::Stem(stem) = &mut node.data {
+        stem.auxin_level *= 1. - degrade_percent;
+    } else if let TreeNodeData::Branch(branch) = &mut node.data {
+        branch.auxin_level *= 1. - degrade_percent;
+    }
+}
+
+fn grow_meristem(node_id: u32, tree: &mut Tree, params: &TreeParams) {
+    let accumulated_transport_cost = tree.get_accumulated_transport_cost(node_id);
+
+    // let parent_node_id = tree.get_node(node_id).unwrap().parent_id;
+
+    let accumulated_transform = tree.get_accumulated_transform(node_id);
+
+    // let parent_node_radius = parent_node.get_radius();
+
+    let node = tree.get_node_mut(node_id).unwrap();
+
+    let meristem = match &mut node.data {
+        TreeNodeData::Meristem(meristem) => meristem,
+        _ => return,
+    };
+
+    let meristem_length_growth =
+        params.transport_cost_factor / accumulated_transport_cost * meristem.auxin_production_level;
+    let area_growth = meristem_length_growth * params.length_to_area_growth_factor;
+
+    meristem.length += meristem_length_growth;
+    update_angle(meristem, accumulated_transform, params);
+
+    let mut auxin_to_add = meristem.auxin_production_level;
+
+    if meristem.length > params.bud_spacing_distance {
+        let phyllotactic_angle = meristem.phyllotactic_angle;
+        let mut new_meristem = meristem.clone();
+        new_meristem.length = meristem.length - params.bud_spacing_distance;
+        let phyllotactic_angle_jitter = rand::rng().random::<f32>() * params.phyllotactic_angle_jitter;
+        new_meristem.phyllotactic_angle = (phyllotactic_angle + params.phyllotactic_angle + phyllotactic_angle_jitter) % TAU;
+
+        node.data = TreeNodeData::Stem(Stem {
+            radius: 0., // TODO
+            length: params.bud_spacing_distance,
+            inclination_angle: meristem.inclination_angle,
+            azimuthal_angle: meristem.azimuthal_angle,
+            age: 0.,
+            auxin_level: meristem.auxin_production_level,
+        });
+
+        add_node(tree, node_id, phyllotactic_angle, new_meristem, params);
+    }
+
+    let mut growth_node_id = node_id;
+    loop {
+        let next_auxin_level = tree
+            .get_parent_node(growth_node_id)
+            .unwrap()
+            .get_auxin_level();
+        let node = tree.get_node_mut(growth_node_id).unwrap();
+
+        let auxin_level = node.get_auxin_level();
+        let mut new_auxin_level = (auxin_level + auxin_to_add).min(1.);
+        let auxin_to_flow = ((new_auxin_level - next_auxin_level) / 2.).max(0.);
+        new_auxin_level -= auxin_to_flow;
+        auxin_to_add = auxin_to_flow;
+
+        if let TreeNodeData::Stem(stem) = &mut node.data {
+            stem.auxin_level = new_auxin_level;
+        } else if let TreeNodeData::Branch(branch) = &mut node.data {
+            branch.auxin_level = new_auxin_level;
+        }
+
+        node.add_area(area_growth);
+
+        if node.is_base() {
+            break;
+        }
+        growth_node_id = node.parent_id;
+    }
+}
+
+fn update_bud(node_id: u32, tree: &mut Tree, params: &TreeParams) {
+    let parent_auxin_level = tree.get_parent_node(node_id).unwrap().get_auxin_level();
+
+    let node = tree.get_node_mut(node_id).unwrap();
+
+    let bud = match &mut node.data {
+        TreeNodeData::Bud(bud) => bud,
+        _ => return,
+    };
+
+    let auxin_difference = bud.auxin_production_level - parent_auxin_level + 0.5;
+    // bud.auxin_production_level = (bud.auxin_production_level + auxin_difference * params.bud_auxin_ramp).max(0.);
+    if auxin_difference > 0. {
+        bud.auxin_production_level = (bud.auxin_production_level + params.bud_auxin_ramp).min(1.);
+    } else {
+        bud.auxin_production_level = (bud.auxin_production_level - params.bud_auxin_ramp).max(0.);
+    }
+
+    if auxin_difference > params.canalization_threshold {
+        node.data = TreeNodeData::Branch(Branch {
+            radius: params.initial_branch_radius, // TODO
+            inclination_angle: PI - params.lateral_gsa,
+            azimuthal_angle: bud.azimuthal_angle,
+            auxin_level: 0.5, // TODO
+        });
+        let new_angle = rand::rng().random::<f32>() * TAU;
+        tree.add_node(
+            node_id,
+            Meristem {
+                kind: MeristemKind::Lateral,
+                length: 0.0001,
+                inclination_angle: 0.,
+                azimuthal_angle: 0.,
+                gravitropic_set_angle: params.lateral_gsa,
+                auxin_production_level: 1. / params.apical_dominance, // TODO
+                phyllotactic_angle: new_angle,
+            },
+        );
+    }
+
+    if auxin_difference > 0. {
+        let parent_node = tree.get_node_mut(node_id).unwrap();
+
+        let auxin_to_flow = (auxin_difference / 2.).max(0.);
+
+        if let TreeNodeData::Stem(stem) = &mut parent_node.data {
+            stem.auxin_level -= auxin_to_flow;
+        } else if let TreeNodeData::Branch(branch) = &mut parent_node.data {
+            branch.auxin_level -= auxin_to_flow;
+        }
+    }
+}
+
+fn add_node(
+    tree: &mut Tree,
+    parent_id: u32,
+    phyllotactic_angle: f32,
+    new_meristem: Meristem,
+    params: &TreeParams,
+) {
+    tree.add_node(parent_id, new_meristem);
+
+    let mut rng = rand::rng();
+    // let random_angle = rng.random::<f32>() * TAU;
+    let angle_spacing = TAU / params.buds_per_node as f32;
+    
+    for i in 0..params.buds_per_node {
+        // let phyllotactic_angle_jitter = rng.random::<f32>() * params.phyllotactic_angle_jitter;
+        let bud_angle = (phyllotactic_angle + (i as f32) * angle_spacing) % TAU;
+        tree.add_node(
+            parent_id,
+            Bud {
+                is_active: false,
+                age: 0.,
+                azimuthal_angle: bud_angle,
+                auxin_production_level: 0.,
+            },
+        );
+    }
+}
+
+fn update_angle(meristem: &mut Meristem, accumulated_transform: Transform, params: &TreeParams) {
+    let mut rng = rand::rng();
+
+    let original_vector = spherical_to_cartesian(
+        meristem.length,
+        meristem.inclination_angle,
+        meristem.azimuthal_angle,
+    );
+
+    let random_angle = rng.random::<f32>() * TAU;
+    let random_magnitude = rng.sample::<f32, StandardUniform>(StandardUniform) - 0.5;
+    let random_vector = Vec3::new(random_angle.cos(), 0.0, random_angle.sin())
+        * random_magnitude
+        * params.random_growth_factor;
+
+    let gravity_direction = accumulated_transform.rotation.inverse() * Vec3::NEG_Y;
+    // info!("gravity {} {}", original_vector, gravity_direction);
+    let mut rotation_axis = gravity_direction.cross(original_vector).normalize();
+    if rotation_axis.is_nan() {
+        let random_angle = rng.random::<f32>() * TAU;
+        let arbitrary = if gravity_direction.x.abs() < 1. {
+            Vec3::X
+        } else {
+            Vec3::Z
+        };
+        let perpendicular = gravity_direction.cross(arbitrary).normalize();
+        let rot = Quat::from_axis_angle(gravity_direction, random_angle);
+        rotation_axis = rot * perpendicular;
+    }
+
+    let gravitropic_direction =
+        Quat::from_axis_angle(rotation_axis, meristem.gravitropic_set_angle) * gravity_direction;
+    // info!("grav {} {}", gravitropic_direction, rotation_axis);
+    let towards_gravitropic =
+        (gravitropic_direction - original_vector.normalize()) * params.gravitropism_strength;
+    // info!(
+    //     "gravity {} {} {}",
+    //     gravity_direction,
+    //     gravitropic_direction,
+    //     original_vector.normalize()
+    // );
+    // info!("towards {}", towards_gravitropic);
+
+    let towards_straight = (Vec3::Y - original_vector) * params.straight_growth_factor;
+
+    let new_vector =
+        (original_vector + random_vector + towards_gravitropic + towards_straight).normalize();
+
+    let (_, inclination, azimuth) = cartesian_to_spherical(new_vector);
+
+    meristem.inclination_angle = inclination;
+    meristem.azimuthal_angle = azimuth;
+}
+
+fn spherical_to_cartesian(r: f32, inclination: f32, azimuth: f32) -> Vec3 {
+    let sin_inclination = inclination.sin();
+    Vec3::new(
+        r * sin_inclination * azimuth.sin(),
+        r * inclination.cos(),
+        r * sin_inclination * azimuth.cos(),
+    )
+}
+
+fn cartesian_to_spherical(vec: Vec3) -> (f32, f32, f32) {
+    let r = vec.length();
+    let inclination = (vec.y / r).acos();
+    let azimuth = vec.x.atan2(vec.z);
+    (r, inclination, azimuth)
+}

src/tree_rendering.rs

@@ -47,10 +47,6 @@ pub fn generate_tree(tree: &Tree) -> Mesh {
         let mut curr_node = node;
         loop {
             branches.push(curr_node.clone());
-            if curr_node.id == curr_node.parent_id {
-                break;
-            }
-
             if curr_node.is_branch_begin() {
                 break;
             }
@@ -77,7 +73,9 @@ fn add_extension(
     // prev_node: Option<&TreeNode>,
     transform: &mut Transform,
 ) {
-    *transform = transform.mul_transform(node.get_rotation());
+    *transform = transform.mul_transform(node.get_transform());
+    let current_circle_start_index = mesh_data.positions.len();
+    let prev_circle_start_index = current_circle_start_index - (RESOLUTION + 1);
     add_circle_points(
         mesh_data,
         RESOLUTION,
@@ -86,19 +84,20 @@ fn add_extension(
         Some(0.), // TODO: fix normal
     );
 
-    let next_circle_start_index = mesh_data.positions.len();
+    // let next_circle_start_index = mesh_data.positions.len();
-    let current_circle_start_index = next_circle_start_index - (RESOLUTION + 1);
     connect_circles(
         mesh_data,
         RESOLUTION,
+        prev_circle_start_index as u32,
         current_circle_start_index as u32,
-        next_circle_start_index as u32,
     );
-    *transform = transform.mul_transform(node.get_translation());
+    // *transform = transform.mul_transform(node.get_translation());
 }
 
 fn add_end(mesh_data: &mut TreeMeshData, node: &TreeNode, transform: &mut Transform) {
-    *transform = transform.mul_transform(node.get_rotation());
+    *transform = transform.mul_transform(node.get_transform());
+    let current_circle_start_index = mesh_data.positions.len();
+    let prev_circle_start_index = current_circle_start_index - (RESOLUTION + 1);
     add_circle_points(
         mesh_data,
         RESOLUTION,
@@ -106,36 +105,44 @@ fn add_end(mesh_data: &mut TreeMeshData, node: &TreeNode, transform: &mut Transf
         node.radius,
         Some(0.), // TODO: fix normal
     );
-    if matches!(node.kind, TreeNodeKind::EndCut) {
-        add_circle_points(mesh_data, RESOLUTION, *transform, node.radius, None);
-    } else if node.offset > node.radius {
-        let new_radius =
-            node.radius - node.radius * 0.5 * (node.offset - node.radius) / node.offset;
-        let new_length = node.offset - new_radius;
-        *transform = transform.mul_transform(Transform::from_translation(Vec3::Y * new_length));
-        let next_circle_start_index = mesh_data.positions.len();
-        add_circle_points(
-            mesh_data,
-            RESOLUTION,
-            *transform,
-            new_radius,
-            Some(0.), // TODO: fix normal
-        );
-        let current_circle_start_index = next_circle_start_index - (RESOLUTION + 1);
     connect_circles(
         mesh_data,
         RESOLUTION,
+        prev_circle_start_index as u32,
         current_circle_start_index as u32,
-            next_circle_start_index as u32,
     );
-        *transform = transform.mul_transform(Transform::from_translation(
-            Vec3::Y * (node.offset - new_length),
-        ));
-    } else {
-        *transform = transform.mul_transform(node.get_translation());
-    }
 
-    let circle_index_start = (mesh_data.positions.len() - (RESOLUTION + 1)) as u32;
+    if matches!(node.kind, TreeNodeKind::Cut) {
+        add_circle_points(mesh_data, RESOLUTION, *transform, node.radius, None);
+    } 
+    // else if node.offset > node.radius {
+    //     let new_radius =
+    //         node.radius - node.radius * 0.5 * (node.offset - node.radius) / node.offset;
+    //     let new_length = node.offset - new_radius;
+    //     // *transform = transform.mul_transform(Transform::from_translation(Vec3::Y * new_length));
+    //     let next_circle_start_index = mesh_data.positions.len();
+    //     add_circle_points(
+    //         mesh_data,
+    //         RESOLUTION,
+    //         *transform,
+    //         new_radius,
+    //         Some(0.), // TODO: fix normal
+    //     );
+    //     let current_circle_start_index = next_circle_start_index - (RESOLUTION + 1);
+    //     connect_circles(
+    //         mesh_data,
+    //         RESOLUTION,
+    //         current_circle_start_index as u32,
+    //         next_circle_start_index as u32,
+    //     );
+    //     // *transform = transform.mul_transform(Transform::from_translation(
+    //     //     Vec3::Y * (node.offset - new_length),
+    //     // ));
+    // } else {
+    //     // *transform = transform.mul_transform(node.get_translation());
+    // }
+
+    // let circle_index_start = (mesh_data.positions.len() - (RESOLUTION + 1)) as u32;
 
     let center_point = transform.transform_point(Vec3::ZERO);
     let center_point_index = mesh_data.positions.len() as u32;
@@ -147,11 +154,11 @@ fn add_end(mesh_data: &mut TreeMeshData, node: &TreeNode, transform: &mut Transf
     mesh_data.uvs.push([0.5, 0.5]); // TODO: UVs
 
     for j in 0..RESOLUTION {
-        let i0 = circle_index_start + j as u32;
+        let i0 = current_circle_start_index + j;
-        let i1 = circle_index_start + ((j + 1) % RESOLUTION) as u32;
+        let i1 = current_circle_start_index + ((j + 1) % RESOLUTION);
         mesh_data
             .indices
-            .extend_from_slice(&[center_point_index, i1, i0]);
+            .extend_from_slice(&[center_point_index, i1 as u32, i0 as u32]);
     }
 }
 
@@ -162,7 +169,16 @@ pub fn generate_branch_mesh(nodes: &[TreeNode]) -> Mesh {
 
     for (i, node) in nodes.iter().enumerate() {
         let next_node = nodes.get(i + 1);
-        if node.is_extension() {
+        if node.is_branch_begin() {
+            transform = transform.mul_transform(node.get_translation());
+            add_circle_points(
+                &mut mesh_data,
+                RESOLUTION,
+                transform,
+                node.radius,
+                Some(0.), // TODO: fix normal
+            );
+        }else if node.is_extension() {
             let next_node = next_node.unwrap(); // TODO: unwrap
             add_extension(&mut mesh_data, &node, &next_node, &mut transform);
         } else if node.is_end() {

src/tree_rendering2.rs

@@ -0,0 +1,215 @@
+use crate::tree2::{Tree, TreeNode, TreeNodeData};
+use bevy::{
+    asset::RenderAssetUsages,
+    prelude::*,
+    render::mesh::{Indices, PrimitiveTopology},
+};
+use std::f32::consts::PI;
+
+const RESOLUTION: usize = 10;
+
+struct TreeMeshData {
+    positions: Vec<[f32; 3]>,
+    normals: Vec<[f32; 3]>,
+    uvs: Vec<[f32; 2]>,
+    indices: Vec<u32>,
+    colors: Vec<[f32; 4]>,
+}
+
+impl Default for TreeMeshData {
+    fn default() -> Self {
+        TreeMeshData {
+            positions: Vec::new(),
+            normals: Vec::new(),
+            uvs: Vec::new(),
+            indices: Vec::new(),
+            colors: Vec::new(),
+        }
+    }
+}
+
+impl TreeMeshData {
+    fn to_mesh(self) -> Mesh {
+        Mesh::new(
+            PrimitiveTopology::TriangleList,
+            RenderAssetUsages::MAIN_WORLD | RenderAssetUsages::RENDER_WORLD,
+        )
+        .with_inserted_attribute(Mesh::ATTRIBUTE_POSITION, self.positions)
+        .with_inserted_attribute(Mesh::ATTRIBUTE_NORMAL, self.normals)
+        .with_inserted_attribute(Mesh::ATTRIBUTE_UV_0, self.uvs)
+        .with_inserted_attribute(Mesh::ATTRIBUTE_COLOR, self.colors)
+        .with_inserted_indices(Indices::U32(self.indices))
+    }
+}
+
+pub fn generate_tree(tree: &Tree) -> Mesh {
+    let mut mesh = TreeMeshData::default().to_mesh();
+
+    let mut branches = Vec::new();
+    for (&id, node) in tree.nodes.iter().filter(|(_, n)| n.is_branch_end()) {
+        let mut curr_node = node;
+        let mut last_node_id = id;
+        loop {
+            branches.push(curr_node.clone());
+            if curr_node.is_branch_begin() {
+                break;
+            }
+            last_node_id = curr_node.parent_id;
+            curr_node = &tree.nodes[&curr_node.parent_id];
+        }
+        branches.reverse();
+        let mut branch_mesh = generate_branch_mesh(branches.as_slice());
+        let transform = tree.get_accumulated_transform(last_node_id);
+        branch_mesh.transform_by(transform);
+        mesh.merge(&branch_mesh).unwrap(); // TODO: unwrap
+        branches.clear();
+    }
+
+    mesh
+}
+
+pub fn generate_branch_mesh(nodes: &[TreeNode]) -> Mesh {
+    let mut mesh_data = TreeMeshData::default();
+
+    let mut transform = Transform::IDENTITY;
+
+    for (i, node) in nodes.iter().enumerate() {
+        transform = transform.mul_transform(node.get_transform());
+        match &node.data {
+            TreeNodeData::Base(base) => {
+                add_circle_points(
+                    &mut mesh_data,
+                    RESOLUTION,
+                    transform,
+                    base.radius,
+                    Some(0.), // TODO: fix normal
+                    [1., 0., 0., 1.],
+                );
+            }
+            TreeNodeData::Bud(bud) => todo!(),
+            TreeNodeData::Meristem(meristem) => {
+                add_extension(
+                    &mut mesh_data,
+                    0.,
+                    &transform,
+                    [
+                        1.,
+                        1. - meristem.auxin_production_level,
+                        1. - meristem.auxin_production_level,
+                        1.,
+                    ],
+                );
+            }
+            TreeNodeData::Stem(stem) => {
+                add_extension(
+                    &mut mesh_data,
+                    stem.radius,
+                    &transform,
+                    [
+                        1.,
+                        (1. - stem.auxin_level).powi(2),
+                        (1. - stem.auxin_level).powi(2),
+                        1.,
+                    ],
+                );
+            }
+            TreeNodeData::CutStem(cut_stem) => {
+                add_extension(
+                    &mut mesh_data,
+                    cut_stem.radius,
+                    &transform,
+                    [1., 1., 1., 1.],
+                );
+                // TODO: add end
+            }
+            TreeNodeData::Branch(branch) => {
+                add_circle_points(
+                    &mut mesh_data,
+                    RESOLUTION,
+                    transform,
+                    branch.radius,
+                    Some(0.), // TODO: fix normal
+                    [
+                        1.,
+                        (1. - branch.auxin_level).powi(2),
+                        (1. - branch.auxin_level).powi(2),
+                        1.,
+                    ],
+                );
+            }
+        }
+    }
+
+    mesh_data.to_mesh()
+}
+
+fn add_extension(
+    mesh_data: &mut TreeMeshData,
+    radius: f32,
+    transform: &Transform,
+    color: [f32; 4],
+) {
+    let current_circle_start_index = mesh_data.positions.len();
+    let prev_circle_start_index = current_circle_start_index - (RESOLUTION + 1);
+    add_circle_points(
+        mesh_data,
+        RESOLUTION,
+        *transform,
+        radius,
+        Some(0.), // TODO: fix normal
+        color,
+    );
+
+    connect_circles(
+        mesh_data,
+        RESOLUTION,
+        prev_circle_start_index as u32,
+        current_circle_start_index as u32,
+    );
+}
+
+fn add_circle_points(
+    mesh_data: &mut TreeMeshData,
+    resolution: usize,
+    transform: Transform,
+    radius: f32,
+    local_y_normal: Option<f32>,
+    color: [f32; 4],
+) {
+    for i in 0..=resolution {
+        let theta = i as f32 / RESOLUTION as f32 * 2.0 * PI;
+        let (cos, sin) = (theta.cos(), theta.sin());
+
+        let normal = if let Some(local_y_normal) = local_y_normal {
+            Vec3::new(cos, local_y_normal, sin).normalize()
+        } else {
+            Vec3::Y
+        };
+        let normal = transform.rotation * normal;
+        let position = transform.transform_point(Vec3::new(radius * cos, 0., radius * sin));
+
+        mesh_data.positions.push(position.to_array());
+        mesh_data.normals.push(normal.into());
+        mesh_data.colors.push(color);
+        // TODO: UVs
+        mesh_data.uvs.push([0.5, 0.5]);
+    }
+}
+
+fn connect_circles(
+    mesh_data: &mut TreeMeshData,
+    resolution: usize,
+    index_circle1: u32,
+    index_circle2: u32,
+) {
+    for i in 0..resolution {
+        let curr = index_circle1 + i as u32;
+        let next = curr + 1;
+        let above = index_circle2 + i as u32;
+        let above_next = above + 1;
+
+        mesh_data
+            .indices
+            .extend_from_slice(&[curr, above, next, next, above, above_next]);
+    }
+}