BERATools 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

beratools/core/algo_line_grouping.py

@@ -12,6 +12,7 @@ Description:
 
     This file hosts code to deal with line grouping and merging, cleanups.
 """
+
 import enum
 from collections import defaultdict
 from dataclasses import dataclass, field
@@ -30,6 +31,7 @@ import beratools.core.constants as bt_const
 TRIMMING_DISTANCE = 75  # meters
 SMALL_BUFFER = 1
 
+
 @enum.unique
 class VertexClass(enum.IntEnum):
     """Enum class for vertex class."""
@@ -107,7 +109,7 @@ def get_angle(line, end_index):
 @dataclass
 class SingleLine:
     """Class to store line and its simplified line."""
-    
+
     line_id: int = field(default=0)
     line: Union[sh_geom.LineString, sh_geom.MultiLineString] = field(default=None)
     sim_line: Union[sh_geom.LineString, sh_geom.MultiLineString] = field(default=None)
@@ -171,7 +173,7 @@ class VertexNode:
 
     def get_line_geom(self, line_id):
         return self.get_line_obj(line_id).line
-            
+
     def get_all_line_ids(self):
         all_line_ids = {i.line_id for i in self.line_list}
         return all_line_ids
@@ -188,7 +190,7 @@ class VertexNode:
     def get_trim_transect(self, poly, line_indices):
         if not poly:
             return None
-        
+
         internal_line = None
         for line_idx in line_indices:
             line = self.get_line_obj(line_idx)
@@ -197,13 +199,13 @@ class VertexNode:
 
         if not internal_line:
             # print("No line is retrieved")
-            return
+            return None
         return internal_line.end_transect()
-    
+
     def _trim_polygon(self, poly, trim_transect):
         if not poly or not trim_transect:
-            return
-        
+            return None
+
         split_poly = shapely.ops.split(poly, trim_transect)
 
         if len(split_poly.geoms) != 2:
@@ -216,7 +218,7 @@ class VertexNode:
                 none_poly = True
 
         if none_poly:
-            return
+            return None
 
         # only two polygons in split_poly
         if split_poly.geoms[0].area > split_poly.geoms[1].area:
@@ -242,15 +244,15 @@ class VertexNode:
                 new_polys.append((idx, out_poly))
 
         return new_polys
-    
+
     def trim_end(self, poly):
         transect = self.get_trim_transect(poly, self.line_not_connected)
         if not transect:
-            return
-        
+            return None
+
         poly = self._trim_polygon(poly, transect)
         return poly
-            # Helper to get the neighbor coordinate based on vertex_index.
+        # Helper to get the neighbor coordinate based on vertex_index.
 
     @staticmethod
     def get_vertex(line_obj, index):
@@ -260,6 +262,7 @@ class VertexNode:
             index += len(coords)
         if 0 <= index < len(coords):
             return sh_geom.Point(coords[index])
+        return None
 
     @staticmethod
     def get_neighbor(line_obj):
@@ -269,9 +272,9 @@ class VertexNode:
             index = 1
         elif line_obj.vertex_index == -1:
             index = -2
-        
+
         return VertexNode.get_vertex(line_obj, index)
-    
+
     @staticmethod
     def parallel_line_centered(p1, p2, center, length):
         """Generate a parallel line."""
@@ -295,7 +298,7 @@ class VertexNode:
         new_p2 = sh_geom.Point(center.x + half_dx, center.y + half_dy)
 
         return sh_geom.LineString([new_p1, new_p2])
-        
+
     def get_transect_for_primary(self):
         """
         Get a transect line from two primary connected lines.
