Merge pull request #21 from natcap/feature/summary-results

Feature/summary results

src/rangeland_production/forage.py

@@ -6,6 +6,7 @@
 from builtins import range
 import re
 import math
+import pickle
 
 import numpy
 import pandas
@@ -627,6 +628,9 @@ def execute(args):
             plant functional type index and each state variable listed in the
             following table:
             https://docs.google.com/spreadsheets/d/1TGCDOJS4nNsJpzTWdiWed390NmbhQFB2uUoMs9oTTYo/edit?usp=sharing
+        args['save_sv_rasters'] (boolean): optional input, default false.
+            Should rasters containing all state variables be saved for each
+            model time step?
 
     Returns:
         None.
@@ -636,6 +640,11 @@ def execute(args):
     starting_month = int(args['starting_month'])
     starting_year = int(args['starting_year'])
     n_months = int(args['n_months'])
+    try:
+        delete_sv_folders = not args['save_sv_rasters']
+    except KeyError:
+        delete_sv_folders = True
+
     # this set will build up the integer months that are used so we can index
     # the mwith temperature later
     temperature_month_set = set()
@@ -1170,9 +1179,32 @@ def execute(args):
             aligned_inputs, provisional_sv_reg, sv_reg, month_reg, pft_id_set,
             current_year, current_month, output_dir)
 
+    # summary results
+    summary_output_dir = os.path.join(output_dir, 'summary_results')
+    os.makedirs(summary_output_dir)
+    summary_shp_path = os.path.join(
+        summary_output_dir, 'grazing_areas_results_rpm.shp')
+    create_vector_copy(
+        args['animal_grazing_areas_path'], summary_shp_path)
+
+    field_pickle_map, field_header_order_list = aggregate_and_pickle_results(
+        output_dir, summary_shp_path)
+    _add_fields_to_shapefile(
+        field_pickle_map, field_header_order_list, summary_shp_path)
+
+    mean_animal_density_path = os.path.join(
+        summary_output_dir, 'mean_animal_density.tif')
+    calc_mean_animal_density(output_dir, mean_animal_density_path)
+
     # clean up
     shutil.rmtree(persist_param_dir)
     shutil.rmtree(PROCESSING_DIR)
