Add 'plot_heatmaps.py'

plot_heatmaps.py

@@ -0,0 +1,142 @@
+"""Plot heatmaps using ArcGIS.
+
+This script must be run from the version of python installed with ArcGIS, e.g.
+C:/Python27/ArcGIS10.4/python plot_heatmaps.py
+
+Note: arcpy does not play nicely with relative paths. This is because arcpy
+changes the python working directory when an mxd is loaded, and during many
+mapping operations. Relative paths may be used as inputs to this script, but
+they will be converted to absolute paths internally.
+"""
+
+import os
+import re
+import sys
+import argparse
+from glob import glob
+
+import arcpy
+
+# Checkout Spatial Analyst toolbox to prevent licence errors
+arcpy.CheckOutExtension('Spatial')
+
+# Overwrite duplicate files
+arcpy.env.overwriteOutput = True
+
+
+def parse_yaml(yaml_name):
+    """Parse yaml file manually (when 'yaml' module is unavailable)
+
+    Args:
+        yaml_name:  path to input yaml file
+
+    Returns:
+        dict of objects from yaml file
+
+    """
+    with open(yaml_name, 'r') as f:
+        yaml = f.read()
+
+    # Fit wrapped strings on one line
+    yaml = re.sub('\n\s', '', yaml)
+
+    # Parse yaml file
+    params = {}
+    for line in yaml.split('\n'):
+        if line:
+            key, val = line.split(':', 1)
+            params[key] = val.strip()
+
+    return params
+
+
+def process(yaml_name):
+    params = parse_yaml(yaml_name)
+    beach = params['BEACH']
+    base_name = os.path.splitext(os.path.basename(params['INPUT LAS']))[0]
+
+    # Make all input and output paths absolute
+    mxd_name = os.path.abspath(params['MAP DOCUMENT'])
+    lyr_name = os.path.abspath(params['SYMBOLOGY LAYER'])
+    input_tif_dir = os.path.abspath(params['TIF HEATMAP FOLDER'])
+    output_jpg_dir = os.path.abspath(params['JPG HEATMAP FOLDER'])
+
+    # Check if previous survey date was provided
+    previous_date = params['PREVIOUS SURVEY DATE']
+    try:
+        int(previous_date)
+    except ValueError:
+        raise ValueError('No previous survey date provided')
+
+    # Set paths for raster files
+    heatmap_raster = os.path.join(input_tif_dir, base_name + '.tif')
+
+    print('processing {}'.format(beach))
+    mxd = arcpy.mapping.MapDocument(mxd_name)
+    df = arcpy.mapping.ListDataFrames(mxd, 'Layers')[0]
+
+    # Load erosion heatmap symbology
+    symbology_layer = arcpy.mapping.Layer(lyr_name)
+
+    # Add heatmap for current beach
+    heatmap_layer = arcpy.mapping.Layer(heatmap_raster)
+    arcpy.mapping.AddLayer(df, heatmap_layer, 'AUTO_ARRANGE')
+
+    # Apply symbology to heatmap
+    update_layer = arcpy.mapping.ListLayers(mxd, heatmap_layer, df)[0]
+    arcpy.mapping.UpdateLayer(df, update_layer, symbology_layer, 'TRUE')
+    arcpy.RefreshTOC()
+
+    # Activate data driven pages
+    ddp = mxd.dataDrivenPages
+
+    for page_num in range(1, ddp.pageCount + 1):
+        ddp.currentPageID = page_num
+
+        # Export beaches that are covered by current survey
+        if ddp.pageRow.Parent == beach.lower():
+            page_name = ddp.pageRow.Beach
+            jpg_name = os.path.join(output_jpg_dir, page_name + '.jpg')
+
+            # Export to jpg
+            arcpy.mapping.ExportToJPEG(mxd, jpg_name)
+
+
+def main():
+    example_text = """examples:
+
+    # Process single survey at specific beach
+    C:/Python27/ArcGIS10.4/python plot_heatmaps.py survey-1-avoca.yaml
+
+    # Process single survey at multiple beaches
+    C:/Python27/ArcGIS10.4/python plot_heatmaps.py survey-1-avoca.yaml survey-1-pearl.yaml
+
+    # Process all surveys at specific beach
+    C:/Python27/ArcGIS10.4/python plot_heatmaps.py *avoca.yaml
+
+    # Process all beaches for specific survey date
+    C:/Python27/ArcGIS10.4/python plot_heatmaps.py survey-1*.yaml
+    """
+
+    # Set up command line arguments
+    parser = argparse.ArgumentParser(
+        epilog=example_text,
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('input', help='path to yaml file(s)', nargs='*')
+
+    # Print usage if no arguments are provided
+    if len(sys.argv) == 1:
+        parser.print_help(sys.stderr)
+        sys.exit(1)
+
+    # Parse arguments
+    args = parser.parse_args()
+    yaml_files = []
+    [yaml_files.extend(glob(f)) for f in args.input]
+
+    for yaml_file in yaml_files:
+        process(yaml_file)
+
+
+if __name__ == '__main__':
+    main()