ad2.py

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from pathlib import Path
import matplotlib.colors as mcolors

# --- Configuration ---
# Define file paths using pathlib
data_dir = Path("data2")
output_dir = Path("output")
data_dir.mkdir(exist_ok=True, parents=True)
output_dir.mkdir(exist_ok=True, parents=True)

precip_path = data_dir / "chirps_MarchJune.nc"
precip_jan_path = data_dir / "chirps_Jan.nc"

# Define precipitation thresholds and pentad labels
rate_precip = [3, 4, 5, 6, 7]
pentads = [
    "March 4", "March 9", "March 14", "March 19", "March 24", "March 29",
    "April 3", "April 8", "April 13", "April 18", "April 23", "April 28",
    "May 3", "May 8", "May 13", "May 18", "May 23", "May 28",
    "Jun 2", "Jun 7", "Jun 12", "Jun 17", "Jun 22", "Jun 27"
]

# --- Data Loading ---
try:
    ds = xr.open_dataset(precip_path)
    ds_Jan = xr.open_dataset(precip_jan_path)
except FileNotFoundError:
    print("Error: Input data files not found. Please ensure 'data/chirps_MarchJune.nc' and 'data/chirps_Jan.nc' exist.")
    exit()


# Extract dimensions and precipitation data
lon = ds["longitude"].values
lat = ds["latitude"].values
precip_total = ds["precip"].values
precip_Jan = ds_Jan["precip"].values
lon2d, lat2d = np.meshgrid(lon, lat)

# --- Data Processing and Onset Calculation ---
# Calculate the mean of January precipitation
mean_Jan = np.mean(precip_Jan, axis=0)

# Create a January mean array that matches the shape of the March-June data
mean_Jan_matchMarchJune = np.tile(mean_Jan[np.newaxis, :, :], (precip_total.shape[0], 1, 1))

# Subtract the January mean to get corrected precipitation
precip_corr = precip_total - mean_Jan_matchMarchJune

# Identify valid points (with sufficient non-NaN data)
varsize = precip_corr.shape
valid_m = np.full((varsize[1], varsize[2]), np.nan)
valid_count = 0
for i in range(varsize[1]):
    for j in range(varsize[2]):
        if np.count_nonzero(~np.isnan(precip_total[:, i, j])) >= 20:
            valid_m[i, j] = 1
            valid_count += 1

# Prepare data for FFT by selecting only valid points
mydata = np.full((varsize[0], valid_count), np.nan)
l = 0
for i in range(varsize[1]):
    for j in range(varsize[2]):
        if not np.isnan(valid_m[i, j]):
            mydata[:, l] = precip_corr[:, i, j]
            l += 1

# --- Main Loop for Onset Calculation and Plotting ---
for mm in rate_precip:
    print(f"Processing for threshold: {mm}mm")
    onset1 = np.full(valid_count, np.nan)

    # Calculate onset for each valid point
    for j in range(valid_count):
        if np.isnan(mydata[:, j]).any():
            continue

        fft1 = np.fft.fft(mydata[:, j])
        fft1[6:67] = 0
        ifft1 = np.fft.ifft(fft1).real
        mean1 = np.mean(ifft1[0:6])

        for i in range(6, len(ifft1)):
            if (ifft1[i] - mm) > mean1:
                onset1[j] = i + 1
                break

    # Reshape onset data back to a 2D grid
    onset_grid = np.full((varsize[1], varsize[2]), np.nan)
    l = 0
    for i in range(varsize[1]):
        for j in range(varsize[2]):
            if not np.isnan(valid_m[i, j]):
                onset_grid[i, j] = onset1[l]
                l += 1

    # --- Plotting ---
    # Plot 1: Rain Intensity at Onset (Continuous Colorbar)
    precip_at_onset = np.full_like(onset_grid, np.nan)
    for i in range(varsize[1]):
        for j in range(varsize[2]):
            if not np.isnan(onset_grid[i, j]):
                onset_pentad_index = int(onset_grid[i, j]) - 1
                if 0 <= onset_pentad_index < varsize[0]:
                    precip_at_onset[i, j] = precip_corr[onset_pentad_index, i, j]

    fig = plt.figure(figsize=(10, 8))
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([40, 160, 0, 60], crs=ccrs.PlateCarree())
    ax.coastlines(resolution="10m")
    ax.add_feature(cfeature.BORDERS.with_scale('10m'), linewidth=0.7)
    gl = ax.gridlines(draw_labels=True, linewidth=0.3, color='gray', linestyle='--')
    gl.top_labels = False
    gl.right_labels = False

    im1 = ax.pcolormesh(
        lon, lat, precip_at_onset,
        transform=ccrs.PlateCarree(),
        cmap="viridis",
        shading="auto"
    )
    cbar1 = fig.colorbar(im1, orientation='horizontal', pad=0.05, fraction=0.03, aspect=30)
    cbar1.set_label("Rain Intensity at Onset (mm)")
    ax.set_title(f"Rain Intensity at Onset (Threshold: {mm}mm)", fontsize=11)
    plt.savefig(output_dir / f"rain_intensity_onset_{mm}mm.png", dpi=300, bbox_inches="tight")
    plt.close(fig)

    # Plot 2: Onset Timing (Discrete Colorbar)
    fig = plt.figure(figsize=(10, 8))
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([40, 160, 0, 60], crs=ccrs.PlateCarree())
    ax.coastlines(resolution="10m")
    ax.add_feature(cfeature.BORDERS.with_scale('10m'), linewidth=0.7)
    gl = ax.gridlines(draw_labels=True, linewidth=0.3, color='gray', linestyle='--')
    gl.top_labels = False
    gl.right_labels = False

    cmap_discrete = plt.cm.get_cmap('turbo', len(pentads))
    norm = mcolors.BoundaryNorm(np.arange(0.5, len(pentads) + 1.5), cmap_discrete.N)
    im2 = ax.pcolormesh(
        lon, lat, onset_grid,
        transform=ccrs.PlateCarree(),
        cmap=cmap_discrete, norm=norm,
        shading="auto"
    )

    cbar2 = fig.colorbar(im2, orientation='horizontal', pad=0.05, fraction=0.03, aspect=30, ticks=np.arange(1, len(pentads) + 1))
    cbar2.set_ticklabels(pentads)
    cbar2.ax.tick_params(rotation=90)
    cbar2.set_label("Monsoon Onset Pentad")
    ax.set_title(f"Monsoon Onset Timing (Threshold: {mm}mm)", fontsize=11)
    plt.savefig(output_dir / f"onset_timing_{mm}mm.png", dpi=300, bbox_inches="tight")
    plt.close(fig)

ds.close()
ds_Jan.close()
print("\nProcessing complete. Plots saved in the 'output' directory.")