@@ -309,7 +312,7 @@ class VertexNode:
         """
         if not self.line_connected or len(self.line_connected[0]) != 2:
             return None
-        
+
         # Retrieve the two connected line objects from the first connectivity group.
         line_ids = self.line_connected[0]
         pt1 = None
@@ -330,20 +333,18 @@ class VertexNode:
         if pt1 is None or pt2 is None:
             return None
 
-        transect = algo_common.generate_perpendicular_line_precise(
-            [pt1, self.vertex, pt2], offset=40
-        )
+        transect = algo_common.generate_perpendicular_line_precise([pt1, self.vertex, pt2], offset=40)
         return transect
-    
+
     def get_transect_for_primary_second(self):
         """
         Get a transect line from the second primary connected line.
-        
-        For the second primary line, this method retrieves the neighbor point from 
-        two lines in the second connectivity group, creates a reference line through the 
-        vertex by mirroring the neighbor point about the vertex, and then generates a 
+
+        For the second primary line, this method retrieves the neighbor point from
+        two lines in the second connectivity group, creates a reference line through the
+        vertex by mirroring the neighbor point about the vertex, and then generates a
         parallel line centered at the vertex.
-        
+
         Returns:
             A LineString representing the transect if available, otherwise None.
 
@@ -378,8 +379,8 @@ class VertexNode:
 
         """
         if len(self.line_connected) == 0:
-            return
-        
+            return None
+
         new_polys = []
         line = self.line_connected[0]
 
@@ -415,23 +416,29 @@ class VertexNode:
             new_polys.append([idx_2, poly_2])
 
         return new_polys
-    
+
     def trim_intersection(self, polys, merge_group=True):
-        """Trim intersection of lines and polygons."""
+        """
+        Trim intersection of lines and polygons.
+
+        TODO: there are polygons of 0 zero.
+
+        """
+
         def get_poly_with_info(line, polys):
             if polys.empty:
                 return None, None, None
-            
+
             for idx, row in polys.iterrows():
                 poly = row.geometry
-                if not poly: # TODO: no polygon
+                if not poly:  # TODO: no polygon
                     continue
 
                 if poly.buffer(SMALL_BUFFER).contains(line):
-                    return idx, poly, row['max_width']
-            
+                    return idx, poly, row["max_width"]
+
             return None, None, None
-                
+
         poly_trim_list = []
         primary_lines = []
         p_primary_list = []
@@ -458,7 +465,8 @@ class VertexNode:
             for line_idx in line_idx_to_trim:
                 line = self.get_line_geom(line_idx)
                 poly_idx, poly, max_width = get_poly_with_info(line, polys)
-                poly_list.append((line_idx, poly_idx, max_width))
+                if poly_idx:
+                    poly_list.append((line_idx, poly_idx, max_width))
 
             poly_list = sorted(poly_list, key=lambda x: x[2])
 
@@ -466,7 +474,7 @@ class VertexNode:
         for i, indices in enumerate(poly_list):
             line_idx = indices[0]
             poly_idx = indices[1]
-            line_cleanup=self.get_line(line_idx)
+            line_cleanup = self.get_line(line_idx)
             poly_cleanup = polys.loc[poly_idx].geometry
             poly_trim = PolygonTrimming(
                 line_index=line_idx,
@@ -486,7 +494,7 @@ class VertexNode:
         # poly_list and poly_trim_list have same index
         for i, indices in enumerate(poly_list):
             p_list = []
-            for p in poly_list[i+1:]:
+            for p in poly_list[i + 1 :]:
                 p_list.append(polys.loc[p[1]].geometry)
 
             poly_trim = poly_trim_list[i]
@@ -520,7 +528,7 @@ class VertexNode:
         elif len(self.line_list) == 1:
             self.vertex_class = VertexClass.SINGLE_WAY
 
-    def has_group_attr(self):
+    def all_has_valid_group_attr(self):
         """If all values in group list are valid value, return True."""
         # TODO: if some line has no group, give advice
         for i in self.line_list:
@@ -532,15 +540,19 @@ class VertexNode:
     def need_regrouping(self):
         pass
 
-    def check_connectivity(self):
-        # TODO add regrouping when new lines are added
-        if self.has_group_attr():
-            if self.need_regrouping():
-                self.group_regroup()
-            else:
-                self.group_line_by_attribute()
-        else:
+    def check_connectivity(self, use_angle_grouping=True):
+        # Fill missing group with -1
+        for line in self.line_list:
+            if line.group is None:
+                line.group = -1
+
+        if self.need_regrouping():
+            self.group_regroup()
+
+        if use_angle_grouping:
             self.group_line_by_angle()
+        else:
+            self.update_connectivity_by_group()
 
         # record line not connected
         all_line_ids = self.get_all_line_ids()
@@ -551,7 +563,7 @@ class VertexNode:
     def group_regroup(self):
         pass
 
-    def group_line_by_attribute(self):
+    def update_connectivity_by_group(self):
         group_line = defaultdict(list)
         for i in self.line_list:
             group_line[i.group].append(i.line_id)
@@ -573,6 +585,7 @@ class VertexNode:
             angle_diff = abs(new_angles[0] - new_angles[1])
             angle_diff = angle_diff if angle_diff <= np.pi else angle_diff - np.pi
 
+            # TODO: check if need to connect when turn angle < TURN_ANGLE_TOLERANCE
             # if angle_diff >= TURN_ANGLE_TOLERANCE:
             self.line_connected.append(
                 (
@@ -587,14 +600,10 @@ class VertexNode:
             for j, angle_2 in enumerate(new_angles[i + 1 :]):
                 if not angle_visited[i + j + 1]:
                     angle_diff = abs(angle_1 - angle_2)
-                    angle_diff = (
-                        angle_diff if angle_diff <= np.pi else angle_diff - np.pi
-                    )
+                    angle_diff = angle_diff if angle_diff <= np.pi else angle_diff - np.pi
                     if (
                         angle_diff < ANGLE_TOLERANCE
-                        or np.pi - ANGLE_TOLERANCE
-                        < abs(angle_1 - angle_2)
-                        < np.pi + ANGLE_TOLERANCE
+                        or np.pi - ANGLE_TOLERANCE < abs(angle_1 - angle_2) < np.pi + ANGLE_TOLERANCE
                     ):
                         angle_visited[j + i + 1] = True  # tenth of PI
                         self.line_connected.append(
@@ -607,19 +616,15 @@ class VertexNode:
 
 class LineGrouping:
     """Class to group lines and merge them."""
-    
-    def __init__(self, in_line_gdf, merge_group=True) -> None:
-        # remove empty and null geometry
-        # self.lines = in_line_gdf.copy()
-        # self.lines = self.lines[
-        #     ~self.lines.geometry.isna() & ~self.lines.geometry.is_empty
-        # ]
+
+    def __init__(self, in_line_gdf, merge_group=True, use_angle_grouping=True) -> None:
         if in_line_gdf is None:
             raise ValueError("Line GeoDataFrame cannot be None")
+        self.use_angle_grouping = use_angle_grouping
 
         if in_line_gdf.empty:
             raise ValueError("Line GeoDataFrame cannot be empty")
-        
+
         self.lines = algo_common.clean_line_geometries(in_line_gdf)
         self.lines.reset_index(inplace=True, drop=True)
         self.merge_group = merge_group
@@ -650,17 +655,19 @@ class LineGrouping:
         if bt_const.BT_GROUP in self.lines.keys():
             self.groups = self.lines[bt_const.BT_GROUP]
             self.has_group_attr = True
-            if self.groups.hasnans:
+            if self.groups.hasnans:  # Todo: check for other invalid values
                 self.need_regrouping = True
 
-        for idx, s_geom, geom, group in zip(
-            *zip(*self.sim_geom.items()), self.lines.geometry, self.groups
-        ):
+        for idx, s_geom, geom, group in zip(*zip(*self.sim_geom.items()), self.lines.geometry, self.groups):
             self.vertex_list.append(VertexNode(idx, geom, s_geom, 0, group))
             self.vertex_list.append(VertexNode(idx, geom, s_geom, -1, group))
 
         v_points = []
         for i in self.vertex_list:
+            if i.vertex is None:
+                print("Vertex is None, skipping.")
+                continue
+
             v_points.append(i.vertex.buffer(SMALL_BUFFER))  # small polygon
 
         # Spatial index of all vertices
@@ -677,10 +684,7 @@ class LineGrouping:
                 for j in s_list:
                     if j != i:
                         # some short line will be very close to each other
-                        if (
-                            vertex.vertex.distance(self.vertex_list[j].vertex)
-                            > bt_const.SMALL_BUFFER
-                        ):
+                        if vertex.vertex.distance(self.vertex_list[j].vertex) > bt_const.SMALL_BUFFER:
                             continue
 
                         vertex.merge(self.vertex_list[j])
@@ -690,7 +694,7 @@ class LineGrouping:
             vertex_visited[i] = True
 
         for i in self.merged_vertex_list:
-            i.check_connectivity()
+            i.check_connectivity(self.use_angle_grouping)
 
         for i in self.merged_vertex_list:
             if i.line_connected:
@@ -721,9 +725,7 @@ class LineGrouping:
         return algo_merge_lines.run_line_merge(self.lines, self.merge_group)
 
     def find_vertex_for_poly_trimming(self):
-        self.vertex_of_concern = [
-            i for i in self.merged_vertex_list if i.vertex_class in CONCERN_CLASSES
-        ]
+        self.vertex_of_concern = [i for i in self.merged_vertex_list if i.vertex_class in CONCERN_CLASSES]
 
     def line_and_poly_cleanup(self):
         sindex_poly = self.polys.sindex
@@ -732,20 +734,21 @@ class LineGrouping:
             s_idx = sindex_poly.query(vertex.vertex, predicate="within")
             if len(s_idx) == 0:
                 continue
-            
+
             #  Trim intersections of primary lines
             polys = self.polys.loc[s_idx].geometry
             if not self.merge_group:
-                if (vertex.vertex_class == VertexClass.FIVE_WAY_TWO_PRIMARY_LINE
+                if (
+                    vertex.vertex_class == VertexClass.FIVE_WAY_TWO_PRIMARY_LINE
                     or vertex.vertex_class == VertexClass.FIVE_WAY_ONE_PRIMARY_LINE
                     or vertex.vertex_class == VertexClass.FOUR_WAY_ONE_PRIMARY_LINE
                     or vertex.vertex_class == VertexClass.FOUR_WAY_TWO_PRIMARY_LINE
-                    or vertex.vertex_class == VertexClass.THREE_WAY_ONE_PRIMARY_LINE):
-
+                    or vertex.vertex_class == VertexClass.THREE_WAY_ONE_PRIMARY_LINE
+                ):
                     out_polys = vertex.trim_primary_end(polys)
                     if len(out_polys) == 0:
                         continue
-                    
+
                     # update polygon DataFrame
                     for idx, out_poly in out_polys:
                         if out_poly:
@@ -766,7 +769,7 @@ class LineGrouping:
                 out_polys = vertex.trim_end_all(polys)
                 if len(out_polys) == 0:
                     continue
-                
+
                 # update polygon DataFrame
                 for idx, out_poly in out_polys:
                     self.polys.at[idx, "geometry"] = out_poly
@@ -780,9 +783,7 @@ class LineGrouping:
                     self.lines.at[p_trim.line_index, "geometry"] = p_trim.line_cleanup
 
                     # update VertexNode's line
-                    self.update_line_in_vertex_node(
-                        p_trim.line_index, p_trim.line_cleanup
-                    )
+                    self.update_line_in_vertex_node(p_trim.line_index, p_trim.line_cleanup)
 
     def get_merged_lines_original(self):
         return self.lines.dissolve(by=bt_const.BT_GROUP)
@@ -793,7 +794,7 @@ class LineGrouping:
             self.group_lines()
 
         self.find_vertex_for_poly_trimming()
-        self.lines["group"] = self.groups  # assign group attribute
+        self.lines[bt_const.BT_GROUP] = self.groups  # assign group attribute
 
     def run_regrouping(self):
         """
