Examples

Here, we apply neutralocean to various ocean models using different horizontal grid structures.

OCCA (latitude-longitude)

OCCA is an ocean atlas on a latitude-longitude grid. The run_OCCA.py example script (below, and on github) shows more of neutralocean’s capabilities. It downloads the dataset, loads the data into memory, selects an equation of state, calculates various approximately neutral surfaces, and more.

# In[Imports]

# Functions to make the Equation of State
from neutralocean.eos import make_eos, make_eos_s_t

# Functions to compute various approximately neutral surfaces
from neutralocean.surface import potential_surf, anomaly_surf, omega_surf

# Functions to load OCCA data
from neutralocean.examples.load_OCCA import load_OCCA

from neutralocean.grid.rectilinear import build_grid, edgedata_to_maps

# In[Load OCCA data]

g, S, T = load_OCCA()  # S, T arranged as (Longitude, Latitude, Depth)
ni, nj, nk = S.shape
Z = -g["RC"]  # Depth vector (note positive and increasing down)

# Select pinning cast in the equatorial Pacific.
# When these are used for 'pin_cast' and 'pin_p', the following surfaces will
# intersect cast (i0,j0) at a depth of z0.
i0, j0 = 220, 80
z0 = 1500.0

# Select vertical interpolation method. Options are "linear" or "pchip"
interp_name = "linear"

# make Boussinesq version of the Jackett and McDougall (1995) equation of state
# --- which is what OCCA used --- and its partial derivatives
eos = make_eos("jmd95", g["grav"], g["rho_c"])
eos_s_t = make_eos_s_t("jmd95", g["grav"], g["rho_c"])

# Build grid adjacency and distance information for neutralocean functions
grid = build_grid(
    (ni, nj), g["wrap"], g["DXCvec"], g["DYCsc"], g["DYGsc"], g["DXGvec"]
)

# Prepare some default options for potential_surf, anomaly_surf, and omega_surf
opts = {}
opts["grid"] = grid
opts["interp"] = interp_name
opts["vert_dim"] = "Depth_c"
opts["eos"] = (eos, eos_s_t)

# In[Potential Density surfaces]

