wolfhece 2.1.13__py3-none-any.whl → 2.1.14__py3-none-any.whl

wolfhece/hydrology/PyWatershed.py

@@ -6,9 +6,11 @@ from matplotlib.axes import Axes
 import matplotlib.pyplot as plt
 import logging
 from tqdm import tqdm
+import numpy as np
 
 from ..PyTranslate import _
 from ..PyVertex import wolfvertex, cloud_vertices
+from ..PyVertexvectors import Zones, zone, vector
 from ..wolf_array import *
 from ..PyCrosssections import crosssections as CrossSections
 from ..GraphNotebook import PlotNotebook
@@ -57,6 +59,11 @@ class Node_Watershed:
 
     flatindex:int = -1  # index de la zone de plat
 
+    def __init__(self):
+        self.cums=0.
+        self.up=None
+        self.down=None
+
     def incr_curvi(self):
         """Incrémentation de la longueur curviligne"""
 
@@ -64,8 +71,7 @@ class Node_Watershed:
             self.cums=0.
         else:
             self.cums = self.down.cums+self.incrs
-
-        for curup in self.up:
+        for curup in self.up:            
             curup.incr_curvi()
 
     def mean_slope_up(self, threshold:float)-> float:
@@ -119,6 +125,95 @@ class Node_Watershed:
         """
         self.strahler = strahler
 
+    def distance(self, x:float, y:float) -> float:
+        """ Distance euclidienne """
+
+        return np.sqrt(pow(self.x-x,2)+pow(self.y-y,2))
+    
+    def get_up_nodes_same_sub(self, excluded_node:list["Node_Watershed"]=[]):
+        """
+        Get all upstream nodes in the same sub-basin
+        """
+
+        all_up = [self]
+        all_rivers = [self] if self.river else []
+        all_runoff = [] if self.river else [self]
+        
+        for curup in self.up:
+
+            if curup in excluded_node:
+                continue
+
+            added = False
+            if curup.sub == self.sub:
+                all_up.append(curup)
+                added = True
+            
+            if curup.river:
+                all_rivers.append(curup)
+            else:
+                all_runoff.append(curup)
+
+            if added:
+                up, river, runoff = curup.get_up_nodes_same_sub(excluded_node)
+                all_up.extend(up)
+                all_rivers.extend(river)
+                all_runoff.extend(runoff)
+        
+        return all_up, all_rivers, all_runoff
+    
+    def get_up_runoff_nodes_same_sub(self):
+        """
+        Get all upstream runoff nodes in the same sub-basin
+        """
+
+        all_up = []
+        for curup in self.up:
+            if curup.sub == self.sub and not curup.river:
+                all_up.append(curup)
+
+                all_up += curup.get_up_runoff_nodes_same_sub()
+        
+        return all_up
+
+    def get_up_rivernodes_same_sub(self):
+        """
+        Get all upstream river nodes in the same sub-basin
+        """
+
+        all_up = []
+        for curup in self.upriver:
+            if curup.sub == self.sub:
+                all_up.append(curup)
+
+                all_up += curup.get_up_rivernodes_same_sub()
+        
+        return all_up
+    
+    def get_up_reaches_same_sub(self) -> list[int]:
+        """
+        Get all upstream reaches in the same sub-basin
+        """
+
+        all_up = [self.reach]
+        for curup in self.upriver:
+            if curup.sub == self.sub:
+                all_up += curup.get_up_reaches_same_sub()
+
+        return np.unique(all_up).tolist()
+    
+    def get_down_reaches_same_sub(self) -> list[int]:
+        """
+        Get all downstream reaches in the same sub-basin
+        """
+
+        all_down = [self.reach]
+        if self.down is not None:
+            if self.down.sub == self.sub:
+                all_down += self.down.get_down_reaches_same_sub()
+
+        return np.unique(all_down).tolist() 
+
 
 class RiverSystem:
     """
