Global IMERG daily consecutive dry days data extraction

IMERG_cdd_5mm.py

# -*- coding: utf-8 -*-
"""
NAME
    imerg_cdd_5mm.py
    Global IMERG daily indices data extraction
DESCRIPTION
    Input data for this script will use IMERG daily data generated by imerg_nc2tif.py
    This script can do Consecutive Dry Days (CDD) calculation with threshold 5mm of rainfall as a rainy
PROCESS
    (i)   Extract IMERG's daily rainfall with value greater/less than 5mm (threshold for a day categoried as rainy day)
    (ii)  If Rainfall < 5 = Yes, then assign 1 otherwise 0. This is for dry condition, for wet are the opposite.
    (iii) For number of consecutive information, it will accumulate to next data calculation result,
          if the value = 1. If not, start from 0 again.
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 imerg_cdd_1mm.py
NOTES
    This script is designed to work with global IMERG data (Final or Late Run)
    If using other data, some adjustment are required: parsing filename, directory, threshold
    All IMERG data and products are available at s3://wbgdecinternal-ntl/climate/
CONTACT
    Benny Istanto
    Climate Geographer
    GOST, The World Bank
LICENSE
    This script is in the public domain, free from copyrights or restrictions.
VERSION
    $Id$
TODO
    xx
"""
import arcpy
from arcpy.sa import *
import os, sys, traceback
import re
from datetime import date, timedelta


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

# Check precipitation data in GeoTIFF format
def create_daily_List(_tif_folder):
    print("start reading list of daily Rainfall data....")
    print("looking for file with naming wld_cli_precips_YYYYMMDD.tif")
    
    Daily_Date_List=[]
    
    for Daily_Data in os.listdir(_tif_folder):
        
        if Daily_Data.endswith(".tif"):
            print("found " + Daily_Data+ " in the daily Rainfall folder")
            
            i_imerg = Daily_Data.index('imerg_')
            ymd = Daily_Data[i_imerg + 6:i_imerg+6+8]
            
            Daily_Data_Date = ymd            
            Daily_Date_List.append(Daily_Data_Date)
    
    return sorted(Daily_Date_List)


# Check if there is Consecutive DRY Days data in output folder
def create_CDD_List(_CDD_folder):
    print("start reading existing Consecutive DRY Days Dataset....")
    print("looking for file with naming wld_cli_cdd_5mm_imerg_YYYYMMDD")
    
    CDD_Date_List=[]
    
    for CDD_Data in os.listdir(_CDD_folder):
        if CDD_Data.endswith(".tif"):
            print("found " + CDD_Data + " in the Consecutive DRY Days folder")
            
            parse_String = CDD_Data.split('_')
            CDD_Data_Date = parse_String[5]
            CDD_Date_List.append(CDD_Data_Date)
    
    return CDD_Date_List


# Execute first Dry condition
def execute_first_CDD(_tiffolder, _CDDFolder, threshold):
    # Spatial reference WGS-84
    sr = arcpy.SpatialReference(4326)
    print("looking at the first daily Rainfall data in tif folder...")
    
    daily_list = create_daily_List(_tiffolder)
    first_date = min(daily_list)
    print("execute first Rainfall data from date "+first_date)
    
    first_data_name = 'wld_cli_precip_1d_imerg_{0}{1}{2}.tif'.format(first_date[0:4], first_date[4:6], first_date[6:8])
    first_daily_data = os.path.join(_tiffolder, first_data_name)
    daily_Date = date(int(first_date[0:4]), int(first_date[4:6]), int(first_date[6:8]))
    dry_date = daily_Date #+ timedelta(days=1)
    print("creating dry data "+str(dry_date)+ " from daily Rainfall data from "+str(daily_Date))
    
    CDDyear = str(dry_date.year)
    CDDmonth = str(dry_date.month)
    CDDday = str(dry_date.day)
    print(str(dry_date))
    
    CDDFilename = 'wld_cli_cdd_5mm_imerg_{0}{1}{2}.tif'.format(CDDyear.zfill(4), CDDmonth.zfill(2), CDDday.zfill(2))
    print("Processing "+CDDFilename)
    
    arcpy.CheckOutExtension("spatial")
    
    outCon = Con(Raster(first_daily_data) < Float(threshold), Float(1), Float(0))
    outCon.save(os.path.join(_CDDFolder, CDDFilename))
    arcpy.DefineProjection_management(os.path.join(_CDDFolder, CDDFilename), sr)
    
    arcpy.CheckInExtension("spatial")
    print("file " + CDDFilename + " is created")


