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]

import neutralocean as no


# In[Load OCCA data]

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

# make Boussinesq version of the Jackett and McDougall (1995) equation of state
# --- which is what OCCA used --- and its partial derivatives
eos = no.load_eos("jmd95", "", grav, rho_c)
eos_s_t = no.load_eos("jmd95", "_s_t", grav, rho_c)

# Perturb S, T to ensure static stability everywhere
print("Begin stabilization of hydrographic casts ...")
from time import time

tic = time()
no.stabilize_ST(S, T, Z, eos=eos, min_dLRPDdz=1e-6, verbose=False)  # about 90 sec
print(f"... done in {time() - tic:.2f} sec")

# Select pinning cast in the equatorial Pacific.
# When these are used for 'pin_cast' and 'pin_p' or 'p_init', 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"

# Build grid adjacency and distance information for neutralocean functions.
# Note OCCA uses a latitude by longitude grid, which is rectangular in memory, hence we
# select the `rectilinear` submodule of the `grid` subpackage.
grid = no.grid.rectilinear.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
opts["eos_s_t"] = 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 = no.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 = no.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.
args = opts.copy()
args["pin_cast"] = {"Longitude_t": 220.5, "Latitude_t": 0.5}
args["pin_p"] = z0
s, t, z, d = no.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 = no.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 = no.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 = no.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) in-situ density anomaly surface.
# By default, this 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["p_init"] = z_delta
args["p_ml"] = {"bottle_index": 1, "ref_p": 0.0}  # see `mld` for info
s, t, z, d = no.omega_surf(S, T, Z, **args)
print(
    f" ** The omega-surface"
    f" initialized from the given in-situ density anomaly surface"
    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 by iteratively making Neutral Tangent Plane links
# outwards from the pinning cast.
# Specify higher than default tolerances: Quit iterations when the
# depth change from one iteration to the next has a root-mean-square value less
# than 1e-6 m, and use a tolerance of 1e-7 m in the "vertical solve" that
# updates the vertical position of the surface during each iteration.
# Also, keep wetting through all the iterations.
args = opts.copy()
args["pin_cast"] = (i0, j0)
args["p_init"] = z0
args["TOL_P_CHANGE_RMS"] = 1e-6
args["TOL_LRPD_MAV"] = 0.0  # deactivate this exit threshold
args["TOL_P_SOLVER"] = 1e-7
args["ITER_STOP_WETTING"] = 99  # > ITER_MAX, so don't stop wetting
s, t, z, d = no.omega_surf(S, T, Z, **args)
s_omega, t_omega, z_omega = s, t, z  # save for later
print(
    f" ** The omega-surface"
    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"
)
args_omega = args.copy()  # save for later


# 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 converged (to the default
# tolerances), this omega surface should basically match the above one.
# Indeed, it will do one iteration, find it is below the tolerances, and exit.
i1, j1 = 315, 110  # North Atlantic
args["pin_cast"] = (i1, j1)
args["p_init"] = z_omega  # set init surface from above
s, t, z, d = no.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

# from neutralocean.lib import _process_casts, xr_to_np


# 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 = no.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 = (no.lib.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)
    """
    z, z2 = (no.lib.xr_to_np(x) for x in (z, z2))
    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 = (no.lib.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 be at a depth `ztiny` deeper at the reference cast.
    # This is slightly better than the above, unless lrpd_z is very near zero
    # in some casts. 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

args["p_init"] = z_init
s, t, z, d = no.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)} and {np.nanmax(b)}."
)

# 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 = no.veronis(z0, S_ref_cast, T_ref_cast, Z, eos=eos, eos_s_t=eos_s_t)
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_ = no.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_ = no.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 = no.mld(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.


def mixedlayer2nan(S, T, Z, z_ml):
    # Mutates S and T!
    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["pin_cast"] = (i0, j0)
args["p_init"] = z0
args["interp"] = "pchip"
args["ITER_MAX"] = 10
args["TOL_P_SOLVER"] = 1e-5
s, t, z, d = no.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 = no.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 = no.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 = no.grid.rectilinear.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 = no.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 = no.lib._process_casts(S, T, Z, "Depth_c")

s, t, z, d = no.anomaly_surf(
    Snp,
    Tnp,
    Znp,
    grid=grid,
    eos=eos,
    eos_s_t=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 = no.load_eos("jmd95", "_p", grav, rho_c)  # for scalar inputs
eos_z = no.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 = no.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).

# Show basic use of neutralocean on a cube-sphere grid using ECCOv4r4 data.
#
# For more advanced usage of neutralocean, see the run_OCCA.py example.

# In[Imports]

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

import neutralocean as no
import neutralocean.grid.xgcm as nogrid

# 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 = nogrid.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 = no.load_eos("jmd95", "", grav, rho_c)
eos_s_t = no.load_eos("jmd95", "_s_t", grav, rho_c)


# Stabilize hydrographic casts
print("Begin stabilization of hydrographic casts ...")
from time import time

tic = time()
no.stabilize_ST(S, T, Z, eos=eos, verbose=False)  # about 140 sec
print(f"... done in {time() - tic:.2f} sec")

# 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, through depth z0 at cast pin_cast; the
# reference depth is z0 also.
# 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, _ = no.potential_surf(
    S, T, Z, eos=eos, vert_dim="k", pin_p=z0, pin_cast=pin_cast, 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 (which requires the eos_s_t input also).
s0, t0 = 34.5, 4.0
s, t, z, d = no.anomaly_surf(
    S,
    T,
    Z,
    grid=grid,
    eos=eos,
    eos_s_t=eos_s_t,
    vert_dim="k",
    ref=(s0, t0),
    isoval=0.0,
)

# Build an omega surface that is initialized by iteratively making Neutral
# Tangent Plane links outwards from the pinning cast.
s, t, z, d = no.omega_surf(
    S,
    T,
    Z,
    grid=grid,
    pin_cast=pin_cast,
    p_init=z0,
    vert_dim="k",
    eos=eos,
    eos_s_t=eos_s_t,
    interp="pchip",
)
z_omega = z

# In[Calculate neutrality error]

# Calculate ϵ neutrality errors on the latest surface, between all pairs of adjacent water columns
e = no.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 = nogrid.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 = nogrid.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. The run_NEMO_ORCA1.py example (below and on github) 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
import neutralocean as no
import neutralocean.grid.tripolar as nogrid

# 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.sizes["j"]  # number of grid points in the meridional y direction
ni = ds.sizes["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 = nogrid.build_grid((nj, ni), e1u, e2v, e2u, e1v)

# In[Approximately Neutral Surfaces]

grav = 9.80665  # gravitational acceleration [m s-2]
rho_c = 1035.0  # Boussinesq reference density [kg m-3]
eos = no.load_eos("jmd95", "", grav, rho_c)
eos_s_t = no.load_eos("jmd95", "_s_t", grav, rho_c)

# Provide reference pressure (actually depth, in Boussinesq) and isovalue
s, t, z, d = no.potential_surf(
    S,
    T,
    Z,
    grid=grid,
    eos=eos,
    eos_s_t=eos_s_t,
    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 = no.omega_surf(
    S, T, Z, grid, pin_cast=(j0, i0), p_init=z0, vert_dim="lev", eos=eos, eos_s_t=eos_s_t
)
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]

# Calculate ϵ neutrality errors on all pairs of adjacent water columns
e = no.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 = nogrid.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()