@@ -129,8 +224,10 @@ class RiverSystem:
     #  reaches
     #  |__['reaches']
     #  |  |__[idx]
-    #  |     |__['upstream']
-    #  |     |__['baselist']
+    #  |     |__['upstream']        # all reaches in upstream
+    #  |     |__['baselist']        # list of nodes in the reach
+    #  |     |__['up']              # **if upstream** node in upstream
+    #  |     |__['fromuptodown']    # **if upstream** list of nodes from upstream to downstream
     #  |__['indexed']
     #  |__['strahler']
 
@@ -205,7 +302,7 @@ class RiverSystem:
         xy = [[curnode.x, curnode.y]  for curnode in nodes]
         self.kdtree = KDTree(xy)
 
-    def get_nearest_nodes(self, xy:np.ndarray, nb=int) -> tuple[np.ndarray, list[Node_Watershed]]:
+    def get_nearest_nodes(self, xy:np.ndarray | vector, nb:int = 1) -> tuple[np.ndarray | float, list[Node_Watershed] | Node_Watershed]:
         """
         Return the distance and the  nearest Node_Watershed
 
@@ -214,10 +311,105 @@ class RiverSystem:
 
         return
         """
+        
+        if isinstance(xy, vector):
+            centroid = xy.asshapely_pol().centroid
+            xy = np.array([[centroid.x, centroid.y]])
+        
         dd, ii = self.kdtree.query(xy, nb)
 
-        return dd, [self.all_nodes[curi] for curi in ii]
+        if isinstance(ii, int | np.int64):
+            return dd, self.all_nodes[ii]
+        elif isinstance(ii, np.ndarray | list):
+            if len(ii) == 1:
+                return dd[0], self.all_nodes[ii[0]]
+            else:
+                return dd, [self.all_nodes[curi] for curi in ii]
+    
+    def get_nodes_in_reaches(self, reaches:list[int])->list[Node_Watershed]:
+        """
+        Get nodes in a reaches
+        """
+        all_nodes = []
+        for cur_reach in reaches:
+            all_nodes.extend(self.reaches['reaches'][cur_reach]['baselist'])
+        
+        return all_nodes
+    
+    def get_downstream_node_in_reach(self, reach:int)->Node_Watershed:
+        """
+        Get downstream node in a reach
+        """
+
+        return self.reaches['reaches'][reach]['baselist'][0]
+    
+    def get_upstream_node_in_reach(self, reach:int)->Node_Watershed:
+        """
+        Get upstream node in a reach
+        """
+
+        return self.reaches['reaches'][reach]['baselist'][-1]
+    
+    def get_downstream_reaches(self, node:Node_Watershed)->list[int]:
+        """
+        Get index of downstream reaches
+        """
 
