Iceland3 model solution overview

import subprocess
import os
import pandas as pd
import netCDF4
import numpy as np
import glob
import time
import matplotlib.pyplot as plt
import copy
import xarray as xr
from datetime import datetime, timedelta 
import dask
from scipy.interpolate import griddata
#from ocean_c_lab_tools import *
#from celluloid import Camera 
#import PyCO2SYS as csys
import seawater as sw
from roms_regrid import *

HAFRO_path='/home/x-uheede/R/HAFRO/Hafro_cruises.xls'
model_grid_path="/home/x-uheede/S/Iceland3_MARBL_2024_60m/P_INPUT/Iceland3_grid_MAT.nc"
# Grid parameters, only modify these if grid is made in MATLAB
vert_levels=60
theta_s_model=5
theta_b_model=2
hc_model=300
model_data_path="/anvil/scratch/x-uheede/Iceland3_MARBL_2024/Iceland3_MARBL_2024_his.2024071???????.nc"
months_analysis=[4] # enter the months you want to analyze for the model
# enter the dates you want to analyze for the observations
target_depth_levels=[10] # Specify depth levels of interest
thinner=1 # specify the temporal frequency of data being read (i.e. no need to read in hourly data)

from roms_tools import Grid, ROMSOutput
grid = Grid.from_file(
    model_grid_path
)

#Only run this cell if grid is made in MATLAB
grid.update_vertical_coordinate(N=vert_levels, theta_s=theta_s_model, theta_b=theta_b_model, hc=hc_model, verbose=False)

import xarray as xr
import numpy as np
target_depth_levels=[1,2,3,4,5,7,9,10,12,14,15,16,18,20,26,30,36,40,50,80] # Specify depth levels of interest
# Load ROMS output using your pattern
roms_output = ROMSOutput(
    grid=grid,
    path=[
        model_data_path,
    ],
    use_dask=True,
)

ds = roms_output.regrid(var_names=["u", "v"],depth_levels=target_depth_levels)

u_rg=ds['u'].mean('time').load()
v_rg=ds['v'].mean('time').load()

import cartopy.crs as ccrs
import cartopy.mpl.ticker as cticker
import matplotlib.pyplot as plt
import numpy as np

# -------------------------------------------------------
#  Select depths and prepare surface/20 m vector plotting
# -------------------------------------------------------
surface_depth = (slice(0,2))
depth20 = 20

u_surf = u_rg.sel(depth=surface_depth).mean('depth')
v_surf = v_rg.sel(depth=surface_depth).mean('depth')

u_20 = u_rg.sel(depth=depth20, method="nearest")
v_20 = v_rg.sel(depth=depth20, method="nearest")

lon = u_rg.lon
lat = u_rg.lat

# -------------------------------------------------------
# 1) Compute full-resolution land mask (before thinning)
# -------------------------------------------------------
mask_surf_full = np.isnan(u_surf) | np.isnan(v_surf)
mask_200_full  = np.isnan(u_20)   | np.isnan(v_20)

# -------------------------------------------------------
# 2) Thin vectors
# -------------------------------------------------------
step = 10
lon_thin = lon[::step]
lat_thin = lat[::step]

u_surf_thin = u_surf[::step, ::step]
v_surf_thin = v_surf[::step, ::step]

u_200_thin = u_20[::step, ::step]
v_200_thin = v_20[::step, ::step]

# -------------------------------------------------------
#  Data for vertical sections
# -------------------------------------------------------
lonA = 360 - 21.8
lonB = 360 - 21.6

uA = u_rg.sel(lon=lonA, method="nearest")
if "time" in uA.dims:
    uA = uA.isel(time=0)
uA = uA.transpose("depth", "lat")

uB = u_rg.sel(lon=lonB, method="nearest")
if "time" in uB.dims:
    uB = uB.isel(time=0)
uB = uB.transpose("depth", "lat")

depth_vals = uA.depth.values
lat_vals_A = uA.lat.values
lat_vals_B = uB.lat.values

extent_hval = [-22.05, -21.3469394652069, 
               64.2, 64.42]

# -------------------------------------------------------
# Plotting using GridSpec
# -------------------------------------------------------
fig = plt.figure(figsize=(14, 10))
gs = fig.add_gridspec(2, 2, height_ratios=[1, 1])

# ===========================
# Panel 1 — Surface vectors
# ===========================
ax1 = fig.add_subplot(gs[0, 0], projection=ccrs.Mercator())
ax1.set_extent(extent_hval)

ax1.text(
    0.02, 0.98, "a)",
    transform=ax1.transAxes,
    fontsize=14,
    fontweight="bold",
    va="top",
    ha="left"
)


# Mask
ax1.contourf(
    lon, lat, mask_surf_full,
    levels=[0.5, 1.5],
    colors=["lightgray"],
    transform=ccrs.PlateCarree(),
    zorder=1
)

