Repository: pyswmm/swmmio

Issue:
Updates to support read/write of additional INP sections
I have been using swmmio to add stormwater/wastewater analysis to WNTR.  The ability to read, modify, and write INP files using pandas DataFrames has been very helpful!  I found a few issues when comparing simulation results with and without the read/write of INP files ([XSECTIONS], [INFLOWS], and [CURVES]). I also needed to support a few additional sections in Model.inp ([CONTROLS] and [PATTERNS]).  I stored controls and patterns in DataFrames with 1 line per entry.  I'm using INP files from SWMM and swmmio to test functionality.  Let me know if you're interested in a PR!

Discussion:
@aerispaha: @kaklise, I'm so happy that swmmio has been helpful for you. 100% we're interested in a PR! Please feel free to submit a PR at your convenience, then we can continue the conversation there :) 

Solution patch:
diff --git a/requirements.txt b/requirements.txt
index 8e7f5a0..74f6f03 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
 # Build dependencies
 pytest
-Pillow==10.2.0
+Pillow==10.3.0
 numpy>=1.16.4
 pandas>=0.24.2
 pyshp==2.1.0
diff --git a/swmmio/core.py b/swmmio/core.py
index f0f8aa0..2da7c5b 100644
--- a/swmmio/core.py
+++ b/swmmio/core.py
@@ -600,6 +600,8 @@ class inp(SWMMIOFile):
         self._streets_df = None
         self._inlets_df = None
         self._inlet_usage_df = None
+        self._patterns_df = None
+        self._controls_df = None
 
         SWMMIOFile.__init__(self, file_path)  # run the superclass init
 
@@ -644,6 +646,8 @@ class inp(SWMMIOFile):
             '[STREETS]',
             '[INLETS]',
             '[INLET_USAGE]',
+            '[PATTERNS]',
+            '[CONTROLS]',
         ]
 
     def save(self, target_path=None):
@@ -1452,7 +1456,7 @@ class inp(SWMMIOFile):
         if self._inflows_df is not None:
             return self._inflows_df
         inf = dataframe_from_inp(self.path, 'INFLOWS', quote_replace='_!!!!_')
-        self._inflows_df = inf.replace('_!!!!_', np.nan)
+        self._inflows_df = inf.replace('_!!!!_', '""') # revert quote replace
         return self._inflows_df
 
     @inflows.setter
@@ -1508,6 +1512,98 @@ class inp(SWMMIOFile):
         """Set inp.timeseries DataFrame."""
         self._timeseries_df = df
 