@@ -821,19 +822,14 @@ class LineGrouping:
         #  remove null geometry
         # TODO make sure lines and polygons match in pairs
         # they should have same amount and spatial coverage
-        self.valid_polys = self.polys[
-            ~self.polys.geometry.isna() & ~self.polys.geometry.is_empty
-        ]
+        self.valid_polys = self.polys[~self.polys.geometry.isna() & ~self.polys.geometry.is_empty]
 
         # save sh_geom.MultiLineString and sh_geom.MultiPolygon
-        self.invalid_polys = self.polys[
-            (self.polys.geometry.geom_type == "MultiPolygon")
-        ]
+        self.invalid_polys = self.polys[(self.polys.geometry.geom_type == "MultiPolygon")]
 
         # check lines
         self.valid_lines = self.merged_lines_trimmed[
-            ~self.merged_lines_trimmed.geometry.isna()
-            & ~self.merged_lines_trimmed.geometry.is_empty
+            ~self.merged_lines_trimmed.geometry.isna() & ~self.merged_lines_trimmed.geometry.is_empty
         ]
         self.valid_lines.reset_index(inplace=True, drop=True)
 
@@ -842,7 +838,7 @@ class LineGrouping:
         ]
         self.invalid_lines.reset_index(inplace=True, drop=True)
 
