Add target_projection and target_pixelsize to ModelSpec and integrate into AWY, UMH

src/natcap/invest/annual_water_yield/annual_water_yield.py

@@ -162,31 +162,41 @@
     about=_model_description,
     input_field_order=[
         ["workspace_dir", "results_suffix"],
+        ["aoi_vector_path"],
         ["precipitation_path", "eto_path", "depth_to_root_rest_layer_path", "pawc_path"],
         ["lulc_path", "biophysical_table_path", "seasonality_constant"],
         ["watersheds_path", "sub_watersheds_path"],
-        ["demand_table_path", "valuation_table_path"]
+        ["demand_table_path", "valuation_table_path"],
+        ["target_projection", "target_pixelsize"]
     ],
     validate_spatial_overlap=True,
-    different_projections_ok=False,
+    different_projections_ok=True,
     aliases=("hwy", "awy"),
     module_name=__name__,
     inputs=[
         spec.WORKSPACE,
         spec.SUFFIX,
         spec.N_WORKERS,
+        spec.AOI.model_copy(update=dict(
+            id="aoi_vector_path",
+            about=gettext("Map of the region over which to run the model."),
+            projected=True,
+            required=False
+        )),
         spec.SingleBandRasterInput(
             id="lulc_path",
             name=gettext("land use/land cover"),
             about=gettext(
                 "Map of land use/land cover codes. Each land use/land cover"
                 " type must be assigned a unique integer code. All values in"
                 " this raster must have corresponding entries in the"
-                " Biophysical Table."
+                " Biophysical Table. This input defines the default target"
+                " projection and alignment for all other spatial data."
             ),
             data_type=int,
             units=None,
-            projected=True
+            is_default_projection=True,
+            is_default_pixelsize=True
         ),
         spec.SingleBandRasterInput(
             id="depth_to_root_rest_layer_path",
@@ -198,15 +208,13 @@
             ),
             data_type=float,
             units=u.millimeter,
-            projected=True
         ),
         spec.SingleBandRasterInput(
             id="precipitation_path",
             name=gettext("precipitation"),
             about=gettext("Map of average annual precipitation."),
             data_type=float,
             units=u.millimeter / u.year,
-            projected=True
         ),
         spec.SingleBandRasterInput(
             id="pawc_path",
@@ -218,11 +226,9 @@
             ),
             data_type=float,
             units=None,
-            projected=True
         ),
         spec.SingleBandRasterInput(
             id="eto_path",
-            projected=True,
             name=gettext("reference evapotranspiration"),
             about=gettext("Map of reference evapotranspiration values."),
             data_type=float,
@@ -243,7 +249,6 @@
                     about=gettext("Unique identifier for each watershed.")
                 )
             ],
-            projected=True
         ),
         spec.VectorInput(
             id="sub_watersheds_path",
@@ -260,7 +265,6 @@
                     about=gettext("Unique identifier for each subwatershed.")
                 )
             ],
-            projected=True
         ),
         spec.CSVInput(
             id="biophysical_table_path",
@@ -415,15 +419,17 @@
                 )
             ],
             index_col="ws_id"
