Changed variable names to use junction, updated test

Author: kaklise

.gitignore

@@ -5,6 +5,7 @@ build/
 wntr.egg-info/
 dist/
 docker/
+/.vscode
 
 *.pyd
 *.pyc

documentation/morph.rst

@@ -82,6 +82,7 @@ This initial set of operations can generate new branch pipes, pipes in series, a
 This cycle repeats until the network can no longer be reduced.  
 The user can specify if branch trimming, series pipe merge, and/or parallel pipe merge should be included in the skeletonization operations.  
 The user can also specify a maximum number of cycles to include in the process. 
+The user can also specify a list of junctions and pipes which should be excluded from skeletonization operations.
 
 .. only:: latex
 

wntr/epanet/io.py

@@ -236,39 +236,30 @@ def __init__(self):
         self.top_comments = []
         self.curves = OrderedDict()
 
-    def read(self, inp_files, wn=None):
+    @staticmethod
+    def read_sections(inp_files:list) -> tuple:
         """
-        Read an EPANET INP file and load data into a water network model object.
-        Both EPANET 2.0 and EPANET 2.2 INP file options are recognized and handled.
+        Read an EPANET INP file and return a dictionary of sections.
 
         Parameters
         ----------
         inp_files : str or list
             An EPANET INP input file or list of INP files to be combined
-        wn : WaterNetworkModel, optional
-            An optional network model to append onto; by default a new model is created.
 
         Returns
         -------
-        WaterNetworkModel
-            A water network model object
+        tuple
+            A tuple of two elements: a dictionary (OrderedDict) of sections, and a list of top comments
 
         """
-        if wn is None:
-            wn = WaterNetworkModel()
-        self.wn = wn
         if not isinstance(inp_files, list):
             inp_files = [inp_files]
-        wn.name = inp_files[0]
 
-        self.curves = OrderedDict()
-        self.top_comments = []
-        self.sections = OrderedDict()
+        sections = OrderedDict()
         for sec in _INP_SECTIONS:
-            self.sections[sec] = []
-        self.mass_units = None
-        self.flow_units = None
+            sections[sec] = []
 
+        top_comments = []
         for filename in inp_files:
           section = None
           lnum = 0
@@ -306,13 +297,55 @@ def read(self, inp_files, wn=None):
                     else:
                         raise ENSyntaxError(201, line_num=lnum, line = line)
                 elif section is None and line.startswith(';'):
-                    self.top_comments.append(line[1:])
+                    top_comments.append(line[1:])
                     continue
                 elif section is None:
                     logger.debug('Found confusing line: %s', repr(line))
                     raise ENSyntaxError(201, line_num=lnum, line=line)
                 # We have text, and we are in a section
-                self.sections[section].append((lnum, line))
+                sections[section].append((lnum, line))
+        return sections, top_comments
+
+    def read(self, inp_files, wn=None, sections=None, top_comments=None):
+        """
+        Read an EPANET INP file and load data into a water network model object.
+        Both EPANET 2.0 and EPANET 2.2 INP file options are recognized and handled.
+
+        Parameters
+        ----------
+        inp_files : str or list
+            An EPANET INP input file or list of INP files to be combined
+        wn : WaterNetworkModel, optional
+            An optional network model to append onto; by default a new model is created.
+
+        Returns
+        -------
+        :class:`~wntr.network.model.WaterNetworkModel`
+            A water network model object
+
+        """
+        if wn is None:
+            wn = WaterNetworkModel()
+        self.wn = wn
+        if not isinstance(inp_files, list):
+            inp_files = [inp_files]
+        wn.name = inp_files[0]
+
+        self.curves = OrderedDict()
+
+        self.mass_units = None
+        self.flow_units = None
+
+        if sections is None:
+            self.sections, self.top_comments = self.read_sections(inp_files)
+        else:
+            self.sections = sections
+            self.top_comments = top_comments
+            if top_comments is None:
+                self.top_comments = []
+            
+
+        assert len([key for key in self.sections.keys()]) == len(_INP_SECTIONS), 'Missing sections in INP file'
 
         # Parse each of the sections
         # The order of operations is important as certain things require prior knowledge

wntr/morph/skel.py

