TIW DCL Paper Figures

import ncar_jobqueue

if "client" in locals():
    client.close()
    del client
# if "cluster" in locals():
#    cluster.close()

# env = {"OMP_NUM_THREADS": "3", "NUMBA_NUM_THREADS": "3"}

# cluster = distributed.LocalCluster(
#    n_workers=8,
#    threads_per_worker=1,
#    env=env
# )

if "cluster" in locals():
    del cluster

# cluster = ncar_jobqueue.NCARCluster(
#    project="NCGD0011",
#    scheduler_options=dict(dashboard_address=":9797"),
# )
# cluster = dask_jobqueue.PBSCluster(
#    cores=9, processes=9, memory="108GB", walltime="02:00:00", project="NCGD0043",
#    env_extra=env,
# )

import dask_jobqueue

cluster = dask_jobqueue.PBSCluster(
    cores=1,  # The number of cores you want
    memory="23GB",  # Amount of memory
    processes=1,  # How many processes
    queue="casper",  # The type of queue to utilize (/glade/u/apps/dav/opt/usr/bin/execcasper)
    local_directory="$TMPDIR",  # Use your local directory
    resource_spec="select=1:ncpus=1:mem=23GB",  # Specify resources
    project="ncgd0011",  # Input your project ID here
    walltime="02:00:00",  # Amount of wall time
    interface="ib0",  # Interface to use
)
cluster.scale(jobs=6)

# %load_ext autoreload
# %autoreload 2

%matplotlib inline

import cf_xarray
import dask
import dcpy
import distributed
import holoviews as hv
import hvplot.xarray
import matplotlib as mpl
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import matplotlib.units as munits
import numpy as np
import pandas as pd
import pump
from pump.calc import ddx, ddy, ddz
import seawater as sw
from holoviews import opts

import xarray as xr

# import hvplot.xarray


# import facetgrid

# mpl.rcParams["savefig.dpi"] = 300
# mpl.rcParams["savefig.bbox"] = "tight"
# mpl.rcParams["figure.dpi"] = 140

munits.registry[np.datetime64] = mdates.ConciseDateConverter()

xr.set_options(keep_attrs=True)

from dask.cache import Cache

cache = Cache(5e9, 0.4)  # Leverage two gigabytes of memory
cache.register()  # Turn cache on globally

hv.extension("bokeh")
# hv.opts.defaults(opts.Image(fontscale=1.5), opts.Curve(fontscale=1.5))


print(dask.__version__)
print(distributed.__version__)
print(xr.__version__)


def unpack(d):
    val = []
    for k, v in d.items():
        if isinstance(v, dict):
            val.extend(unpack(v))
        else:
            val.append(v)
    return val


xr.DataArray([1.0])

2.30.0
2.30.0
0.16.3.dev150+g37522e991

<xarray.DataArray (dim_0: 1)>
array([1.])
Dimensions without coordinates: dim_0

client = distributed.Client(cluster)
client

Client Scheduler: tcp://10.12.1.3:41722 Dashboard: https://jupyterhub.ucar.edu/dav/user/dcherian/proxy/8787/status

Cluster Workers: 0 Cores: 0 Memory: 0 B

def slice_like(obj, other):
    indexes = set(obj.indexes) & set(other.indexes)
    return obj.sel({k: slice(other[k][0].values, other[k][-1].values) for k in indexes})

