Global IMERG daily consecutive wet days data extraction

IMERG_cwd_5mm.py

# -*- coding: utf-8 -*-
"""
NAME
    imerg_cwd_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 Wet Days (CWD) calculation with threshold 5mm of rainfall as a rainy day
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_cwd_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 WET Days data in output folder
def create_CWD_List(_CWD_folder):
    print("start reading existing Consecutive WET Days Dataset....")
    print("looking for file with naming wld_cli_cwd_5mm_imerg_YYYYMMDD")
    
    CWD_Date_List=[]
    
    for CWD_Data in os.listdir(_CWD_folder):
        if CWD_Data.endswith(".tif"):
            print("found " + CWD_Data + " in the Consecutive WET Days folder")
            
            parse_String = CWD_Data.split('_')
            CWD_Data_Date = parse_String[5]
            CWD_Date_List.append(CWD_Data_Date)
    
    return CWD_Date_List


# Execute first Wet condition
def execute_first_CWD(_tiffolder, _CWDFolder, 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]))
    wet_date = daily_Date #+ timedelta(days=1)
    print("creating wet data "+str(wet_date)+ " from daily Rainfall data from "+str(daily_Date))
    
    CWDyear = str(wet_date.year)
    CWDmonth = str(wet_date.month)
    CWDday = str(wet_date.day)
    print(str(wet_date))
    
    CWDFilename = 'wld_cli_cwd_5mm_imerg_{0}{1}{2}.tif'.format(CWDyear.zfill(4), CWDmonth.zfill(2), CWDday.zfill(2))
    print("Processing "+CWDFilename)
    
    arcpy.CheckOutExtension("spatial")
    
    outCon = Con(Raster(first_daily_data) < Float(threshold), Float(0), Float(1))
    outCon.save(os.path.join(_CWDFolder, CWDFilename))
    arcpy.DefineProjection_management(os.path.join(_CWDFolder, CWDFilename), sr)
    
    arcpy.CheckInExtension("spatial")
    print("file " + CWDFilename + " is created")


# Execute next Consecutive DRY Days data
def execute_CWD(_lastdate, _tiffolder, _CWD_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_wetname = 'wld_cli_cwd_5mm_imerg_{0}'.format(_lastdate)
    last_wetfile = os.path.join(_CWD_folder, last_wetname)
    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 WET Days...")
        
        new_wet_date = date_formatted + timedelta(days=1)
        CWDyear1 = str(new_wet_date.year)
        CWDmonth1 = str(new_wet_date.month)
        CWDday1 = str(new_wet_date.day)
        new_wet_name = 'wld_cli_cwd_5mm_imerg_{0}{1}{2}.tif'.format(CWDyear1.zfill(4), CWDmonth1.zfill(2), CWDday1.zfill(2))
        print("Processing Consecutive WET Days from "+last_wetfile+" and "+next_dailydata)
        
        arcpy.CheckOutExtension("spatial")
        
        outCWDCon = Con(Raster(next_dailydata) < Float(threshold), Float(0), Raster(last_wetfile)+Float(1))
        outCWDCon.save(os.path.join(_CWD_folder, new_wet_name))
        arcpy.DefineProjection_management(os.path.join(_CWD_folder, new_wet_name), sr)
        
        arcpy.CheckInExtension("spatial")
        print("Consecutive WET Days File "+new_wet_name+" is created")
    
    else:
        print("next daily data is not available. Exit...")


# Run the script
def create_CWD(_CWD_folder, _tiffolder, threshold):

    CWD_Date_List = create_CWD_List(_CWD_folder)
    Daily_list = create_daily_List(_tiffolder)
    
    # if there is no WET data, creating new WET data
    if len(CWD_Date_List)==0:
        print("No Consecutive WET Days data found...")
        print("Creating first Consecutive WET Days data...")
        
        execute_first_CWD(_tiffolder, _CWD_folder, threshold)
        CWD_Date_List = create_CWD_List(_CWD_folder)
    
    # if there is Consecutive WET Days data
    print("Consecutive WET Days data found. Looking for latest Consecutive WET Days data...")
    
    #Check last Consecutive WET Days available
    last_date = max(CWD_Date_List)

    #Check last daily data availabke
    max_daily_date = max(Daily_list)
    last_CWD_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 WET Days to every daily data available after last Consecutive WET Days data
    while last_daily_date + timedelta(days=1) > last_CWD_date:
    #while last_daily_date > last_CWD_date:

        execute_CWD(last_date, _tiffolder, _CWD_folder, threshold)

        last_CWD_date=last_CWD_date+timedelta(days=1)
        CWDyear2 = str(last_CWD_date.year)
        CWDmonth2 = str(last_CWD_date.month)
        CWDday2 = str(last_CWD_date.day)
        last_date='{0}{1}{2}.tif'.format(CWDyear2.zfill(4), CWDmonth2.zfill(2), CWDday2.zfill(2))

    print("All Consecutive WET 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_CWD('X:\\Temp\\imerg\\products\\cwd\\cwd_5mm_temp',\
            'X:\\Temp\\imerg\\data\\geotiff\\rainfall_1days', 5)