ssr_lpi.py

#!/usr/bin/env python
############################################################################
#
# MODULE:       ssr_lpi.py
# AUTHOR:       Collin Bode, UC Berkeley
#
# PURPOSE:
#   1. Calculate point density neighborhood sums using weighted box kernels.
#   2. Calculate LPI as the ratio of filtered to unfiltered point density.
#   3. Assign monthly LPI rasters based on solar geometry weighting.
#
# COPYRIGHT:    (c) 2011 Collin Bode
#               This program is free software under the GNU General Public
#               License (>=v2). Read the file COPYING that comes with GRASS
#               for details.
#
#############################################################################

import os
import sys
import shutil
import traceback
import datetime as dt
import numpy as np
from scipy.ndimage import convolve

from ssr_params import *
from ssr_utilities import (printout, ensure_dir, raster_path, raster_exists,
                            read_raster, save_raster, get_path, setup_workspace)


def apply_weight_kernel(data, weight_file):
    """Apply a neighborhood sum using a custom weight kernel from file.

    The weight file is a tab-separated grid matching GRASS r.neighbors format.
    NaN/nodata input cells are treated as zero before convolution.
    """
    kernel = np.loadtxt(weight_file)
    safe_data = np.where(np.isnan(data), 0.0, data)
    result = convolve(safe_data.astype(np.float64), kernel, mode='constant', cval=0.0)
    return result.astype(np.float32)


def main():
    setup_workspace()

    tlog = dt.datetime.strftime(dt.datetime.now(), "%Y-%m-%d_h%H")
    log_path = os.path.join(workspace, dir_logs, 'ssr_' + tlog + '_lpi.log')
    lf = open(log_path, 'a')

    ow = lpi_run - 1   # 0 = no overwrite, 1 = overwrite

    printout("STARTING LPI RUN", lf)
    printout("LPI prefix: " + lpipref, lf)
    printout("Boxsize: " + boxsize, lf)
    printout("Overwrite: " + str(ow), lf)
    printout('_________________________________', lf)

    if lpi_run > 0:
        L2start = dt.datetime.now()
        printout('START LPI at ' + dt.datetime.strftime(L2start, "%m-%d %H:%M:%S"), lf)

        lpi_dir = os.path.join(workspace, dir_lpi)
        ensure_dir(lpi_dir)
        scriptPath = get_path()

        weight_num = 4
        weight_name = 'lpi_box18x18_weight'
        # Month -> weight index mapping based on solar geometry
        month_weights = {1: 1, 2: 2, 3: 2, 4: 3, 5: 4, 6: 4,
                         7: 4, 8: 3, 9: 3, 10: 2, 11: 2, 12: 1}

        # Point density raster paths (from ssr_lidar output)
        pdensityfilt_path = raster_path(pdensitypref + LidarPoints[0], dir_lidar)
        pdensityunf_path = raster_path(pdensitypref + LidarPoints[1], dir_lidar)

        printout("Reading filtered density: " + pdensityfilt_path, lf)
        filt_data, filt_profile = read_raster(pdensityfilt_path)
        printout("Reading unfiltered density: " + pdensityunf_path, lf)
        unf_data, _ = read_raster(pdensityunf_path)

        # Treat NoData as zero (r.null equivalent) for filtered density
        filt_data = np.where(np.isnan(filt_data), 0.0, filt_data)

        printout("Running LPI for " + str(weight_num) + " weight kernels", lf)

        for weight_idx in range(1, weight_num + 1):
            w = str(weight_idx)
            weight_file = os.path.join(scriptPath, weight_name + w + '.txt')
            lpi_name = lpipref + 'w' + w
            lpi_out = raster_path(lpi_name, dir_lpi)

            if not ow and os.path.exists(lpi_out):
                printout("Skipping (exists): " + lpi_name, lf)
                continue

            printout("Neighborhood sum, weight " + w, lf)
            pneighfilt = apply_weight_kernel(filt_data, weight_file)
            pneighunf = apply_weight_kernel(unf_data, weight_file)

            # Avoid division by zero
            with np.errstate(invalid='ignore', divide='ignore'):
                if year == 'ym4':
                    lpi_data = 3.71 * (pneighfilt / pneighunf) ** 1.3455
                elif year == 'yr4':
                    lpi_data = 6.5 * (pneighfilt / pneighunf) ** 1.57 + 0.005
                else:
                    lpi_data = pneighfilt / pneighunf

            lpi_data = np.where(np.isinf(lpi_data) | np.isnan(pneighunf) | (pneighunf == 0),
                                np.nan, lpi_data)
            # Cap at 1.0
            lpi_data = np.minimum(lpi_data, 1.0)
            lpi_data = lpi_data.astype(np.float32)

            save_raster(lpi_data, filt_profile, lpi_out)
            printout("Created: " + lpi_name, lf)

        # Copy weight rasters to monthly LPI names
        printout("Assigning monthly LPI rasters", lf)
        for month, weight_idx in month_weights.items():
            wlpi_path = raster_path(lpipref + 'w' + str(weight_idx), dir_lpi)
            monthlpi_name = lpipref + 'm' + str(month).zfill(2)
            monthlpi_path = raster_path(monthlpi_name, dir_lpi)
            if ow or not os.path.exists(monthlpi_path):
                shutil.copy2(wlpi_path, monthlpi_path)
                printout("Copied w" + str(weight_idx) + " -> " + monthlpi_name, lf)

        L2end = dt.datetime.now()
        printout('END at ' + dt.datetime.strftime(L2end, "%m-%d %H:%M:%S") +
                 ', time: ' + str(L2end - L2start), lf)
        printout("DONE with LPI Calculations", lf)
        printout("--------------------------------------", lf)

    lf.close()
    sys.exit("FINISHED.")


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        sys.exit(1)