# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Import packages
# OceanSpy is documented here [readthedocs](https://oceanspy.readthedocs.io/en/latest/) and here [paper](https://joss.theoj.org/papers/10.21105/joss.01506).

# +
# Import OceanSpy
import oceanspy as ospy

# Additional imports
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from cartopy.crs import PlateCarree
import scipy.io as sio
import os

# Just to make the notebook cleaner for the symposium.
# DO NOT USE IT OTHERWISE!
import warnings

warnings.filterwarnings("ignore")
# -

# # Download data
#
# For this illustration a small cutout of the full model fields is downloaded, plus some observations from the real ocean.

data_dir = "IDIES_symposium"
if not os.path.isdir(data_dir):
    # Download data
    import subprocess

    print("Downloading and uncompressing get_started data...")
    print("...it might take a couple of minutes.")
    commands = [
        "wget -v -O IDIES_symposium.tar.gz -L "
        "https://livejohnshopkins-my.sharepoint.com/"
        ":u:/g/personal/malmans2_jh_edu/"
        "ETCLiwVHdN9d4DNKA_zoMI8BkaiNQY2BGjZpLbMBrMGUSQ?"
        "download=1",
        "tar xvzf IDIES_symposium.tar.gz",
        "rm -f IDIES_symposium.tar.gz",
    ]
    subprocess.call("&&".join(commands), shell=True)
    subprocess.call("&&".join(commands), shell=True)

# # Set parameters for plotting

# +
# Coordinates of the sections
X_Latrabjarg = [-29.8, -24.6]
Y_Latrabjarg = [66.8, 65.5]
map_extent = [-39.9, -19, 63, 67.75]
ylim = (-700, 0)

# Colors
land_col = (253 / 255, 180 / 255, 108 / 255)
bottom_col = "grey"
cmap_vel = "RdBu_r"
cmap_hyd = "rainbow"
cmap_N2 = "PuBuGn_r"

# Level for contours (commonly used by observational oceanographers)
lev_vel = np.arange(-0.5, 0.51, 0.05)
lev_Temp = [
    -1,
    -0.25,
    0,
    0.25,
    0.5,
    0.75,
    1,
    1.5,
    2,
    2.5,
    3,
    3.5,
    4,
    4.5,
    5,
    5.5,
    6,
    7,
    8,
]
lev_S = [
    31,
    33,
    33.5,
    34,
    34.2,
    34.4,
    34.5,
    34.6,
    34.7,
    34.75,
    34.8,
    34.85,
    34.9,
    34.92,
    34.96,
    35,
    35.05,
    35.1,
    35.2,
]
lev_Sigma0 = [26.5, 27, 27.25, 27.5, 27.6, 27.7, 27.8, 27.9, 27.95, 28]
lev_N2 = [0, 1.0e-6, 2.0e-6, 5.0e-6, 1.0e-5, 2.0e-5, 5.0e-5, 1.0e-4, 5.0e-4]

# Plot parameters
mpl.rcParams.update({"font.size": 22})
clabel_size = 15
# -

# # Open datasets

# +
# Open model dataset
od_vage2011 = ospy.open_oceandataset.from_zarr(
    os.path.join(data_dir, "vage2011")
)
od_mastropole2017 = ospy.open_oceandataset.from_zarr(
    os.path.join(data_dir, "mastropole2017")
)

# Open observations
vage2011 = sio.loadmat(
    os.path.join(data_dir, "observations/Velocities_VageEtAl2011.mat")
)
mastropole2017 = sio.loadmat(
    os.path.join(data_dir, "observations/meanFields_MastrpoleEtAl2017.mat")
)
bathyobs = sio.loadmat(os.path.join(data_dir, "observations/LatBat.mat"))
# -

# # Show datasets
# Full metadata are included in the numerical model output.

print(
    "Here is the oceandataset to compare with Våge et al., 2011 - Nature Geoscience:\n"
)
print(od_vage2011)
print("\n\nAvailable variables:")
print(od_vage2011.dataset)

print(
    "Here is the oceandataset to compare with Mastropole et al., 2017 - Journal of Physical Oceanography:\n"
)
print(od_mastropole2017)
print("\n\nAvailable variables:")
print(od_mastropole2017.dataset)

