convert O(n2) to O(n) for basin function

changliao1025 · Mar 1, 2024 · d2e73bb · d2e73bb
1 parent e56f143
commit d2e73bb
Show file tree

Hide file tree

Showing 11 changed files with 165 additions and 132 deletions.
diff --git a/pyflowline/algorithms/auxiliary/check_head_water.py b/pyflowline/algorithms/auxiliary/check_head_water.py
@@ -8,7 +8,13 @@ def check_head_water(aFlowline_in, pVertex_start_in):
 
     Returns:
         [int]: [0: not headwater; 1: is headwater]
-    """
+    """    
+    start_vertices = {flowline.pVertex_start for flowline in aFlowline_in}
+    end_vertices = {flowline.pVertex_end for flowline in aFlowline_in}
+    # Check if the vertex is a headwater
+    is_headwater = pVertex_start_in in start_vertices and pVertex_start_in not in end_vertices
+    return int(is_headwater)
+
     #nFlowline = len(aFlowline_in)
     #iFlag_head_water = -1
     #iCount = 0
@@ -28,11 +34,4 @@ def check_head_water(aFlowline_in, pVertex_start_in):
     #    
     #return iFlag_head_water
     #Create sets of all start and end vertices
-    start_vertices = {flowline.pVertex_start for flowline in aFlowline_in}
-    end_vertices = {flowline.pVertex_end for flowline in aFlowline_in}
-
-    # Check if the vertex is a headwater
-    is_headwater = pVertex_start_in in start_vertices and pVertex_start_in not in end_vertices
-
-    return int(is_headwater)
 
diff --git a/pyflowline/algorithms/direction/correct_flowline_direction.py b/pyflowline/algorithms/direction/correct_flowline_direction.py
@@ -18,33 +18,41 @@ def correct_flowline_direction(aFlowline_in, pVertex_outlet_in):
             dDiatance_min = dDiatance
             lIndex_outlet = i  
 
-        unfinished_flowlines.add(aFlowline_in[i])
-         # Update the dictionary
-        vertex_to_flowlines.setdefault(pFlowline.pVertex_start, []).append(pFlowline)
-        vertex_to_flowlines.setdefault(pFlowline.pVertex_end, []).append(pFlowline)
+    for i in range(nFlowline):        
+        if i != lIndex_outlet:
+            pFlowline = aFlowline_in[i]   
+            unfinished_flowlines.add(aFlowline_in[i])
+            # Update the dictionary
+            vertex_to_flowlines.setdefault(pFlowline.pVertex_start, []).append(pFlowline)
+            vertex_to_flowlines.setdefault(pFlowline.pVertex_end, []).append(pFlowline)
 
-
-    unfinished_flowlines.remove(aFlowline_in[lIndex_outlet])
     aVertex_downslope_table = [aFlowline_in[lIndex_outlet].pVertex_start]    
     aFlowline_out= [aFlowline_in[lIndex_outlet]]    
     while unfinished_flowlines:
-        aVertex_downslope_current= []        
+        aVertex_downslope_current= []  
+        iCount = 0      
         for pVertex_dummy in aVertex_downslope_table:       
             to_remove = set()       
-            #for pFlowline in unfinished_flowlines:     
+            #for pFlowline in unfinished_flowlines:              
             for pFlowline in vertex_to_flowlines.get(pVertex_dummy, []):  # Use the dictionary here
-                if pFlowline.pVertex_end == pVertex_dummy :
-                    aVertex_downslope_current.append(pFlowline.pVertex_start)
-                    to_remove.add(pFlowline)
-                    aFlowline_out.append(pFlowline)                    
-                else:
-                    if pFlowline.pVertex_start ==  pVertex_dummy :
-                        pFlowline.reverse()
-                        aVertex_downslope_current.append(pFlowline.pVertex_start)                            
+                if pFlowline in unfinished_flowlines:
+                    if pFlowline.pVertex_end == pVertex_dummy :
+                        aVertex_downslope_current.append(pFlowline.pVertex_start)
                         to_remove.add(pFlowline)