# Provide reference pressure (actually depth, in Boussinesq) and isovalue.
args = opts.copy()
args["ref"] = 0.0
args["isoval"] = 1027.5
s, t, z, d = potential_surf(S, T, Z, **args)
print(
    f" ** The potential density surface (referenced to {d['ref']}m)"
    f" with isovalue = {d['isoval']}kg m-3"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# Provide pin_cast and pin_p: the reference location and depth that the surface intersects
args = opts.copy()
args["ref"] = 0.0
args["pin_cast"] = (i0, j0)
args["pin_p"] = z0
s, t, z, d = potential_surf(S, T, Z, **args)
print(
    f" ** The potential density surface (referenced to {d['ref']}m)"
    f" intersecting the cast indexed by {(i0,j0)} at depth {z0}m"
    f" (isovalue = {d['isoval']}kg m-3)"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# Provide just the location to intersect `(pin_cast, pin_p)`.
# This takes the reference depth `ref` to match `pin_p`.
# Also illustrate using xarray coordinates for pin_cast.
# Also show how to just give the EOS name and the necessary parameters (g and rho_c)
# for its Boussinesq version, rather than using the pre-made EOS's as above.
args = opts.copy()
args["pin_cast"] = {"Longitude_t": 220.5, "Latitude_t": 0.5}
args["pin_p"] = z0
args["eos"] = "jmd95"
args["grav"] = g["grav"]
args["rho_c"] = g["rho_c"]
s, t, z, d = potential_surf(S, T, Z, **args)
assert z.values[i0, j0] == z0  # check pin_cast was indeed (i0,j0)
z_sigma = z  # save for later
print(
    f" ** The potential density surface (referenced to {d['ref']}m)"
    f" intersecting the cast at (180.5 E, 0.5 N) at depth {z0}m"
    f" (isovalue = {d['isoval']}kg m-3)"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# In[Delta surfaces]

# Provide reference salinity and potential temperature values
s0, t0 = 34.5, 4.0

args = opts.copy()
args["ref"] = (s0, t0)
args["isoval"] = 0.0
s, t, z, d = anomaly_surf(S, T, Z, **args)
print(
    f" ** The in-situ density anomaly surface (referenced to {d['ref']})"
    f" with isovalue = {d['isoval']}kg m-3"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# Provide pin_cast and pin_p: the reference location and depth that the surface intersects
args = opts.copy()
args["ref"] = (s0, t0)
args["pin_cast"] = (i0, j0)
args["pin_p"] = z0
s, t, z, d = anomaly_surf(S, T, Z, **args)
print(
    f" ** The in-situ density anomaly surface (referenced to {d['ref']})"
    f" intersecting the cast indexed by {(i0,j0)} at depth {z0}m"
    f" (isovalue = {d['isoval']}kg m-3)"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# Provide just the location to intersect: depth `pin_p` on cast `pin_cast`
# This takes the reference S and T values from that location.
args = opts.copy()
args["pin_cast"] = (i0, j0)
args["pin_p"] = z0
s, t, z, d = anomaly_surf(S, T, Z, **args)
z_delta = z  # save for later
print(
    f" ** The in-situ density anomaly surface (referenced to {d['ref']})"
    f" intersecting the cast indexed by {(i0,j0)} at depth {z0}m"
    f" (isovalue = {d['isoval']}kg m-3)"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# In[Omega surfaces]

# Initialize omega surface with a (locally referenced) potential density surface.
# Provide grid distances. By default, does as many iterations as needed to get
# the root-mean-square change in the locally referenced potential density from
# one iteration to the next to drop below 10^-7 kg m-3, or a maximum of 10
# iterations.
args = opts.copy()
args["pin_cast"] = (i0, j0)
args["pin_p"] = z0
s, t, z, d = omega_surf(S, T, Z, **args)
print(
    f" ** The omega-surface"
    f" initialized from a potential density surface (referenced to {z0}m)"
    f" intersecting the cast indexed by {(i0,j0)} at depth {z0}m"
    f" has root-mean-square ϵ neutrality error {d['e_RMS'][-1]} kg m-4"
)

# Initialize omega surface with a (locally referenced) in-situ density anomaly surface.
#   see `mixed_layer` for details on these parameters.
# Specify a higher max number of iterations, and quit iterations when the
# depth change from one iteration to the next has a root-mean-square value less
# than 1e-7 m.
args["ref"] = (None, None)  # init surface is locally referenced anomaly_surf
args["p_ml"] = {"bottle_index": 1, "ref_p": 0.0}
args["ITER_MAX"] = 20
args["TOL_P_CHANGE_RMS"] = 1e-6
args["TOL_LRPD_MAV"] = 0.0  # deactivate this exit threshold
args["TOL_P_SOLVER"] = 1e-7
s, t, z, d = omega_surf(S, T, Z, **args)
s_omega, t_omega, z_omega = s, t, z  # save for later
args_omega = args.copy()  # save for later
print(
    f" ** The omega-surface"
    f" initialized from an in-situ density anomaly surface (referenced locally to cast {(i0,j0)} at {z0}m)"
    f" intersecting the cast indexed by {(i0,j0)} at depth {z0}m"
    f" has root-mean-square ϵ neutrality error {d['e_RMS'][-1]} kg m-4"
)

# Initialize omega surface with a pre-computed surface.
# In this case, let's continue from the above omega surface, but change the
# pinning location. Since the above omega surface is very nearly converged,
# this omega surface should basically match the above one. Indeed, it will
# do one iteration and find it is below the RMS depth change tolerance, and exit.
i1, j1 = 315, 110  # North Atlantic
args.pop("pin_p", None)
args.pop("ref", None)
args["pin_cast"] = (i1, j1)
args["p_init"] = z_omega  # set init surface from above
args["ITER_START_WETTING"] = 99  # greater than ITER_MAX, so no wetting
s, t, z, d = omega_surf(S, T, Z, **args)
print(
    f" ** The omega-surface"
    f" initialized from an omega surface and pinned to the cast {(i1,j1)} at {z0}m)"
    f" has root-mean-square ϵ neutrality error {d['e_RMS'][-1]} kg m-4"
)


# In[Begin showing more advanced features]

import numpy as np
import numba as nb
from neutralocean.mixed_layer import mixed_layer
from neutralocean.ntp import ntp_epsilon_errors, ntp_epsilon_errors_norms
from neutralocean.label import veronis
from neutralocean.lib import _process_casts, xr_to_np
from neutralocean.ppinterp import make_pp
from neutralocean.eos import make_eos_p, vectorize_eos
from neutralocean.traj import ntp_bottle_to_cast, neutral_trajectory

# In[Show vertical interpolation]

# Build interpolation function using same interpolation kernel ("linear") as
# omega surface we saved above.
# Allow this to operate across all water columns with one call (kind="u").
# Evaluate the interpolant rather than build its coefficients (out="interp").
# Allow that there could be NaNs in the data (nans=True).
# There will be two numpy arrays of dependent data (S, T), sharing the same
# independent data Z (num_dep_vars=2).
interp_two = make_pp(
    interp_name, kind="u", out="interp", nans=True, num_dep_vars=2
)

# Apply interpolation function to interpolate salinity and temperature onto the
# depth of the surface.  This requires working with numpy arrays, not xarrays.
# Evaluate the interpolant, not any of its derivatives (d=0).
s_, t_ = interp_two(z_omega.values, Z, S.values, T.values, 0)


# Check that the results of the above interpolation match (to machine precision)
# the results returned from omega_surf above.
s_check = np.allclose(s_omega.values, s_, atol=1e-15, equal_nan=True)
t_check = np.allclose(t_omega.values, t_, atol=1e-15, equal_nan=True)
if not (s_check and t_check):
    print(
        "Something's wrong; should be able to reconstruct salinity and "
        "temperature on surface by interpolating to surface's depth"
    )


# In[Approximate b, the integrating factor, from a pair of omega surfaces]


def calc_lrpd_z(z):
    """
    Calculate ∂(Locally Referenced Potential Density)/dz on the surface `z`,
    lrpd_z = (ρ_S ∂S/∂z + ρ_Θ ∂Θ/∂z).
    """
    z, Snp, Tnp, Znp = (xr_to_np(x) for x in (z, S, T, Z))
    s, t = interp_two(z, Znp, Snp, Tnp, 0)
    ds, dt = interp_two(z, Znp, Snp, Tnp, 1)
    rs, rt = eos_s_t(s, t, z)
    lrpd_z = rs * ds + rt * dt
    return lrpd_z


def calc_b(z, z2, I0):
    """
    Estimate the integrating factor, b, from a pair of surfaces.

    This estimates b from the vertical component of the neutral relation
    ∇γ = b N, which is approximately true for highly accurate approximately
    neutral surfaces such as omega surfaces, where N = ρ_S ∇S + ρ_Θ ∇Θ is the
    dianeutral vector.
    Thus, we calculate b as
      b = (∂γ/∂z) / Nz,
    where Nz = (ρ_S ∂S/∂z + ρ_Θ ∂Θ/∂z)
    and with the approximation
      (∂γ/∂z) = Δγ / Δz = Δγ / (z2 - z),
    where Δγ is the γ density label on the `z2` surface minus that on the `z` surface.
    We determine Δγ by imposing b = 1 at the reference cast indexed by `I0`, so
      b = (z2[I0] - z[I0]) * Nz[I0] / ((z2 - z) * Nz)
    """

    dz = z2 - z
    lrpd_z = calc_lrpd_z(z)
    b = (lrpd_z[I0] * dz[I0]) / (lrpd_z * dz)
    return b


def same_valid_casts(z, z2):
    z, z2 = (xr_to_np(x) for x in (z, z2))
    return np.all(np.isnan(z) == np.isnan(z2))


# Calculate 2nd omega surface that's `ztiny` below the previous one at `pin_cast`.
# Use the same options as for the previous omega surface, but turn off wetting
# to help ensure the two omega surfaces will have exactly the same wet casts.
# Also, using highly converged omega surfaces (eg small TOL_P_CHANGE_RMS) helps
# the b estimate be smooth.
ztiny = 1e-2
if True:
    # Heave the omega surface rigidly down by an amount `ztiny`
    z_init = z_omega + ztiny
else:
    # Heave the omega surface rigidly if LRPD were the vertical coordinate, so
    # as to land at a depth `ztiny` deeper at the reference cast.
    # This is slightly better than the above, but either way, omega_surf will
    # converge to the same surface
    lrpd_z = calc_lrpd_z(z_omega)
    z_init = z_omega + (ztiny * lrpd_z[i0, j0]) / lrpd_z

args = args_omega.copy()
args.pop("pin_p", None)
args["p_init"] = z_init
args["ITER_START_WETTING"] = 99  # greater than ITER_MAX, so no wetting
s, t, z, d = omega_surf(S, T, Z, **args)
z_omega_pair = z

if not same_valid_casts(z_omega, z_omega_pair):
    print(
        "Warning: pair of surfaces has different valid casts, "
        "so the estimate of b cannot be trusted. "
        "Make sure wetting is turned off for the second surface, and try "
        "reducing the depth difference between the pair of surfaces."
    )

b = calc_b(z_omega, z_omega_pair, (i0, j0))
print(
    "The integrating factor b for the omega surface varies between "
    f"{np.nanmin(b.values)} and {np.nanmax(b.values)}."
)

# In[Veronis Density, used to label an approx neutral surface]
S_ref_cast = S.values[i0, j0]
T_ref_cast = T.values[i0, j0]
rho_v = veronis(z0, S_ref_cast, T_ref_cast, Z, eos="jmd95")
print(
    f"A surface through the cast indexed by {(i0,j0)} at depth {z0}m"
    f" has Veronis density {rho_v} kg m-3"
)

# In[Neutral Trajectory]

# neutral trajectory depth along meridional section
znt = np.full(nj, np.nan)

# Build a neutral trajectory going northwards from cast at (i0, j0), starting
# at depth of z0. Put result into `znt`
Snt = S[i0, j0:nj, :]
Tnt = T[i0, j0:nj, :]
s_, t_, z_ = neutral_trajectory(Snt, Tnt, Z, z0, eos=eos)
znt[j0:nj] = z_

# As above, but go southwards.
Snt = S[i0, j0:0:-1, :]
Tnt = T[i0, j0:0:-1, :]
s_, t_, z_ = neutral_trajectory(Snt, Tnt, Z, z0, eos=eos)
znt[j0:0:-1] = z_

# Consider three other trajectories along this meridional section, following ANSs
zom = z_omega[i0, :].values  # omega
zpd = z_sigma[i0, :].values  # potential density
zda = z_delta[i0, :].values  # density anomaly

# Find index to furthest south cast that has all four trajectories
j_south = np.argmax(
    np.all(
        np.isfinite(np.array((znt, zom, zpd, zda))),
        axis=0,
    )
)
print(
    f"From (i,j)={(i0,j0)} at {z0}m moving south to j={j_south}, ...\n"
    f"Neutral Trajectory              ends at {znt[j_south]:8.4f}m depth,\n"
    f"Omega Surface                   ends at {zom[j_south]:8.4f}m depth,\n"
    f"Potential Density Surface       ends at {zpd[j_south]:8.4f}m depth,\n"
    f"In-situ density anomaly surface ends at {zda[j_south]:8.4f}m depth."
)

print(
    "Uncomment below to plot neutral trajectory and ANSs along meridional section"
)
# import matplotlib.pyplot as plt
# lat = g["YCvec"]  # latitudes
# fig, ax = plt.subplots()
# ax.plot(lat, -zom, label="omega")
# ax.plot(lat, -zpd, label="potential density")
# ax.plot(lat, -zda, label="in-situ density anomaly")
# ax.plot(lat, -znt, "--k", label="neutral trajectory")
# ax.scatter(lat[j0], -z0)  # reference latitude and depth.
# ax.legend()

# In[Remove mixed layer from an omega surface]

# Pre-compute depth of the mixed layer
z_ml = mixed_layer(S, T, Z, eos)

# The correct way to remove the mixed layer from the omega surface algorithm
# is to pass the p_ml parameter, e.g. pass `p_ml=z_ml` in the call below.
# However, here we'll use a more drastic approach: setting our 3D hydrographic
# data to NaN above the mixed layer.
# This serves as a test of the new interpolation methods that allow for NaN data
# in the top of the water column followed by valid (non-NaN) data in the middle
# where the interpolation is done, followed by more NaN data at the bottom.
# This is to allow for models or data with ice shelves.
# Note, this drastic NaN'ing of the input (S,T) data is not expected to give
# exactly the same results as by passing the p_ml parameter, since
# (a) with p_ml, the mixed layer removal is not applied until after the first
#     iteration, and
# (b) with PCHIPs, the interpolant just below the mixed layer is affected by
#     a few of the lower bottles in the mixed layer.
# (c) with p_ml, the surface is set to NaN only if it goes above p_ml which
#     is any real number, whereas with the data-NaN'ing approach, discrete
#     bottles are set to NaN so the surface will only be valid if it is below
#     the first valid bottle below what was NaN'ed out.


@nb.njit
def mixedlayer2nan(S, T, Z, z_ml):
    for n in np.ndindex(z_ml.shape):
        z = z_ml[n]
        for k in range(nk):
            if Z[k] < z:
                S[(*n, k)] = np.nan
                T[(*n, k)] = np.nan
            else:
                break


mixedlayer2nan(S.values, T.values, Z, z_ml)

args = opts.copy()
args["ref"] = (None, None)
args["pin_cast"] = (i0, j0)
args["pin_p"] = z0
args["interp"] = "pchip"
args["ITER_MAX"] = 10
args["ITER_START_WETTING"] = 1
args["TOL_P_SOLVER"] = 1e-5
s, t, z, d = omega_surf(S, T, Z, **args)

# potential and anomaly surfaces don't take `p_ml` as an argument, since they
# aren't iterative algorithms, so we can just remove it manually, e.g.
# z[z < z_ml] = np.nan

# In[Neutrality errors on a surface]
e_RMS, e_MAV = ntp_epsilon_errors_norms(s, t, z_omega, grid, eos_s_t)
print(f"RMS of ϵ is {e_RMS : 4e} [kg m-4])")

# Calculate ϵ neutrality errors on all pairs of adjacent water columns
e = ntp_epsilon_errors(s, t, z, grid, eos_s_t)

# Convert ϵ above into two 2D maps, one for zonal ϵ errors and one for meridional ϵ errors
ex, ey = edgedata_to_maps(e, (ni, nj), g["wrap"])
# These can then be mapped...

# In[Neutral Tangent Plane bottle to cast]

sB, tB, zB = 35.0, 16.0, 500.0  # Thermodynamic properties of a given Bottle
S1 = S.values[180, 80, :]
T1 = T.values[180, 80, :]
s1, t1, z1 = ntp_bottle_to_cast(sB, tB, zB, S1, T1, Z)

# In[Work with Numpy arrays instead of xarrays]

# Convert S and T from xarray to numpy ndarrays, and make vertical dimension
# contiguous in memory.  If not done here in advance, this will be done each
# time potential_surf, anomaly_surf, or omega_surf is called.  Hence, this is
# the recommended method if many approx neutral surfaces will be calculated.
Snp, Tnp, Znp = _process_casts(S, T, Z, "Depth_c")

s, t, z, d = anomaly_surf(
    Snp,
    Tnp,
    Znp,
    grid=grid,
    eos=(eos, eos_s_t),
    vert_dim=-1,
    ref=(s0, t0),
    isoval=0.0,
)


# In[Calculate a large-scale potential vorticity on the above specvol anomaly surface]

# Create function for partial deriv of equation of state with respect to depth z
eos_z = make_eos_p("jmd95", g["grav"], g["rho_c"])  # for scalar inputs
eos_z = vectorize_eos(eos_z)  # for nd inputs

# Build linear interpolation function, operating over a whole dataset (kind="u"),
# evaluating the interpolant on the fly (out="interp"), for two dependent
# variables (num_dep_vars=2).
interp_two = make_pp("linear", kind="u", out="interp", num_dep_vars=2)

# Earth sidereal day period [s]
Earth_day = 86164

# Coriolis param [s-1] on tracer grid
f = 2 * (2 * np.pi / Earth_day) * np.sin(g["YCvec"] * (np.pi / 180))

# evaluate first derivative (d=1) of S(z) and T(z) at depth z.  That is, eval
# ∂S/∂Z and ∂T/∂Z, on the surface.
sz, tz = interp_two(z, Z, S.values, T.values, 1)
rs, rt = eos_s_t(s, t, z)  # ∂ρ/∂S and ∂ρ/∂T, on the surface

# ∂δ/∂z on the surface, where δ is the in-situ density anomaly
δz = rs * sz + rt * tz + (eos_z(s, t, z) - eos_z(s0, t0, z))

# large-scale potential vorticity [m-1 s-1] defined via in-situ density anomaly, on the surface
q = f * δz

ECCOv4 (lat-lon-cap)

ECCOv4 uses a “lat-lon-cap” tiled rectilinear grid. Whereas the OCCA script downloaded the dataset for you via FTP, for ECCOv4r4 you must register for a free Earthdata Login then download the dataset manually. Follow the instructions in the print lines in the run_ECCOv4r4.py script (below and on github).

# In[Imports]

import numpy as np
import xarray as xr
from os.path import expanduser

# Functions to make the Equation of State
from neutralocean.eos import make_eos, make_eos_s_t

# Functions to compute various approximately neutral surfaces
from neutralocean.surface import potential_surf, anomaly_surf, omega_surf

from neutralocean.grid.xgcm import build_grid, edgedata_to_maps
from neutralocean.ntp import ntp_epsilon_errors

# In[Load data]

print(
    """To get started, we need some ECCOv4r4 data. 
Make an Earthdata Login at
    https://urs.earthdata.nasa.gov/
and install "podaac data downloader" following the instructions at
    https://github.com/podaac/data-subscriber'.
Run the following commands to download the ECCOv4r4 grid information and
one day of data:
    podaac-data-downloader -c ECCO_L4_GEOMETRY_LLC0090GRID_V4R4 -d ./data --start-date 1992-01-01T00:00:00Z --end-date 1992-01-08T00:00:00Z -e ""
    podaac-data-downloader -c ECCO_L4_TEMP_SALINITY_LLC0090GRID_DAILY_V4R4 -d ./data --start-date 2002-12-23T00:00:01Z --end-date 2002-12-23T23:59:59Z -e ""
Finally, edit the variable `folder_ecco4` below to point to the directory where
you saved these `*.nc` files.
"""
)
# To get the first and last day of data, also run the following 2 lines:
# podaac-data-downloader -c ECCO_L4_TEMP_SALINITY_LLC0090GRID_DAILY_V4R4 -d ./data --start-date 1992-01-01T00:00:01Z --end-date 1992-01-01T23:59:59Z -e ""
# podaac-data-downloader -c ECCO_L4_TEMP_SALINITY_LLC0090GRID_DAILY_V4R4 -d ./data --start-date 2017-12-31T00:00:01Z --end-date 2017-12-31T23:59:59Z -e ""

folder_ecco4 = expanduser("~/work/data/ECCOv4r4/data")  # << EDIT AS NEEDED >>
file_grid = folder_ecco4 + "GRID_GEOMETRY_ECCO_V4r4_native_llc0090.nc"
date = "2002-12-23"
file_ST = (
    folder_ecco4
    + "OCEAN_TEMPERATURE_SALINITY_day_mean_"
    + date
    + "_ECCO_V4r4_native_llc0090.nc"
)

# Load horizontal grid information
with xr.open_dataset(file_grid) as ds:
    XC, YC, dxC, dyC, dyG, dxG = (
        ds[x].load() for x in ("XC", "YC", "dxC", "dyC", "dyG", "dxG")
    )

# Load hydrographic data
with xr.open_dataset(file_ST) as ds:
    S, T = (ds[x].squeeze().load() for x in ("SALT", "THETA"))

    # Create depth xarray.
    Z = -ds.Z.load()  # Make Z > 0 and increasing down
    Z.attrs.update({"positive": "down"})  # Update attrs to match.

    n = len(ds.i)  # size of each horizontal dimension in each square tile

# Get order of non-vertical dimensions
dims = tuple(x for x in S.dims if x != "k")

# Reorder dimensions to ensure individual water columns are float64 and contiguous in memory
S, T = (x.transpose(*dims, "k").astype(np.float64, order="C") for x in (S, T))

# define the connectivity between faces for the ECCOv4 LLC grid:
# fmt: off
face_connections = {'tile': {
    0: {'X':  ((12, 'Y', False), (3, 'X', False)),
        'Y':  (None,             (1, 'Y', False))},
    1: {'X':  ((11, 'Y', False), (4, 'X', False)),
        'Y':  ((0, 'Y', False),  (2, 'Y', False))},
    2: {'X':  ((10, 'Y', False), (5, 'X', False)),
        'Y':  ((1, 'Y', False),  (6, 'X', False))},
    3: {'X':  ((0, 'X', False),  (9, 'Y', False)),
        'Y':  (None,             (4, 'Y', False))},
    4: {'X':  ((1, 'X', False),  (8, 'Y', False)),
        'Y':  ((3, 'Y', False),  (5, 'Y', False))},
    5: {'X':  ((2, 'X', False),  (7, 'Y', False)),
        'Y':  ((4, 'Y', False),  (6, 'Y', False))},
    6: {'X':  ((2, 'Y', False),  (7, 'X', False)),
        'Y':  ((5, 'Y', False),  (10, 'X', False))},
    7: {'X':  ((6, 'X', False),  (8, 'X', False)),
        'Y':  ((5, 'X', False),  (10, 'Y', False))},
    8: {'X':  ((7, 'X', False),  (9, 'X', False)),
        'Y':  ((4, 'X', False),  (11, 'Y', False))},
    9: {'X':  ((8, 'X', False),  None),
        'Y':  ((3, 'X', False),  (12, 'Y', False))},
    10:{'X': ((6, 'Y', False),  (11, 'X', False)),
        'Y': ((7, 'Y', False),  (2, 'X', False))},
    11:{'X': ((10, 'X', False), (12, 'X', False)),
        'Y': ((8, 'Y', False),  (1, 'X', False))},
    12:{'X': ((11, 'X', False), None),
        'Y': ((9, 'Y', False),  (0, 'X', False))}}}
# fmt: on

# Number of Faces -- len(faces_connections only value)
nf = len(next(iter(face_connections.values())))

# Build list of adjacent water columns and distances between those water column pairs
xsh = ysh = "left"
grid = build_grid(n, face_connections, dims, xsh, ysh, dxC, dyC, dyG, dxG)


# Make Boussinesq version of the Jackett and McDougall (1995) equation of state
#  and its partial derivatives.
# TODO: is this what ECCOv4r4 used?
grav, rho_c = 9.81, 1027.5
eos = make_eos("jmd95", grav, rho_c)
eos_s_t = make_eos_s_t("jmd95", grav, rho_c)

# Select pinning cast, picking the cast closest to (x0,y0)
x0, y0 = (-172, -4)  # longitude, latitude -- Pacific equatorial ocean
pin_cast = np.unravel_index(
    ((XC.values - x0) ** 2 + (YC.values - y0) ** 2).argmin(), XC.shape
)
z0 = 1500.0  # pinning depth

# In[Build approximately neutral surfaces]

# Build potential density surface, with given reference pressure (actually depth,
# for Boussinesq) and given isovalue.  No diagnostics requested, so info about
# the grid is not needed (no `edges` and `geoemtry` provided), and also eos_s_t
# is not needed.
s, t, z, _ = potential_surf(
    S, T, Z, eos=eos, vert_dim="k", ref=0.0, isoval=1027.5, diags=False
)
z_sigma = z

# Build in-situ density anomaly surface with given reference salinity and
# potential temperature values and an isovalue of 0, which means the surface
# will intersect any point where the local (S,T) equals the reference values.
# Also return diagnostics.
s0, t0 = 34.5, 4.0
s, t, z, d = anomaly_surf(
    S,
    T,
    Z,
    grid=grid,
    eos=(eos, eos_s_t),
    vert_dim="k",
    ref=(s0, t0),
    isoval=0.0,
)

# Build an omega surface that is initialized from the above potential density
# surface and is pinned at the cast `pin_cast` (i.e. the omega surface will have
# the same depth as the initializing potential density surface at this cast).
s, t, z, d = omega_surf(
    S,
    T,
    Z,
    grid=grid,
    vert_dim="k",
    p_init=z_sigma,
    pin_cast=pin_cast,
    eos=(eos, eos_s_t),
    interp="pchip",
    ITER_MAX=10,
    ITER_START_WETTING=1,
)
z_omega = z

# In[Calculate neutrality error]

# Calculate ϵ neutrality errors on the latest surface, between all pairs of adjacent water columns
e = ntp_epsilon_errors(s, t, z, grid, eos_s_t)

# Convert the 1D array of ϵ values into two maps of ϵ neutrality errors, one
# for the errors in each of the two lateral ('i' and 'j') dimensions.  These
# neutrality errors can then be mapped or further analyzed.
ei, ej = edgedata_to_maps(e, n, face_connections, dims, xsh, ysh)

# In[Map depth difference between omega and potential]

# I've installed ecco_v4_py by cloning their git repository to ~/work/python/ECCOv4-py
# Follow https://ecco-v4-python-tutorial.readthedocs.io/Installing_Python_and_Python_Packages.html#option-1-clone-into-the-repository-using-git-recommended
# and edit path below as needed.
import sys
import matplotlib.pyplot as plt

sys.path.append("/home/stanley/work/python/ECCOv4-py/")
import ecco_v4_py as ecco

ecco.plot_tiles(
    z_omega - z_sigma,
    cmin=-200,
    cmax=200,
    fig_size=9,
    layout="latlon",
    rotate_to_latlon=True,
    Arctic_cap_tile_location=10,
    show_tile_labels=False,
    show_colorbar=True,
    show_cbar_label=True,
    cbar_label="Depth difference [m]",
)
plt.suptitle("z_omega - z_sigma")

plt.savefig("depth omega - depth sigma.png", bbox_inches="tight")

# In[Double check neutralocean's grid differencing]
from xmitgcm import open_mdsdataset
import xgcm

# To check that the grid works, we'll calculate the backwards differences of the
# sea-surface temperature in both horizontal directions. Begin by extracting SST.
k = 0  # index in vertical dimension
Tnp = T.values  # extract raw numpy data
SST = Tnp[..., k]  # slice the shallowest data from each water column

# Calculate difference between all pairs of adjacent casts
a, b = grid["edges"]
ΔT = SST.reshape(-1)[a] - SST.reshape(-1)[b]

# Decompose ΔT (a 1D array) into two 2D arrays, one for differences in each horizontal dimension
SSTx, SSTy = edgedata_to_maps(ΔT, n, face_connections, dims, xsh, ysh)


# Next, we'll repeat the above differencing using ECCO's methods.
# See https://ecco-v4-python-tutorial.readthedocs.io/VectorCalculus_ECCO_barotropicVorticity.html
print(
    "Download the grid's binary data files from "
    "< https://ndownloader.figshare.com/files/6494721 > "
    f" then exctract it so that {folder_ecco4} contains a folder 'global_oce_llc90'"
)
ds_llc = open_mdsdataset(
    folder_ecco4 + "global_oce_llc90/",
    iters=0,
    geometry="llc",
)

# rename 'face' to 'tile' to match T from the netCDF file
ds_llc = ds_llc.rename({"face": "tile"})

# Make the xgcm grid object to handle the differencing.
grid_llc = xgcm.Grid(
    ds_llc,
    periodic=False,
    face_connections=face_connections,
    coords={
        "X": {"center": "i", "left": "i_g"},
        "Y": {"center": "j", "left": "j_g"},
    },
)

# Backward differences in both horizontal directions
SSTx_ = grid_llc.diff(T[..., k], "X").values
SSTy_ = grid_llc.diff(T[..., k], "Y").values


# Begin checks.

# First, check the two differencing methods are equal, everywhere
assert np.array_equal(SSTx, SSTx_, equal_nan=True)
assert np.array_equal(SSTy, SSTy_, equal_nan=True)
print(
    "Test for equivalence between neutralocean grid differencing and xgcm grid differencing: passed."
)

# Now, check that differencing makes sense at an interior point
(t, j, i) = (9, 40, 40)
assert SSTx[t, j, i] == (Tnp[t, j, i, k] - Tnp[t, j, i - 1, k])
assert SSTy[t, j, i] == (Tnp[t, j, i, k] - Tnp[t, j - 1, i, k])
print(
    f"SSTx[{t}, {j}, {i}] == (T[{t}, {j}, {i}, {k}] - T[{t}, {j}, {i - 1}, {k}])"
)
print(
    f"SSTy[{t}, {j}, {i}] == (T[{t}, {j}, {i}, {k}] - T[{t}, {j - 1}, {i}, {k}])"
)

# Check that differencing makes sense across a boundary.  Find the second point
# that is involved in the difference across the boundary.  Verify by hand that
# this second point is where it ought to be.
# In this case, we're taking a meridional difference from the southern boundary
# of tile 4, and the second point is along the northern boundary of tile 3,
# which makes sense.
t, j, i = 4, 0, 40
idx = np.nonzero(SSTy[t, j, i] == (Tnp[t, j, i, k] - Tnp[..., k]))
t_, j_, i_ = idx[0][0], idx[1][0], idx[2][0]
assert SSTy[t, j, i] == Tnp[t, j, i, k] - Tnp[t_, j_, i_, k]
print(f"SSTy[{t},{j},{i}] == T[{t},{j},{i},{k}] - T[{t_},{j_},{i_},{k}]")
# # Output:
# SSTy[4,0,40] == T[4,0,40,0] - T[3,89,40,0]

ORCA (tripolar)

The ORCA tripolar grid used by NEMO is similar to a latitude-longitude grid, but it splits the north pole singularity into two pieces and places these over land. This example applies neutralocean to CanESM5 data in CMIP6.

# Show basic use of neutralocean on a tripolar grid using CMIP6 archived
# CanESM5 data, which runs NEMO 3.4 on the ORCA1 grid.
#
# For more advanced usage of neutralocean, see the run_OCCA.py example.

# In[Imports]

import xarray as xr
import pooch

from neutralocean.grid.tripolar import build_grid, edgedata_to_maps
from neutralocean.surface import potential_surf, omega_surf

# In[Load data]

# Get salt and potential temperature from CMIP6 archives (734 MB and 963 MB)
url_salt = "http://crd-esgf-drc.ec.gc.ca/thredds/fileServer/esgC_dataroot/AR6/CMIP6/CMIP/CCCma/CanESM5/esm-piControl/r1i1p1f1/Omon/so/gn/v20190429/so_Omon_CanESM5_esm-piControl_r1i1p1f1_gn_629101-630012.nc"
url_theta = "http://crd-esgf-drc.ec.gc.ca/thredds/fileServer/esgC_dataroot/AR6/CMIP6/CMIP/CCCma/CanESM5/esm-piControl/r1i1p1f1/Omon/thetao/gn/v20190429/thetao_Omon_CanESM5_esm-piControl_r1i1p1f1_gn_629101-630012.nc"

# Get ORCA1 horizontal grid distances from a Zenodo archive that is
# unassociated with the CMIP6 CanESM5 run, I suspect is the same as that which
# would have been used by CanESM5. This example is for illustrative purposes
# only! (405 MB)
url_mesh = "https://zenodo.org/record/4432892/files/ORCA1_mesh_mask.nc"

# Download data, using friendly pooch to fetch it if not already downloaded.
hash_salt = "93ba818ac7661c48686ae3af497e832667d0e72eb1968f375d9191c24f11db94"
hash_theta = "5d4085ab355c3a5526fa62d5f545a2322050e08dc24383189d458b413db99a79"
hash_mesh = "3e4560590c338dfd18de5db580f2354b11acb8ae95e13101c66461683385e777"
file_salt = pooch.retrieve(url=url_salt, known_hash=hash_salt)
file_theta = pooch.retrieve(url=url_theta, known_hash=hash_theta)
file_mesh = pooch.retrieve(url=url_mesh, known_hash=hash_mesh)


# Extract data
timestep = 0
ds = xr.open_dataset(file_salt)
nj = ds.dims["j"]  # number of grid points in the meridional y direction
ni = ds.dims["i"]  # number of grid points in the zonal x direction
S = ds["so"].isel({"time": timestep})  # 3D salinity
Z = ds["lev"]  # 1D depth at center of tracer cells

ds = xr.open_dataset(file_theta)
T = ds["thetao"].isel({"time": timestep})  # potential temperature

mesh = xr.open_dataset(file_mesh)  # contains horizontal grid information

# Squeeze out singleton dimension in horizontal grid distances and convert to numpy array
e1u, e2v, e2u, e1v = (mesh[x].data.squeeze() for x in ("e1u", "e2v", "e2u", "e1v"))

# Note: The ORCA1 tripolar horizontal grid is of size (nj, ni) == (291, 360).
# The second dimension is periodic, handling longitude's periodic nature.
# The first dimension has two different boundary conditions. The south
# (the first row) is non-periodic since Antarctica covers the South Pole and
# the ocean grid only goes to -78.3935°S. The north (the last row) is periodic
# with a flipped version of itself, due to ORCA's tripolar grid. Specifically,
# the cell at [nj - 1, i] is adjacent to the cell at [nj - 1, ni - i - 1].
# The grid metrics e1u, e2v, e2u, e1v all have size (nj+1, ni+2) == (292, 362):
# they employ padding to ease applying boundary conditions.
# See documentation in `neutralocean.grid.tripolar` for more information.

# Build the edges of the grid graph, and the metrics associated with each edge
grid = build_grid((nj, ni), e1u, e2v, e2u, e1v)

# In[Approximately Neutral Surfaces]

grav = 9.80665  # gravitational acceleration [m s-2]
rau0 = 1035.0  # Boussinesq reference density [kg m-3]

# Provide reference pressure (actually depth, in Boussinesq) and isovalue
s, t, z, d = potential_surf(
    S,
    T,
    Z,
    grid=grid,
    eos="jmd95",
    grav=grav,
    rho_c=rau0,
    vert_dim="lev",
    ref=0.0,
    isoval=1027.5,
)
s_sigma, t_sigma, z_sigma = s, t, z  # save for later
print(
    f" ** The potential density surface (referenced to {d['ref']}m)"
    f" with isovalue = {d['isoval']}kg m-3"
    f" has root-mean-square ϵ neutrality error {d['e_RMS']} kg m-4"
)

# Pin omega surface to the cast (j0, i0) in the equatorial Pacific at depth 1500m.
j0, i0 = 100, 150
z0 = 1500.0

# Initialize omega surface with a (locally referenced) potential density surface.
# Provide grid distances. By default, does as many iterations as needed to get
# the root-mean-square change in the locally referenced potential density from
# one iteration to the next to drop below 10^-7 kg m-3, or a maximum of 10
# iterations.
s, t, z, d = omega_surf(
    S,
    T,
    Z,
    grid,
    vert_dim="lev",
    pin_cast=(j0, i0),
    pin_p=z0,
    eos="jmd95",
    grav=grav,
    rho_c=rau0,
)
s_omega, t_omega, z_omega = s, t, z  # save for later
print(
    f" ** The omega-surface"
    f" initialized from a potential density surface (referenced to {z0}m)"
    f" intersecting the cast indexed by {(j0,i0)} at depth {z0}m"
    f" has root-mean-square ϵ neutrality error {d['e_RMS'][-1]} kg m-4"
)


# In[Neutrality errors on a surface]
from neutralocean.ntp import ntp_epsilon_errors
from neutralocean.eos import make_eos_s_t

# Prepare function for S and T derivatives of the equation of state
eos_s_t = make_eos_s_t("jmd95", grav, rau0)

# Calculate ϵ neutrality errors on all pairs of adjacent water columns
e = ntp_epsilon_errors(s, t, z, grid, eos_s_t)

# Convert ϵ above into two 2D maps, one each of the "x" and "y" dimensions
ex, ey = edgedata_to_maps(e, (nj, ni))

# These can then be mapped, e.g.:
# import matplotlib.pyplot as plt
# plt.imshow(ex, origin="lower", vmin=-1e-9, vmax=1e-9)
# plt.colorbar()