-    def save_file(self, out_file):
+    def save_file(self, out_file, out_layer="ground_footprint"):
         if not self.valid_lines.empty:
             self.valid_lines["length"] = self.valid_lines.length
             self.valid_lines.to_file(out_file, layer="merged_lines")
@@ -850,9 +846,10 @@ class LineGrouping:
         if not self.valid_polys.empty:
             if "length" in self.valid_polys.columns:
                 self.valid_polys.drop(columns=["length"], inplace=True)
-                
+
             self.valid_polys["area"] = self.valid_polys.area
-            self.valid_polys.to_file(out_file, layer="clean_footprint")
+            layer_name = out_layer
+            self.valid_polys.to_file(out_file, layer=layer_name)
 
         if not self.invalid_lines.empty:
             self.invalid_lines.to_file(out_file, layer="invalid_lines")
@@ -860,6 +857,7 @@ class LineGrouping:
         if not self.invalid_polys.empty:
             self.invalid_polys.to_file(out_file, layer="invalid_polygons")
 
+
 @dataclass
 class PolygonTrimming:
     """Store polygon and line to trim. Primary polygon is used to trim both."""
@@ -878,18 +876,20 @@ class PolygonTrimming:
         if self.line_cleanup.length < 100.0:
             trim_distance = 50.0
 