+        curnode = node
+        downreaches = []
+        while not curnode is None:
+            downreaches.append(curnode.reach)
+            curnode = curnode.down
+
+        return list(np.unique(downreaches))
+    
+    def get_kdtree_downstream(self, node:Node_Watershed)-> tuple[list[Node_Watershed], KDTree]:
+        """
+        Get KDTree of downstream reaches
+        """
+
+        downreaches = self.get_downstream_reaches(node)
+        return self.get_kdtree_from_reaches(downreaches)
+    
+    def get_kdtree_from_reaches(self, reaches:list[int])->tuple[list[Node_Watershed], KDTree]:
+        """
+        Get KDTree from a list of reaches
+        """
+
+        nodes = self.get_nodes_in_reaches(reaches)
+        xy = [[curnode.x, curnode.y]  for curnode in nodes]
+        return nodes, KDTree(xy)
+    
+    def get_downstream_reaches_excluded(self, node:Node_Watershed, excluded:list[int])->list[int]:
+        """
+        Get index of downstream reaches, excepted the excluded ones
+        """
+
+        list_reaches = self.get_downstream_reaches(node)
+
+        for cur in excluded:
+            if cur in list_reaches:
+                list_reaches.remove(cur)
+
+        return list_reaches
+
+    def go_downstream_until_reach_found(self, node:Node_Watershed, reach:int | list[int])->Node_Watershed:
+        """ Go downstream until a reach is found """
+
+        curnode = node
+        if isinstance(reach, int):
+            while not curnode is None:
+                if curnode.reach == reach:
+                    break
+                curnode = curnode.down
+        elif isinstance(reach, list):
+            while not curnode is None:
+                if curnode.reach in reach:
+                    break
+                curnode = curnode.down
+
+        return curnode
+    
     def get_cums(self, whichreach:int=None, whichup:int=None):
         """
         Récupération de la position curvi
@@ -298,6 +490,25 @@ class RiverSystem:
 
         return slope
 
+    def get_upstreams_coords(self):
+        """
+        Récupération des coordonnées des amonts
+        """
+
+        xy = [[curnode.x, curnode.y]  for curnode in self.upstreams['list']]
+        return np.array(xy)
+    
+    def get_nearest_upstream(self, xy:np.ndarray, nb:int) -> tuple[np.ndarray, list[Node_Watershed]]:
+        """
+        Recherche des amonts les plus proches
+        """
+
+        xy_up = self.get_upstreams_coords()
+        loc_kd = KDTree(xy_up)
+        dd, ii =loc_kd.query(xy, nb)
+
+        return dd, [self.upstreams['list'][curi] for curi in ii]
+    
     def create_index(self):
         """
         Incrément d'index depuis l'amont jusque l'exutoire final
@@ -1132,19 +1343,222 @@ class SubWatershed:
                  runoff:list[Node_Watershed],
                  rivers:list[Node_Watershed]) -> None:
 
-        self.parent = parent
-        self.index  = idx
-        self.name   = name
-        self.mask   = mask
-        self.nodes  = nodes
-        self.rivers = rivers
-        self.runoff = runoff
+        self.parent:"Watershed" = parent
+        self.index:int          = idx  # index of subwatershed - **NOT** sorted like in the array
+        self.name:str           = name # name of the subwatershed
+        self.mask:WolfArray     = mask # WolfArray of the subwatershed -- All nodes are masked except the subwatershed
+        self.mask.count()
+        
+        self.nodes:list[Node_Watershed]  = nodes    # all nodes in the subwatershed
+        self.rivers:list[Node_Watershed] = rivers   # only rivers - sorted by dem value --> outlet is the first one
+        self.runoff:list[Node_Watershed] = runoff
+        
         self.idx_reaches = np.unique(np.asarray([x.reach for x in rivers]))
 
+        self._index_sorted = idx
+
+        self._is_virtual:bool        = False
+        self._src_sub:"SubWatershed" = None
+
+    @property
+    def is_virtual(self) -> bool:
+        """ Vérification si le sous-bassin est virtuel """
+
+        return self._is_virtual
+    
     @property
     def surface(self) -> float:
+        """ Surface du bassin versant en m² """
         return self.mask.nbnotnull * self.mask.dx * self.mask.dy
