Add color mapping and fix make_n_trees.py

Author: abhinavj98

__init__.py

@@ -0,0 +1,5 @@
+"""LPY Tree Simulation package."""
+
+from .color_manager import ColorManager
+
+__all__ = ["ColorManager"]

base_lpy.lpy

@@ -23,35 +23,27 @@ from dataclasses import dataclass
 # Override these externs via Lsystem(..., variables) to point at a specific
 # architecture without editing this shared file.
 extern(
+    color_manager = None,
     prototype_dict_path = "examples.Envy.Envy_prototypes.basicwood_prototypes",
     trunk_class_path = "examples.Envy.Envy_prototypes.Trunk",
     simulation_config_class_path = "examples.Envy.Envy_simulation.EnvySimulationConfig",
-    point_generator_path = "examples.Envy.Envy_simulation.generate_points_v_trellis",
-    energy_function_path = "examples.Envy.Envy_simulation.get_energy_mat",
-    decide_function_path = "examples.Envy.Envy_simulation.decide_guide",
-    tie_function_path = "examples.Envy.Envy_simulation.tie",
-    prune_function_path = "examples.Envy.Envy_simulation.prune",
+    simulation_class_path = "examples.Envy.Envy_simulation.EnvySimulation",
     axiom_pitch = 0.0,
     axiom_yaw = 0.0,
-)
-
-# Resolve extern-provided dotted paths so the rest of the script can operate on
+)# Resolve extern-provided dotted paths so the rest of the script can operate on
 # concrete classes/functions exactly like the Envy/UFO drivers do.
 basicwood_prototypes = resolve_attr(prototype_dict_path)
 Trunk = resolve_attr(trunk_class_path)
 SimulationConfigClass = resolve_attr(simulation_config_class_path)
-point_generator_fn = resolve_attr(point_generator_path)
-get_energy_mat = resolve_attr(energy_function_path)
-decide_guide = resolve_attr(decide_function_path)
-tie = resolve_attr(tie_function_path)
-prune = resolve_attr(prune_function_path)
+SimulationClass = resolve_attr(simulation_class_path)
 
 simulation_config = SimulationConfigClass()
+simulation = SimulationClass(simulation_config)
 main_trunk = Trunk(copy_from = basicwood_prototypes['trunk'])
 
 # Build trellis geometry and cadence numbers up front so callbacks stay simple.
 trellis_support = Support(
-    point_generator_fn(simulation_config),
+  simulation.generate_points(),
     simulation_config.support_num_wires,
     simulation_config.support_spacing_wires,
     simulation_config.support_trunk_wire_point,
@@ -62,7 +54,30 @@ main_trunk.tying.guide_target = trellis_support.trunk_wire
 
 # Track parent/child relationships for tying/pruning decisions.
 branch_hierarchy = {main_trunk.name: []}
-enable_color_labeling = simulation_config.label
+enable_semantic_labeling = simulation_config.semantic_label
+enable_instance_labeling = simulation_config.instance_label
+enable_per_cylinder_labeling = simulation_config.per_cylinder_label
+
+if (enable_instance_labeling or enable_per_cylinder_labeling) and color_manager is None:
+  raise ValueError(
+    "Instance labeling is enabled but no color_manager extern was provided."
+  )
+
+
+def _get_energy_mat(branches, arch, _config):
+  return simulation.get_energy_mat(branches, arch)
+
+
+def _decide_guide(energy_matrix, branches, arch, _config):
+  return simulation.decide_guide(energy_matrix, branches, arch)
+
+
+def _tie(lstring, _config):
+  return simulation.tie(lstring)
+
+
+def _prune(lstring, _config):
+  return simulation.prune(lstring)
 
 # L-Py callbacks delegate to the Python helpers so architectures share logic.
 def StartEach(lstring):
@@ -81,12 +96,12 @@ def EndEach(lstring):
         simulation_config,
         main_trunk,
         getIterationNb,
-        get_energy_mat,
-        decide_guide,
-        tie,
-        prune,
+        _get_energy_mat,
+        _decide_guide,
+        _tie,
+        _prune,
     )