# # Plot map

# +
# Mask land
Depth = od_vage2011.dataset["Depth"]
Depth = Depth.where(Depth > 0)
od_vage2011 = od_vage2011.merge_into_oceandataset(Depth.rename("masked_Depth"))

# Plot Depth
fig = plt.figure(figsize=(10, 5))
ax = od_vage2011.plot.horizontal_section(varName="masked_Depth", cmap="bone_r")
line = ax.plot(
    X_Latrabjarg, Y_Latrabjarg, "m-", transform=PlateCarree(), linewidth=5
)
ax.background_patch.set_facecolor(land_col)
ax.set_extent(map_extent)
# -

# # Survey
# Sample the model like the ship sampled the real ocean. This is called "survey" mode. See: [here](https://oceanspy.readthedocs.io/en/latest/generated/oceanspy.subsample.survey_stations.html#oceanspy.subsample.survey_stations).

# Extract transect
od_vage2011_LL = od_vage2011.subsample.survey_stations(
    Xsurv=X_Latrabjarg, Ysurv=Y_Latrabjarg, delta=2
)
print()
od_mastropole2017_LL = od_mastropole2017.subsample.survey_stations(
    Xsurv=X_Latrabjarg, Ysurv=Y_Latrabjarg, delta=2
)

# # Use OceanSpy to compute density and $N^2$
# Density is computed using the seawater equation of state $\rho = \rho (T, S, p)$ and $\sigma_0 = \rho(T, S, 0) -1000$kgm$^{-3}$. See the [docs](https://oceanspy.readthedocs.io/en/latest/generated/oceanspy.compute.potential_density_anomaly.html#oceanspy.compute.potential_density_anomaly).
#
# The Brunt-Vaisala frequency $N$ measures the vertical density stratification:
# \begin{equation}
# N^2 = -\frac{g}{\rho_0}\frac{\partial\sigma_0}{\partial z} .
# \end{equation}
# See the [docs](https://oceanspy.readthedocs.io/en/latest/generated/oceanspy.compute.Brunt_Vaisala_frequency.html).

# Compute density
od_mastropole2017_LL = od_mastropole2017_LL.compute.potential_density_anomaly()
od_mastropole2017_LL = od_mastropole2017_LL.compute.Brunt_Vaisala_frequency()

# # Plot hydrography
# The model fields closely resemble the field data. There are some difference, as discussed in [Almansi et al. (2017)](https://journals.ametsoc.org/doi/full/10.1175/JPO-D-17-0129.1), but they're likely inconsequential.

# +
# Same as before:
# All metadata for the model are handled by OceanSpy and xarray.
# Need to be added manually for observations.

# Initialize figure
xlim = (1, 225)
fig, axes = plt.subplots(3, 2, figsize=(20, 15), sharex=True, sharey=True)
for a, ax in enumerate(axes.flatten()):
    # Pick variable
    if a in [0, 1]:
        var_mod = "Temp"
        var_obs = "THE"
        cmap = cmap_hyd
        levels = lev_Temp
        fmt_cbar = "%.2f"
    elif a in [2, 3]:
        var_mod = "S"
        var_obs = "SAL"
        cmap = cmap_hyd
        levels = lev_S
        fmt_cbar = "%.2f"
    elif a in [4, 5]:
        var_mod = "N2"
        var_obs = "NSQ"
        cmap = cmap_N2
        levels = lev_N2
        fmt_cbar = "%.0e"

    # Isopycnals properties
    contour_kwargs = {
        "levels": lev_Sigma0,
        "colors": ["k" if (lev != 27.8) else "m" for lev in lev_Sigma0],
        "linewidths": [2 if (lev != 27.8) else 3 for lev in lev_Sigma0],
    }
    clabel_kwargs = {"fontsize": clabel_size, "fmt": "%2.2f"}

    if a % 2 == 0:
        # Observations
        cols = mpl.colors.BoundaryNorm(levels, 256)
        contf = ax.contourf(
            mastropole2017["X"],
            -mastropole2017["Y"].astype("float"),
            mastropole2017[var_obs],
            cmap=cmap,
            levels=levels,
            norm=cols,
        )
        cont = ax.contour(
            mastropole2017["X"],
            -mastropole2017["Y"].astype("float"),
            mastropole2017["SIG"],
            **contour_kwargs,
        )
        plt.colorbar(contf, ax=ax, format=fmt_cbar)
        ax.clabel(cont, **clabel_kwargs)
        ax.fill_between(
            np.squeeze(bathyobs["regdist"]),
            -np.squeeze(bathyobs["regbat"]),
            -np.squeeze(bathyobs["regbat"]) * 0 - 700,
            color=bottom_col,
            zorder=10,
        )

    else:
        # Model
        ax = od_mastropole2017_LL.plot.vertical_section(
            varName=var_mod,
            contourName="Sigma0",
            plotType="contourf",
            levels=levels,
            ax=ax,
            xlim=xlim,
            cmap=cmap,
            contour_kwargs=contour_kwargs,
            clabel_kwargs=clabel_kwargs,
            cbar_kwargs={"format": fmt_cbar},
        )
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_facecolor(bottom_col)
    if a == 0:
        ax.set_title("Observations")
    elif a == 1:
        ax.set_title("Model")
    else:
        ax.set_title(None)
