Fix bug when to_gis() is called on a network with a leak

wntr/epanet/io.py

@@ -3053,36 +3053,50 @@ def _read_control_line(line, wn, flow_units, control_name):
     current = line.split()
     if current == []:
         return
-    link_name = current[1]
-    link = wn.get_link(link_name)
+
+    element_name = current[1]
+
+    # For the case of leak demands
+    if current[0] == 'JUNCTION':
+        element = wn.get_node(element_name)
+    else:
+        element = wn.get_link(element_name)
+
     if current[5].upper() != 'TIME' and current[5].upper() != 'CLOCKTIME':
         node_name = current[5]
     current = [i.upper() for i in current]
-    current[1] = link_name # don't capitalize the link name
+    current[1] = element_name # don't capitalize the link name
 
     # Create the control action object
 
     status = current[2].upper()
     if status == 'OPEN' or status == 'OPENED' or status == 'CLOSED' or status == 'ACTIVE':
         setting = LinkStatus[status].value
-        action_obj = wntr.network.ControlAction(link, 'status', setting)
+        action_obj = wntr.network.ControlAction(element, 'status', setting)
     else:
-        if isinstance(link, wntr.network.Pump):
-            action_obj = wntr.network.ControlAction(link, 'base_speed', float(current[2]))
-        elif isinstance(link, wntr.network.Valve):
-            if link.valve_type == 'PRV' or link.valve_type == 'PSV' or link.valve_type == 'PBV':
+        if isinstance(element, wntr.network.Pump):
+            action_obj = wntr.network.ControlAction(element, 'base_speed', float(current[2]))
+        elif isinstance(element, wntr.network.Valve):
+            if element.valve_type == 'PRV' or element.valve_type == 'PSV' or element.valve_type == 'PBV':
                 setting = to_si(flow_units, float(current[2]), HydParam.Pressure)
-            elif link.valve_type == 'FCV':
+            elif element.valve_type == 'FCV':
                 setting = to_si(flow_units, float(current[2]), HydParam.Flow)
-            elif link.valve_type == 'TCV':
+            elif element.valve_type == 'TCV':
                 setting = float(current[2])
-            elif link.valve_type == 'GPV':
+            elif element.valve_type == 'GPV':
                 setting = current[2]
             else:
-                raise ValueError('Unrecognized valve type {0} while parsing control {1}'.format(link.valve_type, line))
-            action_obj = wntr.network.ControlAction(link, 'setting', setting)
+                raise ValueError('Unrecognized valve type {0} while parsing control {1}'.format(element.valve_type, line))
+            action_obj = wntr.network.ControlAction(element, 'setting', setting)
+        elif isinstance(element, wntr.network.Node):
+            if status=="TRUE":
+                setting = True
+            elif status=="FALSE":
+                setting = False
+            action_obj = wntr.network.ControlAction(element, 'leak_status', setting)
+
         else:
-            raise RuntimeError(('Links of type {0} can only have controls that change\n'.format(type(link))+
+            raise RuntimeError(('Links of type {0} can only have controls that change\n'.format(type(element))+
                                 'the link status. Control: {0}'.format(line)))
 
     # Create the control object

wntr/gis/network.py

@@ -114,6 +114,11 @@ def _extract_geodataframe(df, crs=None, valid_base_names=None,
                 if 'node_type' in df.columns:
                     geom = [Point((x,y)) for x,y in df['coordinates']]
                     del df['node_type']
+
+                    # do not carry over leak attributes to dataframe. 
+                    del df['leak']
+                    del df['leak_area']
+                    del df['leak_discharge_coeff']
                 elif 'link_type' in df.columns:
                     geom = []
                     for link_name in df['name']:

wntr/network/base.py

@@ -192,6 +192,11 @@ def leak_status(self):
     # @leak_status.setter
     # def leak_status(self, value):
     #     self._leak_status = value
+
+    @property
+    def leak(self):
+        """float: (read-only) the current simulation leak area at the node"""
+        return self._leak
 
     @property
     def leak_area(self):
@@ -257,8 +262,8 @@ def to_dict(self):
         d['node_type'] = self.node_type
         for k in dir(self):
             if not k.startswith('_') and \
-              k not in ['demand', 'head', 'leak_area', 'leak_demand',
-                        'leak_discharge_coeff', 'leak_status', 'level', 'pressure', 'quality', 'vol_curve', 'head_timeseries']:
+              k not in ['demand', 'head', 'leak_demand', 'leak_status', 
+                        'level', 'pressure', 'quality', 'vol_curve', 'head_timeseries']:
                 try:
                     val = getattr(self, k)
                     if not isinstance(val, types.MethodType):

wntr/network/controls.py

@@ -12,7 +12,7 @@
 import abc
 from wntr.utils.ordered_set import OrderedSet
 from collections import OrderedDict
-from .elements import Tank, Junction, Valve, Pump, Reservoir, Pipe
+from .elements import Tank, Junction, Valve, Pump, Reservoir, Pipe, Link
 from wntr.utils.doc_inheritor import DocInheritor
 import warnings
 from typing import Hashable, Dict, Any, Tuple, MutableSet, Iterable
@@ -1753,7 +1753,9 @@ def __repr__(self):
         return '<ControlAction: {}, {}, {}>'.format(str(self._target_obj), str(self._attribute), str(self._repr_value()))
 
     def __str__(self):
-        return "{} {} {} IS {}".format(self._target_obj.link_type.upper(),
+        target_obj_type = (self._target_obj.link_type if isinstance(self._target_obj, Link) else 
+                           self._target_obj.node_type)
+        return "{} {} {} IS {}".format(target_obj_type.upper(),
                                        self._target_obj.name,
                                        self._attribute.upper(),
                                        self._repr_value())

wntr/network/io.py

@@ -137,6 +137,11 @@ def from_dict(d: dict, append=None):
                 j.pressure_exponent = node.setdefault("pressure_exponent")
                 j.required_pressure = node.setdefault("required_pressure")
                 j.tag = node.setdefault("tag")
+
+                j._leak = node.setdefault("leak", False)
+                j._leak_area = node.setdefault("leak_area", 0.0)
+                j._leak_discharge_coeff = node.setdefault("leak_discharge_coeff", 0.0)
+
                 # custom additional attributes
                 for attr in list(set(node.keys()) - set(dir(j))):
                     setattr( j, attr, node[attr] )

wntr/tests/test_network.py

@@ -953,7 +953,7 @@ def setUpClass(self):
 
         self.wntr = wntr
 
-        inp_file = join(ex_datadir, "Net6.inp")
+        self.inp_file = join(ex_datadir, "Net3.inp")
         self.inp_files = [join(ex_datadir, f) for f in ["Net1.inp", "Net2.inp", "Net3.inp", "Net6.inp"]]
 
     @classmethod
@@ -979,6 +979,29 @@ def test_json_pattern_dump(self):
         wn = wntr.network.WaterNetworkModel()
         wn.add_pattern('pat0', [0,1,0,1,0,1,0])
         self.wntr.network.write_json(wn, f'temp.json')
+
+    def test_dict_with_leak(self):
+        # This covers a bug where writing controls to a dictionary broke if a leak was added
+        wn = wntr.network.WaterNetworkModel(self.inp_file)
+        junction_name = wn.junction_name_list[0]
+        junction = wn.get_node(junction_name)
+        junction.add_leak(wn, area=1, start_time=0, end_time=3600)
+        wn_dict = wn.to_dict()
+        wn2 = wntr.network.from_dict(wn_dict)
+
+        # check leak controls and node properties
+        start_control = wn.get_control(wn.control_name_list[18])
+        start_control2 = wn2.get_control(wn2.control_name_list[18])
+        assert str(start_control) == str(start_control2)
+
+        end_control = wn.get_control(wn.control_name_list[19])
+        end_control2 = wn2.get_control(wn2.control_name_list[19])
+        assert str(end_control) == str(end_control2)
+
+        junction2 = wn2.get_node(junction_name)
+        assert junction2._leak
+        assert junction2._leak_area == 1
+        assert junction2._leak_discharge_coeff == 0.75
 
 
 @unittest.skipIf(not has_geopandas,