-
+generalized_cylinder = getattr(simulation_config, "use_generalized_cylinder", False)
 # =============================================================================
 # L-SYSTEM GRAMMAR DEFINITION
 # =============================================================================
@@ -149,18 +164,28 @@ grow_object(plant_segment) :
     #Update internal state of the plant segment
     plant_segment.grow_one()
     #Update physical representation
-    if enable_color_labeling:
+    if enable_semantic_labeling:
       # Add color visualization if labeling is enabled -- Can this be moved to the start of building the segment?
       r, g, b = plant_segment.info.color
+      plant_part = plant_segment.name.split("_")[0] # Get the part type (e.g., "trunk", "branch")
+      color_manager.set_unique_color((r,g,b), plant_part)  # Ensure part type has a color assig
+      nproduce SetColor(r,g,b)
+    if enable_instance_labeling:
+        # Instance-level labeling (unique color per branch instance)
+      r, g, b = color_manager.get_unique_color(plant_segment.name)
       nproduce SetColor(r,g,b)
+
     #Produce internode segments (n cylinders per growth step)
     n_cylinders = int(plant_segment.growth.growth_length / plant_segment.growth.cylinder_length)
     for i in range(n_cylinders):
+      if enable_per_cylinder_labeling:
+        r, g, b = color_manager.get_unique_color(plant_segment.name, if_exists=False)
+        nproduce SetColor(r,g,b)  
       nproduce I(plant_segment.growth.cylinder_length, plant_segment.growth.thickness, plant_segment)
     #Produce bud (after all cylinders in this growth step)
     if plant_segment.pre_bud_rule(plant_segment, simulation_config):
       for generated in plant_segment.post_bud_rule(plant_segment, simulation_config):
-              nproduce new(generated[0], *generated[1])
+        nproduce new(generated[0], *generated[1])
 
     if should_bud(plant_segment, simulation_config):
         # Age-based bud production for lateral branching
@@ -190,15 +215,18 @@ bud(bud_parameters) :
    bud_parameters.num_buds+=1
    bud_parameters.type.info.num_branches+=1
 
