plotting.py

import cartopy
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
import cartopy.crs as ccrs

from tools import t_test


# Configuration for cartopy plots
cartopy.config["data_dir"] = (
    "/discover/nobackup/projects/jh_tutorials/JH_examples/JH_datafiles/Cartopy"
)
cartopy.config["pre_existing_data_dir"] = (
    "/discover/nobackup/projects/jh_tutorials/JH_examples/JH_datafiles/Cartopy"
)


def plot_training_loss(all_train_losses, EPOCHS):
    plt.plot(range(EPOCHS), all_train_losses[0], label="Training Loss")
    plt.xlabel("Epoch")
    plt.ylabel("L1 Loss")
    plt.title("Training Loss")
    plt.grid()
    plt.show()


def global_timeseries_plot(Y_pred, Y_train, Y_test, VARIABLE):

    # Gather weights
    weights = np.cos(np.deg2rad(Y_pred.lat))
    weights.name = "weights"
    weights.mean()

    # Initialize a flag to track if the label has been added
    label_added = False

    # Adds traininng data plots
    for i in range(len(Y_train)):
        if not label_added:
            Y_train[i][VARIABLE[0]].weighted(weights).mean(dim=["lat", "lon"]).plot(
                color="grey", label="Training Data"
            )
            label_added = True  # Set the flag to True after adding the label
        else:
            Y_train[i][VARIABLE[0]].weighted(weights).mean(dim=["lat", "lon"]).plot(
                color="grey"
            )

    # Plot model predictions
    Y_pred[VARIABLE[0]].weighted(weights).mean(dim=["lat", "lon"]).plot(
        color="black", label="Model Predictions"
    )

    # Plot test output
    Y_test[VARIABLE[0]].weighted(weights).mean(dim=["lat", "lon"]).plot(
        color="green", label="Test Data (SSP245)"
    )

    # Formatting
    plt.title("Global Average Timeseries")
    plt.xlabel("Year")
    plt.ylabel("Precipitation (mm/day)")
    plt.legend()
    plt.grid()


def global_anomaly_plot(Y_pred, Y_test, p_value, VARIABLE):

    # Extract data for the average of 2080-2100
    average_of_runs = Y_pred
    prediction_tsurf = average_of_runs
    prediction_tsurf_2080_2100 = prediction_tsurf.sel(year=slice(2080, 2100)).mean(
        dim="year"
    )[VARIABLE[0]]
    validation_tsurf_2080_2100 = (
        Y_test[VARIABLE[0]].sel(year=slice(2080, 2100)).mean(dim="year")
    )
    diff = prediction_tsurf_2080_2100 - validation_tsurf_2080_2100

    # Create figure
    fig = plt.figure(figsize=(10, 6))
    ax = plt.subplot(111, projection=ccrs.Robinson())

    # Create custom colormap settings
    cmap = plt.cm.coolwarm
    vmin, vmax = -1, 1

    # Plot with custom colorbar but suppress automatic display
    im = diff.plot(
        ax=ax,
        transform=ccrs.PlateCarree(),
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        extend="both",
        cbar_kwargs={
            "label": "kg/m²/s",
            "shrink": 0.75,
            "extendfrac": 0.05,
            "extendrect": False,
            "aspect": 30,
            "format": "%.1f",
            "ticks": [-1, -0.5, 0, 0.5, 1],
        },
    )

    # Add coastlines
    ax.coastlines()

    # Formatting
    plt.title(
        "Precipitation Error (Y_hat - Y_test)\nEmulator vs. NASA Global Climate Model\n2080-2100"
    )
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")

    # Display the plot once
    plt.show()