+    
+    @property
+    def area(self) -> float:
+        """ Surface du bassin versant en km² """
+        return self.surface / 1.e6
+    
+    @property
+    def area_outlet(self) -> float:
+        """ Surface du bassin à l'exutoire """
+
+        return self.outlet.uparea
+    
+    @property
+    def outlet(self) -> Node_Watershed:
+        """ Outlet of the subbasin """
+
+        return self.rivers[0]
+
+    def is_reach_in_sub(self, idx_reach:int) -> bool:
+        """ Vérification si un bief est dans le sous-bassin """
+
+        return idx_reach in self.idx_reaches
+    
+    def is_in_rivers(self, node:Node_Watershed) -> bool:
+        """ Vérification si un noeud est dans les rivières """
+
+        return node in self.rivers
+
+    def get_list_nodes_river(self, idx_reach:int) -> list[Node_Watershed]:
+        """
+        Récupération des noeuds d'une rivière
+        """
+
+        return [x for x in self.rivers if x.reach==idx_reach]
+    
+    def get_nearest_river(self, x, y) -> Node_Watershed:
+        """
+        Récupération du noeud de rivière le plus proche
+        """
+
+        return min(self.rivers, key=lambda x: x.distance(x,y))
+    
+    def get_max_area_in_reach(self, idx_reach:int) -> float:
+        """
+        Récupération de la surface maximale dans un bief
+        """
+
+        return max([x.uparea for x in self.get_list_nodes_river(idx_reach)])
+        
+    def get_min_area_in_reach(self, idx_reach:int) -> float:
+        """
+        Récupération de la surface minimale dans un bief
+        """
+
+        return min([x.uparea for x in self.get_list_nodes_river(idx_reach)])
+    
+    def get_min_area_along_reaches(self, reaches:list[int], starting_node:Node_Watershed = None) -> float:    
+        """ Aire drainée à la limite amont du ss-bassin """
+        
+        if starting_node is None:
+            starting_node = self.outlet
+        
+        upriver = starting_node.upriver
+        
+        area_min = []
+
+        idx_reach = 0
+        for curnode in upriver:
+            if curnode.reach in reaches:
+                idx_reach = curnode.reach
+                area_min.append(self.get_min_area_in_reach(idx_reach))
+
+        if len(area_min) == 0:
+            return None
+        else:
+            return min(area_min)
+        
+    def get_up_rivernode_outside_sub(self, starting_node:Node_Watershed, reaches:list[int]) -> Node_Watershed:
+        """ Récupération du noeud de rivière en amont du sous-bassin """
+        
+        def up_in_reaches(node:Node_Watershed, reaches:list[int]) -> Node_Watershed:            
+        
+            if len(node.upriver) == 0:
+                # No upstream node
+                return node
+            else:
+                # Iterate over upriver nodes
+                for curup in node.upriver:
+                    if curup.reach in reaches:
+                        # If the reach is in the list, return the node
+                        return curup        
+                
+                # No node found in the list of reaches
+                return node
+
+        if self._is_virtual:
+            return self._src_sub.get_up_rivernode_outside_sub(starting_node, reaches)
+
+        up = up_in_reaches(starting_node, reaches)
+        loc_up = None
+        while up is not starting_node and up.sub == starting_node.sub and up is not loc_up:
+            # bouclage parfois utile en fonction de la superposition ou non du tracé
+            # vectoriel de lit mineur vis-à-vis de la discrétsation hydrologique
+            loc_up = up
+            up = up_in_reaches(up, reaches)
+
+        return up if up is not loc_up else None
+
+    def get_area_outside_sub_if_exists(self, starting_node:Node_Watershed, reaches:list[int]) -> float:
+        """ Aire drainée en amont du sous-bassin """
+        
+        up_outside = self.get_up_rivernode_outside_sub(starting_node, reaches)
+        
+        if up_outside is None:
+            return 0.
+        else:
+            return up_outside.uparea        
+    
+    def get_river_nodes_from_upareas(self, min_area:float, max_area:float) -> list[Node_Watershed]:
+        """
+        Récupération des noeuds de rivière entre deux surfacesde BV.
+
+        Les surfaces sont exprimées en km².
+        
+        Les bornes sont incluses.
+
+        :param min_area: surface minimale
+        :param max_area: surface maximale
+        """
 