-                        aFlowline_out.append(pFlowline)
+                        aFlowline_out.append(pFlowline)  
+                        iCount = iCount + 1                  
+                    else:
+                        if pFlowline.pVertex_start ==  pVertex_dummy :
+                            pFlowline.reverse()
+                            aVertex_downslope_current.append(pFlowline.pVertex_start)                            
+                            to_remove.add(pFlowline)
+                            aFlowline_out.append(pFlowline)
+                            iCount = iCount + 1     
 
-            unfinished_flowlines -= to_remove                    
+            unfinished_flowlines -= to_remove  
+
+        if iCount == 0:            
+           break
 
         if len(unfinished_flowlines)==0:
            break

diff --git a/pyflowline/algorithms/index/define_stream_order.py b/pyflowline/algorithms/index/define_stream_order.py
@@ -11,19 +11,28 @@
     pass
 
 def update_head_water_stream_order(aFlowline_in):
-    nFlowline = len(aFlowline_in)
-    aFlowline_out = list()
-    for i in range(nFlowline):
-        pFlowline = aFlowline_in[i]          
-        pVertex_start = pFlowline.pVertex_start  
-        if check_head_water(aFlowline_in, pVertex_start)==1:
-            pFlowline.iStream_order = 1
-        else:
-            pFlowline.iStream_order = -1
-
-        aFlowline_out.append(pFlowline)
-
-    return aFlowline_out
+    start_vertices = {flowline.pVertex_start for flowline in aFlowline_in}
+    end_vertices = {flowline.pVertex_end for flowline in aFlowline_in}
+
+    for flowline in aFlowline_in:
+        pVertex_start = flowline.pVertex_start
+        is_headwater = pVertex_start in start_vertices and pVertex_start not in end_vertices
+        flowline.iStream_order = 1 if is_headwater else -1
+
+    return aFlowline_in
+
+    #nFlowline = len(aFlowline_in)
+    #aFlowline_out = list()
+    #for i in range(nFlowline):
+    #    pFlowline = aFlowline_in[i]          
+    #    pVertex_start = pFlowline.pVertex_start  
+    #    if check_head_water(aFlowline_in, pVertex_start)==1:
+    #        pFlowline.iStream_order = 1
+    #    else:
+    #        pFlowline.iStream_order = -1
+    #    
+    #    aFlowline_out.append(pFlowline)
+    #return aFlowline_out
 
 
 def define_stream_order(aFlowline_in, aConfluence_in):
@@ -53,19 +62,7 @@ def define_stream_order(aFlowline_in, aConfluence_in):
             nConfleunce = len(aConfluence_in)
             aFlag_confluence_treated = np.full(nConfleunce, 0, dtype=int)
             #build rtree for confluence
-            index_confluence = RTree( max_cap=5, min_cap=2)
-            #for i in range(nConfleunce):
-            #    lID = i 
-            #    pVertex_confluence = aConfluence_in[i].pVertex_confluence 
-            #    x = pVertex_confluence.dLongitude_degree
-            #    y = pVertex_confluence.dLatitude_degree   
-            #    left =   x - 1E-5
-            #    right =  x + 1E-5
-            #    bottom = y - 1E-5
-            #    top =    y + 1E-5
-            #    pBound= (left, bottom, right, top)                
-            #    index_confluence.insert(lID, pBound)  #
-
+            index_confluence = RTree( max_cap=5, min_cap=2)          
             for i, confluence in enumerate(aConfluence_in):
                 pVertex_confluence = confluence.pVertex_confluence 
                 x, y = pVertex_confluence.dLongitude_degree, pVertex_confluence.dLatitude_degree
@@ -101,13 +98,7 @@ def define_stream_order(aFlowline_in, aConfluence_in):
                         aFlag_confluence_treated[i] = 1
                         #now we can process the downstream
                         #get unique value