+   nproduce [
+   if generalized_cylinder:
+     nproduce @Gc
    if hasattr(new_branch, 'curve_x_range'):
      # Curved branches: set up custom growth guide curve
-     curve = create_bezier_curve(x_range=new_branch.curve_x_range, y_range=new_branch.curve_y_range, z_range=new_branch.curve_z_range, seed_value=time.time())
-     nproduce[@GcSetGuide(curve, new_branch.growth.max_length)
-   else:
-     # Straight branches: use default orientation
-     nproduce [@Gc
+     curve = create_bezier_curve(x_range=new_branch.curve_x_range, y_range=new_branch.curve_y_range, z_range=new_branch.curve_z_range, seed_value=rd.randint(0,1000))
+     nproduceSetGuide(curve, new_branch.growth.max_length)
    # Produce new branch with random orientation and growth object
-   nproduce @RGetPos(new_branch.location.start)WoodStart(ParameterSet(type = new_branch))/(rd.random()*360)&(rd.random()*360)grow_object(new_branch)GetPos(new_branch.location.end)@Ge]bud(bud_parameters)
+   nproduce @RGetPos(new_branch.location.start)WoodStart(ParameterSet(type = new_branch))/(rd.random()*360)&(rd.random()*360)grow_object(new_branch)GetPos(new_branch.location.end)
+   if generalized_cylinder:
+     nproduce @Ge
+   produce ]bud(bud_parameters)
 
 # -----------------------------------------------------------------------------
 # GEOMETRIC INTERPRETATION (HOMOMORPHISM)

color_manager.py

@@ -0,0 +1,53 @@
+"""
+Color management utilities for L-Py tree simulation system.
+
+This module provides utilities for assigning unique colors to tree segments
+for visualization and labeling purposes.
+"""
+
+import itertools
+import json
+
+
+class ColorManager:
+    """Manages assignment of unique colors to named entities."""
+    
+    def __init__(self):
+        self.color_to_name = {}
+        self.name_to_color = {}
+        # Permute all possible colors to make a list
+        self.all_colors = list(itertools.product(range(256), repeat=3))  # 0-255 inclusive
+        self.color_pointer = 0
+
+    def get_unique_color(self, name, if_exists=True):
+        """Get a unique RGB color tuple for the given name."""
+        if name in self.name_to_color and if_exists:
+            return self.name_to_color[name]
+        
+        if self.color_pointer >= len(self.all_colors):
+            raise ValueError("Ran out of unique colors!")
+        
+        unique_color = self.all_colors[self.color_pointer]
+        self.color_pointer += 1
+
+        # Assign and save mapping
+        self.name_to_color[name] = unique_color
+        self.color_to_name[unique_color] = name
+
+        return unique_color
+
+    def export_mapping(self, filename):
+        """Export color -> name mapping to JSON"""
+        export_dict = {}
+        for color, name in self.color_to_name.items():
+            export_dict[str(color)] = name
+
+        with open(filename, "w") as f:
+            json.dump(export_dict, f, indent=4)
+
+        print(f"Exported color mappings to {filename}")
+
+    def set_unique_color(self, color, name):
+        """Set a specific color for a given name."""
+        self.name_to_color[name] = color
+        self.color_to_name[color] = name

examples/Envy/Envy.lpy

@@ -2,23 +2,25 @@ from pathlib import Path
 from openalea.lpy import Lsystem
 from openalea.lpy import *
 from openalea.plantgl.all import *
+from lpy_treesim import ColorManager
 
 BASE_LPY_PATH = Path(__file__).resolve().parents[2] / "base_lpy.lpy"
 
+color_manager = ColorManager()
+
 vars = {
     "prototype_dict_path": "examples.Envy.Envy_prototypes.basicwood_prototypes",
     "trunk_class_path": "examples.Envy.Envy_prototypes.Trunk",
     "simulation_config_class_path": "examples.Envy.Envy_simulation.EnvySimulationConfig",
-    "point_generator_path": "examples.Envy.Envy_simulation.generate_points_v_trellis",
-    "energy_function_path": "examples.Envy.Envy_simulation.get_energy_mat",
-    "decide_function_path": "examples.Envy.Envy_simulation.decide_guide",
-    "tie_function_path": "examples.Envy.Envy_simulation.tie",
-    "prune_function_path": "examples.Envy.Envy_simulation.prune",
+    "simulation_class_path": "examples.Envy.Envy_simulation.EnvySimulation",
+    "color_manager": color_manager,
     "axiom_pitch": 0.0,
     "axiom_yaw": 0.0,
 }
 
 lsystem = Lsystem(str(BASE_LPY_PATH), vars)
 
 for lstring in lsystem:
-    lsystem.plot(lstring)
+    lsystem.plot(lstring)
+
+# color_manager.export_mapping("envy_color_mapping.json")

examples/Envy/Envy_simulation.py

@@ -25,8 +25,9 @@ class EnvySimulationConfig(SimulationConfig):
     trellis_z_spacing: float = 0.45
 
     # Envy-specific Growth Parameters
-    growth_length: float = 0.1
-    bud_spacing_age: int = 2
+    semantic_label: bool = True
+    instance_label: bool = False
+    per_cylinder_label: bool = False
 
 
 class EnvySimulation(TreeSimulationBase):
@@ -60,31 +61,4 @@ def generate_points(self):
             pts.append((x[i], y[i], z[i]))
         return pts
 
-
-# Backwards compatibility: provide standalone functions that use the class
-def generate_points_v_trellis(simulation_config):
-    """Backward compatibility wrapper for generate_points."""
-    sim = EnvySimulation(simulation_config)
-    return sim.generate_points()
-
-def get_energy_mat(branches, arch, simulation_config):
-    """Backward compatibility wrapper for get_energy_mat."""
-    sim = EnvySimulation(simulation_config)
-    return sim.get_energy_mat(branches, arch)
-
-def decide_guide(energy_matrix, branches, arch, simulation_config):
-    """Backward compatibility wrapper for decide_guide."""
-    sim = EnvySimulation(simulation_config)
-    return sim.decide_guide(energy_matrix, branches, arch)
-
-def prune(lstring, simulation_config):
-    """Backward compatibility wrapper for prune."""
-    sim = EnvySimulation(simulation_config)
-    return sim.prune(lstring)
-
-def tie(lstring, simulation_config):
-    """Backward compatibility wrapper for tie."""
-    sim = EnvySimulation(simulation_config)
-    return sim.tie(lstring)
-
 

examples/UFO/UFO.lpy

@@ -2,18 +2,18 @@ from pathlib import Path
 from openalea.lpy import Lsystem
 from openalea.lpy import *
 from openalea.plantgl.all import *
+from lpy_treesim import ColorManager
 
 BASE_LPY_PATH = Path(__file__).resolve().parents[2] / "base_lpy.lpy"
 
+color_manager = ColorManager()
+
 vars = {
     "prototype_dict_path": "examples.UFO.UFO_prototypes.basicwood_prototypes",
     "trunk_class_path": "examples.UFO.UFO_prototypes.Trunk",
     "simulation_config_class_path": "examples.UFO.UFO_simulation.UFOSimulationConfig",
-    "point_generator_path": "examples.UFO.UFO_simulation.generate_points_ufo",
-    "energy_function_path": "examples.UFO.UFO_simulation.get_energy_mat",
-    "decide_function_path": "examples.UFO.UFO_simulation.decide_guide",
-    "tie_function_path": "examples.UFO.UFO_simulation.tie",
-    "prune_function_path": "examples.UFO.UFO_simulation.prune",
+    "simulation_class_path": "examples.UFO.UFO_simulation.UFOSimulation",
+    "color_manager": color_manager,
     "axiom_pitch": 270.0,
     "axiom_yaw": 0.0,
 }
@@ -22,3 +22,5 @@ lsystem = Lsystem(str(BASE_LPY_PATH), vars)
 
 for lstring in lsystem:
     lsystem.plot(lstring)
+
+# color_manager.export_mapping("ufo_color_mapping.json")

examples/UFO/UFO_simulation.py

@@ -26,6 +26,9 @@ class UFOSimulationConfig(SimulationConfig):
 
     # UFO-specific Growth Parameters
     thickness_multiplier: float = 1.2  # Multiplier for internode thickness
+    semantic_label: bool = True
+    instance_label: bool = False
+    per_cylinder_label: bool = False
 
 
 class UFOSimulation(TreeSimulationBase):
@@ -61,31 +64,3 @@ def generate_points(self):
             wire_attachment_points.append((x[point_index], y[point_index], z[point_index]))
 
         return wire_attachment_points
-
-
-# Backwards compatibility: provide standalone functions that use the class
-# Backwards compatibility: provide standalone functions that use the class
-def generate_points_ufo(simulation_config):
-    """Backward compatibility wrapper for generate_points."""
-    sim = UFOSimulation(simulation_config)
-    return sim.generate_points()
-
-def get_energy_mat(branches, arch, simulation_config):
-    """Backward compatibility wrapper for get_energy_mat."""
-    sim = UFOSimulation(simulation_config)
-    return sim.get_energy_mat(branches, arch)
-
-def decide_guide(energy_matrix, branches, arch, simulation_config):
-    """Backward compatibility wrapper for decide_guide."""
-    sim = UFOSimulation(simulation_config)
-    return sim.decide_guide(energy_matrix, branches, arch)
-
-def prune(lstring, simulation_config):
-    """Backward compatibility wrapper for prune."""
-    sim = UFOSimulation(simulation_config)
-    return sim.prune(lstring)
-
-def tie(lstring, simulation_config):
-    """Backward compatibility wrapper for tie."""
-    sim = UFOSimulation(simulation_config)
-    return sim.tie(lstring)