-        poly_primary = poly_primary.intersection(
-            vertex.buffer(trim_distance)
-        )
-            
+        poly_primary = poly_primary.intersection(vertex.buffer(trim_distance))
+
         self.poly_primary = poly_primary
-        
+
         # TODO: check why there is such cases
         if self.poly_cleanup is None:
             print("No polygon to trim.")
             return
-        
+
         midpoint = self.line_cleanup.interpolate(0.5, normalized=True)
+        if self.poly_primary is None or self.poly_primary.is_empty:
+            # print("Warning: No valid primary polygon for trimming; skipping difference operation.")
+            # TODO: handle this case
+            return
         diff = self.poly_cleanup.difference(self.poly_primary)
         if diff.geom_type == "Polygon":
             self.poly_cleanup = diff

beratools/core/algo_merge_lines.py

@@ -10,27 +10,57 @@ Description:
     This script is part of the BERA Tools.
     Webpage: https://github.com/appliedgrg/beratools
 
-    This file is intended to be hosting algorithms and utility functions/classes 
+    This file is intended to be hosting algorithms and utility functions/classes
     for merging lines.
 """
+
 from itertools import pairwise
 from operator import itemgetter
 
 import networkit as nk
 import shapely.geometry as sh_geom
+from shapely.geometry import GeometryCollection, LineString, MultiLineString
+from shapely.ops import linemerge
 
 import beratools.core.algo_common as algo_common
 import beratools.core.constants as bt_const
 
 
+def safe_linemerge(geom):
+    if isinstance(geom, (MultiLineString, GeometryCollection)):
+        return linemerge(geom)
+    elif isinstance(geom, LineString):
+        return geom
+    else:
+        return geom
+
+def custom_line_merge(geom):
+    if geom.geom_type == "MultiLineString":
+        # First try shapely's linemerge (fast)
+        merged = linemerge(geom)
+        # If still MultiLineString, use MergeLines for complex cases
+        if isinstance(merged, sh_geom.MultiLineString):
+            worker = MergeLines(merged)
+            merged = worker.merge_all_lines()
+            return merged if merged else geom
+        else:
+            return merged
+    elif geom.geom_type == "LineString":
+        return geom
+    else:
+        return geom
+
+
 def run_line_merge(in_line_gdf, merge_group):
     out_line_gdf = in_line_gdf
     if merge_group:
+        if bt_const.BT_GROUP not in in_line_gdf.columns:
+            in_line_gdf = in_line_gdf.copy()
+            in_line_gdf[bt_const.BT_GROUP] = range(1, len(in_line_gdf) + 1)
         out_line_gdf = in_line_gdf.dissolve(by=bt_const.BT_GROUP, as_index=False)
 
-    out_line_gdf.geometry = out_line_gdf.line_merge()
-
-    for row in out_line_gdf.itertuples():
+    out_line_gdf.geometry = out_line_gdf.geometry.apply(safe_linemerge)
+    for i, row in enumerate(out_line_gdf.itertuples()):
         if isinstance(row.geometry, sh_geom.MultiLineString):
             worker = MergeLines(row.geometry)
             merged_line = worker.merge_all_lines()
@@ -39,9 +69,10 @@ def run_line_merge(in_line_gdf, merge_group):
 
     out_line_gdf = algo_common.clean_line_geometries(out_line_gdf)
     out_line_gdf.reset_index(inplace=True, drop=True)
-    out_line_gdf['length'] = out_line_gdf.geometry.length
+    out_line_gdf["length"] = out_line_gdf.geometry.length
     return out_line_gdf
 
+
 class MergeLines:
     """Merge line segments in MultiLineString."""
 
@@ -61,7 +92,7 @@ class MergeLines:
         self.line_segs = [line for line in self.line_segs if line.length > 1e-3]
         self.multi_line = sh_geom.MultiLineString(self.line_segs)
         m = sh_geom.mapping(self.multi_line)
-        self.end = [(i[0], i[-1]) for i in m['coordinates']]
+        self.end = [(i[0], i[-1]) for i in m["coordinates"]]
 
         self.G = nk.Graph(edgesIndexed=True)
         self.G.addNodes(2)
@@ -112,12 +143,13 @@ class MergeLines:
                 single_path = False
 
         return single_path
+
     def get_merged_line_for_component(self, component):
         sub = nk.graphtools.subgraphFromNodes(self.G, component)
         lines = None
         if nk.graphtools.maxDegree(sub) >= 3:  # not simple path
             edges = [self.G.edgeId(i[0], i[1]) for i in list(sub.iterEdges())]
-            lines =  itemgetter(*edges)(self.line_segs)
+            lines = itemgetter(*edges)(self.line_segs)
         elif nk.graphtools.maxDegree(sub) == 2:
             lines = self.merge_single_line(component)
 
@@ -180,7 +212,7 @@ class MergeLines:
             pair = pairs[i]
             poly_t = self.node_poly[pair[0]]
             point_t = sh_geom.Point(self.end[id][0])
-            if  poly_t.contains(point_t):
+            if poly_t.contains(point_t):
                 line = self.line_segs[id]
             else:
                 # line = reverse(self.line_segs[id])