+    @property
+    def patterns(self):
+        """
+        Get/set patterns section of the model
+        
+        :return: dataframe of patterns
+        
+        >>> from swmmio.examples import pump_control
+        >>> # NOTE, only the first 5 columns are shown in the following example
+        >>> pump_control.inp.patterns.iloc[:,0:5]  #doctest: +NORMALIZE_WHITESPACE
+                Type  Factor1  Factor2  Factor3  Factor4
+        Name                                            
+        DWF   HOURLY   0.0151  0.01373  0.01812  0.01098
+        """
+        
+        if self._patterns_df is not None:
+            return self._patterns_df
+        self._patterns_df = dataframe_from_inp(self.path, '[PATTERNS]')
+
+        if self._patterns_df.shape[0] > 0:
+            # reformat, 1 row per pattern
+            pattern_entry_list = []
+            for name, pattern in self._patterns_df.groupby('Name'):
+                pattern_entry = {}
+                pattern_entry['Name'] = name
+                pattern_entry['Type'] = pattern['Type'].iloc[0]
+                if pattern.shape[0] > 1:
+                    # shift pattern values to the right
+                    pattern.iloc[1::, 1::] = pattern.iloc[1::, 0:-1].values
+                    pattern['Factors'] = pattern['Factors'].astype(float)
+                values = pattern.iloc[:, 1:].values.flatten()
+                for i in range(len(values)):
+                    pattern_entry['Factor'+str(i+1)] = values[i]
+                pattern_entry_list.append(pattern_entry)
+
+            self._patterns_df = pd.DataFrame(pattern_entry_list)
+            self._patterns_df.set_index('Name', inplace=True)
+
+        return self._patterns_df
+
+    @patterns.setter
+    def patterns(self, df):
+        """Set inp.patterns DataFrame."""
+        self._patterns_df = df
+
+    @property
+    def controls(self):
+        """
+        Get/set controls section of the model
+        
+        :return: dataframe of controls
+        
+        >>> from swmmio.examples import pump_control
+        >>> pump_control.inp.controls  #doctest: +NORMALIZE_WHITESPACE
+                                                                            Control
+        Name                                                                       
+        RULE PUMP1A  IF NODE SU1 DEPTH >= 4 THEN PUMP PUMP1 status = ON PRIORITY 1 
+        RULE PUMP1B  IF NODE SU1 DEPTH < 1 THEN PUMP PUMP1 status = OFF PRIORITY 1 
+        """
+        
+        if self._controls_df is None:
+            self._controls_df = dataframe_from_inp(self.path, "[CONTROLS]")
+        
+        if self._controls_df.shape[0] > 0:
+            # reformat, 1 row per control
+            control_entry_list = []
+            control_entry = {}
+            controls = self._controls_df['[CONTROLS]']
+            # make sure the first entry starts with RULE
+            assert controls[0][0:5] == "RULE "
+            for row in controls:
+                if row[0:5] == 'RULE ': # new control
+                    if len(control_entry) > 0: # add control to the list
+                        control_entry_list.append(control_entry)
+                    control_entry = {}
+                    control_entry['Name'] = row.rstrip() # remove white space
+                    control_entry['Control'] = ''
+                else:
+                    control_entry['Control'] = control_entry['Control'] + row + ' '
+            if len(control_entry) > 0: # add last control to the list
+                control_entry_list.append(control_entry)
+            
+            self._controls_df = pd.DataFrame(control_entry_list)
+            self._controls_df.set_index('Name', inplace=True)
+                
+        return self._controls_df
+
+    @controls.setter
+    def controls(self, df):
+        """Set inp.controls DataFrame."""
+        self._controls_df = df
+        
     @property
     def tags(self):
         """
diff --git a/swmmio/defs/inp_sections.yml b/swmmio/defs/inp_sections.yml
index ab898ed..715c0a7 100644
--- a/swmmio/defs/inp_sections.yml
+++ b/swmmio/defs/inp_sections.yml
@@ -78,8 +78,8 @@ inp_file_objects:
   STREETS: [Name, Tcrown, Hcurb, Sx, nRoad, a, W, Sides, Tback, Sback, nBack]
   INLETS: [Name, Type, Param1, Param2, Param3, Param4, Param5]
   INLET_USAGE: [Link, Inlet, Node, Number, "%Clogged", Qmax, aLocal, wLocal, Placement]
-
-
+  PATTERNS: [Name, Type, Factors]
+  CONTROLS: [blob]
 
 inp_section_tags:
   ['[TITLE', '[OPTION', '[FILE', '[RAINGAGES', '[TEMPERATURE', '[EVAP',
diff --git a/swmmio/defs/section_headers.yml b/swmmio/defs/section_headers.yml
index f52adb0..f93658a 100644
--- a/swmmio/defs/section_headers.yml
+++ b/swmmio/defs/section_headers.yml
@@ -132,6 +132,16 @@ rpt_sections:
     - SlopePerc
     - RainGage
     - Outlet
+  Pumping Summary:
+    - PercentUtilized
+    - NumberOfStartUps
+    - MinFlowCFS
+    - AvgFlowCFS
+    - MaxFlowCFS
+    - TotalVolume(MG)
+    - PowerUsage(kW-hr)
+    - PercentTimeOffPumpCurveLow
+    - PercentTimeOffPumpCurveHigh
 
 swmm5_version:
   1.13:
diff --git a/swmmio/examples.py b/swmmio/examples.py
index 9a4b5da..53c1445 100644
--- a/swmmio/examples.py
+++ b/swmmio/examples.py
@@ -1,7 +1,8 @@
 from swmmio import Model
 from swmmio.tests.data import (MODEL_A_PATH, MODEL_EX_1, MODEL_FULL_FEATURES_XY,
                                MODEL_FULL_FEATURES__NET_PATH, MODEL_FULL_FEATURES_XY_B,
-                               MODEL_GREEN_AMPT, MODEL_TEST_INLET_DRAINS, MODEL_GROUNDWATER)
+                               MODEL_GREEN_AMPT, MODEL_TEST_INLET_DRAINS, MODEL_GROUNDWATER,
+                               MODEL_PUMP_CONTROL)
 
 # example models
 philly = Model(MODEL_A_PATH, crs="+init=EPSG:2817")
@@ -12,3 +13,4 @@ walnut = Model(MODEL_EX_1)
 green = Model(MODEL_GREEN_AMPT)
 streets = Model(MODEL_TEST_INLET_DRAINS)
 groundwater = Model(MODEL_GROUNDWATER)
+pump_control = Model(MODEL_PUMP_CONTROL)
diff --git a/swmmio/utils/dataframes.py b/swmmio/utils/dataframes.py
index d7e1b29..fa32b09 100644
--- a/swmmio/utils/dataframes.py
+++ b/swmmio/utils/dataframes.py
@@ -56,7 +56,7 @@ def create_dataframe_multi_index(inp_path, section='CURVES'):
 
     df = pd.DataFrame(data=data, columns=cols)
     if sect == 'CURVES':
-        df = df.set_index(['Name', 'Type'])
+        df = df.set_index(['Name']) #, 'Type']) # changed index to Name to correct INP write
     elif sect == 'TIMESERIES':
         df = df.set_index(['Name'])
 
@@ -92,6 +92,7 @@ def dataframe_from_rpt(rpt_path, section, element_id=None):
         # and get the list of columns to use for parsing this section
         end_strings = list(headers.keys())
         end_strings.append('***********')
+        end_strings.append('Analysis begun on')
         start_strings = [section, '-'*20, '-'*20]
     cols = headers[section]['columns']
 
@@ -142,16 +143,22 @@ def dataframe_from_inp(inp_path, section, additional_cols=None, quote_replace='
     # replace occurrences of double quotes ""
     s = s.replace('""', quote_replace)
 
-    # count tokens in first non-empty line, after the header, ignoring comments
+    # find max tokens after the header, ignoring comments
     # if zero tokens counted (i.e. empty line), fall back to headers dict
-    n_tokens = len(re.sub(r"(\n)\1+", r"\1", s).split('\n')[1].split(';')[0].split())
+    n_tokens = 0
+    lines = re.sub(r"(\n)\1+", r"\1", s).split('\n')
+    for line in lines[1::]:
+        n_tokens = max([n_tokens, len(line.split(';')[0].split())])
     n_tokens = len(headers[sect]['columns']) if n_tokens == 0 else n_tokens
 
     # and get the list of columns to use for parsing this section
     # add any additional columns needed for special cases (build instructions)
     additional_cols = [] if additional_cols is None else additional_cols
-    cols = headers[sect]['columns'][:n_tokens] + additional_cols
-
+    cols = headers[sect]['columns'][:n_tokens]
+    if n_tokens > len(cols):
+        cols = cols + ["col"+str(len(cols)+i) for i in range(n_tokens-len(cols))]
+    cols = cols + additional_cols
+    
     if headers[sect]['columns'][0] == 'blob':
         # return the whole row, without specific col headers
         return pd.read_csv(StringIO(s), delim_whitespace=False)
diff --git a/swmmio/version_control/utils.py b/swmmio/version_control/utils.py
index 6cd2c1a..119950d 100644
--- a/swmmio/version_control/utils.py
+++ b/swmmio/version_control/utils.py
@@ -81,6 +81,11 @@ def write_inp_section(file_object, allheaders, sectionheader, section_data, pad_
                 formatters=objectformatter  # {'Comment':formatter}
             )
 
+        # Deliminate control string using keywords
+        if sectionheader == '[CONTROLS]':
+            for sep in [' IF ', ' THEN ', ' PRIORITY ', ' AND ',  ' OR ', ' ELSE ']:
+                add_str = add_str.replace(sep, '\n'+sep)
+
         # write the dataframe as a string
         f.write(add_str + '\n\n')
 


Test patch:
diff --git a/swmmio/tests/data/__init__.py b/swmmio/tests/data/__init__.py
index 5682893..d1f81ed 100644
--- a/swmmio/tests/data/__init__.py
+++ b/swmmio/tests/data/__init__.py
@@ -48,3 +48,6 @@ BUILD_INSTR_01 = os.path.join(DATA_PATH, 'test_build_instructions_01.txt')
 
 df_test_coordinates_csv = os.path.join(DATA_PATH, 'df_test_coordinates.csv')
 OUTFALLS_MODIFIED = os.path.join(DATA_PATH, 'outfalls_modified_10.csv')
+
+# SWMM example model
+MODEL_PUMP_CONTROL = os.path.join(DATA_PATH, 'Pump_Control_Model.inp')
diff --git a/swmmio/tests/test_dataframes.py b/swmmio/tests/test_dataframes.py
index 05b010c..42c21b8 100644
--- a/swmmio/tests/test_dataframes.py
+++ b/swmmio/tests/test_dataframes.py
@@ -1,12 +1,16 @@
 from io import StringIO
+from pandas.testing import assert_series_equal
 
 from swmmio.elements import Links
 from swmmio.tests.data import (MODEL_FULL_FEATURES_PATH, MODEL_FULL_FEATURES__NET_PATH,
                                BUILD_INSTR_01, MODEL_XSECTION_ALT_01, df_test_coordinates_csv,
                                MODEL_FULL_FEATURES_XY, DATA_PATH, MODEL_XSECTION_ALT_03,
                                MODEL_CURVE_NUMBER, MODEL_MOD_HORTON, MODEL_GREEN_AMPT, MODEL_MOD_GREEN_AMPT,
-                               MODEL_INFILTRAION_PARSE_FAILURE, OWA_RPT_EXAMPLE)
+                               MODEL_INFILTRAION_PARSE_FAILURE, OWA_RPT_EXAMPLE,
+                               MODEL_TEST_INLET_DRAINS, MODEL_PUMP_CONTROL)
 from swmmio.utils.dataframes import (dataframe_from_rpt, dataframe_from_inp, dataframe_from_bi)
+from swmmio.utils.text import get_inp_sections_details
+from swmmio.run_models.run import run_simple
 import swmmio
 
 import pandas as pd
@@ -173,7 +177,7 @@ def test_inflow_dwf_dataframe():
 
     inf = m.inp.inflows
     assert (inf.loc['dummy_node2', 'Time Series'] == 'my_time_series')
-    assert (pd.isna(inf.loc['dummy_node6', 'Time Series']))
+    assert inf.loc['dummy_node6', 'Time Series'] == '""'
 
 
 def test_coordinates():
@@ -258,3 +262,71 @@ def test_polygons(test_model_02):
     assert poly1.equals(test_model_02.inp.polygons)
 
     # print()
+
+def test_inp_sections():
+    
+    # Additional models could be added to this test, or additional features 
+    # could be added to the full feature model
+    inpfiles = [
+                MODEL_FULL_FEATURES_PATH,
+                MODEL_CURVE_NUMBER,
+                MODEL_MOD_HORTON,
+                MODEL_GREEN_AMPT,
+                MODEL_TEST_INLET_DRAINS,
+                MODEL_PUMP_CONTROL,
+                ]
+    
+    from swmmio.defs import INP_OBJECTS
+    all_sections = set(sec.upper() for sec in INP_OBJECTS.keys())
+    unsupported_sections = set()
+    tested_sections = set()
+    
+    for inpfile in inpfiles:
+        print(inpfile)
+        
+        # Run SWMM with the original INP file
+        headers = get_inp_sections_details(inpfile, include_brackets=False)
+        run_simple(inpfile)
+        model = swmmio.Model(inpfile, include_rpt=True)
+        links = model.links()
+        
+        # Check that the simulation results loaded into the swmmio model
+        assert 'MaxQ' in links.columns
+    
+        # Create a swmmio model, reassign each section to force the use
+        # of "replace_inp_section", save the INP file, and then run SWMM
+        temp_inpfile = 'temp.inp'
+        model = swmmio.Model(inpfile, include_rpt=False)
+        empty_sections = []
+        for sec in headers.keys():
+            try:
+                df = getattr(model.inp, sec.lower())
+                if not df.empty:
+                    tested_sections.add(sec.upper())
+                else:
+                    empty_sections.append(sec.upper())
+                setattr(model.inp, sec, df.copy())
+            except:
+                unsupported_sections.add(sec.upper())
+        model.inp.save(temp_inpfile)
+        headers2 = get_inp_sections_details(temp_inpfile, 
+                                            include_brackets=False)
+        run_simple(temp_inpfile)
+        model2 = swmmio.Model(temp_inpfile, include_rpt=True)
+        links2 = model2.links()
+        
+        # Check that the INP file headers are the same
+        header_set = set(headers.keys() - empty_sections)
+        header2_set = set(headers2.keys() - empty_sections)
+        assert header_set == header2_set
+        # Check that the simulation results loaded into the swmmio model
+        assert 'MaxQ' in links2.columns
+        # Check that results are the same
+        assert_series_equal(links['MaxQ'], links2['MaxQ']) 
+        
+    # Print empty and unsupported INP file sections
+    print('Unsupported sections', unsupported_sections)
+    print('Untested sections', all_sections - unsupported_sections - tested_sections)
+    
+if __name__ == "__main__":
+    test_inp_sections()
\ No newline at end of file