-        )
+        ),
+        spec.TARGET_PROJECTION,
+        spec.TARGET_PIXELSIZE
     ],
     outputs=[
         spec.VectorOutput(
             id="watershed_results_wyield",
             path="output/watershed_results_wyield.shp",
             about=gettext(
                 "Shapefile containing biophysical output values per"
-                " watershed."
+                " watershed. Watershed reprojected to match target_projection."
             ),
             geometry_types={"POLYGON"},
             fields=WATERSHED_OUTPUT_FIELDS
@@ -443,7 +449,7 @@
             path="output/subwatershed_results_wyield.shp",
             about=gettext(
                 "Shapefile containing biophysical output values per"
-                " subwatershed."
+                " subwatershed. Subwatershed reprojected to match target_projection."
             ),
             geometry_types={"POLYGON"},
             fields=SUBWATERSHED_OUTPUT_FIELDS
@@ -630,6 +636,10 @@ def execute(args):
         args['workspace_dir'] (string): a path to the directory that will write
             output and other temporary files during calculation. (required)
 
+        args['aoi_path'] (str): (optional) Path to a polygon vector of the area
+            over which the model should be run. Must be projected in a
+            coordinate system.
+
         args['lulc_path'] (string): a path to a land use/land cover raster
             whose LULC indexes correspond to indexes in the biophysical table
             input. Used for determining soil retention and other biophysical
@@ -691,6 +701,16 @@ def execute(args):
             valuation: 'ws_id', 'time_span', 'discount', 'efficiency',
             'fraction', 'cost', 'height', 'kw_price'
 
+        args['target_projection'] (string): (optional) if a non-empty string,
+            id of spatial input that defines the target projection. If that
+            spatial input is a raster, rather than a vector, it will also
+            represent the target alignment for other spatial inputs.
+
+        args['target_pixelsize'] (string): (optional) if a non-empty string,
+            id of spatial input that defines the target pixel size for other
+            spatial inputs. If ``target_projection`` is a vector, this spatial
+            input will also represent the target alignment.
+
         args['n_workers'] (int): (optional) The number of worker processes to
             use for processing this model.  If omitted, computation will take
             place in the current process.
@@ -729,16 +749,34 @@ def execute(args):
                 'valuation table to see if they are missing: '
                 f'"{", ".join(str(x) for x in sorted(missing_ws_ids))}"')
 
+    # reproject watersheds_path to target_projection
+    target_projection_wkt = utils.get_raster_or_vector_projection(
+        args[args['target_projection']])
+    # Reproject watersheds_path even if it has the `target_projection` to
+    # create a copy so we don't modify the original when doing zonal stats 
+    reproject_watersheds_task = graph.add_task(
+        pygeoprocessing.reproject_vector,
+        args=(args['watersheds_path'], target_projection_wkt,
+              file_registry['watershed_results_wyield']),
+        target_path_list=[file_registry['watershed_results_wyield']],
+        task_name='reproject_watersheds')
     watershed_paths_list = [(
-        args['watersheds_path'], 'ws_id',
-        file_registry['watershed_results_wyield'],
+        'ws_id', file_registry['watershed_results_wyield'],
         file_registry['watershed_results_wyield_csv'])]
+    dependent_tasks_for_watersheds_list = [reproject_watersheds_task]
 
     if args['sub_watersheds_path']:
+        reproject_sub_watersheds_task = graph.add_task(
+            pygeoprocessing.reproject_vector,
+            args=(args['sub_watersheds_path'], target_projection_wkt,
+                  file_registry['subwatershed_results_wyield']),
+            target_path_list=[file_registry['subwatershed_results_wyield']],
+            task_name='reproject_sub_watersheds')
         watershed_paths_list.append((
-            args['sub_watersheds_path'], 'subws_id',
-            file_registry['subwatershed_results_wyield'],
+            'subws_id', file_registry['subwatershed_results_wyield'],
             file_registry['subwatershed_results_wyield_csv']))
+        dependent_tasks_for_watersheds_list.append(
+            reproject_sub_watersheds_task)
 
     base_raster_path_list = [
         args['eto_path'],
@@ -754,17 +792,31 @@ def execute(args):
         file_registry['pawc'],
         file_registry['clipped_lulc']]
 
-    target_pixel_size = pygeoprocessing.get_raster_info(
-        args['lulc_path'])['pixel_size']
+    if pygeoprocessing.get_gis_type(
+            args[args['target_projection']]) == pygeoprocessing.RASTER_TYPE:
+        raster_align_index = base_raster_path_list.index(
+            args[args['target_projection']])
+    else:
+        # fallback to aligning everything to the arg with default pixel size
+        raster_align_index = base_raster_path_list.index(
+            args[args['target_pixelsize']])
+    target_pixel_size = utils.get_raster_pixel_size_in_target_proj_units(
+        args[args['target_pixelsize']], target_projection_wkt)
+    base_vector_path_list = [file_registry['watershed_results_wyield']]
+    if args['aoi_vector_path']:
+        base_vector_path_list.append(args['aoi_vector_path'])
+
     align_raster_stack_task = graph.add_task(
         pygeoprocessing.align_and_resize_raster_stack,
         args=(base_raster_path_list, aligned_raster_path_list,
               ['near'] * len(base_raster_path_list),
               target_pixel_size, 'intersection'),
-        kwargs={'raster_align_index': 4,
-                'base_vector_path_list': [args['watersheds_path']]},
+        kwargs={'raster_align_index': raster_align_index,
+                'base_vector_path_list': base_vector_path_list,
+                'target_projection_wkt': target_projection_wkt},
         target_path_list=aligned_raster_path_list,
-        task_name='align_raster_stack')
+        task_name='align_raster_stack',
+        dependent_task_list=[reproject_watersheds_task])
     # Joining now since this task will always be the root node
     # and it's useful to have the raster info available.
     align_raster_stack_task.join()