# -

# # Use OceanSpy to rotate velocities
# We want to plot the velocity component perpendicular to the survey section, like in the observations.

# Compute aligned velocities
od_vage2011_LL = od_vage2011_LL.compute.survey_aligned_velocities()

# # Plot velocities
# Again, the model velocity field is realistic. The differences between the measurements and the model results are likely due to different chaotic fluctuations (geostrophic turbulence), not systematic model bias. We have not attempted to match the phases of the chaotic fluctuations in this experiment.

# +
# All metadata for the model are handled by OceanSpy and xarray.
# Need to be added manually for observations.

# Plot
xlim = (0, 175)
fig, axes = plt.subplots(2, 1, figsize=(10, 10), sharex=True)

# Observations
ax = axes[0]
contf = ax.contourf(
    vage2011["X"],
    -vage2011["Y"].astype("float"),
    -vage2011["OrtVel"] * 1.0e-2,
    levels=lev_vel,
    cmap=cmap_vel,
    extend="both",
)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
plt.colorbar(contf, ax=ax)
_ = ax.set_facecolor(bottom_col)
_ = ax.set_title("Observations")

# Model
ax = axes[1]
ax = od_vage2011_LL.plot.vertical_section(
    varName="ort_Vel",
    plotType="contourf",
    levels=lev_vel,
    ax=ax,
    xlim=xlim,
    cmap=cmap_vel,
)
ax.set_ylim(ylim)
_ = ax.set_facecolor(bottom_col)
_ = ax.set_title("Model")
# -

# # Compute Ertel Potential Vorticity
#
# Ertel PV is computed as
# \begin{equation}
# Q =
# - \frac{\omega \cdot \nabla \rho}{\rho} =
# (f + \zeta)\frac{N^2}{g}
# + \frac{\left(\zeta_H+e\hat{\mathbf{y}}\right)
# \cdot\nabla_H\rho}{\rho_0}
# \end{equation}
#
# This is the single most important variable in dynamical oceanography, yet it's almost impossible to measure in the real ocean. Using the ultra-high resolution numerical model it's easily computed (see: [here](https://oceanspy.readthedocs.io/en/latest/generated/oceanspy.compute.Ertel_potential_vorticity.html)).
#
# We compute and plot the average $Q$ between 100 and 150m depth.

# +
# Compute Ertel PV
od_vage2011_surf = od_vage2011.subsample.cutout(ZRange=[-100, -150])
od_vage2011_surf = od_vage2011_surf.compute.Ertel_potential_vorticity()

# Plot Ertel PV
fig = plt.figure(figsize=(10, 5))
od_vage2011_surf.plot.horizontal_section(
    varName="Ertel_PV",
    contourName="Depth",
    meanAxes="Z",
    center=False,
    robust=True,
)
# -



StarcoderdataJupyterScriptsDedupFiltered