aew_lifecycle_composite.py

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# colormap for pride flag
colors = ["#000000", "#613915", "#FFAFC8","#FFFFFF", "#74D7EE", "#E40303","#FF8C00", "#FFED00", "#008026", "#004DFF", "#750787"] 

# create the colormap object, N=number of colors inside the gradient, customize as needed!
pride = LinearSegmentedColormap.from_list("my_custom_map", colors, N=20)

#••••••••••••••••••••••••••••••••••
# settings
#••••••••••••••••••••••••••••••••••
YEARS = range(1950, 2015)

# Model : resolution
MODELS = {
    "LR": "2",
    "MR": "1",
    "HR": "0.5",
    "VHR": "0.25",
}

N_PHASES = 3
titles = ["inital stage", "middle stage", "late stage"]

##••••••••••••••••••••••••••••••••••
# FIGURE
#••••••••••••••••••••••••••••••••••
fig, axes = plt.subplots(
    len(MODELS),
    N_PHASES,
    figsize=(15, 18),
    constrained_layout=True,
)

mappable = None

#••••••••••••••••••••••••••••••••••
# LOOP OVER MODELS
#••••••••••••••••••••••••••••••••••
for row, (MODEL, RESOLUTION) in enumerate(MODELS.items()):

    print(f"\nProcessing {MODEL}")

    #••••••••••••••••••••••••••••••••••
    # Storage across all years
    #••••••••••••••••••••••••••••••••••
    u_phase = [[] for _ in range(N_PHASES)]
    v_phase = [[] for _ in range(N_PHASES)]

    first_wave = None

    #••••••••••••••••••••••••••••••••••
    # LOOP YEARS
    #••••••••••••••••••••••••••••••••••
    for YEAR in YEARS:

        BASE = Path(
            f"/data/kellie.thrower/IPSL_{MODEL}_QTRACK/"
            f"{YEAR}_JunOct_{RESOLUTION}_deg/AEW_wave_files"
        )

        wave_files = sorted(BASE.glob("wave_*.nc"))

        if len(wave_files) == 0:
            print(f"No wave files found: {MODEL} {YEAR}")
            continue

        print(
            f"{MODEL} {YEAR}: {len(wave_files)} waves"
        )

        if first_wave is None:
            first_wave = wave_files[0]

        #••••••••••••••••••••••••••••••••••
        # Read waves
        #••••••••••••••••••••••••••••••••••
        for f in wave_files:

            ds = xr.open_dataset(f)

            T = ds.sizes["time"]

            if T < 3:
                ds.close()
                continue

            edges = np.linspace(
                0,
                T,
                N_PHASES + 1
            ).astype(int)

            for phase in range(N_PHASES):

                t0 = edges[phase]
                t1 = edges[phase + 1]

                if t1 - t0 < 1:
                    continue

                u_phase[phase].append(
                    ds["u700"]
                    .isel(time=slice(t0, t1))
                    .values
                )

                v_phase[phase].append(
                    ds["v700"]
                    .isel(time=slice(t0, t1))
                    .values
                )

            ds.close()

    #••••••••••••••••••••••••••••••••••
    # Check data exists
    #••••••••••••••••••••••••••••••••••
    if first_wave is None:
        print(f"No data for {MODEL}")
        continue

    #••••••••••••••••••••••••••••••••••
    # Composite means
    #••••••••••••••••••••••••••••••••••
    u_mean = []
    v_mean = []

    for phase in range(N_PHASES):

        u = np.concatenate(
            u_phase[phase],
            axis=0
        )

        v = np.concatenate(
            v_phase[phase],
            axis=0
        )

        u_mean.append(
            np.nanmean(u, axis=0)
        )

        v_mean.append(
            np.nanmean(v, axis=0)
        )

    #••••••••••••••••••••••••••••••••••
    # Grid
    #••••••••••••••••••••••••••••••••••
    ds0 = xr.open_dataset(first_wave)

    x = ds0["x"].values
    y = ds0["y"].values

    X, Y = np.meshgrid(
        x,
        y
    )

    ds0.close()

    #••••••••••••••••••••••••••••••••••
    # Plot lifecycle phases
    #••••••••••••••••••••••••••••••••••
    for phase in range(N_PHASES):

        ax = axes[row, phase]

        speed = np.sqrt(
            u_mean[phase]**2 +
            v_mean[phase]**2
        )

        cf = ax.contourf(
            X,
            Y,
            speed,
            levels=20,
            cmap=pride,
            vmin=0,
            vmax=12,
        )

        if mappable is None:
            mappable = cf

        # Streamlines
        ax.streamplot(
            X,
            Y,
            u_mean[phase],
            v_mean[phase],
            color="white",
            density=1,
            linewidth=1.2,
            arrowsize=1.2,
        )

        # Maximum wind location
        iy, ix = np.unravel_index(
            np.nanargmax(speed),
            speed.shape,
        )

        ax.scatter(
            X[iy, ix],
            Y[iy, ix],
            marker="*",
            s=200,
            color="white",
            edgecolor="black",
            zorder=5,
        )

        # Reference axes
        ax.axhline(
            0,
            color="k",
            linewidth=0.8
        )

        ax.axvline(
            0,
            color="k",
            linewidth=0.8
        )

        ax.set_xlim(
            x.min(),
            x.max()
        )

        ax.set_ylim(
            y.min(),
            y.max()
        )

        ax.set_aspect(
            "equal",
            adjustable="box"
        )

        # Column titles
        if row == 0:
            ax.set_title(
                titles[phase],
                fontsize=14,
                weight="bold"
            )

        # Row labels
        if phase == 0:
            ax.set_ylabel(
                f"{MODEL}\nMeridional distance (deg)",
                weight="bold"
            )

        # Bottom labels
        if row == len(MODELS)-1:
            ax.set_xlabel(
                "Zonal distance (deg)",
                weight="bold"
            )


#••••••••••••••••••••••••••••••••••
# COLORBAR
#••••••••••••••••••••••••••••••••••

cbar = fig.colorbar(
    mappable,
    ax=axes,
    orientation="vertical",
    shrink=0.9,
    pad=0.02,
)

cbar.set_label(
    "700-hPa Wind Speed (m s$^{-1}$)",
    fontsize=12,
)


#••••••••••••••••••••••••••••••••••
# title
#••••••••••••••••••••••••••••••••••
plt.suptitle(
    f"AEW Lifecycle Composite Structure "
    f"({YEARS.start}-{YEARS.stop-1})",
    fontsize=18,
    weight="bold",
)


#••••••••••••••••••••••••••••••••••
# save
#••••••••••••••••••••••••••••••••••
plt.savefig(
    f"aew_lifecycle_composite_all_models_"
    f"{YEARS.start}-{YEARS.stop-1}.png",
    dpi=300,
    bbox_inches="tight",
)

plt.show()