correction direction speedup

changliao1025 · Feb 28, 2024 · 4ad153e · 4ad153e
1 parent f4daf79
commit 4ad153e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 28 deletions.
diff --git a/pyflowline/algorithms/direction/correct_flowline_direction.py b/pyflowline/algorithms/direction/correct_flowline_direction.py
@@ -1,7 +1,7 @@
 import sys
-sys.setrecursionlimit(10000)
 import numpy as np 
 from pyflowline.algorithms.auxiliary.check_head_water import check_head_water
+sys.setrecursionlimit(100000)
 lFlowlineIndex=0
 def correct_flowline_direction(aFlowline_in, pVertex_outlet_in):    
     """_summary_ This function should expect the flowline may not be ordered, so the stream order info is not available.
@@ -20,28 +20,16 @@ def correct_flowline_direction(aFlowline_in, pVertex_outlet_in):
     global lFlowlineIndex    
     nFlowline = len(aFlowline_in)
     # Create sets for faster lookup
-    pVertex_start_in_set = {flowline.pVertex_start for flowline in aFlowline_in}
-    pVertex_end_in_set = {flowline.pVertex_end for flowline in aFlowline_in}
-
-    aFlag_process=np.full(nFlowline, 0, dtype =int)
-    iFlag_first = 1
-    for i in range(nFlowline):
-        pFlowline = aFlowline_in[i]
-        iFlag_dam = pFlowline.iFlag_dam
-        iStream_order = pFlowline.iStream_order
-        pVertex_start = pFlowline.pVertex_start
-        pVertex_end = pFlowline.pVertex_end         
-        if iFlag_first == 1:
-            dDiatance_min = pVertex_end.calculate_distance( pVertex_outlet_in )                             
-            lIndex_outlet = i            
-            iFlag_first=0    
-        else:        
-            dDiatance = pFlowline.pVertex_end.calculate_distance(pVertex_outlet_in)              
-            if  dDiatance < dDiatance_min:
-                dDiatance_min = dDiatance
-                #found it
-                lIndex_outlet = i      
-
+    pVertex_start_in_set = {flowline.pVertex_start for flowline in aFlowline_in}    
+    aFlag_process=np.full(nFlowline, 0, dtype =int)   
+    dDiatance_min = float('inf')
+    for i in range(nFlowline):        
+        pFlowline = aFlowline_in[i]                
+        pVertex_end = pFlowline.pVertex_end
+        dDiatance = pVertex_end.calculate_distance(pVertex_outlet_in)
+        if dDiatance < dDiatance_min:
+            dDiatance_min = dDiatance
+            lIndex_outlet = i                  
 
     lFlowlineIndex = 0 
     pFlowline = aFlowline_in[lIndex_outlet]
@@ -52,7 +40,7 @@ def correct_flowline_direction(aFlowline_in, pVertex_outlet_in):
     pVertex_end = pFlowline.pVertex_end
     lFlowlineIndex = lFlowlineIndex + 1    
     #we might find more than 1 upstream    
-    def find_upstream_flowline(pVertex_start_in, pVertex_end_in):
+    def find_upstream_flowline( pVertex_end_in):
         nUpstream = 0 
         aUpstream=list()
         aReverse=list()
@@ -89,9 +77,8 @@ def tag_upstream(pVertex_start_in, pVertex_end_in):
         if(check_head_water(aFlowline_in, pVertex_start_in)==1):            
             pass
         else:
-            nUpstream, aUpstream, aReverse = find_upstream_flowline(pVertex_start_in, pVertex_end_in)
-            if nUpstream > 0:
-
+            nUpstream, aUpstream, aReverse = find_upstream_flowline( pVertex_end_in)
+            if nUpstream > 0:                
                 for j in range(nUpstream):
                     pFlowline = aFlowline_in[ aUpstream[j] ] 
                     if (aReverse[j]==1):             

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
+#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/classes/basin.py b/pyflowline/classes/basin.py
@@ -520,6 +520,7 @@ def basin_flowline_simplification(self):
         try:
             print('Basin ',  self.sBasinID, 'started correction flow direction')
             ptimer.start()
+            sys.stdout.flush()
             nFlowline_before = len(aFlowline_basin_simplified)  
             #this flowline is ordered from downstream to upstream      
             aFlowline_basin_simplified = correct_flowline_direction(aFlowline_basin_simplified,  pVertex_outlet )
@@ -541,6 +542,7 @@ def basin_flowline_simplification(self):
         try:
             print('Basin ',  self.sBasinID, 'started loop removal')
             ptimer.start()
+            sys.stdout.flush()
             aFlowline_basin_simplified = remove_flowline_loop(aFlowline_basin_simplified)   
             ptimer.stop() 
             sys.stdout.flush()