-                        iStream_order = max(aStrord) if len(set(aStrord)) > 1 else aStrord[0] + 1
-                        #dummy = np.array(aStrord)
-                        #dummy1 = np.unique(dummy)
-                        #if len(dummy1) == 1: #all upstreams have the same order
-                        #    iStream_order = aStrord[0] + 1
-                        #else:
-                        #    iStream_order = np.max(dummy)  
+                        iStream_order = max(aStrord) if len(set(aStrord)) > 1 else aStrord[0] + 1                        
 
                         #update
                         pFlowline_downstream.iStream_order = iStream_order
@@ -137,11 +128,7 @@ def define_stream_order(aFlowline_in, aConfluence_in):
                             pass
 
 
-            #for i in range(nFlowline):            
-            #    pFlowline = aFlowline_in[i]      
-            #    pFlowline.iStream_order =    aFlowline_in[i].iStream_order   
-            #    aFlowline_out.append(pFlowline)
-            #    aStream_order[i] = pFlowline.iStream_order
+
             for i, flowline in enumerate(aFlowline_in):                
                 aFlowline_out.append(flowline)
                 aStream_order[i] = flowline.iStream_order

diff --git a/pyflowline/algorithms/index/define_stream_segment_index.py b/pyflowline/algorithms/index/define_stream_segment_index.py
@@ -1,4 +1,4 @@
-#from pyflowline.algorithms.auxiliary.check_head_water import check_head_water
+
 def define_stream_segment_index(aFlowline_in):
     """build stream segment index, because they are ordered from outlet to headwater, so the index is from 1 to n
 

diff --git a/pyflowline/algorithms/intersect/intersect_flowline_with_flowline.py b/pyflowline/algorithms/intersect/intersect_flowline_with_flowline.py
@@ -90,7 +90,7 @@ def intersect_flowline_with_flowline( sFilename_flowline_a_in, sFilename_flowlin
                 pGeometry_flowline_b = pFeature_flowline_b.GetGeometryRef()
 
                 if iFlag_id ==1:
-                    lFlowlineID = pFeature_flowline_b.GetField("id")
+                    lFlowlineID = pFeature_flowline_b.GetField("lineid")
                 else:
                     lFlowlineID = -1
 

diff --git a/pyflowline/algorithms/merge/merge_flowline.py b/pyflowline/algorithms/merge/merge_flowline.py
@@ -37,25 +37,21 @@ def merge_flowline(aFlowline_in,
     aIndex_middle = np.array(aIndex_middle_in)
     aIndex_confluence = np.array(aIndex_confluence_in)
     if aIndex_middle.size == 0:
-        return aFlowline_in
-
-    #aVertex_headwater=aVertex[aIndex_headwater]    
-    #aVertex_middle=aVertex[aIndex_middle]   
-    #if aIndex_middle.size == 0:
-    #    return aFlowline_in
+        return aFlowline_in 
 
     #convert to set
     aVertex_headwater_set = set(aVertex[aIndex_headwater])  
     aVertex_middle_set = set(aVertex[aIndex_middle])  
 
     if aIndex_confluence.size > 0:        
-        iFlag_confluence = 1
-        #aVertex_confluence=aVertex[aIndex_confluence]    
+        iFlag_confluence = 1  
         aVertex_confluence_set = set(aVertex[aIndex_confluence])  
     else:
         iFlag_confluence = 0
-        pass
 
+    # Build the dictionary
+    vertex_to_flowline = {flowline.pVertex_end: flowline for flowline in aFlowline_in}  
+
     def merge_flowline_reach(lIndex_in, pVertex_start_in):
         global lID
         pFlowline = aFlowline_in[lIndex_in]
@@ -64,15 +60,20 @@ def merge_flowline_reach(lIndex_in, pVertex_start_in):
 
         #while (find_vertex_in_list(aVertex_middle.tolist(), pVertex_current)[0] ==1):            
         while (pVertex_current in aVertex_middle_set): 
-            for j in range(0, nFlowline):      
-                pFlowline2 = aFlowline_in[j]                
-                pVertex_start = pFlowline2.pVertex_start
-                pVertex_end = pFlowline2.pVertex_end
-                if pVertex_end == pVertex_current:
-                    pFlowline = pFlowline.merge_upstream(pFlowline2)
-                    pVertex_current = pVertex_start                    
-                    break
-
+            if pVertex_current in vertex_to_flowline:
+                pFlowline2 = vertex_to_flowline[pVertex_current]
+                pFlowline = pFlowline.merge_upstream(pFlowline2)
+                pVertex_current = pFlowline2.pVertex_start
+
+            #old method    
+            #for j in range(0, nFlowline):      
+            #    pFlowline2 = aFlowline_in[j]                
+            #    pVertex_start = pFlowline2.pVertex_start
+            #    pVertex_end = pFlowline2.pVertex_end
+            #    if pVertex_end == pVertex_current:
+            #        pFlowline = pFlowline.merge_upstream(pFlowline2)
+            #        pVertex_current = pVertex_start                    
+            #        break                
 
         #go to next 
         #if find_vertex_in_list(aVertex_headwater.tolist(), pVertex_current)[0] ==1: 
@@ -97,26 +98,8 @@ def merge_flowline_reach(lIndex_in, pVertex_start_in):
                     pVertex_end = pFlowline3.pVertex_end
                     if pVertex_end == pVertex_current:
                         merge_flowline_reach(k, pVertex_start)
-
-
-    #find outlet
-    #iFlag_first=1
-    #for i in range(nFlowline):        
-    #    pFlowline = aFlowline_in[i]                
-    #    pVertex_start = pFlowline.pVertex_start
-    #    pVertex_end = pFlowline.pVertex_end
-    #    dDiatance = pVertex_end.calculate_distance( pVertex_outlet_in)
-    #    if iFlag_first ==1:
-    #        dDiatance_min = dDiatance                
-    #        lIndex_outlet = i            
-    #        iFlag_first=0    
-    #    else:            
-    #        if  dDiatance < dDiatance_min:
-    #            dDiatance_min = dDiatance           
-    #            lIndex_outlet = i                
-    #            pass    
-    #        else:
-    #            pass
+
+
     #find outlet
     dDiatance_min = float('inf')
     for i in range(nFlowline):        
@@ -131,19 +114,23 @@ def merge_flowline_reach(lIndex_in, pVertex_start_in):
     pVertex_start = pFlowline.pVertex_start
     pVertex_end = pFlowline.pVertex_end   
 
+
+
     #now start from outlet
     if iFlag_confluence == 1:
         #check whether outlet is a confluence
         #if  (find_vertex_in_list(aVertex_confluence.tolist(), pVertex_end)[0] ==1):
         if pVertex_end in aVertex_confluence_set:
+            #if pVertex_end in end_vertex_to_flowline:
+            #    merge_flowline_reach(end_vertex_to_flowline[pVertex_end], pVertex_start_dummy)      
             for i in range(nFlowline):
                 pFlowline = aFlowline_in[i]  
                 pVertex_start_dummy = pFlowline.pVertex_start
                 pVertex_end_dummy = pFlowline.pVertex_end
                 if pVertex_end == pVertex_end_dummy:
                     #this is 
                     merge_flowline_reach(i, pVertex_start_dummy)  
-            pass
+
         else:
             merge_flowline_reach(lIndex_outlet, pVertex_start)   
     else:

diff --git a/pyflowline/algorithms/simplification/remove_small_river.py b/pyflowline/algorithms/simplification/remove_small_river.py
@@ -1,4 +1,3 @@
-from pyflowline.algorithms.auxiliary.check_head_water import check_head_water
 def remove_small_river(aFlowline_in, dThreshold_in):
     """Remove small river that meet the threshold and headwater requirement, also dam flowline are reserved
 

