IMERG 2 days or more precipitation running sum

imerg_runningsum_daily.py

# -*- coding: utf-8 -*-
"""
NAME
    imerg_runningsum_daily.py
    Global IMERG daily running sum.
DESCRIPTION
    Input data for this script will use IMERG daily data generated by imerg_nc2tif.py
    This script can do calculation on consecutive 2 days or more precipitation running sum
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_runningsum_daily.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 in nc4 format are available at JNB Server Drive J:\\Data\\GLOBAL\\CLIMATE\\imerg\\nc4\\
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
import multiprocessing

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

# 2 days precip accumulation running sum
def running_sum_2days(data_folder, output2_folder):
    # Spatial reference WGS-84
    sr = arcpy.SpatialReference(4326)
    print("looking at the first daily Rainfall data in tif folder...")
    
    for data_p in os.listdir(data_folder):
        if data_p.endswith(".tif"):

            i_imerg = data_p.index('imerg_')
            ymd = data_p[i_imerg + 6:i_imerg+6+8]

            date_formatted = date(int(ymd[0:4]), int(ymd[4:6]), int(ymd[6:8]))
            date_plus1 = date_formatted + timedelta(days=1)

            filename_2d = 'idn_cli_precip_2d_imerg_{0}{1}{2}.tif'.format(str(date_plus1.year).zfill(4), 
                                                                        str(date_plus1.month).zfill(2), 
                                                                        str(date_plus1.day).zfill(2))
            print("Processing "+filename_2d)

            data_file_1 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_formatted.strftime('%Y%m%d')))
            data_file_2 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus1.strftime('%Y%m%d')))

            if os.path.exists(data_file_1) and os.path.exists(data_file_2):
                print(str(date_formatted)+" next 1 days data exist. Start calculating....")
                arcpy.CheckOutExtension("spatial")
                
                data_combined = arcpy.sa.CellStatistics([data_file_1, data_file_2], "SUM", "DATA")
                data_combined.save(os.path.join(output2_folder, filename_2d))
                arcpy.DefineProjection_management(os.path.join(output2_folder, filename_2d), sr)
                
                print(filename_2d+ ' is succesfully created')
                arcpy.CheckInExtension("spatial")

            else:
                print(str(date_formatted)+" does not have complete 1 following data")

# 3 days precip accumulation running sum
def running_sum_3days(data_folder, output3_folder):
    # Spatial reference WGS-84
    sr = arcpy.SpatialReference(4326)
    print("looking at the first daily Rainfall data in tif folder...")
    
    for data_p in os.listdir(data_folder):
        if data_p.endswith(".tif"):

            i_imerg = data_p.index('imerg_')
            ymd = data_p[i_imerg + 6:i_imerg+6+8]

            date_formatted = date(int(ymd[0:4]), int(ymd[4:6]), int(ymd[6:8]))
            date_plus1 = date_formatted + timedelta(days=1)
            date_plus2 = date_formatted + timedelta(days=2)

            filename_3d = 'idn_cli_precip_3d_imerg_{0}{1}{2}.tif'.format(str(date_plus2.year).zfill(4), 
                                                                        str(date_plus2.month).zfill(2), 
                                                                        str(date_plus2.day).zfill(2))
            print("Processing "+filename_3d)

            data_file_1 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_formatted.strftime('%Y%m%d')))
            data_file_2 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus1.strftime('%Y%m%d')))
            data_file_3 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus2.strftime('%Y%m%d')))

            if os.path.exists(data_file_1) and os.path.exists(data_file_2) and os.path.exists(data_file_3):
                print(str(date_formatted)+" next 2 days data exist. Start calculating....")
                arcpy.CheckOutExtension("spatial")
                
                data_combined = arcpy.sa.CellStatistics([data_file_1, data_file_2, data_file_3], "SUM", "DATA")
                data_combined.save(os.path.join(output3_folder, filename_3d))
                arcpy.DefineProjection_management(os.path.join(output3_folder, filename_3d), sr)
                
                print(filename_3d+ ' is succesfully created')
                arcpy.CheckInExtension("spatial")

            else:
                print(str(date_formatted)+" does not have complete 2 following data")

# 4 days precip accumulation running sum
def running_sum_4days(data_folder, output4_folder):
    # Spatial reference WGS-84
    sr = arcpy.SpatialReference(4326)
    print("looking at the first daily Rainfall data in tif folder...")
    
    for data_p in os.listdir(data_folder):
        if data_p.endswith(".tif"):

            i_imerg = data_p.index('imerg_')
            ymd = data_p[i_imerg + 6:i_imerg+6+8]

            date_formatted = date(int(ymd[0:4]), int(ymd[4:6]), int(ymd[6:8]))
            date_plus1 = date_formatted + timedelta(days=1)
            date_plus2 = date_formatted + timedelta(days=2)
            date_plus3 = date_formatted + timedelta(days=3)

            filename_4d = 'idn_cli_precip_4d_imerg_{0}{1}{2}.tif'.format(str(date_plus3.year).zfill(4), 
                                                                        str(date_plus3.month).zfill(2), 
                                                                        str(date_plus3.day).zfill(2))
            print("Processing "+filename_4d)

            data_file_1 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_formatted.strftime('%Y%m%d')))
            data_file_2 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus1.strftime('%Y%m%d')))
            data_file_3 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus2.strftime('%Y%m%d')))
            data_file_4 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus3.strftime('%Y%m%d')))

            if os.path.exists(data_file_1) and os.path.exists(data_file_2) and os.path.exists(data_file_3) and os.path.exists(data_file_4):
                print(str(date_formatted)+" next 3 days data exist. Start calculating....")
                arcpy.CheckOutExtension("spatial")
                
                data_combined = arcpy.sa.CellStatistics([data_file_1, data_file_2, data_file_3, data_file_4], "SUM", "DATA")
                data_combined.save(os.path.join(output4_folder, filename_4d))
                arcpy.DefineProjection_management(os.path.join(output4_folder, filename_4d), sr)
                
                print(filename_4d+ ' is succesfully created')
                arcpy.CheckInExtension("spatial")

            else:
                print(str(date_formatted)+" does not have complete 3 following data")

# 5 days precip accumulation running sum
def running_sum_5days(data_folder, output5_folder):
    # Spatial reference WGS-84
    sr = arcpy.SpatialReference(4326)
    print("looking at the first daily Rainfall data in tif folder...")
    
    for data_p in os.listdir(data_folder):
        if data_p.endswith(".tif"):

            i_imerg = data_p.index('imerg_')
            ymd = data_p[i_imerg + 6:i_imerg+6+8]

            date_formatted = date(int(ymd[0:4]), int(ymd[4:6]), int(ymd[6:8]))
            date_plus1 = date_formatted + timedelta(days=1)
            date_plus2 = date_formatted + timedelta(days=2)
            date_plus3 = date_formatted + timedelta(days=3)
            date_plus4 = date_formatted + timedelta(days=4)

            filename_5d = 'idn_cli_precip_5d_imerg_{0}{1}{2}.tif'.format(str(date_plus4.year).zfill(4), 
                                                                        str(date_plus4.month).zfill(2), 
                                                                        str(date_plus4.day).zfill(2))
            print("Processing "+filename_5d)

            data_file_1 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_formatted.strftime('%Y%m%d')))
            data_file_2 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus1.strftime('%Y%m%d')))
            data_file_3 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus2.strftime('%Y%m%d')))
            data_file_4 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus3.strftime('%Y%m%d')))
            data_file_5 = os.path.join(data_folder, 'idn_cli_precip_1d_imerg_{0}.tif'.format(date_plus4.strftime('%Y%m%d')))

            if os.path.exists(data_file_1) and os.path.exists(data_file_2) and os.path.exists(data_file_3) and os.path.exists(data_file_4) and os.path.exists(data_file_5):
                print(str(date_formatted)+" next 4 days data exist. Start calculating....")
                arcpy.CheckOutExtension("spatial")
                
                data_combined = arcpy.sa.CellStatistics([data_file_1, data_file_2, data_file_3, data_file_4, data_file_5], "SUM", "DATA")
                data_combined.save(os.path.join(output5_folder, filename_5d))
                arcpy.DefineProjection_management(os.path.join(output5_folder, filename_5d), sr)
                
                print(filename_5d+ ' is succesfully created')
                arcpy.CheckInExtension("spatial")

            else:
                print(str(date_formatted)+" does not have complete 4 following data")

# 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 (input folder, outputX folder)
        running_sum_2days('X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_1d','X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_2d')
        running_sum_3days('X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_1d','X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_3d')
        running_sum_4days('X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_1d','X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_4d')
        running_sum_5days('X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_1d','X:\\Temp\\imerg\\data\\geotiff\\IDN\\precip_5d')