Add 'generate_heatmaps.py'

generate_heatmaps.py

@@ -0,0 +1,133 @@
+"""Generate heatmaps using ArcGIS.
+
+This script must be run from the version of python installed with ArcGIS, e.g.
+C:/Python27/ArcGIS10.4/python generate_heatmaps.py
+"""
+
+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(':')
+            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]
+
+    # Set all paths relative to current script location
+    pwd = os.getcwd()
+    input_tif_dir = os.path.join(pwd, params['TIF DEM FOLDER'])
+    mxd_name = os.path.join(pwd, params['MAP DOCUMENT'])
+    symb_lyr_name = os.path.join(pwd, params['SYMBOLOGY LAYER'])
+    clip_poly = os.path.join(pwd, params['HEATMAP CROP POLY'])
+    output_tif_dir = os.path.join(pwd, params['TIF HEATMAP FOLDER'])
+    previous_date = params['PREVIOUS SURVEY DATE']
+
+    # Check if previous survey date was provided
+    try:
+        int(previous_date)
+    except ValueError:
+        raise ValueError('No previous survey date provided')
+
+    # Set paths for raster files
+    current_raster = os.path.join(input_tif_dir, base_name + '.tif')
+    previous_base_name = re.sub('\d+', previous_date, base_name)
+    previous_raster = os.path.join(input_tif_dir, previous_base_name + '.tif')
+    heatmap_raster = os.path.join(output_tif_dir, base_name + '.tif')
+
+    print('processing {}'.format(beach))
+
+    # Save heatmap in memory
+    heatmap_memory = 'in_memory/' + base_name
+
+    # Subtract the previous raster from the current raster
+    arcpy.Minus_3d(
+        in_raster_or_constant1=current_raster,
+        in_raster_or_constant2=previous_raster,
+        out_raster=heatmap_memory)
+
+    # Clip the heatmap
+    arcpy.Clip_management(
+        in_raster=heatmap_memory,
+        rectangle='350000 6290000 360000 6300000',
+        out_raster=heatmap_raster,
+        in_template_dataset=clip_poly,
+        nodata_value='-9.999000e+003',
+        clipping_geometry='ClippingGeometry',
+        maintain_clipping_extent='NO_MAINTAIN_EXTENT')
+
+
+def main():
+    example_text = """examples:
+
+    # Process single survey at specific beach
+    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1-avoca.yaml
+
+    # Process single survey at multiple beaches
+    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1-avoca.yaml survey-1-pearl.yaml
+
+    # Process all surveys at specific beach
+    C:/Python27/ArcGIS10.4/python generate_heatmaps.py *avoca.yaml
+
+    # Process all beaches for specific survey date
+    C:/Python27/ArcGIS10.4/python generate_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()