Global TerraClimate's SPEI-based drought indicators blend

spei_monthly_blend.py

# -*- coding: utf-8 -*-
"""
NAME
    spei_monthly_blend.py
    Global TerraClimate's SPEI-based dry/wet indicators blend.
DESCRIPTION
    Input data for this script will use TerraClimate SPEI monthly data in GeoTIFF format
    This experimental SPEI blends integrate several SPEI scales into a single product.
    The combines 3-, 6-, 9-, 12-, and 24-month SPEI to estimate the overall dry/wet conditions.
METHOD
    The SPEI values are weighted according to a normalized portion of the inverse of their
    contributing periods. For example, a 3-month SPEI is weighted to 0.453 of the toral SPEI blend.
    To calculate the SPEI 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
        SPEI for a total of 54-months)
        2. Divide the total contributing months by the SPEI analysis period (for a 3-month SPEI,
        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 spei_monthly_blend.py
NOTES
    This script is designed to work with global TerraClimate SPEI-based product.
    If using other data, some adjustment are required: parsing filename, directory, threshold.
    All TerraClimate 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_spei03 = r"X:\\Temp\\TerraClimate\\SPEI\\spei03"
folder_spei06 = r"X:\\Temp\\TerraClimate\\SPEI\\spei06"
folder_spei09 = r"X:\\Temp\\TerraClimate\\SPEI\\spei09"
folder_spei12 = r"X:\\Temp\\TerraClimate\\SPEI\\spei12"
folder_spei24 = r"X:\\Temp\\TerraClimate\\SPEI\\spei24"
folder_temp = r"X:\\Temp\\TerraClimate\\SPEI\\temp"
# Set the output folder
folder_blend = r"X:\\Temp\\TerraClimate\\SPEI\\spei_blend_03to24"

def spei_blend(folder_spei03, folder_spei06, folder_spei09, folder_spei12, folder_spei24, folder_blend):
    # Loop through the files in the input folder
    for raster in os.listdir(folder_spei03):
        # 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_terraclimate_spei03_19810301.tif
            if "_" in raster:
                year = int(raster[28:32])
                month = int(raster[32:34])
                day = int(raster[34:36])
                date = datetime(year, month, day)

            # Construct the corresponding filenames in the stats folder
            raster_suffixes = ["spei03", "spei06", "spei09", "spei12", "spei24"]
            rasters = [os.path.join(eval("folder_" + suffix), "wld_cli_terraclimate_" + suffix + "_" + date.strftime("%Y%m%d") + ".tif") for suffix in raster_suffixes]

            # Create a dictionary of folder and raster pairs
            raster_folders = {"spei03": folder_spei03, "spei06": folder_spei06, "spei09": folder_spei09, "spei12": folder_spei12, "spei24": folder_spei24}

            # 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_spei03 = arcpy.Raster(os.path.join(folder_spei03, rasters[0]))
                raster_spei06 = arcpy.Raster(os.path.join(folder_spei06, rasters[1]))
                raster_spei09 = arcpy.Raster(os.path.join(folder_spei09, rasters[2]))
                raster_spei12 = arcpy.Raster(os.path.join(folder_spei12, rasters[3]))
                raster_spei24 = arcpy.Raster(os.path.join(folder_spei24, rasters[4]))
            except Exception:
                continue

            # Apply the weight on each raster and create virtual raster objects
            spei03_weighted = arcpy.sa.Times(raster_spei03, 0.453)
            spei03_virtual = arcpy.sa.Raster(spei03_weighted)
            spei06_weighted = arcpy.sa.Times(raster_spei06, 0.226)
            spei06_virtual = arcpy.sa.Raster(spei06_weighted)
            spei09_weighted = arcpy.sa.Times(raster_spei09, 0.151)
            spei09_virtual = arcpy.sa.Raster(spei09_weighted)
            spei12_weighted = arcpy.sa.Times(raster_spei12, 0.113)
            spei12_virtual = arcpy.sa.Raster(spei12_weighted)
            spei24_weighted = arcpy.sa.Times(raster_spei24, 0.057)
            spei24_virtual = arcpy.sa.Raster(spei24_weighted)

            # Calculate the sum of the weighted rasters
            output_speiblend = os.path.join(folder_blend, "wld_cli_terraclimate_speiblend_03to24_" + date.strftime("%Y%m%d") + ".tif")
            speiblend = arcpy.sa.CellStatistics([spei03_virtual,spei06_virtual,spei09_virtual,spei12_virtual,spei24_virtual], "SUM", "DATA")
            speiblend.save(output_speiblend)
            print(output_speiblend + " completed")

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

# Main function
def main():
    # Call the spei_blend() function for the input folder
    spei_blend(folder_spei03, folder_spei06, folder_spei09, folder_spei12, folder_spei24, folder_blend)

if __name__ == '__main__':
    main()