Merge pull request natcap#1626 from emilyanndavis/bugfix/1600-ufrm-pr…

…eserve-source-aoi-fields

UFRM: preserve source AOI fields in output vector

HISTORY.rst

@@ -55,6 +55,9 @@ Unreleased Changes
 * Habitat Quality
     * Access raster is now generated from the reprojected access vector.
       (https://github.com/natcap/invest/issues/1615)
+* Urban Flood Risk
+    * Fields present on the input AOI vector are now retained in the output.
+      (https://github.com/natcap/invest/issues/1600)
 
 3.14.2 (2024-05-29)
 -------------------

src/natcap/invest/urban_flood_risk_mitigation.py

@@ -173,14 +173,14 @@
                     "about": "Map of runoff volume.",
                     "bands": {1: {"type": "number", "units": u.meter**3}}
                 },
-                "reprojected_aoi.gpkg": {
+                "reprojected_aoi.shp": {
                     "about": (
                         "Copy of AOI vector reprojected to the same spatial "
                         "reference as the LULC."),
                     "geometries": spec_utils.POLYGONS,
                     "fields": {}
                 },
-                "structures_reprojected.gpkg": {
+                "structures_reprojected.shp": {
                     "about": (
                         "Copy of built infrastructure vector reprojected to "
                         "the same spatial reference as the LULC."),
@@ -408,14 +408,14 @@ def execute(args):
         task_name='calculate service built raster')
 
     reprojected_aoi_path = os.path.join(
-        intermediate_dir, 'reprojected_aoi.gpkg')
+        intermediate_dir, 'reprojected_aoi')
     reprojected_aoi_task = task_graph.add_task(
         func=pygeoprocessing.reproject_vector,
         args=(
             args['aoi_watersheds_path'],
             target_sr_wkt,
             reprojected_aoi_path),
-        kwargs={'driver_name': 'GPKG'},
+        kwargs={'driver_name': 'ESRI Shapefile'},
         target_path_list=[reprojected_aoi_path],
         task_name='reproject aoi/watersheds')
 
@@ -435,7 +435,7 @@ def execute(args):
             (runoff_retention_raster_path, 1),
             reprojected_aoi_path),
         store_result=True,
-        dependent_task_list=[runoff_retention_task],
+        dependent_task_list=[runoff_retention_task, reprojected_aoi_task],
         task_name='zonal_statistics over runoff_retention raster')
 
     runoff_retention_volume_stats_task = task_graph.add_task(
@@ -444,7 +444,7 @@ def execute(args):
             (runoff_retention_vol_raster_path, 1),
             reprojected_aoi_path),
         store_result=True,
-        dependent_task_list=[runoff_retention_vol_task],
+        dependent_task_list=[runoff_retention_vol_task, reprojected_aoi_task],
         task_name='zonal_statistics over runoff_retention_volume raster')
 
     damage_per_aoi_stats = None
@@ -457,13 +457,13 @@ def execute(args):
             args['built_infrastructure_vector_path'] not in ('', None)):
         # Reproject the built infrastructure vector to the target SRS.
         reprojected_structures_path = os.path.join(
-            intermediate_dir, 'structures_reprojected.gpkg')
+            intermediate_dir, 'structures_reprojected')
         reproject_built_infrastructure_task = task_graph.add_task(
             func=pygeoprocessing.reproject_vector,
             args=(args['built_infrastructure_vector_path'],
                   target_sr_wkt,
                   reprojected_structures_path),
-            kwargs={'driver_name': 'GPKG'},
+            kwargs={'driver_name': 'ESRI Shapefile'},
             target_path_list=[reprojected_structures_path],
             task_name='reproject built infrastructure to target SRS')
 
@@ -512,7 +512,7 @@ def _write_summary_vector(
         runoff_ret_vol_stats, flood_volume_stats, damage_per_aoi_stats=None):
     """Write a vector with summary statistics.
 
-    This vector will always contain two fields::
+    This vector will always contain three fields::
 
         * ``'flood_vol'``: The volume of flood (runoff), in m3, per watershed.
         * ``'rnf_rt_idx'``: Average of runoff retention values per watershed
@@ -553,21 +553,11 @@ def _write_summary_vector(
         ``None``
     """
     source_aoi_vector = gdal.OpenEx(source_aoi_vector_path, gdal.OF_VECTOR)
-    source_aoi_layer = source_aoi_vector.GetLayer()
-    source_geom_type = source_aoi_layer.GetGeomType()
-    source_srs_wkt = pygeoprocessing.get_vector_info(
-        source_aoi_vector_path)['projection_wkt']
-    source_srs = osr.SpatialReference()
-    source_srs.ImportFromWkt(source_srs_wkt)
-
     esri_driver = gdal.GetDriverByName('ESRI Shapefile')