# Execute next Consecutive DRY Days data
def execute_CDD(_lastdate, _tiffolder, _CDD_folder, threshold):
    # Spatial reference WGS-84
    sr = arcpy.SpatialReference(4326)
    
    date_formatted = date(int(_lastdate[0:4]), int(_lastdate[4:6]), int(_lastdate[6:8]))
    last_dryname = 'wld_cli_cdd_5mm_imerg_{0}'.format(_lastdate)
    last_dryfile = os.path.join(_CDD_folder, last_dryname)
    next_daily_date = date_formatted + timedelta(days=1)
    next_dailyname = 'wld_cli_precip_1d_imerg_{0}.tif'.format(next_daily_date.strftime('%Y%m%d'))
    #next_dailyname = 'wld_cli_precip_1d_imerg_{0}{1}{2}.tif'.format(_lastdate[0:4], _lastdate[4:6], _lastdate[6:8])
    next_dailydata = os.path.join(_tiffolder, next_dailyname)
    
    if arcpy.Exists(next_dailydata):
        print("next daily data is available...")
        print("start processing next Consecutive DRY Days...")
        
        new_dry_date = date_formatted + timedelta(days=1)
        CDDyear1 = str(new_dry_date.year)
        CDDmonth1 = str(new_dry_date.month)
        CDDday1 = str(new_dry_date.day)
        new_dry_name = 'wld_cli_cdd_5mm_imerg_{0}{1}{2}.tif'.format(CDDyear1.zfill(4), CDDmonth1.zfill(2), CDDday1.zfill(2))
        print("Processing Consecutive DRY Days from "+last_dryfile+" and "+next_dailydata)
        
        arcpy.CheckOutExtension("spatial")
        
        outCDDCon = Con(Raster(next_dailydata) < Float(threshold), Raster(last_dryfile)+Float(1), Float(0))
        outCDDCon.save(os.path.join(_CDD_folder, new_dry_name))
        arcpy.DefineProjection_management(os.path.join(_CDD_folder, new_dry_name), sr)
        
        arcpy.CheckInExtension("spatial")
        print("Consecutive DRY Days File "+new_dry_name+" is created")
    
    else:
        print("next daily data is not available. Exit...")


# Run the script
def create_CDD(_CDD_folder, _tiffolder, threshold):

    CDD_Date_List = create_CDD_List(_CDD_folder)
    Daily_list = create_daily_List(_tiffolder)
    
    # if there is no DRY data, creating new DRY data
    if len(CDD_Date_List)==0:
        print("No Consecutive DRY Days data found...")
        print("Creating first Consecutive DRY Days data...")
        
        execute_first_CDD(_tiffolder, _CDD_folder, threshold)
        CDD_Date_List = create_CDD_List(_CDD_folder)
    
    # if there is Consecutive DRY Days data
    print("Consecutive DRY Days data found. Looking for latest Consecutive DRY Days data...")
    
    #Check last Consecutive DRY Days available
    last_date = max(CDD_Date_List)

    #Check last daily data availabke
    max_daily_date = max(Daily_list)
    last_CDD_date = date(int(last_date[0:4]), int(last_date[4:6]), int(last_date[6:8]))
    last_daily_date = date(int(max_daily_date[0:4]), int(max_daily_date[4:6]), int(max_daily_date[6:8]))
    
    # process Consecutive DRY Days to every daily data available after last Consecutive DRY Days data
    while last_daily_date + timedelta(days=1) > last_CDD_date:
    #while last_daily_date > last_CDD_date:

        execute_CDD(last_date, _tiffolder, _CDD_folder, threshold)

        last_CDD_date=last_CDD_date+timedelta(days=1)
        CDDyear2 = str(last_CDD_date.year)
        CDDmonth2 = str(last_CDD_date.month)
        CDDday2 = str(last_CDD_date.day)
        last_date='{0}{1}{2}.tif'.format(CDDyear2.zfill(4), CDDmonth2.zfill(2), CDDday2.zfill(2))

    print("All Consecutive DRY Days data is available")


# Let's go!
if __name__ == '__main__':
    # Global Environment settings
    with arcpy.EnvManager(scratchWorkspace=r"X:\ArcGIS_TEMP\Scratch.gdb", \
        workspace=r"X:\ArcGIS_TEMP\Default.gdb"):
        # Run the function (output folder, input folder, threshold)
        create_CDD('X:\\Temp\\imerg\\products\\cdd\\cdd_5mm_temp',\
            'X:\\Temp\\imerg\\data\\geotiff\\rainfall_1days', 5)