diff --git a/pyflowline/classes/basin.py b/pyflowline/classes/basin.py
@@ -359,7 +359,8 @@ def basin_flowline_simplification(self):
         ptimer = pytimer()
 
         if self.iFlag_dam == 1:
-            sFilename_flowline_filter = self.sFilename_flowline_filter
+            #sFilename_flowline_filter = self.sFilename_flowline_filter
+            sFilename_flowline_filter = self.sFilename_flowline_filter_geojson
             aFlowline_basin_filtered_raw, pSpatial_reference = read_flowline_geojson( sFilename_flowline_filter )
             aVertex_filtered = find_flowline_vertex(aFlowline_basin_filtered_raw)
 
@@ -460,8 +461,8 @@ def basin_flowline_simplification(self):
             print('Basin ',  self.sBasinID, ' has dam', nFlowline_before, nFlowline_after)
             ptimer.stop()
         else:
-            print('Basin ',  self.sBasinID, ' has no dam')
-            sFilename_flowline_filter = self.sFilename_flowline_filter
+            print('Basin ',  self.sBasinID, ' has no dam')            
+            sFilename_flowline_filter = self.sFilename_flowline_filter_geojson #sFilename_flowline_filter = self.sFilename_flowline_filter
             aFlowline_basin_filtered, pSpatial_reference = read_flowline_geojson( sFilename_flowline_filter )
             #aVertex_filtered = find_flowline_vertex(aFlowline_basin_filtered)
 