@@ -14,7 +14,8 @@
 logger = logging.getLogger(__name__)
 
 def skeletonize(wn, pipe_diameter_threshold, branch_trim=True, series_pipe_merge=True, 
-                parallel_pipe_merge=True, max_cycles=None, use_epanet=True, 
+                parallel_pipe_merge=True, max_cycles=None, use_epanet=True,
+                pipes_to_exclude:list=[], junctions_to_exclude:list=[],
                 return_map=False, return_copy=True):
     """
     Perform network skeletonization using branch trimming, series pipe merge, 
@@ -43,6 +44,10 @@ def skeletonize(wn, pipe_diameter_threshold, branch_trim=True, series_pipe_merge
     use_epanet: bool, optional
         If True, use the EpanetSimulator to compute headloss in pipes.  
         If False, use the WNTRSimulator to compute headloss in pipes.
+    pipes_to_exclude: list, optional
+        List of pipe names to exclude from skeletonization
+    junctions_to_exclude: list, optional
+        List of junction names to exclude from skeletonization
     return_map: bool, optional
         If True, return a skeletonization map. The map is a dictionary 
         that includes original nodes as keys and a list of skeletonized nodes 
@@ -60,7 +65,12 @@ def skeletonize(wn, pipe_diameter_threshold, branch_trim=True, series_pipe_merge
         nodes as keys and a list of skeletonized nodes that were merged into 
         each original node as values.
     """
-    skel = _Skeletonize(wn, use_epanet, return_copy)
+    if len(pipes_to_exclude) > 0:
+        assert len(set(pipes_to_exclude) - set(wn.pipe_name_list)) == 0
+    if len(junctions_to_exclude) > 0:
+        assert len(set(junctions_to_exclude) - set(wn.junction_name_list)) == 0
+    
+    skel = _Skeletonize(wn, use_epanet, return_copy, pipes_to_exclude, junctions_to_exclude)
 
     skel.run(pipe_diameter_threshold, branch_trim, series_pipe_merge, 
              parallel_pipe_merge, max_cycles)
@@ -73,7 +83,7 @@ def skeletonize(wn, pipe_diameter_threshold, branch_trim=True, series_pipe_merge
 
 class _Skeletonize(object):
 
-    def __init__(self, wn, use_epanet, return_copy):
+    def __init__(self, wn, use_epanet, return_copy, pipes_to_exclude, junctions_to_exclude):
 
         if return_copy:
             # Get a copy of the WaterNetworkModel
@@ -93,6 +103,7 @@ def __init__(self, wn, use_epanet, return_copy):
         self.skeleton_map = skel_map
 
         # Get a list of junction and pipe names that are associated with controls
+        # Add them to junctions and pipes to exclude
         junc_with_controls = []
         pipe_with_controls = []
         for name, control in self.wn.controls():
@@ -101,8 +112,10 @@ def __init__(self, wn, use_epanet, return_copy):
                     junc_with_controls.append(req.name)
                 elif isinstance(req, Pipe):
                     pipe_with_controls.append(req.name)
-        self.junc_with_controls = list(set(junc_with_controls))
-        self.pipe_with_controls = list(set(pipe_with_controls))
+        self.junc_to_exclude = list(set(junc_with_controls))
+        self.junc_to_exclude.extend(junctions_to_exclude)
+        self.pipe_to_exclude = list(set(pipe_with_controls))
+        self.pipe_to_exclude.extend(pipes_to_exclude)
 
         # Calculate pipe headloss using a single period EPANET simulation
         duration = self.wn.options.time.duration
@@ -163,7 +176,7 @@ def branch_trim(self, pipe_threshold):
         patterns) to the neighboring junction.
         """
         for junc_name in self.wn.junction_name_list:
-            if junc_name in self.junc_with_controls:
+            if junc_name in self.junc_to_exclude:
                 continue
             neighbors = list(nx.neighbors(self.G,junc_name))
             if len(neighbors) > 1:
@@ -181,7 +194,7 @@ def branch_trim(self, pipe_threshold):
             pipe = self.wn.get_link(pipe_name)
             if not ((isinstance(pipe, Pipe)) and \
                 (pipe.diameter <= pipe_threshold) and \
-                pipe_name not in self.pipe_with_controls):
+                pipe_name not in self.pipe_to_exclude):
                 continue
 
             logger.info('Branch trim: '+ str(junc_name) + str(neighbors))
@@ -215,7 +228,7 @@ def series_pipe_merge(self, pipe_threshold):
         to the nearest junction.
         """
         for junc_name in self.wn.junction_name_list:
-            if junc_name in self.junc_with_controls:
+            if junc_name in self.junc_to_exclude:
                 continue
             neighbors = list(nx.neighbors(self.G,junc_name))
             if not (len(neighbors) == 2):
@@ -239,8 +252,8 @@ def series_pipe_merge(self, pipe_threshold):
                 (isinstance(pipe1, Pipe)) and \
                 ((pipe0.diameter <= pipe_threshold) and \
                 (pipe1.diameter <= pipe_threshold)) and \
-                pipe_name0 not in self.pipe_with_controls and \
-                pipe_name1 not in self.pipe_with_controls):
+                pipe_name0 not in self.pipe_to_exclude and \
+                pipe_name1 not in self.pipe_to_exclude):
                 continue
             # Find closest neighbor junction
             if (isinstance(neigh_junc0, Junction)) and \
@@ -305,7 +318,7 @@ def parallel_pipe_merge(self, pipe_threshold):
         """
 
         for junc_name in self.wn.junction_name_list:
-            if junc_name in self.junc_with_controls:
+            if junc_name in self.junc_to_exclude:
                 continue
             neighbors = nx.neighbors(self.G,junc_name)
             for neighbor in [n for n in neighbors]:
@@ -322,8 +335,8 @@ def parallel_pipe_merge(self, pipe_threshold):
                        (isinstance(pipe1, Pipe)) and \
                         ((pipe0.diameter <= pipe_threshold) and \
                         (pipe1.diameter <= pipe_threshold)) and \
-                        pipe_name0 not in self.pipe_with_controls and \
-                        pipe_name1 not in self.pipe_with_controls):
+                        pipe_name0 not in self.pipe_to_exclude and \
+                        pipe_name1 not in self.pipe_to_exclude):
                         continue
 
                     logger.info('Parallel pipe merge: '+ str(junc_name) + str((pipe_name0, pipe_name1)))

wntr/tests/test_morph.py

@@ -328,6 +328,42 @@ def test_skeletonize_with_controls(self):
         self.assertEqual(skel_wn.num_nodes, wn.num_nodes - 17)
         self.assertEqual(skel_wn.num_links, wn.num_links - 22)
 
+    def test_skeletonize_with_excluding_nodes_and_pipes(self):
+
+        inp_file = join(datadir, "skeletonize.inp")
+        wn = wntr.network.WaterNetworkModel(inp_file)
+        
+        # Run skeletonization without excluding junctions or pipes
+        skel_wn = wntr.morph.skeletonize(wn, 12.0 * 0.0254, use_epanet=False)
+        # Junction 13 and Pipe 60 are not in the skeletonized model
+        assert "13" not in skel_wn.junction_name_list
+        assert "60" not in skel_wn.pipe_name_list
+        
+        # Run skeletonization excluding Junction 13 and Pipe 60
+        skel_wn = wntr.morph.skeletonize(wn, 12.0 * 0.0254, use_epanet=False, 
+                                         junctions_to_exclude=["13"], 
+                                         pipes_to_exclude=["60"])
+        # Junction 13 and Pipe 60 are in the skeletonized model
+        assert "13" in skel_wn.junction_name_list
+        assert "60" in skel_wn.pipe_name_list
+        self.assertEqual(skel_wn.num_nodes, wn.num_nodes - 17)
+        self.assertEqual(skel_wn.num_links, wn.num_links - 22)
+
+        # Change diameter of link 60 one link connected to Junction 13 to be 
+        # greater than 12, should get some results as above
+        # Note, link 11 is connected to Junction 13
+        link = wn.get_link("60")
+        link.diameter = 16 * 0.0254
+        link_connected_to_13 = wn.get_links_for_node('13')[0]
+        link = wn.get_link(link_connected_to_13)
+        link.diameter = 16 * 0.0254
+
+        skel_wn = wntr.morph.skeletonize(wn, 12.0 * 0.0254, use_epanet=False)
+        assert "13" in skel_wn.junction_name_list
+        assert "60" in skel_wn.pipe_name_list
+        self.assertEqual(skel_wn.num_nodes, wn.num_nodes - 17)
+        self.assertEqual(skel_wn.num_links, wn.num_links - 22)
+
     def test_series_merge_properties(self):
 
         wn = wntr.network.WaterNetworkModel()