Refactor. Trees and Forests.

python_scripts/add_trees_to_open_street_map/add_trees.py

@@ -19,7 +19,7 @@ from shapely import LineString, geometry, wkt
 from shapely.ops import transform
 
 from tree import Forest
-from import_existing_trees import get_existing_trees
+from import_existing_trees import get_existing_forest
 
 # TODO: Use Args
 # TODO: Document
@@ -99,17 +99,9 @@ def set_plot(bounds, to_local_coordinates):
     return ax
 
 
-def place_trees(existing_trees_coords, ways, region, to_local, tree_distance, min_distance_2):
+def place_trees(forest, ways, region, to_local, tree_distance, min_distance_2):
     local_region = transform(to_local.transform, region)
 
-    existing_trees = kdtree.create(existing_trees_coords or [(0, 0)], dimensions=2)
-
-    tree_xs = []
-    tree_ys = []
-    for x, y in existing_trees_coords:
-        tree_xs.append(x)
-        tree_ys.append(y)
-
     for way in ways:
         width = float(way.tags.get("width", 0))
         highway = way.tags.get("highway")
@@ -145,12 +137,9 @@ def place_trees(existing_trees_coords, ways, region, to_local, tree_distance, mi
             x = potential_tree.x
             y = potential_tree.y
             if local_region.contains(geometry.Point(x, y)):
-                _nearest_tree, distance_2 = existing_trees.search_nn((x, y))
-                if distance_2 > min_distance_2:
-                    existing_trees.add((x, y))
-                    tree_xs.append(x)
-                    tree_ys.append(y)
-    return tree_xs, tree_ys
+                forest.add_tree_if_possible(min_distance_2, x, y)
+
+    return forest.xs_ys
 
 
 def plot_trees(bounds, tree_xs, tree_ys, tree_distance):
@@ -214,15 +203,17 @@ def main(wkt_polygon, epsg_id, tree_distance, min_distance, import_tree_shp):
     ways = get_osm_roads(bounds)
 
     if import_tree_shp:
-        existing_trees = get_existing_trees(import_tree_shp)
+        existing_forest = get_existing_forest(import_tree_shp)
     else:
-        existing_trees = Forest()
+        existing_forest = Forest()
+
+    print(existing_forest)
 
     to_local = Transformer.from_crs("EPSG:4326", f"EPSG:{epsg_id}", always_xy=True)
 
     set_plot(bounds, to_local)
-    tree_xs, tree_ys = place_trees(existing_trees, ways, region,
-                                   to_local, tree_distance, min_distance**2)
+    tree_xs, tree_ys = place_trees(existing_forest, ways, region, to_local, tree_distance, min_distance**2)
+    print(existing_forest)
 
     plot_trees(bounds, tree_xs, tree_ys, tree_distance)
     export_map(bounds, tree_xs, tree_ys, epsg_id)

python_scripts/add_trees_to_open_street_map/import_existing_trees.py

 import geopandas as gpd
 from tree import Tree, Forest
 
-def get_existing_trees(shp_input):
+def get_existing_forest(shp_input):
     print(f"Importing {shp_input}")
     df = gpd.read_file(shp_input)
     trees = [Tree(p.x, p.y) for p in df.geometry]
     return Forest(trees)
 
 if __name__ == "__main__":
-    print(repr(get_existing_trees('existing_trees/Trees_ideal_2_20240227.shp')))
+    print(repr(get_existing_forest('existing_trees/Trees_ideal_2_20240227.shp')))

python_scripts/add_trees_to_open_street_map/tree.py

 from dataclasses import dataclass
 import kdtree
 
+
 @dataclass
 class Tree:
     x: float
@@ -12,7 +13,7 @@ class Tree:
     type: str = None
 
     def __len__(self):
-        return 2 # x & y
+        return 2  # x & y
 
     def __getitem__(self, i):
         return [self.x, self.y][i]
@@ -27,9 +28,22 @@ class Forest:
 
         self.trees = existing_trees
 
+    def add_tree_if_possible(self, min_distance_2, x, y, *params):
+        _nearest_tree, distance_2 = self.kd_tree.search_nn((x, y))
+        if distance_2 > min_distance_2:
+            self.kd_tree.add((x, y))
+            self.trees.append(Tree(x, y, *params))
+
+    @property
+    def xs_ys(self):
+        xs, ys = [], []
+        for tree in self.trees:
+            xs.append(tree.x)
+            ys.append(tree.y)
+        return xs, ys
+
     def __str__(self):
         return f"Forest with {len(self.trees)} trees."
 
     def __repr__(self):
         return "\n".join([str(self)] + [str(tree) for tree in self.trees])
-