Global CHIRPS's SPI-based dry/wet indicators blend

spi_monthly_blend.py

# -*- coding: utf-8 -*-
"""
NAME
    spi_monthly_blend.py
    Global CHIRPS's SPI-based dry/wet indicators blend.
DESCRIPTION
    Input data for this script will use CHIRPS SPI monthly data in GeoTIFF format
    This experimental SPI blends integrate several SPI scales into a single product.
    The combines 3-, 6-, 9-, 12-, and 24-month SPI to estimate the overall dry/wet conditions.
METHOD
    The SPI values are weighted according to a normalized portion of the inverse of their
    contributing periods. For example, a 3-month SPI is weighted to 0.453 of the toral SPI blend.
    To calculate the SPI blend weighting factor, can follow below procedure:
        1. Sum the total number of contributing months (GOST runs a 3-, 6-, 9-, 12-, and 24-months
        SPI for a total of 54-months)
        2. Divide the total contributing months by the SPI analysis period (for a 3-month SPI,
        54/3 = 18)
        3. Divide the fractional contribution by the sum of all of the fractional contributions
        (ensures that total weights always sum to 1.00)
REQUIREMENT
    ArcGIS must installed before using this script, as it required arcpy module.
EXAMPLES
    C:\\Program Files\\ArcGIS\\Pro\\bin\\Python\\envs\\arcgispro-py3\\python spi_monthly_blend.py
NOTES
    This script is designed to work with global CHIRPS SPI-based product.
    If using other data, some adjustment are required: parsing filename, directory, threshold.
    All CHIRPS data and products are available at s3://wbgdecinternal-ntl/climate/
CONTACT
    Benny Istanto
    Climate Geographer
    GOST/DECAT/DEC Data Group, The World Bank
LICENSE
    This script is in the public domain, free from copyrights or restrictions.
VERSION
    $Id$
TODO
    xx
"""
import os
import arcpy
from datetime import datetime
from arcpy.sa import *

# Check out the Spatial Analyst extension
arcpy.CheckOutExtension("spatial")

# To avoid overwriting outputs, change overwriteOutput option to False.
arcpy.env.overwriteOutput = True

# Set the input folders
folder_spi03 = r"X:\\Temp\\CHIRPS\\SPI\\spi03"
folder_spi06 = r"X:\\Temp\\CHIRPS\\SPI\\spi06"
folder_spi09 = r"X:\\Temp\\CHIRPS\\SPI\\spi09"
folder_spi12 = r"X:\\Temp\\CHIRPS\\SPI\\spi12"
folder_spi24 = r"X:\\Temp\\CHIRPS\\SPI\\spi24"
folder_temp = r"X:\\Temp\\CHIRPS\\SPI\\temp"
# Set the output folder
folder_blend = r"X:\\Temp\\CHIRPS\\SPI\\spi_blend_03to24"

def spi_blend(folder_spi03, folder_spi06, folder_spi09, folder_spi12, folder_spi24, folder_blend):
    # Loop through the files in the input folder
    for raster in os.listdir(folder_spi03):
        # Check if the file is a TIFF file
        if raster.endswith(".tif") or raster.endswith(".tiff"):
            # Get the year, month, and day or day of year from the raster name
            # wld_cli_chirps_spi03_19810301.tif
            if "_" in raster:
                year = int(raster[21:25])
                month = int(raster[25:27])
                day = int(raster[27:29])
                date = datetime(year, month, day)

            # Construct the corresponding filenames in the stats folder
            raster_suffixes = ["spi03", "spi06", "spi09", "spi12", "spi24"]
            rasters = [os.path.join(eval("folder_" + suffix), "wld_cli_chirps_" + suffix + "_" + date.strftime("%Y%m%d") + ".tif") for suffix in raster_suffixes]

            # Create a dictionary of folder and raster pairs
            raster_folders = {"spi03": folder_spi03, "spi06": folder_spi06, "spi09": folder_spi09, "spi12": folder_spi12, "spi24": folder_spi24}

            # Iterate over the dictionary and check if each raster exists in its corresponding folder
            for folder, rasters_folder in raster_folders.items():
                for raster in rasters:
                    try:
                        if not arcpy.Exists(os.path.join(rasters_folder, raster)):
                            print("Error: {} not found in {}".format(raster, rasters_folder))
                            raise Exception("Input raster missing")
                    except Exception:
                        continue

            # Prepare the input
            try:
                raster_spi03 = arcpy.Raster(os.path.join(folder_spi03, rasters[0]))
                raster_spi06 = arcpy.Raster(os.path.join(folder_spi06, rasters[1]))
                raster_spi09 = arcpy.Raster(os.path.join(folder_spi09, rasters[2]))
                raster_spi12 = arcpy.Raster(os.path.join(folder_spi12, rasters[3]))
                raster_spi24 = arcpy.Raster(os.path.join(folder_spi24, rasters[4]))
            except Exception:
                continue

            # Apply the weight on each raster and create virtual raster objects
            spi03_weighted = arcpy.sa.Times(raster_spi03, 0.453)
            spi03_virtual = arcpy.sa.Raster(spi03_weighted)
            spi06_weighted = arcpy.sa.Times(raster_spi06, 0.226)
            spi06_virtual = arcpy.sa.Raster(spi06_weighted)
            spi09_weighted = arcpy.sa.Times(raster_spi09, 0.151)
            spi09_virtual = arcpy.sa.Raster(spi09_weighted)
            spi12_weighted = arcpy.sa.Times(raster_spi12, 0.113)
            spi12_virtual = arcpy.sa.Raster(spi12_weighted)
            spi24_weighted = arcpy.sa.Times(raster_spi24, 0.057)
            spi24_virtual = arcpy.sa.Raster(spi24_weighted)

            # Calculate the sum of the weighted rasters
            output_spiblend = os.path.join(folder_blend, "wld_cli_chirps_spiblend_03to24_" + date.strftime("%Y%m%d") + ".tif")
            spiblend = arcpy.sa.CellStatistics([spi03_virtual,spi06_virtual,spi09_virtual,spi12_virtual,spi24_virtual], "SUM", "DATA")
            spiblend.save(output_spiblend)
            print(output_spiblend + " completed")

            # Check out the Spatial Analyst extension
            arcpy.CheckInExtension("spatial")

# Main function
def main():
    # Call the spi_blend() function for the input folder
    spi_blend(folder_spi03, folder_spi06, folder_spi09, folder_spi12, folder_spi24, folder_blend)

if __name__ == '__main__':
    main()