+    if delete_sv_folders:
+        for month_index in range(-1, n_months):
+            shutil.rmtree(
+                os.path.join(
+                    args['workspace_dir'],
+                    'state_variables_m%d' % month_index))
 
 
 def raster_multiplication(
@@ -13483,16 +13515,11 @@ def _estimate_animal_density(
         sv_reg (dict): map of key, path pairs giving paths to state variables,
             including carbon in aboveground biomass for each plant functional
             type
-        obs_biomass_path (string): path to output location where observed
-            biomass for this timestep should be saved as geotiff
         month_reg (dict): map of key, path pairs giving paths to intermediate
             calculated values that are shared between submodels, including
             estimated animal density
 
     Side effects:
-        creates a geotiff raster of observed biomass calculated from remotely
-            sensed vegetation index at the location indicated by
-            `obs_biomass_path`
         creates or modifies the raster indicated by month_reg['animal_density']
 
     """
@@ -13789,6 +13816,183 @@ def copy_intermediate_sv(pft_id_set, sv_reg, intermediate_sv_dir):
     return intermediate_sv_reg
 
 
+def create_vector_copy(base_vector_path, target_vector_path):
+    """Create a copy of base vector."""
+    if os.path.isfile(target_vector_path):
+        os.remove(target_vector_path)
+    base_vector = gdal.OpenEx(base_vector_path, gdal.OF_VECTOR)
+    driver = gdal.GetDriverByName('ESRI Shapefile')
+    target_vector = driver.CreateCopy(
+        target_vector_path, base_vector)
+    target_vector = None  # seemingly uncessary but gdal seems to like it.
+
+
+def aggregate_and_pickle_results(output_dir, grazing_areas_path):
+    """Calculate aggregated biomass and diet sufficiency per grazing area.
+
+    Parameters:
+        output_dir (string): path to directory containing model outputs: diet
+            sufficiency, potential biomass, standing biomass rasters per time
+            step
+        grazing_areas_path (string): path to shapefile giving the location of
+            grazing animals. Zonal mean values are calculated per feature in
+            this dataset.
+
+    Side effects:
+        creates pickle files for mean diet sufficiency, mean potential biomass,
+            and mean standing biomass
+
+    Returns:
+        a dictionary with the following keys, where values contain paths to
+            pickled zonal statistics for the given result:
+            'potbiom_m' (mean total potential biomass)
+            'stdbiom_m' (mean total standing biomass)
+            'dietsuff_m' (mean diet sufficiency)
+
+    """
+    def dump_zonal_stats(raster_prefix, pickle_path):
+        """Calculate zonal mean and save as pickle."""
+        raster_path_list = [
+            os.path.join(output_dir, f) for f in os.listdir(output_dir) if
+            f.startswith(raster_prefix) and f.endswith('.tif')]
+        df_list = []
+        for raster_path in raster_path_list:
+            zonal_stat_dict = pygeoprocessing.zonal_statistics(
+                (raster_path, 1), grazing_areas_path)
+            zonal_df = pandas.DataFrame(
+                {
+                    'fid': [
+                        key for key, value in sorted(
+                            zonal_stat_dict.items())],
+                    'sum': [
+                        value['sum'] for key, value in
+                        sorted(zonal_stat_dict.items())],
+                    'count': [
+                        value['count'] for key, value in
+                        sorted(zonal_stat_dict.items())]
+                })
+            df_list.append(zonal_df)
+        summary_df = pandas.concat(df_list)
+        df_means = summary_df.groupby('fid').mean()
+        df_means['mean'] = df_means['sum'] / df_means['count']
+        df_means = df_means[['mean']].to_dict(orient='index')
+        with open(pickle_path, 'wb') as target_pickle_file:
+            pickle.dump(df_means, target_pickle_file)
+
+
+    potential_biomass_pickle_path = os.path.join(
+        PROCESSING_DIR, 'potential_biomass.pickle')
+    dump_zonal_stats('potential_biomass', potential_biomass_pickle_path)
+    diet_suff_pickle_path = os.path.join(
+        PROCESSING_DIR, 'diet_sufficiency.pickle')
+    dump_zonal_stats('diet_sufficiency', diet_suff_pickle_path)
+    standing_biomass_pickle_path = os.path.join(
+        PROCESSING_DIR, 'standing_biomass.pickle')
+    dump_zonal_stats('standing_biomass', standing_biomass_pickle_path)
+
+    field_pickle_map = {
+        'potbiom_m': potential_biomass_pickle_path,
+        'stdbiom_m': standing_biomass_pickle_path,
+        'dietsuff_m': diet_suff_pickle_path,
+    }
+    field_header_order_list = ['potbiom_m', 'stdbiom_m', 'dietsuff_m']
+    return field_pickle_map, field_header_order_list
+
+
+def _add_fields_to_shapefile(
+        field_pickle_map, field_header_order, target_vector_path):
+    """Add fields and values to an OGR layer open for writing.
+
+    Parameters:
+        field_pickle_map (dict): maps field name to a pickle file that is a
+            result of pygeoprocessing.zonal_stats with FIDs that match
+            `target_vector_path`.
+        field_header_order (list of string): a list of field headers in the
+            order to appear in the output table.
+        target_vector_path (string): path to target vector file.
+
+    Returns:
+        None.
+
+    """
+    target_vector = gdal.OpenEx(
+        target_vector_path, gdal.OF_VECTOR | gdal.GA_Update)
+    target_layer = target_vector.GetLayer()
+    field_summaries = {}
+    for field_name in field_header_order:
+        field_def = ogr.FieldDefn(field_name, ogr.OFTReal)
+        field_def.SetWidth(24)
+        field_def.SetPrecision(11)
+        target_layer.CreateField(field_def)
+        with open(field_pickle_map[field_name], 'rb') as pickle_file:
+            field_summaries[field_name] = pickle.load(pickle_file)
+
+    for feature in target_layer:
+        fid = feature.GetFID()
+        for field_name in field_header_order:
+            feature.SetField(
+                field_name, float(field_summaries[field_name][fid]['mean']))
+        # Save back to datasource
+        target_layer.SetFeature(feature)
+    target_layer = None
+    target_vector = None
+
+
+def calc_mean_animal_density(output_dir, target_path):
+    """Calculate mean animal density inside each pixel during the model run.
+
+    Calculate mean animal density per pixel across months of the simulation.
+    Calculate the mean for all pixels that have >0 valid value during the
+    months of the simulation, discarding months without valid values in the
+    calculation of the average.
+
+    Parameters:
+        output_dir (string): path to directory containing output rasters
+            produced during the model run, including animal density rasters
+        target_path (string): path to location where raster with mean animal
+            density should be saved
+
+    Side effects:
+        creates or modifies the raster at `target_path` to contain mean animal
+            density (animals/ha) across model timesteps
+
+    Returns:
+        None
+
+    """
+    def raster_mean_op(*raster_list):
+        """Calculate the mean value pixel-wise from rasters in raster_list."""
+        valid_mask = numpy.any(
+            ~numpy.isclose(numpy.array(raster_list), _TARGET_NODATA), axis=0)
+        # get number of valid observations per pixel
+        num_observations = numpy.count_nonzero(
+            ~numpy.isclose(numpy.array(raster_list), _TARGET_NODATA), axis=0)
+        for r in raster_list:
+            numpy.place(r, numpy.isclose(r, _TARGET_NODATA), [0])
+        sum_of_rasters = numpy.sum(raster_list, axis=0)
+
+        divide_mask = (
+            (sum_of_rasters > 0) &
+            (num_observations > 0) &
+            valid_mask)
+
+        mean_of_rasters = numpy.empty(
+            sum_of_rasters.shape, dtype=numpy.float32)
+        mean_of_rasters[:] = _TARGET_NODATA
+        mean_of_rasters[valid_mask] = 0
+        mean_of_rasters[divide_mask] = numpy.divide(
+            sum_of_rasters[divide_mask], num_observations[divide_mask])
+        return mean_of_rasters
+
+    raster_path_list = [
+        os.path.join(output_dir, f) for f in os.listdir(output_dir) if
+        f.startswith('animal_density_') and f.endswith('.tif')]
+
+    pygeoprocessing.raster_calculator(
+        [(path, 1) for path in raster_path_list], raster_mean_op,
+            target_path, gdal.GDT_Float32, _TARGET_NODATA)
+
+
 @validation.invest_validator
 def validate(args, limit_to=None):
     """Validate args to ensure they conform to ``execute``'s contract.

src/rangeland_production/ui/forage.py

@@ -311,6 +311,12 @@ def __init__(self):
             label=u'Initial Conditions Table: PFT State Variables',
             validator=self.validator)
         self.add_input(self.pft_initial_table)
+        self.save_sv_rasters = inputs.Checkbox(
+            args_key=u'save_sv_rasters',
+            helptext=(u"Should rasters for each state variable, for each "
+                "timestep, be saved?"),
+            label=u'Save State Variable Rasters')
+        self.add_input(self.save_sv_rasters)
 
     def assemble_args(self):
         args = {
@@ -352,6 +358,7 @@ def assemble_args(self):
                 self.initial_conditions_dir.value(),
             self.site_initial_table.args_key: self.site_initial_table.value(),
             self.pft_initial_table.args_key: self.pft_initial_table.value(),
+            self.save_sv_rasters.args_key: self.save_sv_rasters.value(),
         }
 
         return args