-    target_watershed_vector = esri_driver.Create(
-        target_vector_path, 0, 0, 0, gdal.GDT_Unknown)
-    layer_name = os.path.splitext(os.path.basename(
-        target_vector_path))[0]
-    LOGGER.debug(f"creating layer {layer_name}")
-    target_watershed_layer = target_watershed_vector.CreateLayer(
-        layer_name, source_srs, source_geom_type)
+    esri_driver.CreateCopy(target_vector_path, source_aoi_vector)
+    target_watershed_vector = gdal.OpenEx(target_vector_path,
+                                          gdal.OF_VECTOR | gdal.GA_Update)
+    target_watershed_layer = target_watershed_vector.GetLayer()
 
     target_fields = ['rnf_rt_idx', 'rnf_rt_m3', 'flood_vol']
     if damage_per_aoi_stats is not None:
@@ -579,13 +569,9 @@ def _write_summary_vector(
         field_def.SetPrecision(11)
         target_watershed_layer.CreateField(field_def)
 
-    target_layer_defn = target_watershed_layer.GetLayerDefn()
-    for base_feature in source_aoi_layer:
-        feature_id = base_feature.GetFID()
-        target_feature = ogr.Feature(target_layer_defn)
-        base_geom_ref = base_feature.GetGeometryRef()
-        target_feature.SetGeometry(base_geom_ref.Clone())
-        base_geom_ref = None
+    target_watershed_layer.ResetReading()
+    for target_feature in target_watershed_layer:
+        feature_id = target_feature.GetFID()
 
         pixel_count = runoff_ret_stats[feature_id]['count']
         if pixel_count > 0:
@@ -611,7 +597,8 @@ def _write_summary_vector(
         target_feature.SetField(
             'flood_vol', float(flood_volume_stats[feature_id]['sum']))
 
-        target_watershed_layer.CreateFeature(target_feature)
+        target_watershed_layer.SetFeature(target_feature)
+
     target_watershed_layer.SyncToDisk()
     target_watershed_layer = None
     target_watershed_vector = None

tests/test_ufrm.py

@@ -59,20 +59,27 @@ def test_ufrm_regression(self):
         """UFRM: regression test."""
         from natcap.invest import urban_flood_risk_mitigation
         args = self._make_args()
+        input_vector = gdal.OpenEx(args['aoi_watersheds_path'],
+                                   gdal.OF_VECTOR)
+        input_layer = input_vector.GetLayer()
+        input_fields = [field.GetName() for field in input_layer.schema]
+
         urban_flood_risk_mitigation.execute(args)
 
         result_vector = gdal.OpenEx(os.path.join(
             args['workspace_dir'], 'flood_risk_service_Test1.shp'),
             gdal.OF_VECTOR)
         result_layer = result_vector.GetLayer()
 
-        # Check that all four expected fields are there.
+        # Check that all expected fields are there.
+        output_fields = ['aff_bld', 'serv_blt', 'rnf_rt_idx',
+                         'rnf_rt_m3', 'flood_vol']
+        output_fields += input_fields
         self.assertEqual(
-            set(('aff_bld', 'serv_blt', 'rnf_rt_idx', 'rnf_rt_m3',
-                 'flood_vol')),
+            set(output_fields),
             set(field.GetName() for field in result_layer.schema))
 
-        result_feature = result_layer.GetFeature(0)
+        result_feature = result_layer.GetNextFeature()
         for fieldname, expected_value in (
                 ('aff_bld', 187010830.32202843),
                 ('serv_blt', 13253546667257.65),
@@ -85,6 +92,11 @@ def test_ufrm_regression(self):
             self.assertAlmostEqual(
                 result_val, expected_value, places=-places_to_round)
 
+        input_feature = input_layer.GetNextFeature()
+        for fieldname in input_fields:
+            self.assertEqual(result_feature.GetField(fieldname),
+                             input_feature.GetField(fieldname))
+
         result_feature = None
         result_layer = None
         result_vector = None
@@ -94,6 +106,11 @@ def test_ufrm_regression_no_infrastructure(self):
         from natcap.invest import urban_flood_risk_mitigation
         args = self._make_args()
         del args['built_infrastructure_vector_path']
+        input_vector = gdal.OpenEx(args['aoi_watersheds_path'],
+                                   gdal.OF_VECTOR)
+        input_layer = input_vector.GetLayer()
+        input_fields = [field.GetName() for field in input_layer.schema]
+
         urban_flood_risk_mitigation.execute(args)
 
         result_raster = gdal.OpenEx(os.path.join(
@@ -115,9 +132,11 @@ def test_ufrm_regression_no_infrastructure(self):
         result_layer = result_vector.GetLayer()
         result_feature = result_layer.GetFeature(0)
 
-        # Check that only the two expected fields are there.
+        # Check that only the expected fields are there.
+        output_fields = ['rnf_rt_idx', 'rnf_rt_m3', 'flood_vol']
+        output_fields += input_fields
         self.assertEqual(
-            set(('rnf_rt_idx', 'rnf_rt_m3', 'flood_vol')),
+            set(output_fields),
             set(field.GetName() for field in result_layer.schema))
 
         for fieldname, expected_value in (
@@ -218,7 +237,6 @@ def test_ufrm_explicit_zeros_in_table(self):
         except ValueError:
             self.fail('unexpected ValueError when testing curve number row with all zeros')
 
-
     def test_ufrm_string_damage_to_infrastructure(self):
         """UFRM: handle str(int) structure indices.