Correction of rlen/rdz

Interface.py

@@ -184,6 +184,26 @@ def get_next(jj1, ii1, trip, grid):
         k+=1
     return ii,jj
 
+@jit(nopython=True)
+def is_connected(jj1, ii1,jj2,ii2, trip):
+    jj = jj1; ii = ii1
+    while (trip[jj,ii] < 9) and (trip[jj,ii] > 0):
+        t = trip[jj,ii]
+        if (t == 6) or (t == 5) or (t == 4):
+            jj+= 1
+        elif (t == 8) or (t == 1) or (t == 2):
+            jj -= 1
+        if (t == 2) or (t == 3) or (t == 4):
+            ii += 1
+        elif (t == 8) or (t == 7) or (t == 6):
+            ii -= 1
+        if (ii<0) or (jj<0) or (ii>=trip.shape[1]) or (jj>=trip.shape[0]):
+            break
+        if (ii==ii2) and (jj == jj2): 
+            return True
+    return False 
+
+
 #
 class HydroSuper :
     def __init__(self, nbvmax, hydrodata, hydrooverlap, nbasmax, part, comm) :
@@ -223,7 +243,7 @@ class HydroSuper :
         #
         self.basin_topoindex_stream = self.basin_topoindex_stream[:,:self.nwbas]
         # Set the number of inflows per basin. For the moment twice the maximum number of basins.
-        self.inflowmax = 50#max(10, self.nwbas*2)
+        self.inflowmax = max(10, self.nwbas*2)
         print("Maximum number of basin created : {0}".format(self.nwbas))
         ijdim=[]
         for i in range(self.nbpt) :
@@ -442,6 +462,61 @@ class HydroSuper :
 
         return
 
+    def correct_mainstream(self, part, hydrogrid):
+        istr, iend, jstr, jend = hydrogrid.box[:]
+        lon =  hydrogrid.ncfile.variables["nav_lon"][0,istr:iend]
+        lat =  hydrogrid.ncfile.variables["nav_lat"][jstr:jend,0]
+        trip = hydrogrid.ncfile.variables["trip"][jstr:jend,istr:iend].astype(np.int32)
+        #
+        grid = np.zeros(trip.shape, dtype = np.int32)
+        htu = np.zeros(trip.shape, dtype = np.int32)
+        for pt in range(self.nbpt):
+            for bs in range(self.basin_count[pt]):
+                latp, lonp = self.basin_outcoor[pt,bs,:]
+                # MERIT and HydoSHEDS are regular - quicker that way
+                i = np.where(lon==lonp)[0][0]
+                j = np.where(lat==latp)[0][0]
+                grid[j,i] = pt+1
+                htu[j,i] = bs+1
+
+
+        orog = hydrogrid.ncfile.variables["orog"][jstr:jend,istr:iend]
+        disto = hydrogrid.ncfile.variables["disto"][jstr:jend,istr:iend]
+        
+        ok = 0; nook = 0; bad = 0
+        for pt in part.landcorelist:
+            for bs in range(self.basin_count[pt]):
+                olat, olon = self.basin_outcoor[pt,bs,:]
+                if self.inflow_number[pt,bs] > 0:
+                   fac = [self.basin_fac[int(self.inflow_grid[pt,bs,i])-1, int(self.inflow_basin[pt,bs,i])-1] for i in range(int(self.inflow_number[pt,bs]))]
+                   ind = np.argmax(fac)
+                   ig = int(self.inflow_grid[pt,bs,ind])-1; ib = int(self.inflow_basin[pt,bs,ind])-1
+                   ilat, ilon = self.basin_outcoor[ig,ib,:]
+                   i0 = np.where(lon==ilon)[0][0]
+                   j0 = np.where(lat==ilat)[0][0]
+                   i1 = np.where(lon==olon)[0][0]
+                   j1 = np.where(lat==olat)[0][0]
+                   
+                   ii,jj = get_next(j0, i0, trip, grid)
+                   if (ii == i1) and (jj == j1):
+                       # then correct
+                       dorog = max(orog[j0,i0] - orog[j1,i1],0.1)
+                       dlen = disto[j0,i0] - disto[j1,i1]
+                       topo = np.sqrt(dlen**3/dorog)/1000
+                       self.basin_rdz[pt,bs] = dorog
+                       self.basin_rlen[pt,bs] = dlen
+                       self.basin_topoindex_stream[pt,bs] = topo
+                       ok+=1
+                   else:
+                       print("error - not directly connected")
+                       conn = is_connected(j0,i0,j1,i1,trip) 
+                       print("Is connected : ", conn)
+                       nook+=1
+                       if not conn:
+                          bad += 1
+        del orog; del disto; del grid; del htu; del trip; del lon; del lat
+        print(f"Proc {part.rank}, num = {ok}, error = {nook}, bad = {bad}")
+        return 
 
     def killbas(self, tokill, totakeover, numops):
         ops = tokill.shape[1]

RoutingPreProc.py

@@ -124,6 +124,12 @@ gc.collect()
 comm.Barrier()
 #
 hsuper.correct_outflows(part)
+
+##################
+
+hsuper.correct_mainstream(part, hydrogrid)
+
+
 #
 INFO("=================== Compute fetch ====================")
 t = time.time()