+        return [x for x in self.rivers if x.uparea>=min_area and x.uparea<=max_area]
+
+    def is_in_subwatershed(self, vec:vector) -> bool:
+        """ Vérification si un vecteur est dans le sous-bassin """
+
+        centroid = vec.asshapely_pol().centroid
+
+        i, j = self.mask.get_ij_from_xy(centroid.x, centroid.y)
+
+        return self.mask.array.mask[i,j] == False
+    
+    def filter_zones(self, zones_to_filter:Zones, force_virtual_if_any:bool = False) -> list[vector]:
+        """
+        Filtrage des zones pour ne garder que celle se trouvant dans le sous-bassin
+        """
+
+        if self._is_virtual and not force_virtual_if_any:
+            return self._src_sub.filter_zones(zones_to_filter)
+
+        return [curvec for curzone in zones_to_filter.myzones for curvec in curzone.myvectors if self.is_in_subwatershed(curvec)]
+    
+    def get_virtual_subwatershed(self, outlet:Node_Watershed, excluded_nodes:list[Node_Watershed] = []) -> "SubWatershed":
+        """
+        Création d'un sous-bassin virtuel sur base d'une maille rivière aval
+        """
+
+        if not outlet.river:
+            logging.error(_('The outlet should be a river node'))
+            return None
+
+        mymask = WolfArray(mold = self.mask)
+        mymask.array.mask[:,:] = True
+
+        all, river, runoff = outlet.get_up_nodes_same_sub(excluded_nodes)
+
+        for curnode in all:
+            mymask.array.mask[curnode.i,curnode.j] = False
+
+        newsub = SubWatershed(self.parent, 
+                            self.name + '_virtual', 
+                            self.parent.nb_subs + 1, 
+                            mymask, 
+                            all,
+                            runoff,
+                            river, 
+                            )
+        
+        newsub._is_virtual = True
+        newsub._src_sub = self
+        
+        return newsub
+    
+    def get_downstream_node_in_reach(self, reach:int) -> Node_Watershed:
+        """
+        Récupération du noeud aval dans un bief
+        """
+
+        # rivers are sorted by dem value, so the first one is the outlet
+        return self.get_list_nodes_river(reach)[0]
 class Watershed:
     """Classe bassin versant"""
 
@@ -1164,6 +1578,7 @@ class Watershed:
     couplednodes:list # forced exchanges
 
     subcatchments: list[SubWatershed]
+    virtualcatchments : list[SubWatershed]
     statisticss: dict
 
     couplednodesxy:list[float,float,float,float]
@@ -1188,6 +1603,8 @@ class Watershed:
                  dir_mnt_subpixels:str=None,
                  *args, **kwargs):
 
+        self.virtualcatchments = []
+
         logging.info(_('Read files...'))
 
         self.dir=os.path.normpath(dir)
@@ -1209,7 +1626,11 @@ class Watershed:
 
             if lines[0]=='COORDINATES':
                 for xy in enumerate(lines[1:]):
-                    xup,yup,xdown,ydown=xy[1].split('\t')
+                    xy_split = xy[1].split('\t')
+                    if len(xy_split)==4:
+                        xup,yup,xdown,ydown=xy_split
+                    else:
+                        xup,yup,xdown,ydown=xy_split[:4]
                     self.couplednodesxy.append([float(xup),float(yup),float(xdown),float(ydown)])
                     self.couplednodesij.append([self.subs_array.get_ij_from_xy(float(xup),float(yup)),self.subs_array.get_ij_from_xy(float(xdown),float(ydown))])
 
@@ -1258,6 +1679,71 @@ class Watershed:
     @property
     def resolution(self):
         return self.header.dx