@@ -618,7 +619,7 @@ def basin_flowline_simplification(self):
             pass
 
         #the final vertex info
-        print('Basin ',  self.sBasinID, 'find flowline confluence')
+        print('Basin ', self.sBasinID, 'find flowline confluence')
         sys.stdout.flush()
         ptimer.start()
         aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity, pVertex_outlet = find_flowline_confluence(aFlowline_basin_simplified,  pVertex_outlet)
@@ -941,7 +942,6 @@ def basin_build_confluence(self, aFlowline_basin_in, aVertex_confluence_in):
 
         return aConfluence_basin
 
-
     def basin_analyze(self):
         """
         Analyze the basin results including length, sinuosity, and breaching angle
@@ -1101,7 +1101,7 @@ def tojson(self):
         """
         aSkip = ['aFlowline_basin_filtered', \
                 'aFlowline_basin_simplified','aFlowline_basin_conceptual','aConfluence_basin_simplified',
-                'aConfluence_basin_conceptual','pRTree_flowline', 'pRTree_edge']
+                'aConfluence_basin_conceptual','pRTree_flowline', 'pRTree_edge','aFlowline_basin_edge']
 
         obj = self.__dict__.copy()
         for sKey in aSkip:

diff --git a/pyflowline/formats/convert_flowline_to_geojson.py b/pyflowline/formats/convert_flowline_to_geojson.py
@@ -22,11 +22,15 @@ def convert_flowline_to_geojson(iFlag_type_in, sFilename_geojson_in, sFilename_g
         pass
     else:
         if iFlag_type_in == 1:
-
-            aFlowline_basin, pSpatial_reference = read_flowline_geojson( sFilename_geojson_in )   
+            sFilename_dummy = sFilename_geojson_out
+            sFilename_out = sFilename_dummy.replace('.geojson', '_withID.geojson')            
+            aFlowline_basin, pSpatial_reference = read_flowline_geojson( sFilename_geojson_in , sFilename_out = sFilename_out)   
+            #get lineid
+            aFlowlineID = [pFlowline.lFlowlineID for pFlowline in aFlowline_basin]          
             #convert it
             iFlag_projected = 0       
-            export_flowline_to_geojson(aFlowline_basin, sFilename_geojson_out)
+            export_flowline_to_geojson(aFlowline_basin, sFilename_geojson_out, aAttribute_field=['lineid'],
+                                       aAttribute_dtype=['int'], aAttribute_data=[aFlowlineID])
         else:
             if iFlag_type_in == 2:
                 #polygon
@@ -56,7 +60,6 @@ def convert_shapefile_to_geojson(iFlag_type_in, sFilename_shapefile_in, sFilenam
     else:
         if iFlag_type_in == 1:
             aFlowline_basin, pSpatial_reference = read_flowline_shapefile( sFilename_shapefile_in )    
-            #aFlowline_basin, pSpatial_reference = read_flowline_geojson( sFilename_shapefile_in )   
             #convert it
 
             iFlag_projected = 0