@@ -853,8 +905,6 @@ def execute(args):
         dependent_task_list=[align_raster_stack_task],
         task_name='create_veg_raster')
 
-    dependent_tasks_for_watersheds_list = []
-
     LOGGER.info('Calculate PET from Ref Evap times Kc')
     calculate_pet_task = graph.add_task(
         func=pygeoprocessing.raster_map,
@@ -939,19 +989,11 @@ def execute(args):
 
     # Aggregate results to watershed polygons, and do the optional
     # scarcity and valuation calculations.
-    for base_ws_path, ws_id_name, target_ws_path, target_csv_path in watershed_paths_list:
-        # make a copy so we don't modify the original
-        # do zonal stats with the copy so that FIDS are correct
-        copy_watersheds_vector_task = graph.add_task(
-            func=copy_vector,
-            args=[base_ws_path, target_ws_path],
-            target_path_list=[target_ws_path],
-            task_name='create copy of watersheds vector')
-
+    for ws_id_name, target_ws_path, target_csv_path in watershed_paths_list:
         zonal_stats_task_list = []
         zonal_stats_pickle_list = []
 
-        # Do zonal stats with the input shapefiles provided by the user
+        # Do zonal stats with the input shapefiles provided by the users
         # and store results dictionaries in pickles
         for key_name, rast_path in raster_names_paths_list:
             target_stats_pickle = file_registry[f'{ws_id_name}_{key_name.lower()}']
@@ -960,9 +1002,7 @@ def execute(args):
                 func=zonal_stats_tofile,
                 args=(target_ws_path, rast_path, target_stats_pickle),
                 target_path_list=[target_stats_pickle],
-                dependent_task_list=[
-                    *dependent_tasks_for_watersheds_list,
-                    copy_watersheds_vector_task],
+                dependent_task_list=dependent_tasks_for_watersheds_list,
                 task_name=f'{ws_id_name}_{key_name}_zonalstats'))
 
         # Add the zonal stats data to the output vector's attribute table
@@ -972,8 +1012,7 @@ def execute(args):
             args=(target_ws_path, ws_id_name, zonal_stats_pickle_list,
                   valuation_df),
             target_path_list=[target_ws_path],
-            dependent_task_list=[
-                *zonal_stats_task_list, copy_watersheds_vector_task],
+            dependent_task_list=zonal_stats_task_list,
             task_name=f'create_{ws_id_name}_vector_output')
 
         # Export a CSV with all the fields present in the output vector
@@ -992,22 +1031,6 @@ def execute(args):
 def wyield_op(fractp, precip): return (1 - fractp) * precip
 
 
-def copy_vector(base_vector_path, target_vector_path):
-    """Wrapper around CreateCopy that handles opening & closing the dataset.
-
-    Args:
-        base_vector_path: path to the vector to copy
-        target_vector_path: path to copy the vector to
-
-    Returns:
-        None
-    """
-    esri_shapefile_driver = gdal.GetDriverByName('ESRI Shapefile')
-    base_dataset = gdal.OpenEx(base_vector_path, gdal.OF_VECTOR)
-    esri_shapefile_driver.CreateCopy(target_vector_path, base_dataset)
-    base_dataset = None
-
-
 def write_output_vector_attributes(target_vector_path, ws_id_name,
                                    stats_path_list, valuation_df):
     """Add data attributes to the vector outputs of this model.