# Vectors (store handle)
Q1 = ax1.quiver(
    lon_thin.values,
    lat_thin.values,
    u_surf_thin.values,
    v_surf_thin.values,
    transform=ccrs.PlateCarree(),
    scale=3,
    zorder=2
)

# --- NEW: reference arrow using quiverkey ---
ax1.quiverkey(
    Q1, 
    X=0.12, Y=0.92,  # position within axes (0–1)
    U=0.1,           # reference speed in m/s
    label="0.1 m/s",
    labelpos="E",
    coordinates='axes'
)

# --- vertical slice lines ---
ax1.plot([lonA-360, lonA-360], [lat.min(), lat.max()], 
         color='black', alpha=0.6, linewidth=2, transform=ccrs.PlateCarree())
ax1.plot([lonB-360, lonB-360], [lat.min(), lat.max()], 
         color='black', alpha=0.6, linewidth=2, transform=ccrs.PlateCarree())

ax1.set_title("Surface velocity vectors (thinned)")
...
gl1 = ax1.gridlines(
    crs=ccrs.PlateCarree(),
    draw_labels=True,
    linewidth=0.8,
    color="gray",
    alpha=0.6,
    linestyle="--"
)

gl1.top_labels = False
gl1.right_labels = False
gl1.xlabel_style = {"size": 11}
gl1.ylabel_style = {"size": 11}

gl1.xformatter = cticker.LongitudeFormatter()
gl1.yformatter = cticker.LatitudeFormatter()

# ===========================
# Panel 2 — 20 m vectors
# ===========================
ax2 = fig.add_subplot(gs[0, 1], projection=ccrs.Mercator())
ax2.set_extent(extent_hval)

ax2.text(
    0.02, 0.98, "b)",
    transform=ax2.transAxes,
    fontsize=14,
    fontweight="bold",
    va="top",
    ha="left"
)


# Mask
ax2.contourf(
    lon, lat, mask_200_full,
    levels=[0.5, 1.5],
    colors=["lightgray"],
    transform=ccrs.PlateCarree(),
    zorder=1
)

# Vectors (store handle)
Q2 = ax2.quiver(
    lon_thin.values, lat_thin.values,
    u_200_thin.values, v_200_thin.values,
    color="black",
    transform=ccrs.PlateCarree(),
    zorder=2, scale=3
)

# --- NEW: reference arrow ---
ax2.quiverkey(
    Q2, 
    X=0.12, Y=0.92,
    U=0.1, 
    label="0.1 m/s",
    labelpos="E",
    coordinates='axes'
)

ax2.set_title("20 m velocity vectors (thinned)")
...
gl2 = ax2.gridlines(
    crs=ccrs.PlateCarree(),
    draw_labels=True,
    linewidth=0.8,
    color="gray",
    alpha=0.6,
    linestyle="--"
)

gl2.top_labels = False
gl2.right_labels = False
gl2.xlabel_style = {"size": 11}
gl2.ylabel_style = {"size": 11}

gl2.xformatter = cticker.LongitudeFormatter()
gl2.yformatter = cticker.LatitudeFormatter()


# ===========================
# Panel 3 — Vertical section at A
# ===========================
ax3 = fig.add_subplot(gs[1, 0])
ax3.text(
    0.02, 0.98, "c)",
    transform=ax3.transAxes,
    fontsize=14,
    fontweight="bold",
    va="top",
    ha="left"
)

cf = ax3.contourf(lat_vals_A, depth_vals, uA, levels=np.arange(-0.158,0.158,0.004),cmap='RdBu_r')
ax3.set_title(f"Zonal velocity u(z, lat) at {lonA-360:.2f}°E")

fig.colorbar(cf, ax=ax3, label="u (m/s)")

ax3.set_ylim(0, 40)
ax3.invert_yaxis()
ax3.set_xlim(64.3, 64.34)
ax3.set_xlabel("Latitude", fontsize=11)
ax3.set_ylabel("Depth (m)", fontsize=11)
ax3.tick_params(axis='both', labelsize=11)
# ===========================
# Panel 4 — Vertical section at B
# ===========================
ax4 = fig.add_subplot(gs[1, 1])
ax4.text(
    0.02, 0.98, "d)",
    transform=ax4.transAxes,
    fontsize=14,
    fontweight="bold",
    va="top",
    ha="left"
)

cf = ax4.contourf(lat_vals_B, depth_vals, uB, levels=np.arange(-0.114,0.114,0.004), cmap='RdBu_r')
ax4.set_title(f"Zonal velocity u(z, lat) at {lonB-360:.2f}°E")

fig.colorbar(cf, ax=ax4, label="u (m/s)")
ax4.set_xlabel("Latitude", fontsize=11)
ax4.set_ylabel("Depth (m)", fontsize=11)
ax4.tick_params(axis='both', labelsize=11)

ax4.set_ylim(0, 40)
ax4.invert_yaxis()
ax4.set_xlim(64.36, 64.41)

plt.tight_layout()
plt.show()