+    
+    # def impose_sorted_index_subbasins(self, new_idx=list[int]):
+    #     """
+    #     Tri des sous-bassins
+    #     """
+
+    #     for cursub in self.subcatchments:
+    #         cursub._index_sorted = new_idx[cursub.index]
+
+    def set_names_subbasins(self, new_names:list[tuple[int,str]]):
+        """
+        Renommage des sous-bassins
+        """
+
+        for cursub, curname in new_names:
+            self.get_subwatershed(cursub).name  = curname
+
+    def add_virtual_subwatershed(self, subwater:SubWatershed):
+        """     
+        Ajout d'un sous-bassin virtuel
+        """
+
+        self.virtualcatchments.append(subwater)
+
+        subwater.name += str(len(self.virtualcatchments))
+
+    def create_virtual_subwatershed(self, outlet:Node_Watershed, excluded_nodes:list[Node_Watershed] = []):
+        """
+        Création d'un sous-bassin virtuel
+        """
+
+        newsub = self.get_subwatershed(outlet.sub).get_virtual_subwatershed(outlet, excluded_nodes=excluded_nodes)
+        
+        self.add_virtual_subwatershed(newsub)
+
+        return newsub
+
+    def get_subwatershed(self, idx_sorted_or_name:int | str) -> SubWatershed:
+        """
+        Récupération d'un sous-bassin sur base de l'index trié
+        """
+        
+        if isinstance(idx_sorted_or_name, str):
+            for cur_sub in self.subcatchments:
+                if cur_sub.name == idx_sorted_or_name:
+                    return cur_sub
+                
+            if len(self.virtualcatchments)>0:
+                for cur_sub in self.virtualcatchments:
+                    if cur_sub.name == idx_sorted_or_name:
+                        return cur_sub
+                        
+        elif isinstance(idx_sorted_or_name, int):        
+            for cur_sub in self.subcatchments:
+                if cur_sub._index_sorted+1 == idx_sorted_or_name:
+                    return cur_sub   
+                
+            if len(self.virtualcatchments)>0:
+                for cur_sub in self.virtualcatchments:
+                    if cur_sub._index_sorted+1 == idx_sorted_or_name:
+                        return cur_sub
+        else:
+            logging.error(_('Index must be an integer or a string!'))
+        
+        return None
 
     def get_node_from_ij(self, i:int,j:int):
         """
@@ -1498,6 +1984,8 @@ class Watershed:
 
         self.rivers=list(filter(lambda x: x.river,self.nodes))
         self.rivers.sort(key=lambda x: x.dem['dem_after_corr'])
+
+        # FIXME : Caution the following iterative function can induce a RecursionError in Debug for bigger systems
         sys.setrecursionlimit(len(self.nodes))
         self.outlet.incr_curvi()
 
@@ -1980,3 +2468,53 @@ class Watershed:
             cur_node.time = wolf_time[cur_node.i, cur_node.j]
 
         self.to_update_times = False
+
+    def get_subwatershed_from_ij(self, i:int, j:int) -> SubWatershed:
+        """
+        Récupération d'un sous-bassin sur base des indices (i,j)
+
+        :return: SubWatershed : sous-bassin or None
+        """
+
+        if self.subs_array.array.mask[i,j]:
+            return None
+
+        idx_sub = self.subs_array.array[i,j]
+
+        return self.subcatchments[idx_sub-1]
+    
+    def get_subwatershed_from_xy(self, x:float, y:float) -> SubWatershed:
+        """
+        Récupération d'un sous-bassin sur base des coordonnées (x,y)
+
+        :return: SubWatershed : sous-bassin or None
+        """
+
+        i,j = self.header.get_ij_from_xy(x,y)
+        return self.get_subwatershed_from_ij(i,j)
+    
+    def get_subwatershed_from_vector(self, vec:vector) -> tuple[SubWatershed, bool, list[SubWatershed]]:
+        """
+        Récupération d'un sous-bassin sur base d'un vecteur.
+
+        Recherche sur base du centroid du vecteur
+
+        :param vec: vecteur
+
+        :return: tuple(SubWatershed, bool, list[SubWatershed]) : sous-bassin, entièrement dans le sous-bassin, autres sous-bassins
+        """
+
+        centroid = vec.asshapely_pol().centroid
+
+        sub = self.get_subwatershed_from_xy(centroid.x, centroid.y)
+
+        entirely = True
+        others = []
+        for curvert in vec.myvertices:
+            cursub = self.get_subwatershed_from_xy(curvert.x, curvert.y)
+            if cursub is not None and cursub != sub:
+                logging.warning(_('The vector is not entirely in the same subcatchment'))
+                entirely = False
+                others.append(cursub)
+
+        return sub, entirely, others