Repository: natcap/invest (python)

Issue:
Support zero-padded month numbers in SWY?
A [user on the forums](https://community.naturalcapitalproject.org/t/precipitation-data-not-red/3385/3) wanted to have their SWY month numbers zero-padded, which the model can't currently handle.

So instead of a file named `et0-1.tif`, we could also support `et0-01.tif`.

This is a really interesting idea, though, since padding single-digit months would help to keep things in lexicographic order when listing files in a directory and wouldn't be difficult to implement as an option since we're using a regular expression to find these files anyways.

Pull request description:
#1166 Add support for zero-padded months in SWY
## Description
Updated the regex to allow users to have zero-padded month numbers in the monthly precipitation and evapotranspiration files. Now, filenames like et0_01.tif and precip01.tif are permitted. 

Fixes #1166

## Checklist
- [x] Updated HISTORY.rst and link to any relevant issue (if these changes are user-facing)
- [ ] Updated the user's guide (if needed)
- [ ] Tested the Workbench UI (if relevant)


Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 09455fe2e..c0b54a365 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -86,6 +86,10 @@ Unreleased Changes
       updated the stated dtype of most pollination model outputs to be float32
       instead of the float64 dtype that was being assumed previously.  This
       will result in smaller output filesizes with minimal loss of precision.
+* Seasonal Water Yield
+    * Added support for zero padding in month numbers in ET and precipitation
+      file names (i.e., users can now name their file Precip_01.tif).
+      (https://github.com/natcap/invest/issues/1166)
 * Urban Flood Risk
     * Fields present on the input AOI vector are now retained in the output.
       (https://github.com/natcap/invest/issues/1600)
diff --git a/src/natcap/invest/seasonal_water_yield/seasonal_water_yield.py b/src/natcap/invest/seasonal_water_yield/seasonal_water_yield.py
index 728859a2d..e5c2ba57e 100644
--- a/src/natcap/invest/seasonal_water_yield/seasonal_water_yield.py
+++ b/src/natcap/invest/seasonal_water_yield/seasonal_water_yield.py
@@ -659,33 +659,8 @@ def execute(args):
     output_align_list = [
         file_registry['lulc_aligned_path'], file_registry['dem_aligned_path']]
     if not args['user_defined_local_recharge']:
-        precip_path_list = []
-        et0_path_list = []
-
-        et0_dir_list = [
-            os.path.join(args['et0_dir'], f) for f in os.listdir(
-                args['et0_dir'])]
-        precip_dir_list = [
-            os.path.join(args['precip_dir'], f) for f in os.listdir(
-                args['precip_dir'])]
-
-        for month_index in range(1, N_MONTHS + 1):
-            month_file_match = re.compile(r'.*[^\d]%d\.[^.]+$' % month_index)
-
-            for data_type, dir_list, path_list in [
-                    ('et0', et0_dir_list, et0_path_list),
-                    ('Precip', precip_dir_list, precip_path_list)]:
-                file_list = [
-                    month_file_path for month_file_path in dir_list
-                    if month_file_match.match(month_file_path)]
-                if len(file_list) == 0:
-                    raise ValueError(
-                        "No %s found for month %d" % (data_type, month_index))
-                if len(file_list) > 1:
-                    raise ValueError(
-                        "Ambiguous set of files found for month %d: %s" %
-                        (month_index, file_list))
-                path_list.append(file_list[0])
+        precip_path_list = _get_monthly_file_lists(N_MONTHS, args['precip_dir'])
+        et0_path_list = _get_monthly_file_lists(N_MONTHS, args['et0_dir'])
 
         input_align_list = (
             precip_path_list + [args['soil_group_path']] + et0_path_list +
@@ -1380,6 +1355,42 @@ def _aggregate_recharge(
     aggregate_vector = None
 
 
+def _get_monthly_file_lists(n_months, in_dir):
+    """Create list of monthly files for data type
+
+    Parameters:
+        n_months (int): Number of months to iterate over (should be 12)
+        in_dir (string): Path to directory of monthly files (for specific
+                         data type)
+
+    Raises:
+        ValueError: If no file or multiple files are found for a month
+
+    Returns:
+        list: contains monthly file paths for data type
+    """
+    in_path_list = [os.path.join(in_dir, f) for f in os.listdir(in_dir)]
+    out_path_list = []
+
+    for month_index in range(1, n_months + 1):
+        month_file_pattern = re.compile(r'.*[^\d]0?%d\.[^.]+$' % month_index)
+        file_list = [
+            month_file_path for month_file_path in in_path_list
+            if month_file_pattern.match(month_file_path)]
+        if len(file_list) == 0:
+            raise ValueError(
+                "No files found in %s for month %d. Please ensure that \
+                    filenames end in the month number (e.g., precip_1.tif)."
+                % (in_dir, month_index))
+        if len(file_list) > 1:
+            raise ValueError(
+                "Ambiguous set of files found for month %d: %s" %
+                (month_index, file_list))
+        out_path_list.append(file_list[0])
+
+    return out_path_list
+
+
 @validation.invest_validator
 def validate(args, limit_to=None):
     """Validate args to ensure they conform to `execute`'s contract.


Test patch:
diff --git a/tests/test_seasonal_water_yield_regression.py b/tests/test_seasonal_water_yield_regression.py
index e795d8599..2b4359cba 100644
--- a/tests/test_seasonal_water_yield_regression.py
+++ b/tests/test_seasonal_water_yield_regression.py
@@ -166,6 +166,30 @@ def make_precip_rasters(precip_dir_path):
         make_raster_from_array(precip_array, precip_raster_path)
 
 
+def make_zeropadded_rasters(dir_path, prefix):
+    """Make twelve 1x1 raster files with filenames ending in zero-padded
+    month number.
+
+    Args:
+        dir_path (str): path to the directory for saving the rasters.
+        file_prefix (str): prefix of new files to create.
+
+    Returns:
+        list: monthly raster filenames
+    """
+    size = 1
+    monthly_raster_list = []
+
+    for month in range(1, 13):
+        raster_path = os.path.join(
+            dir_path, prefix + str(month).zfill(2) + '.tif')
+        temp_array = numpy.full((size, size), 1, dtype=numpy.int8)
+        make_raster_from_array(temp_array, raster_path)
+        monthly_raster_list.append(raster_path)
+
+    return monthly_raster_list
+
+
 def make_recharge_raster(recharge_ras_path):
     """Make a 100x100 raster of user defined recharge.
 
@@ -318,6 +342,61 @@ class SeasonalWaterYieldUnusualDataTests(unittest.TestCase):
         """Delete workspace after test is done."""
         shutil.rmtree(self.workspace_dir, ignore_errors=True)
 
+    def test_zeropadded_monthly_filenames(self):
+        """test filenames with zero-padded months in
+        _get_monthly_file_lists function
+        """
+        from natcap.invest.seasonal_water_yield.seasonal_water_yield import _get_monthly_file_lists
+
+        n_months = 12
+
+        # Make directory and file names with zero-padded months
+        test_precip_dir_path = os.path.join(self.workspace_dir,
+                                            'test_0pad_precip_dir')
+        os.makedirs(test_precip_dir_path)
+        precip_file_list = make_zeropadded_rasters(test_precip_dir_path, 'Prcp')
+
+        test_eto_dir_path = os.path.join(self.workspace_dir,
+                                         'test_0pad_eto_dir')
+        os.makedirs(test_eto_dir_path)
+        eto_file_list = make_zeropadded_rasters(test_eto_dir_path, 'et0_')
+
+        # Create list of monthly files for data_type
+        eto_path_list = _get_monthly_file_lists(
+            n_months, test_eto_dir_path)
+        
+        precip_path_list = _get_monthly_file_lists(
+            n_months, test_precip_dir_path)
+        
+        # Verify that the returned lists match the input
+        self.assertEqual(precip_path_list, precip_file_list)
+        self.assertEqual(eto_path_list, eto_file_list)
+
+    def test_nonpadded_monthly_filenames(self):
+        """test filenames without zero-padded months in
+        _get_monthly_file_lists function
+        """
+        from natcap.invest.seasonal_water_yield.seasonal_water_yield import _get_monthly_file_lists
+
+        n_months = 12
+
+        # Make directory and file names with (non-zero-padded) months
+        precip_dir_path = os.path.join(self.workspace_dir, 'precip_dir')
+        os.makedirs(precip_dir_path)
+        make_precip_rasters(precip_dir_path)
+
+        precip_path_list = _get_monthly_file_lists(
+            n_months, precip_dir_path)
+        
+        # Create lists of monthly filenames to which to compare function output
+        # Note this is hardcoded to match the filenames created in make_precip_rasters
+        match_precip = [os.path.join(precip_dir_path,
+                                     "precip_mm_" + str(m) + ".tif")
+                                     for m in range(1, n_months + 1)]
+        
+        # Verify that the returned lists match the input
+        self.assertEqual(precip_path_list, match_precip)
+
     def test_ambiguous_precip_data(self):
         """SWY test case where there are more than 12 precipitation files."""
         from natcap.invest.seasonal_water_yield import seasonal_water_yield