column can be specified

copro/scripts/postprocessing/avg_over_time.py

@@ -10,13 +10,14 @@
 @click.command()
 @click.option('-t0', '--start-year', type=int)
 @click.option('-t1', '--end-year', type=int)
+@click.option('-c', '--column', help='column name', default='chance_of_conflict', type=str)
 @click.option('--geojson/--no-geojson', help='save output to geojson or not', default=False)
 @click.option('--png/--no-png', help='save output to png or not', default=True)
 @click.argument('input_dir', type=click.Path())
 @click.argument('output_dir',type=click.Path())
 @click.argument('selected_polygons', type=click.Path())
 
-def main(input_dir=None, output_dir=None, selected_polygons=None, start_year=None, end_year=None, geojson=None, png=None):
+def main(input_dir=None, output_dir=None, selected_polygons=None, start_year=None, end_year=None, geojson=None, png=None, column=None):
     """Post-processing script to calculate average model output over a user-specifeid period or all output geojson-files.
     """    
 
@@ -80,7 +81,7 @@ def main(input_dir=None, output_dir=None, selected_polygons=None, start_year=Non
     # add number of predicted conflicts over all selected years
     y_out = pd.merge(y_out, y_df.nr_predicted_conflicts.groupby(y_df.ID).sum().to_frame(), on='ID')
     # determine chance of conflict over all selected years
-    y_out['chance_of_conflict'] = y_out.nr_predicted_conflicts / y_out.nr_predictions
+    y_out[column] = y_out.nr_predicted_conflicts / y_out.nr_predictions
     # add geometry informatoin for each polygon
     y_out = pd.merge(y_out, global_df, on='ID', how='left')
 
@@ -99,11 +100,11 @@ def main(input_dir=None, output_dir=None, selected_polygons=None, start_year=Non
     # if specified, save as png-file to output-dir
     if png:
         fig, ax = plt.subplots(1, 1)
-        gdf_out.plot(column='chance_of_conflict', ax=ax, 
+        gdf_out.plot(column=column, ax=ax, 
                         cmap='Reds', 
                         vmin=0, vmax=1,
                         legend=True,
-                        legend_kwds={'label': 'chance_of_conflict', 'orientation': "vertical"})
+                        legend_kwds={'label': column, 'orientation': "vertical"})
 
         click.echo('saving to {}'.format(os.path.abspath(os.path.join(output_dir, 'merged_output_{}.png'.format(suffix)))))
         plt.savefig(os.path.abspath(os.path.join(output_dir, 'merged_output_{}.png'.format(suffix))), dpi=300, bbox_inches='tight')