Regional Ocean: Animations of Surface Fields

Regional Ocean: Animations of Surface Fields#

Hide code cell source

 
%load_ext autoreload
%autoreload 2
import xarray as xr
import os
from IPython.display import HTML

import regional_utils as utils

Hide code cell source

 
case_name = ""  # "/glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/CESM_Output/"
CESM_output_dir = ""  # "CROCODILE_tutorial_nwa12_MARBL"

# As regional domains vary so much in purpose, simulation length, and extent, we don't want to assume a minimum duration
## Thus, we ignore start and end dates and simply reduce/output over the whole time frame for all of the examples given.
start_date = None  # "0001-01-01"
end_date = None  # "0101-01-01

save_figs = False
fig_output_dir = None

lc_kwargs = {}
serial = False

sfc_variables = []  # ['SSH', 'tos', 'sos']
max_frames = None  # 60

# Parameters
case_name = "CROCODILE_tutorial_nwa12_MARBL"
CESM_output_dir = (
    "/glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/CESM_Output/"
)
start_date = ""
end_date = ""
save_figs = True
fig_output_dir = None
lc_kwargs = {"threads_per_worker": 1}
serial = True
sfc_variables = ["SSH", "tos", "sos"]
max_frames = 60
subset_kwargs = {}
product = "/glade/work/ajanney/crocodile_2025/CUPiD/examples/regional_ocean/computed_notebooks//ocn/Regional_Ocean_Animations.ipynb"

Hide code cell source

 
OUTDIR = f"{CESM_output_dir}/{case_name}/ocn/hist/"
print("Output directory is:", OUTDIR)

Output directory is: /glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/CESM_Output//CROCODILE_tutorial_nwa12_MARBL/ocn/hist/

Open sfc_datas and Define Paths#

Hide code cell source

 
case_output_dir = os.path.join(CESM_output_dir, case_name, "ocn", "hist")

# Xarray time decoding things
time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)

## Static data includes hgrid, vgrid, bathymetry, land/sea mask
static_data = xr.open_mfdataset(
    os.path.join(case_output_dir, f"*static.nc"),
    decode_timedelta=True,
    decode_times=time_coder,
)

## Surface Data
sfc_data = xr.open_mfdataset(
    os.path.join(case_output_dir, f"*sfc*.nc"),
    decode_timedelta=True,
    decode_times=time_coder,
)

## Monthly Domain Data
# monthly_data = xr.open_mfdataset(
#     os.path.join(case_output_dir, f"*z*.nc"),
#     decode_timedelta=True,
#     decode_times=time_coder,
# )

## Image/Gif Output Directory
if fig_output_dir is None:
    image_output_dir = os.path.join(
        "/glade/derecho/scratch/",
        os.environ["USER"],
        "archive",
        case_name,
        "ocn",
        "cupid_images",
    )
else:
    image_output_dir = os.path.join(fig_output_dir, case_name, "ocn", "cupid_images")
if not os.path.exists(image_output_dir):
    os.makedirs(image_output_dir)
print("Image output directory is:", image_output_dir)

Image output directory is: /glade/derecho/scratch/ajanney/archive/CROCODILE_tutorial_nwa12_MARBL/ocn/cupid_images

## Select for only the variables we want to analyze
if len(sfc_variables) > 0:
    print("Selecting only the following surface variables:", sfc_variables)
    sfc_data = sfc_data[sfc_variables]

## Apply time boundaries
## if they are the right format
if len(start_date.split("-")) == 3 and len(end_date.split("-")) == 3:
    import cftime

    calendar = sfc_data.time.encoding.get("calendar", "standard")

    calendar_map = {
        "gregorian": cftime.DatetimeProlepticGregorian,
        "noleap": cftime.DatetimeNoLeap,
    }

    CFTime = calendar_map.get(calendar, cftime.DatetimeGregorian)
    y, m, d = [int(i) for i in start_date.split("-")]
    start_date_time = CFTime(y, m, d)
    y, m, d = [int(i) for i in end_date.split("-")]
    end_date_time = CFTime(y, m, d)

    print(
        f"Applying time range from start_date: {start_date_time} and end_date: {end_date_time}."
    )

    sfc_data = sfc_data.sel(time=slice(start_date_time, end_date_time))

sfc_time_bounds = [
    sfc_data["time"].isel(time=0).compute().item(),
    sfc_data["time"].isel(time=-1).compute().item(),
]

print(f"Surface Data Time Bounds: {sfc_time_bounds[0]} to {sfc_time_bounds[-1]}")

Selecting only the following surface variables: ['SSH', 'tos', 'sos']
Surface Data Time Bounds: 1999-12-30 12:00:00 to 2000-11-28 12:00:00

The GIF!#

if sfc_data["time"].size > max_frames:
    sfc_data = sfc_data.isel(time=slice(0, max_frames))

for gif_variable in sfc_variables:

    field = sfc_data[gif_variable]

    coords = utils.chooseGeoCoords(field.dims)
    areacello = utils.chooseAreacello(field.dims)

    anim = utils.create2DFieldAnimation(
        field,
        latitude=static_data[coords["latitude"]],
        longitude=static_data[coords["longitude"]],
        iter_dim="time",
        interval=150,
        save=save_figs,
        save_path=image_output_dir,
    )

from matplotlib import rcParams

rcParams["animation.embed_limit"] = (
    200 * 1024 * 1024
)  # constrains max size of HTML to be displayed in notebook

HTML(anim.to_jshtml())

Not rendered here!#

Contents