gcm1 = pump.model.Model(
    "/glade/u/home/dcherian/pump/TPOS_MITgcm_1_hb/HOLD_NC_LINKS/",
    name="gcm1",
    full=False,
    budget=False,
)
# gcm1.tao.load()
# gcm1.tao = xr.merge([gcm1.tao, pump.calc.get_tiw_phase(gcm1.tao.v)])
# gcm1.surface["zeta"] = xr.open_dataset(gcm1.dirname + "obs_subset/surface_zeta.nc").zeta
# gcm1.surface = gcm1.surface.chunk({"time": 1})
# gcm1.surface["zeta"] = gcm1.surface.v.differentiate("longitude") - gcm1.surface.u.differentiate("latitude")

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-5-17e9045cc9c0> in <module>
----> 1 gcm1 = pump.model.Model(
      2     "/glade/u/home/dcherian/pump/TPOS_MITgcm_1_hb/HOLD_NC_LINKS/",
      3     name="gcm1",
      4     full=False,
      5     budget=False,

~/pump/pump/model.py in __init__(self, dirname, name, kind, full, budget)
    510             self.depth = None
    511 
--> 512         self.update_coords()
    513         self.read_metrics()
    514 

~/pump/pump/model.py in update_coords(self)
    770         elif "latitude" in self.full:
    771             ds = self.full
--> 772         elif "latitude" in self.budget:
    773             ds = self.budget
    774         else:

AttributeError: 'Model' object has no attribute 'budget'

gcm1.tao["mld"] = pump.calc.get_mld(gcm1.tao.dens)
gcm1.tao["dcl"] = pump.calc.get_dcl_base_Ri(gcm1.tao)

Read data

/glade/scratch/dcherian/TPOS_MITgcm_1_hb/110w-period-4-with-jq.nc has in-situ density

# period4 = xr.open_dataset("/glade/scratch/dcherian/TPOS_MITgcm_1_hb/110w-period-4-with-jq.nc")

period4 = xr.open_zarr(
    "/glade/work/dcherian/pump/zarrs/110w-period-4-4.zarr", consolidated=True
).chunk({"time": 50})
period4.latitude.attrs = {"long_name": "Latitude", "units": "°N"}
metrics = (
    pump.model.read_metrics(
        f"/glade/campaign/cgd/oce/people/dcherian/TPOS_MITgcm_1_hb/"
    )
    .reindex_like(period4, method="nearest")
    .load()
)
# from http://mailman.mitgcm.org/pipermail/mitgcm-support/2010-December/006921.html
period4["Um_Impl"] = period4.VISrI_Um.diff("depth", label="lower") / (
    metrics.rAw * metrics.drF * metrics.hFacW
)
period4["Vm_Impl"] = period4.VISrI_Vm.diff("depth", label="lower") / (
    metrics.rAw * metrics.drF * metrics.hFacW
)


# period4["time"] = period4["time"] - pd.Timedelta("7h")
# period4["dens"] = pump.mdjwf.dens(period4.salt, period4.theta, np.array([0.0]))
# period4 = pump.calc.calc_reduced_shear(period4)
# period4["mld"] = pump.calc.get_mld(period4.dens)
# period4["dcl_base"] = pump.calc.get_dcl_base_Ri(period4)
# period4["dcl"] = period4["mld"] - period4["dcl_base"]
# period4["sst"] = period4.theta.isel(depth=0)
# period4["eucmax"] = pump.calc.get_euc_max(
#    period4.u.sel(latitude=0, method="nearest", drop=True)
# )

period4["nonlocal_flux"] = 1035 * 3995 * period4.KPPg_TH / metrics.RAC
# period4.nonlocal_flux.isel(longitude=1).sel(latitude=3.5, method="nearest").plot(y="depth")
snapshot_times = pd.to_datetime(
    [
        "1995-11-18 00:00",
        # "1995-11-24 18:00",
        "1995-11-26 06:00",
        "1995-12-03 18:00",
        # "1995-12-09 18:00",
    ]
)

period4["dcl_mask"] = (
    (period4.depth < period4.mld)
    & (period4.depth > period4.dcl_base)
    & (period4.dcl > 30)
)

import dcpy.interpolate

period4["hRib"] = (
    dcpy.interpolate.pchip_roots_old(
        period4.KPPRi.sortby("depth").chunk({"depth": -1}), dim="depth", target=0.3
    )
    .squeeze()
    .persist()
)

period4 = period4.persist()
# period4.depth.attrs["units"] = "m"
# period4.Jq.attrs["long_name"] = "$J^t_q$"
# period4.Jq.attrs["units"] = "W/m²"
# period4.S2.attrs["long_name"] = "$S²$"
# period4.S2.attrs["units"] = "$s^{-2}$"
# period4.Ri.attrs["long_name"] = "$Ri$"
# period4.N2.attrs["long_name"] = "$N²$"
# period4.N2.attrs["units"] = "$s^{-2}$"

/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/core/indexing.py:1375: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  value = value[(slice(None),) * axis + (subkey,)]
/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/core/indexing.py:1369: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  return self.array[key]
/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/core/indexing.py:1369: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  return self.array[key]
/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/core/indexing.py:1369: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  return self.array[key]

vort = pump.calc.vorticity(period4).isel(longitude=1).compute()
vort["vy"] = pump.calc.ddy(period4.v.isel(longitude=1)).compute()

(vort.f + vort.vx - vort.uy).sel(depth=slice(-60)).mean("depth").plot(
    x="time", robust=True
)

<matplotlib.collections.QuadMesh at 0x2accedcf4ee0>

test viscous term

period4.Um_Impl.isel(longitude=1).sel(latitude=0, method="nearest").rolling(
    depth=4
).mean().plot(x="time", robust=True)

<matplotlib.collections.QuadMesh at 0x2af30fc93c50>

Figure 1: snapshot

gcm1.read_budget()

/glade/u/home/dcherian/python/xarray/xarray/core/indexing.py:1369: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  return self.array[key]
/glade/u/home/dcherian/python/xarray/xarray/core/indexing.py:1369: PerformanceWarning: Slicing is producing a large chunk. To accept the large
chunk and silence this warning, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    ...     array[indexer]

To avoid creating the large chunks, set the option
    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    ...     array[indexer]
  return self.array[key]

gcm1 = pump.model.Model(
    "/glade/u/home/dcherian/pump/TPOS_MITgcm_1_hb/HOLD_NC_LINKS/",
    name="gcm1",
    full=True,
    budget=True,
)
gcm1.full = pump.calc_reduced_shear(gcm1.full)
gcm1.full["mld"] = pump.calc.get_mld(gcm1.full.dens.sel(depth=slice(-200)))
gcm1.full["dcl_base"] = pump.calc.get_dcl_base_Ri(gcm1.full)
gcm1.full

calc uz
calc vz

Exception ignored in: <bound method GCDiagnosis._gc_callback of <distributed.utils_perf.GCDiagnosis object at 0x2b60d609da90>>
Traceback (most recent call last):
  File "/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/distributed/utils_perf.py", line 183, in _gc_callback
    def _gc_callback(self, phase, info):
KeyboardInterrupt: 

calc S2
calc N2

Exception ignored in: <bound method GCDiagnosis._gc_callback of <distributed.utils_perf.GCDiagnosis object at 0x2b60d609da90>>
Traceback (most recent call last):
  File "/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/distributed/utils_perf.py", line 193, in _gc_callback
    self._fractional_timer.start_timing()
  File "/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/distributed/utils_perf.py", line 116, in start_timing
    assert self._cur_start is None
AssertionError: 

calc shred2
Calc Ri

<xarray.Dataset>
Dimensions:          (depth: 345, latitude: 480, longitude: 1500, time: 2947)
Coordinates:
  * latitude         (latitude) float32 -12.0 -11.95 -11.9 ... 11.9 11.95 12.0
  * longitude        (longitude) float32 -170.0 -169.9 -169.9 ... -95.05 -95.0
  * depth            (depth) float32 -0.5 -1.5 -2.5 ... -5.666e+03 -5.766e+03
  * time             (time) datetime64[ns] 1995-09-01 ... 1997-01-04T12:00:00
Data variables:
    theta            (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    u                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    v                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    w                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    salt             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    KPP_diffusivity  (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    dens             (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    uz               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    vz               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    S2               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    shear            (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    N2               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    shred2           (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    Ri               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    mld              (time, latitude, longitude) float32 dask.array<chunksize=(1, 138, 430), meta=np.ndarray>
    dcl_base         (time, latitude, longitude) float32 dask.array<chunksize=(1, 69, 215), meta=np.ndarray>
Attributes:
    title:              daily snapshot from 1/20 degree Equatorial Pacific MI...
    easting:            longitude
    northing:           latitude
    field_julian_date:  400296
    julian_day_unit:    days since 1950-01-01 00:00:00

time = "1995-12-04"

# gcm1.full["time"] = gcm1.full.time - pd.Timedelta("8h")
# gcm1.budget["time"] = gcm1.budget.time - pd.Timedelta("8h")
extract = gcm1.full.sel(time=time)
extract["Jq"] = gcm1.budget.Jq.sel(time=time)
extract["oceQnet"] = gcm1.budget.oceQnet.sel(time=time)
# extract.sel(depth=slice(-200)).to_zarr("/glade/scratch/dcherian/TPOS_MITgcm_1_hb/snapshot.zarr")
extract

<xarray.Dataset>
Dimensions:          (depth: 345, latitude: 480, longitude: 1500, time: 6)
Coordinates:
  * latitude         (latitude) float32 -12.0 -11.95 -11.9 ... 11.9 11.95 12.0
  * longitude        (longitude) float32 -170.0 -169.9 -169.9 ... -95.05 -95.0
  * depth            (depth) float64 -0.5 -1.5 -2.5 ... -5.666e+03 -5.766e+03
  * time             (time) datetime64[ns] 1995-12-04 ... 1995-12-04T20:00:00
Data variables:
    theta            (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    u                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    v                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    w                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    salt             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    KPP_diffusivity  (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    dens             (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    uz               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    vz               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    S2               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    shear            (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    N2               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    shred2           (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    Ri               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    mld              (time, latitude, longitude) float32 dask.array<chunksize=(1, 138, 430), meta=np.ndarray>
    dcl_base         (time, latitude, longitude) float32 dask.array<chunksize=(1, 69, 215), meta=np.ndarray>
    Jq               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    oceQnet          (time, latitude, longitude) float32 dask.array<chunksize=(1, 480, 1500), meta=np.ndarray>
Attributes:
    title:              daily snapshot from 1/20 degree Equatorial Pacific MI...
    easting:            longitude
    northing:           latitude
    field_julian_date:  400296
    julian_day_unit:    days since 1950-01-01 00:00:00

(
    extract.sel(depth=slice(-200), latitude=slice(-10, 10))
    .chunk({"latitude": 138})
    .to_zarr("/glade/work/dcherian/pump/zarrs/snapshot.zarr", mode="w")
)

<xarray.backends.zarr.ZarrStore at 0x2b6298eddfa0>

integrate jq in the dcl over a day

snapshot = xr.open_zarr("/glade/work/dcherian/pump/zarrs/snapshot.zarr")
snapshot["ijq"] = (
    snapshot.Jq.sel(depth=slice(-250))
    .where(
        (snapshot.depth.sel(depth=slice(-250)) < snapshot.mld)
        & (snapshot.depth >= snapshot.dcl_base)
    )
    .sum("depth")
)
snapshot["dcl"] = snapshot["mld"] - snapshot["dcl_base"]
snapshot["dcl"] = snapshot.dcl.where(snapshot.dcl > 5, 0)
snapshot["sst"] = snapshot.theta.isel(depth=0)

snapshot["umean"] = snapshot.u.sel(depth=slice(-50, -250)).mean("depth")

snapshot["sst"].attrs["long_name"] = "SST"
snapshot["sst"].attrs["units"] = "°C"

snapshot["oceQnet"].attrs["long_name"] = "$Q_{net}$"
snapshot["oceQnet"].attrs["units"] = "W/m²"

snapshot["ijq"].attrs["long_name"] = "Avg. $J_q^t$ in low Ri layer"
snapshot["ijq"].attrs["units"] = "W/m²"

snapshot["dcl"].attrs["long_name"] = "Low Ri layer width"
snapshot["dcl"].attrs["units"] = "m"

snapshot["umean"].attrs["long_name"] = "Mean 50-250m $u$"
snapshot["umean"].attrs["units"] = "m/s"

xr.set_options(keep_attrs=True)

extract = snapshot.mean("time").compute()
snapshot

<xarray.Dataset>
Dimensions:          (depth: 200, latitude: 400, longitude: 1500, time: 6)
Coordinates:
  * depth            (depth) float32 -0.5 -1.5 -2.5 ... -197.5 -198.5 -199.5
  * latitude         (latitude) float32 -9.996 -9.946 -9.896 ... 9.946 9.996
  * longitude        (longitude) float32 -170.0 -169.9 -169.9 ... -95.05 -95.0
  * time             (time) datetime64[ns] 1995-12-04 ... 1995-12-04T20:00:00
Data variables: (12/22)
    Jq               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    KPP_diffusivity  (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    N2               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    Ri               (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    S2               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    dcl_base         (time, latitude, longitude) float32 dask.array<chunksize=(1, 138, 215), meta=np.ndarray>
    ...               ...
    vz               (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    w                (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 100, 138, 430), meta=np.ndarray>
    ijq              (time, latitude, longitude) float64 dask.array<chunksize=(1, 138, 215), meta=np.ndarray>
    dcl              (time, latitude, longitude) float32 dask.array<chunksize=(1, 138, 215), meta=np.ndarray>
    sst              (time, latitude, longitude) float32 dask.array<chunksize=(1, 138, 430), meta=np.ndarray>
    umean            (time, latitude, longitude) float32 dask.array<chunksize=(1, 138, 430), meta=np.ndarray>
Attributes:
    easting:            longitude
    field_julian_date:  400296
    julian_day_unit:    days since 1950-01-01 00:00:00
    northing:           latitude
    title:              daily snapshot from 1/20 degree Equatorial Pacific MI...

f, axx = plt.subplots(
    5, 1, constrained_layout=True, sharex=True, sharey=True, squeeze=False
)
ax = dict(zip(["sflx", "sst", "umean", "jqt", "dcl"], axx.flat))
cbar_kwargs = {"label": ""}

(-1 * extract.oceQnet).plot(
    ax=ax["sflx"], robust=True, cmap=mpl.cm.RdBu_r, center=0, cbar_kwargs=cbar_kwargs
)

extract.sst.plot(
    x="longitude",
    robust=True,
    cmap=mpl.cm.RdYlBu_r,
    ax=ax["sst"],
    vmax=27.5,
    cbar_kwargs=cbar_kwargs,
)

extract.umean.plot(
    ax=ax["umean"], robust=True, cmap=mpl.cm.PuOr_r, cbar_kwargs=cbar_kwargs
)

(extract.ijq / 50).plot(
    x="longitude",
    ax=ax["jqt"],
    vmax=0,
    robust=True,
    cmap=mpl.cm.Blues_r,
    cbar_kwargs=cbar_kwargs,
)

extract.dcl.plot(
    x="longitude", ax=ax["dcl"], robust=True, cmap=mpl.cm.BuGn, cbar_kwargs=cbar_kwargs
)

axx[0, 0].set_yticks([-10, -5, 0, 5, 10])
axx[0, 0].set_xticks([-95, -110, -125, -140, -155, -170])

ax["jqt"].plot([-114, -105, -105, -114, -114], [-2, -2, 5, 5, -2], color="k", lw=1)

[dcpy.plots.lat_lon_ticks(aa) for aa in axx.flat]
axx[0, 0].set_ylim([-8, 8])
dcpy.plots.clean_axes(axx)
dcpy.plots.label_subplots(
    axx.squeeze(),
    labels=[
        "Net surface heat flux $Q_{net}$ [W/m²]",
        "SST [°C]",
        "Mean 50-250m $u$ [m/s]",
        "Avg $J_q^t$ in low Ri layer [W/m²]",
        "Low Ri layer thickness [m]",
    ],
    y=0.05,
)
f.set_size_inches((8, 8))
f.savefig("images/snapshot-sflx-sst-jq-dcl.png")

Figure 2: validation

Shows that TIWs and mean state are reasonable

johnson = pump.obs.read_johnson()
tao = pump.obs.read_tao_zarr()
sections = xr.open_zarr("/glade/work/dcherian/pump/zarrs/gcm1-sections-rechunked.zarr")

lon = -110
taosub = (
    tao.sel(longitude=lon, time=slice("1990", "2000"))
    .groupby("time.dayofyear")
    .mean()
    .rename({"dayofyear": "time"})
)

modelsub = (
    sections.sel(longitude=lon, latitude=0, method="nearest").sel(time="1996")
    # .pipe(slice_like, taosub.time)
)


def plot(data, ax=None, **kwargs):
    if ax is None:
        ax = plt.gca()

    mean, std = dask.compute(data.mean("time"), data.std("time"))

    if "depth" in data.dims:
        kwargs.update(dict(y="depth"))
        plotfunc = "fill_betweenx"
        coord = data.depth
    elif "latitude" in data.dims:
        kwargs.update(dict(x="latitude"))
        plotfunc = "fill_between"
        coord = data.latitude
    else:
        raise ValueError(f"Don't know how to handle data with dims: {data.dims}")

    hdl = mean.plot(**kwargs, ax=ax)
    getattr(ax, plotfunc)(
        coord,
        (mean - std).values,
        (mean + std).values,
        color=hdl[0].get_color(),
        alpha=0.3,
    )
    return hdl[0]


f = plt.figure(constrained_layout=True)
gs = f.add_gridspec(4, 2, height_ratios=[1.5, 1, 1, 1])
ax = [0, 1, 2, 3, 4]
ax[0] = f.add_subplot(gs[0, 0])
ax[1] = f.add_subplot(gs[0, 1], sharey=ax[0])
ax[2] = f.add_subplot(gs[1, :])
ax[3] = f.add_subplot(gs[2, :], sharex=ax[2])
ax[4] = f.add_subplot(gs[3, :], sharex=ax[3])

# f, ax = plt.subplots(1, 2, sharey=True, constrained_layout=True)
htao = plot(taosub.T, ax=ax[0])
hmod = plot(modelsub.theta, ax=ax[0])
hjohn = johnson.temp.sel(longitude=lon, latitude=0).plot(y="depth", ax=ax[0])

plot(taosub.u, ax=ax[1])
plot(modelsub.u, ax=ax[1])
johnson.u.sel(longitude=lon, latitude=0).plot(y="depth", ax=ax[1])

plot(sections.sel(longitude=lon, depth=0, method="nearest").theta, ax=ax[3], color="C1")
johnson.temp.sel(longitude=lon, depth=0, method="nearest").plot(
    x="latitude", ax=ax[3], color="C2"
)
# plot(sections.sel(longitude=lon, depth=0, method="nearest").theta, ax=ax[3])

plot(
    sections.sel(longitude=lon, depth=-75, method="nearest").theta, ax=ax[4], color="C1"
)
johnson.temp.sel(longitude=lon, depth=-75, method="nearest").plot(
    x="latitude", ax=ax[4], color="C2"
)

ax[0].legend(
    handles=(htao, hmod, hjohn[0]),
    labels=["TAO", "Model", "Johnson"],
    loc="lower right",
)

(
    -1
    * sections.sel(longitude=lon, time="1996")
    .mean("time")
    .sel(depth=slice(-250))
    .u.integrate("depth")
).plot(ax=ax[2], color="C1")

(
    (-1 * johnson.sel(longitude=lon).sel(depth=slice(-250)).u.integrate("depth")).plot(
        xlim=(-7.5, 7.5), ax=ax[2], color="C2"
    )
)


ax[0].set_xlim((12, None))
ax[0].set_ylim(-250, 0)

dcpy.plots.liney(0, ax=ax[2])
dcpy.plots.clean_axes(ax[:2])
ax[0].set_xlabel("$T$ [°C]")
ax[1].set_xlabel("$u$ [m/s]")
ax[2].set_ylabel("$\int u dz$ [m²/s]")

ax[3].set_ylabel("$T$ [°C]")
ax[4].set_ylabel("$T$ [°C]")
ax[3].set_xlabel("")
ax[2].set_xlabel("")
ax[4].set_xlabel("Latitude [°N]")
ax[4].set_xlim((-7.5, 7.5))

dcpy.plots.clean_axes([[aa] for aa in ax[2:]])

[aa.set_title("") for aa in ax]
dcpy.plots.label_subplots(ax, labels=["", "", "Top 250m", "Surface", "75m"])
f.set_size_inches((dcpy.plots.pub_fig_width("jpo", "two column"), 8))
f.suptitle("(0°N, 110°W)")

f.savefig("images/gcm-20-validation.png")

Task exception was never retrieved
future: <Task finished name='Task-186048' coro=<WebSocketProtocol13.write_message.<locals>.wrapper() done, defined at /glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/tornado/websocket.py:1102> exception=WebSocketClosedError()>
Traceback (most recent call last):
  File "/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/tornado/websocket.py", line 1104, in wrapper
    await fut
tornado.iostream.StreamClosedError: Stream is closed

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/tornado/websocket.py", line 1106, in wrapper
    raise WebSocketClosedError()
tornado.websocket.WebSocketClosedError

Figure 3: 0, 3.5N deep cycle heat flux + Ri profile

f, ax = pump.plot.plot_jq_sst(
    period4.isel(longitude=1),
    lon=-110,
    periods=slice("1995-11-15", "1995-12-15"),
    lat=[0, 3.5],
    time_Ri={
        0: slice("1995-11-17", "1995-11-29"),
        3.5: slice("1995-11-28", "1995-12-09"),
    },
)

axes = list(ax.flat)
del axes[1]
del axes[2]
[aa.set_title("") for aa in ax[:, 0]]
axes[-2].set_xlabel("time")
dcpy.plots.label_subplots(
    axes,
    x=0.01,
    y=0.05,
    labels=[
        "SST [°C]",
        "Avg $J_q^t$ in low Ri layer [W/m²]",
        "$J_q^t$ [W/m²] at 3.5°N",
        "$Ri$",
        "$J_q^t$ [W/m²] at 0°",
        "$Ri$",
    ],
)
f.text(0.7, 0.53, "$J_q^t$ [W/m²]")
f.savefig("../images/sst-deep-cycle-3N.png")

Figure 4: daily cycles at 3.5N

jrafull = pump.obs.read_jra()
station = (
    period4.isel(longitude=1)
    .sel(latitude=3.5, method="nearest")
    .chunk({"depth": -1, "time": 10})
)
forcing = pump.obs.interp_jra_to_station(jrafull, station)
fluxes, coare = pump.calc.coare_fluxes_jra(station, forcing)

/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:466: RuntimeWarning: invalid value encountered in greater
  k = u10 > umax
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:480: RuntimeWarning: invalid value encountered in greater
  charn = np.where(ut > 0, 0.011 + (ut - 10) / (18 - 10) * (0.018 - 0.011), 0.011)
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:481: RuntimeWarning: invalid value encountered in greater
  charn[ut > 18] = 0.018

station = station.merge(fluxes, compat="override")
station = station.merge(coare, compat="override")
station = station.merge(forcing, compat="override")
station = station.isel(time=slice(1, -1))

kpp = pump.calc.calc_kpp_terms(station)

kpp.hbl.plot()
kpp.hunlimit.plot()

%matplotlib ipympl

(
    period4.KPP_diffusivity.rolling(depth=2)
    .mean()
    .isel(longitude=1)
    .sel(latitude=3.5, method="nearest")
    .sel(depth=slice(0, -120))
    .plot(x="time", norm=mpl.colors.LogNorm(1e-5, 1e-2))
)
(-1 * kpp.hunlimit).plot()
(-1 * kpp.hbl).plot()

Warning: Cannot change to a different GUI toolkit: ipympl. Using qt instead.

[<matplotlib.lines.Line2D at 0x2b1a5003e490>]

/glade/u/home/dcherian/miniconda3/envs/dcpy2/lib/python3.7/site-packages/matplotlib/colors.py:1171: RuntimeWarning: invalid value encountered in less_equal
  mask |= resdat <= 0

ds = (
    period4.isel(longitude=1)
    .sel(time=slice("1995-11-25", "1995-12-07"), depth=slice(-100))
    # .sel(time=slice("1995-11-05", "1995-12-15"), depth=slice(-130))
    .sel(latitude=3.5, method="nearest")
    .assign(latitude=3.5)
).compute()
ds["Jq"] = ds.Jq.rolling(depth=2).mean()
ds["Jq"].attrs = period4.Jq.attrs

with plt.rc_context({"font.size": 8}):
    f, _ = pump.plot.plot_daily_cycles(ds)

f.savefig("images/110-deep-cycle-35.png")

Figure 5: DCL everywhere

sections = pump.utils.read_gcm1_zarr_subset(gcm1).chunk({"time": 300, "depth": 100})
sections.to_zarr(
    "/glade/work/dcherian/pump/gcm1-sections-rechunked.zarr",
    consolidated=True,
    encoding={k: {} for k in sections.data_vars},
    mode="w",
)

xarray.Dataset

• Dimensions:
  • depth: 345
  • latitude: 480
  • longitude: 4
  • time: 2947
• Coordinates: (4)
  • depth
    (depth)
    float32
    -0.5 -1.5 ... -5665.922 -5765.922

    array([-5.000000e-01, -1.500000e+00, -2.500000e+00, ..., -5.565922e+03,
           -5.665922e+03, -5.765922e+03], dtype=float32)

  • latitude
    (latitude)
    float32
    -12.0 -11.949896 ... 11.949896 12.0

    array([-12.      , -11.949896, -11.899791, ...,  11.899791,  11.949896,
            12.      ], dtype=float32)

  • longitude
    (longitude)
    int64
    -110 -125 -140 -155

    array([-110, -125, -140, -155])

  • time
    (time)
    datetime64[ns]
    1995-09-01 ... 1997-01-04T12:00:00

    _CoordinateAxisType :

    Time

    axis :

    T

    long_name :

    Time (hours since 1950-01-01)

    standard_name :

    time

    array(['1995-09-01T00:00:00.000000000', '1995-09-01T04:00:00.000000000',
           '1995-09-01T08:00:00.000000000', ..., '1997-01-03T20:00:00.000000000',
           '1997-01-04T00:00:00.000000000', '1997-01-04T12:00:00.000000000'],
          dtype='datetime64[ns]')

• Data variables: (8)
  • Jq
    (time, depth, latitude, longitude)
    float64
    dask.array<chunksize=(300, 100, 240, 1), meta=np.ndarray>

    long_name :

    $J_q^t$

    units :

    W/m$^2$

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     15.62 GB 
     
     57.60 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 240, 1) 
        
     Count 
     141777 Tasks 
     320 Chunks 
        
     Type 
     float64 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • KPP_diffusivity
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • dens
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • salt
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • theta
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • u
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • v
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

  • w
    (time, depth, latitude, longitude)
    float32
    dask.array<chunksize=(300, 100, 138, 1), meta=np.ndarray>

    
    
    
    
      
    
        
     
     Array 
     Chunk 
      
      
    
        
     Bytes 
     7.81 GB 
     
     16.56 MB 
        
     Shape 
     (2947, 345, 480, 4) 
     
     (300, 100, 138, 1) 
        
     Count 
     236401 Tasks 
     640 Chunks 
        
     Type 
     float32 
     numpy.ndarray 
      
    
    
    
    
    
    
    

• Attributes: (5)

  easting :

  longitude

  field_julian_date :

  400296

  julian_day_unit :

  days since 1950-01-01 00:00:00

  northing :

  latitude

  title :

  daily snapshot from 1/20 degree Equatorial Pacific MITgcm simulation

sections = xr.open_zarr(
    "/glade/work/dcherian/pump/gcm1-sections-rechunked.zarr", consolidated=True
)
sections["dens"] = pump.mdjwf.dens(sections.salt, sections.theta, np.array([0.0]))
sections = pump.calc.calc_reduced_shear(sections)
sections["mld"] = pump.calc.get_mld(sections.dens)
sections

medianRi = (
    sections.Ri.where(
        (sections.depth < sections.mld) & (sections.depth > sections.mld - 40)
    )
    .resample(time="D")
    .median(["time", "depth"])
    .compute()
)
medianRi = medianRi.assign_coords(longitude=np.round(medianRi.longitude))
medianRi.attrs["long_name"] = "Median $Ri$ within 40m of the mixed layer base"

sst = sections.theta.isel(depth=0).rolling(time=6).mean().persist()
u = (
    sections.u.sel(depth=slice(-250))
    .rolling(time=4 * 6, min_periods=1)
    .mean()
    .persist()
)
eucmax = pump.calc.get_euc_bounds(u).compute()

(
    sections.u.sel(depth=slice(-150), time=slice("1996-04-01", "1996-07-01"))
    .sel(
        longitude=-125,
        latitude=2,
        method="nearest",
    )
    .hvplot.quadmesh(y="depth")
)

(
    sections.Ri.sel(depth=slice(-50, 90))
    .median("depth")
    .sel(longitude=-125, method="nearest")
    .sel(time=slice("1995-Dec", "1996-Aug-01"))
    .plot(x="time", robust=True, vmax=0.8)
)

<matplotlib.collections.QuadMesh at 0x2ade290ed210>

import itertools

import facetgrid

labels = [lon + "°W" for lon in str(np.abs(medianRi.longitude.values))[1:-1].split(" ")]
cbkwargs = {"orientation": "horizontal", "shrink": 0.8, "aspect": 30}


def plot_euc(ax, loc):
    kwargs = dict(ax=ax, _labels=False, add_legend=False, hue="variable", zorder=10)
    eucmax.sel(**loc).to_array().plot.line(color="w", lw=1.5, **kwargs)
    eucmax.sel(**loc).to_array().plot.line(color="k", lw=0.75, **kwargs)
    ax.set_title("")


with plt.rc_context({"font.size": 10}):
    fg = facetgrid.facetgrid(row=medianRi.longitude, col=["Ri", "SST"])
    fg.map_col(
        "Ri",
        medianRi.sel(latitude=slice(-7, 7)),
        xr.plot.pcolormesh,
        x="time",
        levels=np.arange(0.3, 0.61, 0.1),
        cmap=mpl.cm.Reds_r,
    )
    fg.map_col(
        "SST",
        sst.sel(latitude=slice(-7, 7)),
        xr.plot.pcolormesh,
        robust=True,
        cmap=mpl.cm.Blues_r,
    )
    fg.set_ylabels("latitude")
    # fg.set_row_labels(labels)

    fg.fig.set_size_inches(dcpy.plots.pub_fig_width("jpo", "two column"), 8)
    fg.fig.colorbar(
        fg.handles["Ri"][0],
        ax=fg.axes[:, 0],
        **cbkwargs,
        label="Median $Ri$ within 40m of mixed layer base",
    )
    fg.fig.colorbar(
        fg.handles["SST"][0],
        ax=fg.axes[:, 1],
        **cbkwargs,
        label="SST [°C]",
        extend="both",
    )
    dcpy.plots.label_subplots(
        fg.axes.flat,
        y=0.09,
        labels=list(itertools.chain(*[(v, v) for v in labels])),
        color="w",
        bbox=dict(alpha=0.4, color="black", linewidth=0),
        # fontweight="bold",
        fontsize="small",
    )
    for loc, ax in zip(fg.row_locs, fg.axes):
        [plot_euc(aa, {"longitude": loc}) for aa in ax]
    [dcpy.plots.concise_date_formatter(ax, show_offset=False) for ax in fg.axes[-1, :]]

fg.axes[0, 0].set_xlim((None, "1996-Apr-01"))
fg.fig.savefig("images/Ri-lat-time.png", dpi=300)

/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/plot/plot.py:970: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  primitive = ax.pcolormesh(x, y, z, **kwargs)

labels = [lon + "°W" for lon in str(np.abs(medianRi.longitude.values))[1:-1].split(" ")]

with plt.rc_context({"font.size": 10}):
    f, ax = plt.subplots(4, 2, sharex=True, sharey=True)

    fg = medianRi.sel(latitude=slice(-7, 7)).plot.contourf(
        aspect=3,
        x="time",
        levels=np.arange(0.3, 0.81, 0.1),
        cmap=mpl.cm.Spectral,
        row="longitude",
        add_colorbar=False,
    )

    levels = [24.6, 24.95, 25, 26.25]
    kwargs = dict(_labels=False, add_legend=False, hue="variable", zorder=10)
    for loc, ax, level in zip(fg.name_dicts.flat, fg.axes.flat, levels):
        # hc = sst.sel(**loc, method="nearest").plot.contour(
        #    ax=ax, x="time", levels=[level], colors="k", add_labels=False, linewidths=1,
        # )
        kwargs.update(ax=ax)
        eucmax.sel(**loc).to_array().plot.line(color="w", lw=1.5, **kwargs)
        eucmax.sel(**loc).to_array().plot.line(color="k", lw=0.75, **kwargs)
        ax.set_title("")

    fg.set_xlabels("")
    dcpy.plots.label_subplots(
        fg.axes.flat, y=0.1, labels=labels, fontsize="small", color="w"
    )
    fg.fig.set_size_inches(dcpy.plots.pub_fig_width("jpo", "medium 2"), 8)
    plt.tight_layout(pad=0.7)
    fg.add_colorbar(**{"orientation": "horizontal", "shrink": 0.8, "aspect": 40})

# fg.fig.savefig("images/Ri-lat-time.png", dpi=300)

Figure 6: TIW lat-depth slices

subset = (
    period4.set_coords(["mld", "dcl_base"])
    .sel(latitude=slice(-3, 6))
    .isel(longitude=1)[["sst", "uz", "vz", "mld", "dcl_base", "dcl_mask", "S2", "N2"]]
).compute()

pump.plot.plot_tiw_lat_lon_summary(subset, times=snapshot_times)
plt.gcf().savefig("images/surface-summary.png", dpi=150)

<Figure size 906.6x700 with 0 Axes>

Figure 7: Shred2 depth-lat snapshots

pump.plot.plot_tiw_period_snapshots(
    period4.set_coords(["mld", "dcl_base"]).sel(latitude=slice(-5, 5)),
    lon=-110,
    period=None,
    times=snapshot_times,
)
plt.gcf().savefig("images/dcl-shred2.png", dpi=150)

/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/plot/plot.py:970: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  primitive = ax.pcolormesh(x, y, z, **kwargs)
/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/plot/plot.py:970: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  primitive = ax.pcolormesh(x, y, z, **kwargs)
/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/plot/plot.py:970: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  primitive = ax.pcolormesh(x, y, z, **kwargs)

Figure 7: horizontal vorticity vector

dirname = "/glade/campaign/cgd/oce/people/dcherian/TPOS_MITgcm_1_hb/HOLD"

tiw_all = xr.open_mfdataset(f"{dirname}/File*.nc", parallel=True)

tiw = tiw_all.sel(time="1995-Dec-01 00:00", longitude=slice(-115, -106))
tiw_vort = pump.calc.vorticity(tiw)
tiw_vort

<xarray.Dataset>
Dimensions:    (depth: 280, latitude: 400, longitude: 180)
Coordinates:
  * latitude   (latitude) float32 -10.0 -9.95 -9.9 -9.85 ... 9.85 9.9 9.95 10.0
  * longitude  (longitude) float32 -115.0 -114.9 -114.9 ... -106.1 -106.1 -106.0
  * depth      (depth) float32 -0.5 -1.5 -2.5 -3.5 ... -378.4 -391.7 -406.3
    time       datetime64[ns] 1995-12-01
Data variables:
    x          (depth, latitude, longitude) float32 dask.array<chunksize=(280, 400, 180), meta=np.ndarray>
    y          (depth, latitude, longitude) float32 dask.array<chunksize=(280, 400, 180), meta=np.ndarray>
    z          (depth, latitude, longitude) float32 dask.array<chunksize=(280, 400, 180), meta=np.ndarray>
    vx         (depth, latitude, longitude) float32 dask.array<chunksize=(280, 400, 180), meta=np.ndarray>
    uy         (depth, latitude, longitude) float32 dask.array<chunksize=(280, 400, 180), meta=np.ndarray>
    f          (latitude) float32 -2.526e-05 -2.513e-05 ... 2.513e-05 2.526e-05

tiw["dens"] = pump.mdjwf.dens(tiw.salt, tiw.theta, tiw.depth)
tiw = pump.calc.calc_reduced_shear(tiw)
tiw["mld"] = pump.calc.get_mld(tiw.dens)
tiw["dcl_base"] = pump.calc.get_dcl_base_Ri(tiw)
tiw["dcl"] = tiw["mld"] - tiw["dcl_base"]

calc uz
calc vz
calc S2
calc N2
calc shred2
Calc Ri

f, ax = plt.subplots(2, 1, sharey=True, constrained_layout=True)
pump.plot.vor_streamplot(
    period4, vort, stride=3, refspeed=0.5, ax=ax[0], colorbar=False
)
pump.plot.vor_streamplot(
    tiw, tiw_vort, stride=2, refspeed=0.5, x="longitude", scale=0.4, ax=ax[1]
)
f.set_size_inches((dcpy.plots.pub_fig_width("jpo", "medium 1"), 8))
dcpy.plots.label_subplots(ax, y=1.02, x=0.025, labels=["110°W", "1995-Dec-01 00:00"])
plt.gcf().savefig("../images/110-period-4-vort-vec-rot.png")

[Text(0.025, 1.02, '(a) 110°W'), Text(0.025, 1.02, '(b) 1995-Dec-01 00:00')]

x = "longitude"
refspeed = 0.5

tiw = tiw_all.sel(time="1995-Dec-01 00:00", longitude=slice(-115, -106))

subset = (
    tiw[["u", "v"]]
    .sel(latitude=slice(-1, 5), depth=slice(0, -70))
    .isel(latitude=slice(0, -1, 1))
    .mean("depth")
).compute()


dx = np.arange(0, 0.05 * subset.sizes["longitude"], 0.05)

dy = subset.latitude.data

dx, y, u, v = dask.compute(
    dx,
    dy,
    subset.u.transpose("latitude", x) + refspeed,
    subset.v.transpose("latitude", x),
)

plt.streamplot(
    dx,
    dy,
    u.values,
    v.values,
    color="k",
    linewidth=3,
    arrowstyle="-",
    density=0.8,
)

<matplotlib.streamplot.StreamplotSet at 0x2b0fe5d1d1f0>

pump.plot.vor_streamplot(period4, vort, stride=3, refspeed=0.5)

plt.gcf().savefig("../images/110-period-4-vort-vec-rot.png")

Figure 8: shear evolution + snapshots

duzdt, dvzdt = dask.compute(*pump.calc.estimate_shear_evolution_terms(period4))
duzdt

<xarray.Dataset>
Dimensions:    (depth: 200, latitude: 400, time: 779)
Coordinates:
  * depth      (depth) float64 -0.5 -1.5 -2.5 -3.5 ... -197.5 -198.5 -199.5
  * latitude   (latitude) float32 -10.0 -9.949875 -9.89975 ... 9.949875 10.0
    longitude  float32 -110.00916
  * time       (time) datetime64[ns] 1995-11-14T17:00:00 ... 1995-12-17T03:00:00
Data variables:
    xadv       (time, depth, latitude) float32 2.5738722e-09 ... 2.641062e-10
    yadv       (time, depth, latitude) float32 1.7001174e-10 ... -1.6879712e-10
    str        (time, depth, latitude) float32 1.6083877e-08 ... 1.447781e-10
    tilt       (latitude, time, depth) float32 -2.5338989e-07 ... 2.8220213e-09
    vtilt      (latitude, time, depth) float32 -2.2721191e-07 ... 2.814436e-09
    htilt      (time, depth, latitude) float32 -2.617798e-08 ... 7.58544e-12
    buoy       (time, depth, latitude) float32 1.524855e-08 ... 2.5125453e-09
    fric       (time, depth, latitude) float32 nan nan nan nan ... nan nan nan
Attributes:
    description:  Zonal shear evolution terms

dcl = period4.dcl.resample(time="D").mean().rolling(latitude=3).mean().isel(longitude=1)

mpl.rcParams["hatch.color"] = "lightgray"
dcl.plot(x="time", robust=True)
dcl.where(dcl < 20).plot.contourf(x="time", hatches=".", colors="w")

<matplotlib.contour.QuadContourSet at 0x2b1e48139e48>

f, ax = pump.plot.plot_shear_terms(
    period4,
    duzdt.drop_vars(["xadv", "yadv"]).resample(time="D").mean(),
    dvzdt.drop_vars(["xadv", "yadv"]).resample(time="D").mean(),
    mask_dcl=False,
)

[
    dcpy.plots.linex(snapshot_times, zorder=3, ax=aa, lw=1, ls="--", color="k")
    for aa in unpack(ax)
]

f.savefig("images/shear-evolution.png")

/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/core/nanops.py:142: RuntimeWarning: Mean of empty slice
  return np.nanmean(a, axis=axis, dtype=dtype)

(ddy(period4.u) + ddx(period4.v)).isel(longitude=1).sel(depth=slice(-60)).mean(
    "depth"
).plot(x="time", robust=True)

<matplotlib.collections.QuadMesh at 0x2b09bee68cd0>

S2 = xr.Dataset()
S2["str"] = duzdt["str"] * period4.uz.isel(longitude=1) + dvzdt[
    "str"
] * period4.vz.isel(longitude=1)
S2["htilt"] = (-period4.uz * period4.vz * (ddy(period4.u) + ddx(period4.v))).isel(
    longitude=1
)
S2

<xarray.Dataset>
Dimensions:    (depth: 200, latitude: 400, time: 779)
Coordinates:
  * depth      (depth) float64 -0.5 -1.5 -2.5 -3.5 ... -197.5 -198.5 -199.5
  * latitude   (latitude) float32 -10.0 -9.949875 -9.89975 ... 9.949875 10.0
    longitude  float32 -110.00916
  * time       (time) datetime64[ns] 1995-11-14T17:00:00 ... 1995-12-17T03:00:00
Data variables:
    str        (time, depth, latitude) float32 dask.array<chunksize=(50, 200, 400), meta=np.ndarray>
    htilt      (time, depth, latitude) float32 dask.array<chunksize=(50, 200, 400), meta=np.ndarray>

period4.S2.where(period4.dcl_mask).mean("depth").isel(longitude=1).plot(
    robust=True, x="time", vmin=0, cmap=mpl.cm.Spectral_r
)

<matplotlib.collections.QuadMesh at 0x2b0b0c89c0a0>

S2.to_array().sel(depth=slice(-60)).mean("depth").sum("variable").plot(
    robust=True, x="time", vmax=5e-10, ylim=(-3, 6)
)

<matplotlib.collections.QuadMesh at 0x2acaecbdafd0>

S2.to_array().where(period4.dcl_mask.isel(longitude=1)).sel(depth=slice(-60)).mean(
    "depth"
).plot(col="variable", robust=True, x="time", vmax=5e-10, ylim=(-3, 6))

<xarray.plot.facetgrid.FacetGrid at 0x2b0b007288e0>

ddy(period4.w).isel(longitude=1).sel(depth=slice(-60)).mean("depth").plot(
    x="time", robust=True
)

<matplotlib.collections.QuadMesh at 0x2b09be691c70>

uyvx.isel(longitude=1).sel(depth=slice(-60)).mean("depth").plot(
    y="latitude", robust=True
)

<matplotlib.collections.QuadMesh at 0x2b09bae84ac0>

Figure 9: KPP deep cycle for period4

%matplotlib inline

subset.cf.describe()

Axes:
	X: []
	Y: []
	Z: ['depth']
	T: ['time']

Coordinates:
	longitude: []
	latitude: []
	vertical: []
	time: ['time']

Cell Measures:
	area: []
	volume: []

Standard Names:

subset.cf.axes

{'Z': ['depth'], 'T': ['time']}

jra = pump.obs.read_jra_95()
subset = period4.isel(longitude=1).sel(latitude=3.5, method="nearest").load()

subset.depth.attrs["axis"] = "Z"
forcing = pump.obs.interp_jra_to_station(jra, subset)
fluxes, coare = pump.calc.coare_fluxes_jra(subset.expand_dims("latitude"), forcing)
fluxes["netflux"] *= -1  # sign convention matches Jq
kpp = pump.calc.calc_kpp_terms(subset.merge(fluxes).chunk({"depth": -1}).fillna(0))

/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:777: RuntimeWarning: invalid value encountered in power
  x = (1 - 18 * zet) ** 0.25
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:780: RuntimeWarning: invalid value encountered in power
  x = (1 - 10 * zet) ** 0.3333
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:731: RuntimeWarning: invalid value encountered in sqrt
  x = (1 - 15 * zet) ** 0.5
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:733: RuntimeWarning: invalid value encountered in power
  x = (1 - 34.15 * zet) ** 0.3333
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:755: RuntimeWarning: invalid value encountered in power
  x = (1 - 15 * zet) ** 0.25
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:757: RuntimeWarning: invalid value encountered in power
  x = (1 - 10.15 * zet) ** 0.3333
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:516: RuntimeWarning: invalid value encountered in power
  ug = np.where(Bf > 0, Beta * (Bf * zi) ** 0.333, 0.2)
/glade/u/home/dcherian/pump/pump/calc.py:1147: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.
  h = xr.apply_ufunc(

import hvplot.xarray

(
    (-1 * kpp.hmonob).sel(time=slice("1995-12-03", "1995-12-04")).hvplot.line()
    * (-1 * subset.KPPhbl).sel(time=slice("1995-12-03", "1995-12-04")).hvplot.line()
)

%matplotlib inline

subset2 = xr.merge([subset, fluxes]).squeeze()
f, ax = pump.plot.plot_dcl(
    subset2.sel(time=slice("1995-12-03", "1995-12-04"), depth=slice(-100)),
    kpp_terms=False,
    lw=1.5,
)
ax[0, 0].set_title("(3.5°N, 110°W)")
ax[0, 0].set_ylim(None, 300)
# kpp.hmonob.plot(ax=ax[1, 0], zorder=14, color='g', marker='o')
dcpy.plots.label_subplots(ax.flat)
f.set_size_inches(dcpy.plots.pub_fig_width("jpo", "two column"), 8)
f.savefig("../images/kpp-deep-cycle-35.png")

debugging

budget = xr.open_mfdataset(
    "/glade/campaign/cgd/oce/people/bachman/TPOS_MITgcm_1_hb/HOLD_NC_LINKS/Day_0[4-6]*sf.nc",
    parallel=True,
)
budget["time"] = budget.time - pd.Timedelta("7h")
netflux = (
    budget.oceQnet.sel(time=slice(period4.time[0].values, period4.time[-1].values))
    .interp(longitude=subset.longitude, latitude=subset.latitude)
    .compute()
)

xr.set_options(display_style="text")

<xarray.core.options.set_options at 0x2baf840c0ee0>

%matplotlib widget

plt.figure()
netflux.plot(x="time", marker="x")
fluxes.netflux.assign_coords(time=fluxes.time + pd.Timedelta("24h")).plot(
    x="time", marker="."
)
fluxes.netflux.assign_coords(time=fluxes.time + pd.Timedelta("24h")).interp(
    time=netflux.time
).plot(marker="o", ls="none")
plt.xlim("1995-11-30", "1995-12-01 00:00")

(9464.0, 9465.0)

## binary

file = "/glade/scratch/bachman/TPOS_MITgcm_1_hb/swrad.bin_1995-97"
array = np.memmap(file, dtype=np.float64, mode="r", order="C", shape=(4377, 480, 1500))
swrad = xr.DataArray(
    array,
    dims=["time", "latitude", "longitude"],
    coords={
        "time": jra.time,  # pd.date_range("1995-09-01", end="1997-03-01", freq="3H") - pd.Timedelta("7h"),
        "latitude": budget.latitude,
        "longitude": budget.longitude,
    },
)
binary = swrad.sel(time=slice(period4.time[0].values, period4.time[-1].values)).interp(
    latitude=3.5, longitude=-110
)
binary = binary.copy(data=binary.data.byteswap())

# JRA SW
jra_sw = jra.rsds.sel(
    time=slice(period4.time[0].values, period4.time[-1].values)
).interp(latitude=3.5, longitude=-110)

plt.figure()
# (0.9 * binary).plot()
(0.9 * jra_sw).plot()
# budget.oceQsw.sel(latitude=0, longitude=-110, method="nearest").sel(
#    time=slice(period4.time[0].values, period4.time[-1].values)
# ).fillna(0).plot()
fluxes.short.assign_coords(time=fluxes.time).plot()

[<matplotlib.lines.Line2D at 0x2baf6ec0b730>]

plt.figure()

subset.KPPbo.plot(marker="x")

[<matplotlib.lines.Line2D at 0x2b43adaa4940>]

Figure 10: partition KPP heat flux

Jq = xr.Dataset()
Jq["Total DCL"] = period4.Jq.where(period4.dcl_mask)
Jq["Within KPP BL"] = period4.Jq.where(
    period4.dcl_mask & (period4.depth > -period4.KPPhbl)
)
Jq["Below KPP BL"] = period4.Jq.where(
    period4.dcl_mask & (period4.depth <= -period4.KPPhbl)
)
Jq = Jq.isel(longitude=1).to_array().load()

S2 = (
    1e4
    * period4.S2.where(period4.dcl_mask)
    .isel(longitude=1)
    .mean(["time", "depth"])
    .load()
)
S2.attrs["long_name"] = "Median $S²$ in low Ri layer"
S2.attrs["units"] = "$10^{-4} s^{-2}$"

N2 = (
    1e4
    * period4.N2.where(period4.dcl_mask)
    .isel(longitude=1)
    .mean(["time", "depth"])
    .load()
)
# N2.attrs["long_name"] = "Median $N²$ in low Ri layer"
N2.attrs["units"] = "$10^{-4} s^{-2}$"

kpphbl = (
    period4.KPPhbl.where(period4.dcl_mask.sum("depth") > 5).isel(longitude=1).compute()
)
kpphbl.attrs["long_name"] = "$H_{KPP}$"
kpphbl.attrs["units"] = "m"

zri = (
    np.abs(period4.dcl_base.where(period4.dcl_mask.sum("depth") > 5))
    .isel(longitude=1)
    .compute()
)
zri.attrs["long_name"] = "$z_{MLD} - z_{Ri}$"
zri.attrs["units"] = "m"

mld = (
    np.abs(period4.mld.where(period4.dcl_mask.sum("depth") > 5))
    .isel(longitude=1)
    .compute()
)
mld.attrs["long_name"] = "$z_{MLD} - z_{Ri}$"
mld.attrs["units"] = "m"

f, axx = plt.subplots(1, 3, constrained_layout=1, squeeze=False, sharey=True)

ax = axx[0, 0]
(
    (
        (-1 * Jq).fillna(0).integrate("time", datetime_unit="s").integrate("depth")
        / (86400 * 30 * 50)
    ).plot.line(
        hue="variable",
        y="latitude",
        add_legend=False,
        ax=ax,
    )
)

N2.plot(ax=axx[0, 1], color="k", y="latitude")
S2.plot(ax=axx[0, 1], color="C2", y="latitude")
axx[0, 1].legend(["$N²$", "$S²$"])
axx[0, 1].set_xlabel("Mean $N²,S²$ in low Ri layer \n[$10^{-4}$ s$^{-2}$]")

(
    (mld).mean("time")
    # .rolling(latitude=5, center=True)
    # .mean()
    .plot(ax=axx[0, 2], y="latitude", color="C1")
)
(
    (kpphbl).mean("time")
    # .rolling(latitude=5, center=True)
    # .mean()
    .plot(ax=axx[0, 2], y="latitude", color="k")
)
(
    (zri).mean("time")
    # .rolling(latitude=5, center=True)
    # .mean()
    .plot(ax=axx[0, 2], y="latitude", color="C2")
)

axx[0, 2].legend(["$z_{MLD}$", "$H_{KPP}$", "$z_{Ri}$"])


# (
#    zri.mean("time")
# .rolling(latitude=5, center=True)
# .mean()
#    .plot(ax=axx[0, 2], y="latitude", color="k")
# )

axx[0, 2].set_xlabel("[m]")
ax.legend(Jq["variable"].values, loc="lower right")
ax.set_xlabel("Avg $J_q$ [W/m²]")
# ax.set_title("Parameterized heat flux \nin low Ri layer at 110°W")
dcpy.plots.clean_axes(axx)
dcpy.plots.label_subplots(axx.flat)

for aa in axx.flat:
    aa.set_title("")
    aa.xaxis.set_minor_locator(mpl.ticker.AutoMinorLocator(4))
    aa.grid(True, which="both")
    aa.set_ylim((-3, 5))

for aa in axx.flat[1:]:
    aa.set_xlim((0, None))
axx[0, 0].set_xlim((None, 0))

f.set_size_inches((dcpy.plots.pub_fig_width("jpo", "medium 2"), 5))

f.savefig("images/kpp-deep-cycle-bl-int-heat-flux.png")

N² budget

def avg1(ds, dim):
    return (ds + ds.shift({dim: -1})) / 2

%matplotlib inline

RAC = metrics.sel(latitude=subset.latitude, longitude=subset.longitude).RAC

plt.figure()
subset = period4.isel(longitude=1, time=120).sel(latitude=0, method="nearest")

myJq = (
    1035 * 3999 * (subset.KPPdiffT * subset.theta.diff("depth", label="upper"))
) + subset.nonlocal_flux.fillna(0)

(subset.Jq).plot(y="depth")
myJq.plot(y="depth", ylim=(-120, 0))
dcpy.plots.liney(-1 * subset.KPPhbl)

period4

<xarray.Dataset>
Dimensions:        (depth: 200, latitude: 400, longitude: 3, time: 779)
Coordinates:
  * depth          (depth) float64 -0.5 -1.5 -2.5 -3.5 ... -197.5 -198.5 -199.5
  * latitude       (latitude) float32 -10.0 -9.949875 -9.89975 ... 9.949875 10.0
  * longitude      (longitude) float32 -110.0592 -110.00916 -109.95913
  * time           (time) datetime64[ns] 1995-11-14T17:00:00 ... 1995-12-17T0...
Data variables:
    DFrI_TH        (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    ETAN           (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    Jq             (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    KPPRi          (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    KPPbo          (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    KPPdiffT       (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    KPPg_TH        (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    KPPhbl         (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    KPPviscA       (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    N2             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    Ri             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    S2             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    Um_Diss        (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    VISrI_Um       (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    VISrI_Vm       (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    Vm_Diss        (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    dcl            (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    dcl_base       (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    dens           (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    eucmax         (time, longitude) float64 dask.array<chunksize=(779, 3), meta=np.ndarray>
    mld            (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    salt           (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    shear          (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    shred2         (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    sst            (time, latitude, longitude) float32 dask.array<chunksize=(1, 400, 3), meta=np.ndarray>
    theta          (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    u              (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    uz             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    v              (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    vz             (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    w              (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    Um_Impl        (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    Vm_Impl        (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    nonlocal_flux  (time, depth, latitude, longitude) float64 dask.array<chunksize=(1, 200, 400, 3), meta=np.ndarray>
    dcl_mask       (depth, time, latitude, longitude) bool dask.array<chunksize=(200, 1, 400, 3), meta=np.ndarray>
Attributes:
    easting:            longitude
    field_julian_date:  400296
    julian_day_unit:    days since 1950-01-01 00:00:00
    northing:           latitude
    title:              daily snapshot from 1/20 degree Equatorial Pacific MI...

ds = period4

ds["b"] = -9.81 / 1035 * period4.dens

from pump.calc import ddx, ddy, ddz


dN2dt = xr.Dataset()
dN2dt["xadv"] = -ds.u * ddx(ds.N2)
dN2dt["yadv"] = -ds.v * ddy(ds.N2)
dN2dt["zadv"] = -ds.w * ddz(ds.N2)
dN2dt["uzbx"] = -ds.uz * ddx(ds.b)
dN2dt["vzby"] = -ds.vz * ddy(ds.b)
dN2dt["wzbz"] = -ddz(ds.w) * ddz(ds.b)
dN2dt["mix"] = -ddz(9.81 * 2e-4 * ddz(ds.Jq / 1035 / 3995))

dN2dt["adv"] = dN2dt.xadv + dN2dt.yadv + dN2dt.zadv
dN2dt["tilt"] = dN2dt.uzbx + dN2dt.vzby + dN2dt.wzbz
dN2dt["lag"] = dN2dt.tilt + dN2dt.mix
mask = period4.dcl_mask.where(period4.dcl > 30, False)
dN2dt = dN2dt.persist()

f, axx = plt.subplots(
    4, 1, sharex=True, sharey=True, squeeze=False, constrained_layout=True
)

ax = dict(zip(["N2", "lag", "tilt", "mix"], axx.flat))
hN2 = (
    period4.N2.where(mask)
    .isel(longitude=1)
    .mean("depth")
    .plot(
        x="time",
        vmin=1e-5,
        vmax=2e-4,
        cmap=mpl.cm.GnBu,
        levels=np.arange(0, 2.1e-4, 5e-5),
        ylim=(-3, 5),
        ax=ax["N2"],
        add_colorbar=False,
    )
)

for key in list(ax.keys())[1:]:
    hdN2 = (
        -1 * dN2dt[key].where(mask).isel(longitude=1).fillna(0).integrate("depth")
    ).plot(
        ax=ax[key],
        vmax=3e-8,
        x="time",
        ylim=(-3.5, 5),
        add_colorbar=False,
    )
    ax[key].set_title("")

cbN2 = f.colorbar(
    hN2, ax=ax["N2"], orientation="horizontal", aspect=50, label="$N²$ [s$^{-2}$]"
)
fmt = mpl.ticker.ScalarFormatter()
fmt.set_powerlimits((-1, 1))
cbN2.ax.xaxis.set_major_formatter(fmt)

cbdN2 = f.colorbar(
    hdN2,
    ax=axx.flat[1:],
    orientation="horizontal",
    aspect=50,
    label="$∂N²/∂t$ [s$^{-3}$]",
    extend="both",
)

axx[0, 0].set_title("")

[
    period4.dcl.isel(longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(levels=[30], ax=aa, x="time", colors="w")
    for aa in axx.flat
]
[
    period4.dcl.isel(longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(levels=[30], ax=aa, x="time", colors="k", linewidths=1)
    for aa in axx.flat
]

dcpy.plots.label_subplots(
    axx.flat,
    labels=["$N²$", "TILT + STRETCH + MIX", "TILT + STRETCH", "MIX"],
    y=0.05,
    fontsize="small",
)

dcpy.plots.clean_axes(axx)
dcpy.plots.concise_date_formatter(axx.flat[-1], show_offset=False)
axx[0, 0].set_title("")

f.set_size_inches((dcpy.plots.pub_fig_width("jpo", "single column"), 8))
f.savefig("../images/period4-N2.png", dpi=150)

/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/plot/plot.py:970: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  primitive = ax.pcolormesh(x, y, z, **kwargs)

(
    dN2dt[["lag", "tilt", "mix"]]
    .where(mask)
    .isel(longitude=1)
    # .sel(depth=slice(-20, -50))
    .sum("depth")
    .to_array()
    # .resample(time="D")
    # .sum()
    .plot(vmax=3e-8, x="time", col="variable", ylim=(-3.5, 5))
)

<xarray.plot.facetgrid.FacetGrid at 0x2b7e0d350d90>

(
    dN2dt.where(period4.dcl_mask)
    .isel(longitude=1)
    .sum("depth")[["uzbx", "vzby", "wzbz", "fric"]]
    .to_array()
    # .resample(time="D")
    # .mean()
    .plot(col="variable", col_wrap=2, robust=True, x="time")
)

<xarray.plot.facetgrid.FacetGrid at 0x2ae026275100>

f, ax = plt.subplots(1, 3, sharex=True, sharey=True, constrained_layout=True)

kwargs = dict(
    x="time", robust=True, cbar_kwargs={"orientation": "horizontal"}, ylim=(-3, 5)
)

period4.N2.where(period4.dcl_mask).isel(longitude=1).median("depth").plot(
    **kwargs, ax=ax[0], **pump.plot.cmaps["N2"]
)

(
    dN2dt.lag.where(period4.dcl_mask)
    .isel(longitude=1)
    .mean("depth")
    # .resample(time="D")
    # .mean()
    .plot(**kwargs, ax=ax[1])
)

(
    dN2dt.fric.where(period4.dcl_mask)
    .isel(longitude=1)
    .mean("depth")
    # .resample(time="D")
    # .mean()
    .plot(**kwargs, ax=ax[2])
)

dcpy.plots.clean_axes(ax)
for aa in ax.flat:
    aa.set_title("")
    aa.set_xlabel("")

dcpy.plots.label_subplots(
    ax.flat, labels=["$N^2$ [s$^{-2}$]", "Lagrangian terms", "Mixing"]
)

/glade/u/home/dcherian/miniconda3/envs/dcpy/lib/python3.8/site-packages/xarray/plot/plot.py:970: MatplotlibDeprecationWarning: Passing parameters norm and vmin/vmax simultaneously is deprecated since 3.3 and will become an error two minor releases later. Please pass vmin/vmax directly to the norm when creating it.
  primitive = ax.pcolormesh(x, y, z, **kwargs)

[Text(0.05, 0.9, '(a) $N^2$ [s$^{-2}$]'),
 Text(0.05, 0.9, '(b) Lagrangian terms'),
 Text(0.05, 0.9, '(c) Mixing')]

(
    period4.N2.isel(longitude=1)
    .sel(time="1995-12-05 12:00")
    .plot(y="depth", robust=True, vmax=1e-4, cmap=mpl.cm.Greens)
)

<matplotlib.collections.QuadMesh at 0x2ac09b2fde10>

f, ax = plt.subplots(1, 2, sharex=True, sharey=True, constrained_layout=True)

period4.isel(longitude=1).mld.resample(time="D").min().plot.contourf(
    x="time",
    robust=True,
    levels=[0, -10, -20, -30, -40, -50],
    cmap=mpl.cm.Blues_r,
    ax=ax[0],
    cbar_kwargs={"orientation": "horizontal"},
)
(
    period4.dcl.isel(longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(
        x="time",
        levels=[30, 60],
        add_labels=False,
        colors="k",
        ax=ax[0],
    )
)

period4.isel(longitude=1).dcl_base.resample(time="D").min().plot.contourf(
    x="time",
    robust=True,
    levels=[0, -40, -50, -60, -70, -80, -90],
    cmap=mpl.cm.Greens_r,
    ax=ax[1],
    cbar_kwargs={"orientation": "horizontal"},
)
(
    period4.dcl.isel(longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(
        x="time",
        levels=[30, 60],
        add_labels=False,
        colors="k",
        ax=ax[1],
    )
)

ax[1].set_ylim((-3, 5))

(-3.0, 5.0)

(
    period4.N2.isel(longitude=1)
    .sel(depth=slice(-40, -60), latitude=slice(-2, 6.5))
    .mean("depth")
).plot(robust=True, x="time", levels=np.arange(5e-5, 2e-4, 2.5e-5), cmap=mpl.cm.Greens)
(
    period4.dcl.isel(longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(
        x="time",
        levels=[30, 60],
        add_labels=False,
        colors="k",
    )
)

<matplotlib.contour.QuadContourSet at 0x2ac079403d50>

DCL mask is reasonable

(
    period4.Jq.where(period4.dcl_mask)
    .isel(longitude=1)
    .mean("depth")
    .resample(time="D")
    .mean()
    .plot(x="time", robust=True)
)

<matplotlib.collections.QuadMesh at 0x2ac07d828110>

def subber(ds):
    return (
        ds.sel(time="1995-12-05")
        .isel(time=23, longitude=1)
        .sel(latitude=slice(-3, 7), depth=slice(-120))
    )


subber(dN2dt).fric.rolling(depth=2).mean().plot(
    y="depth", robust=True, cmap=mpl.cm.BrBG_r
)
# (-subber(period4).Jq.rolling(depth=2).mean().compute().pipe(ddz)).plot.contour(
#    y="depth", robust=True, levels=12, colors="k", linewidths=0.25
# )
subber(period4).mld.plot(color="orange")
subber(period4).dcl_base.plot(color="k")
dcpy.plots.liney((-40, -60), zorder=10)

dN2dt.lag.where(period4.dcl_mask).isel(longitude=1).mean("depth").plot(
    x="time", robust=True
)

<matplotlib.collections.QuadMesh at 0x2ae07db30610>

%matplotlib inline

ddt = (
    dN2dt.where(period4.dcl_mask)
    .isel(longitude=1)
    .sel(latitude=slice(-2, 6.5))
    .mean("depth")
    .resample(time="D")
    .mean()
    .to_array()
).compute()

fg = ddt.plot(
    col="variable",
    x="time",
    robust=True,
    col_wrap=3,
    vmax=4e-10,
    cbar_kwargs={"orientation": "horizontal", "label": "∂N²/∂t"},
)

fg.map(
    lambda: period4.dcl.isel(longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(x="time", levels=[30, 60], add_labels=False, colors="k")
)

fg.fig.suptitle("Daily avg(N² budget terms averaged over DCL)", y=1.01)
fg.fig.savefig("images/N2-budget-period4.png")

building intuition for effect of mixing on N²

def subber(ds):
    return (
        ds.sel(time="1995-12-05")
        .isel(time=23, longitude=1)
        .sel(latitude=slice(-3, 7), depth=slice(-120))
    )


(
    subber(period4.Jq)
    .rolling(depth=12, center=True)
    .mean()
    .compute()
    .diff("depth")
    .rolling(depth=12, center=True)
    .mean()
    .diff("depth")
    .rolling(depth=12, center=True)
    .mean()
    .sel(latitude=0, method="nearest")
    .plot()
)

[<matplotlib.lines.Line2D at 0x2ac0ac4eac90>]

z = ds.depth[:60]
J = np.sin(-2 * np.pi * z / 50)
J.plot(y="depth")
(-J.differentiate("depth") * 10).plot(y="depth")
(-J.differentiate("depth").differentiate("depth") * 100).plot(y="depth")

[<matplotlib.lines.Line2D at 0x2af3782eeb90>]

Supplementary Figures

shear advection terms

f, ax = pump.plot.plot_shear_terms(
    period4,
    duzdt[["xadv", "yadv"]].resample(time="D").mean(),
    dvzdt[["xadv", "yadv"]].resample(time="D").mean(),
    mask_dcl=False,
)

[
    dcpy.plots.linex(snapshot_times, zorder=3, ax=aa, lw=1, ls="--", color="k")
    for aa in unpack(ax)
]

f.savefig("images/shear-evolution-advective.png")

KPP at eq

jra = pump.obs.read_jra_95()
subset = period4.isel(longitude=1).sel(latitude=0, method="nearest").load()
forcing = pump.obs.interp_jra_to_station(jra, subset)
fluxes, coare = pump.calc.coare_fluxes_jra(subset, forcing)
kpp = pump.calc.calc_kpp_terms(subset.merge(fluxes).chunk({"depth": -1}).fillna(0))

/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:733: RuntimeWarning: invalid value encountered in power
  x = (1 - 34.15 * zet) ** 0.3333
/glade/u/home/dcherian/python/xcoare/xcoare/coare35vn.py:516: RuntimeWarning: invalid value encountered in power
  ug = np.where(Bf > 0, Beta * (Bf * zi) ** 0.333, 0.2)

%matplotlib inline
for sub in [
    period4.isel(longitude=1).sel(latitude=0, method="nearest"),
    period4.isel(longitude=1).sel(latitude=3.5, method="nearest"),
]:
    (sub.hRib - (-1 * sub.KPPhbl)).plot()

%matplotlib widget


subset2.hRib.plot(color="r", lw=3)
(-1 * subset2.KPPhbl).plot(color="k", lw=1)

[<matplotlib.lines.Line2D at 0x2b4c11699490>]

subset2 = xr.merge([subset, fluxes])
f, ax = pump.plot.plot_dcl(
    subset2.sel(time=slice("1995-11-17", "1995-11-27"), depth=slice(-40)),
    kpp_terms=False,
    lw=1.5,
)