fix: updated the initial_quality to valid gis name

USEPA · Dec 16, 2024 · 97a378d · 97a378d
1 parent effb245
commit 97a378d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/wntr/network/base.py b/wntr/network/base.py
@@ -321,6 +321,7 @@ class Link(six.with_metaclass(abc.ABCMeta, object)):
         end_node_name
         initial_status
         initial_setting
+        initial_quality
         tag
         vertices
 

diff --git a/wntr/network/elements.py b/wntr/network/elements.py
@@ -888,7 +888,8 @@ class Pipe(Link):
                         "minor_loss",
                         "initial_status",
                         "check_valve"]
-    _optional_attributes = ["bulk_coeff",
+    _optional_attributes = ["initial_quality",
+                            "bulk_coeff",
                             "wall_coeff",
                             "vertices",
                             "tag"]
@@ -1044,6 +1045,7 @@ class Pump(Link):
         speed_timeseries
         initial_status
         initial_setting
+        initial_quality
         efficiency
         energy_price
         energy_pattern
@@ -1074,7 +1076,8 @@ class Pump(Link):
                         "base_speed",
                         "speed_pattern_name",
                         "initial_status"]
-    _optional_attributes = ["initial_setting",
+    _optional_attributes = ["initial_quality",
+                            "initial_setting",
                             "efficiency",
                             "energy_pattern",
                             "energy_price",
@@ -1256,6 +1259,7 @@ class HeadPump(Pump):
         speed_timeseries
         initial_status
         initial_setting
+        initial_quality
         pump_type
         pump_curve_name
         efficiency
@@ -1476,6 +1480,7 @@ class PowerPump(Pump):
         speed_timeseries
         initial_status
         initial_setting
+        initial_quality
         pump_type
         power
         efficiency
@@ -1566,6 +1571,7 @@ class Valve(Link):
         valve_type
         initial_status
         initial_setting
+        initial_quality
         vertices
         tag
 
@@ -1592,7 +1598,8 @@ class Valve(Link):
                         "minor_loss",
                         "initial_setting",
                         "initial_status"]
-    _optional_attributes = ["vertices",
+    _optional_attributes = ["initial_quality",
+                            "vertices",
                             "tag"]
 
     def __init__(self, name, start_node_name, end_node_name, wn):

diff --git a/wntr/tests/test_network.py b/wntr/tests/test_network.py
@@ -1128,6 +1128,7 @@ def test_valid_gis_names(self):
             data_columns.append(data.index.name)
 
             # Check that all data columns are valid
+            print(set(data_columns)-set(valid_columns))
             assert len(set(data_columns)-set(valid_columns)) == 0
             # Check that all required columns are in the data
             assert len(set(required_columns)-set(data_columns)) == 0