Calculate POP2 heat budget using xgcm#
In this notebook, we are going to use xgcm with metrics to demonstrate budget closure for the 0.1 degree horizontal resolution version of POP2. Note that the lower resolution has more parameterizations and therefore does not close following this notebook. This notebook was contributed by Anna-Lena Deppenmeier.
This is an image of the POP output structure on the horizontal B-grid courtesy of Yassir Eddebbar.

Import packages and define functions#
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from tqdm import tqdm
import pop_tools
def pop_find_lat_ind(loc, LATDAT):
return np.abs(LATDAT[:, 0].values - loc).argmin()
def pop_find_lon_ind(loc, LONDAT, direction="w"):
if direction.lower() in ["east", "e"]:
value = loc
elif direction.lower() in ["west", "w"]:
value = 360 - loc
else:
print("I do not know which direction.")
return np.nanargmin(np.abs(LONDAT[152, :].values - value))
Load Dataset#
# open sample data
filepath = pop_tools.DATASETS.fetch('Pac_POP0.1_JRA_IAF_1993-12-6-test.nc')
ds = xr.open_dataset(filepath)
# get DZU and DZT, needed for operations later on
filepath_g = pop_tools.DATASETS.fetch("Pac_grid_pbc_1301x305x62.tx01_62l.2013-07-13.nc")
ds_g = xr.open_dataset(filepath_g)
ds
<xarray.Dataset> Dimensions: (nlat: 305, nlon: 1301, time: 1, z_t: 62, z_w: 62, z_w_bot: 62, z_w_top: 62) Coordinates: * time (time) object 0036-12-07 00:00:00 * z_t (z_t) float32 500.0 1.5e+03 2.5e+03 ... 5.625e+05 5.875e+05 * z_w (z_w) float32 0.0 1e+03 2e+03 ... 5.25e+05 5.5e+05 5.75e+05 * z_w_top (z_w_top) float32 0.0 1e+03 2e+03 ... 5.5e+05 5.75e+05 * z_w_bot (z_w_bot) float32 1e+03 2e+03 3e+03 ... 5.75e+05 6e+05 ULONG (nlat, nlon) float64 ... ULAT (nlat, nlon) float64 ... TLONG (nlat, nlon) float64 ... TLAT (nlat, nlon) float64 ... Dimensions without coordinates: nlat, nlon Data variables: (12/26) TEND_TEMP (time, z_t, nlat, nlon) float32 ... UET (time, z_t, nlat, nlon) float32 ... VNT (time, z_t, nlat, nlon) float32 ... WTT (time, z_w_top, nlat, nlon) float32 ... KPP_SRC_TEMP (time, z_t, nlat, nlon) float32 ... UVEL (time, z_t, nlat, nlon) float32 ... ... ... hflux_factor float64 2.439e-05 KMT (nlat, nlon) float64 ... KMU (nlat, nlon) float64 ... DIA_IMPVF_TEMP (time, z_w_bot, nlat, nlon) float32 ... SHF (time, nlat, nlon) float32 ... SHF_QSW (time, nlat, nlon) float32 ... Attributes: title: g.e20.G.TL319_t13.control.001_hfreq history: none Conventions: CF-1.0; http://www.cgd.ucar.edu/cms/eaton/netcdf/CF-cu... time_period_freq: day_5 model_doi_url: https://doi.org/10.5065/D67H1H0V contents: Diagnostic and Prognostic Variables source: CCSM POP2, the CCSM Ocean Component revision: $Id: tavg.F90 89091 2018-04-30 15:58:32Z altuntas@ucar... calendar: All years have exactly 365 days. start_time: This dataset was created on 2018-12-14 at 16:05:58.8 cell_methods: cell_methods = time: mean ==> the variable values are ...
- nlat: 305
- nlon: 1301
- time: 1
- z_t: 62
- z_w: 62
- z_w_bot: 62
- z_w_top: 62
- time(time)object0036-12-07 00:00:00
- long_name :
- time
- bounds :
- time_bound
array([cftime.DatetimeNoLeap(36, 12, 7, 0, 0, 0, 0)], dtype=object)
- z_t(z_t)float32500.0 1.5e+03 ... 5.875e+05
- long_name :
- depth from surface to midpoint of layer
- units :
- centimeters
- positive :
- down
- valid_min :
- 500.0
- valid_max :
- 587499.06
array([5.000000e+02, 1.500000e+03, 2.500000e+03, 3.500000e+03, 4.500000e+03, 5.500000e+03, 6.500000e+03, 7.500000e+03, 8.500000e+03, 9.500000e+03, 1.050000e+04, 1.150000e+04, 1.250000e+04, 1.350000e+04, 1.450000e+04, 1.550000e+04, 1.650984e+04, 1.754790e+04, 1.862913e+04, 1.976603e+04, 2.097114e+04, 2.225783e+04, 2.364088e+04, 2.513702e+04, 2.676542e+04, 2.854837e+04, 3.051192e+04, 3.268680e+04, 3.510935e+04, 3.782276e+04, 4.087846e+04, 4.433777e+04, 4.827367e+04, 5.277280e+04, 5.793729e+04, 6.388626e+04, 7.075633e+04, 7.870025e+04, 8.788252e+04, 9.847059e+04, 1.106204e+05, 1.244567e+05, 1.400497e+05, 1.573946e+05, 1.764003e+05, 1.968944e+05, 2.186457e+05, 2.413972e+05, 2.649001e+05, 2.889385e+05, 3.133405e+05, 3.379793e+05, 3.627670e+05, 3.876452e+05, 4.125768e+05, 4.375392e+05, 4.625190e+05, 4.875083e+05, 5.125028e+05, 5.375000e+05, 5.624991e+05, 5.874991e+05], dtype=float32)
- z_w(z_w)float320.0 1e+03 ... 5.5e+05 5.75e+05
- long_name :
- depth from surface to top of layer
- units :
- centimeters
- positive :
- down
- valid_min :
- 0.0
- valid_max :
- 574999.06
array([ 0. , 1000. , 2000. , 3000. , 4000. , 5000. , 6000. , 7000. , 8000. , 9000. , 10000. , 11000. , 12000. , 13000. , 14000. , 15000. , 16000. , 17019.682, 18076.129, 19182.125, 20349.932, 21592.344, 22923.312, 24358.453, 25915.58 , 27615.26 , 29481.47 , 31542.373, 33831.227, 36387.473, 39258.047, 42498.887, 46176.656, 50370.688, 55174.91 , 60699.668, 67072.86 , 74439.805, 82960.695, 92804.35 , 104136.82 , 117104.016, 131809.36 , 148290.08 , 166499.2 , 186301.44 , 207487.39 , 229803.9 , 252990.4 , 276809.84 , 301067.06 , 325613.84 , 350344.88 , 375189.2 , 400101.16 , 425052.47 , 450026.06 , 475012. , 500004.7 , 525000.94 , 549999.06 , 574999.06 ], dtype=float32)
- z_w_top(z_w_top)float320.0 1e+03 ... 5.5e+05 5.75e+05
- long_name :
- depth from surface to top of layer
- units :
- centimeters
- positive :
- down
- valid_min :
- 0.0
- valid_max :
- 574999.06
array([ 0. , 1000. , 2000. , 3000. , 4000. , 5000. , 6000. , 7000. , 8000. , 9000. , 10000. , 11000. , 12000. , 13000. , 14000. , 15000. , 16000. , 17019.682, 18076.129, 19182.125, 20349.932, 21592.344, 22923.312, 24358.453, 25915.58 , 27615.26 , 29481.47 , 31542.373, 33831.227, 36387.473, 39258.047, 42498.887, 46176.656, 50370.688, 55174.91 , 60699.668, 67072.86 , 74439.805, 82960.695, 92804.35 , 104136.82 , 117104.016, 131809.36 , 148290.08 , 166499.2 , 186301.44 , 207487.39 , 229803.9 , 252990.4 , 276809.84 , 301067.06 , 325613.84 , 350344.88 , 375189.2 , 400101.16 , 425052.47 , 450026.06 , 475012. , 500004.7 , 525000.94 , 549999.06 , 574999.06 ], dtype=float32)
- z_w_bot(z_w_bot)float321e+03 2e+03 ... 5.75e+05 6e+05
- long_name :
- depth from surface to bottom of layer
- units :
- centimeters
- positive :
- down
- valid_min :
- 1000.0
- valid_max :
- 599999.06
array([ 1000. , 2000. , 3000. , 4000. , 5000. , 6000. , 7000. , 8000. , 9000. , 10000. , 11000. , 12000. , 13000. , 14000. , 15000. , 16000. , 17019.682, 18076.129, 19182.125, 20349.932, 21592.344, 22923.312, 24358.453, 25915.58 , 27615.26 , 29481.47 , 31542.373, 33831.227, 36387.473, 39258.047, 42498.887, 46176.656, 50370.688, 55174.91 , 60699.668, 67072.86 , 74439.805, 82960.695, 92804.35 , 104136.82 , 117104.016, 131809.36 , 148290.08 , 166499.2 , 186301.44 , 207487.39 , 229803.9 , 252990.4 , 276809.84 , 301067.06 , 325613.84 , 350344.88 , 375189.2 , 400101.16 , 425052.47 , 450026.06 , 475012. , 500004.7 , 525000.94 , 549999.06 , 574999.06 , 599999.06 ], dtype=float32)
- ULONG(nlat, nlon)float64...
- long_name :
- array of u-grid longitudes
- units :
- degrees_east
[396805 values with dtype=float64]
- ULAT(nlat, nlon)float64...
- long_name :
- array of u-grid latitudes
- units :
- degrees_north
[396805 values with dtype=float64]
- TLONG(nlat, nlon)float64...
- long_name :
- array of t-grid longitudes
- units :
- degrees_east
[396805 values with dtype=float64]
- TLAT(nlat, nlon)float64...
- long_name :
- array of t-grid latitudes
- units :
- degrees_north
[396805 values with dtype=float64]
- TEND_TEMP(time, z_t, nlat, nlon)float32...
- long_name :
- Tendency of Thickness Weighted TEMP
- units :
- degC/s
- grid_loc :
- 3111
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- UET(time, z_t, nlat, nlon)float32...
- long_name :
- Flux of Heat in grid-x direction
- units :
- degC/s
- grid_loc :
- 3211
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- VNT(time, z_t, nlat, nlon)float32...
- long_name :
- Flux of Heat in grid-y direction
- units :
- degC/s
- grid_loc :
- 3121
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- WTT(time, z_w_top, nlat, nlon)float32...
- long_name :
- Heat Flux Across Top Face
- units :
- degC/s
- grid_loc :
- 3112
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- KPP_SRC_TEMP(time, z_t, nlat, nlon)float32...
- long_name :
- TEMP tendency from KPP non local mixing term
- units :
- degC/s
- grid_loc :
- 3111
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- UVEL(time, z_t, nlat, nlon)float32...
- long_name :
- Velocity in grid-x direction
- units :
- centimeter/s
- grid_loc :
- 3221
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- VVEL(time, z_t, nlat, nlon)float32...
- long_name :
- Velocity in grid-y direction
- units :
- centimeter/s
- grid_loc :
- 3221
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- WVEL(time, z_w_top, nlat, nlon)float32...
- long_name :
- Vertical Velocity
- units :
- centimeter/s
- grid_loc :
- 3112
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- TEMP(time, z_t, nlat, nlon)float32...
- long_name :
- Potential Temperature
- units :
- degC
- grid_loc :
- 3111
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- DXT(nlat, nlon)float64...
- long_name :
- x-spacing centered at T points
- units :
- centimeters
[396805 values with dtype=float64]
- DXU(nlat, nlon)float64...
- long_name :
- x-spacing centered at U points
- units :
- centimeters
[396805 values with dtype=float64]
- DYT(nlat, nlon)float64...
- long_name :
- y-spacing centered at T points
- units :
- centimeters
[396805 values with dtype=float64]
- DYU(nlat, nlon)float64...
- long_name :
- y-spacing centered at U points
- units :
- centimeters
[396805 values with dtype=float64]
- TAREA(nlat, nlon)float64...
- long_name :
- area of T cells
- units :
- centimeter^2
[396805 values with dtype=float64]
- UAREA(nlat, nlon)float64...
- long_name :
- area of U cells
- units :
- centimeter^2
[396805 values with dtype=float64]
- QSW_3D(time, z_w_top, nlat, nlon)float32...
- long_name :
- Solar Short-Wave Heat Flux
- units :
- watt/m^2
- grid_loc :
- 3112
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- HDIFE_TEMP(time, z_t, nlat, nlon)float32...
- long_name :
- TEMP Horizontal Diffusive Flux in grid-x direction
- units :
- degC/s
- grid_loc :
- 3211
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- HDIFN_TEMP(time, z_t, nlat, nlon)float32...
- long_name :
- TEMP Horizontal Diffusive Flux in grid-y direction
- units :
- degC/s
- grid_loc :
- 3121
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- dz(z_t)float32...
- long_name :
- thickness of layer k
- units :
- centimeters
array([ 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1019.6808, 1056.4484, 1105.9951, 1167.807 , 1242.4133, 1330.9678, 1435.141 , 1557.1259, 1699.6796, 1866.2124, 2060.9023, 2288.852 , 2556.247 , 2870.575 , 3240.8372, 3677.7725, 4194.031 , 4804.2236, 5524.7544, 6373.192 , 7366.945 , 8520.893 , 9843.658 , 11332.466 , 12967.199 , 14705.344 , 16480.709 , 18209.135 , 19802.234 , 21185.957 , 22316.51 , 23186.494 , 23819.45 , 24257.217 , 24546.78 , 24731.014 , 24844.328 , 24911.975 , 24951.291 , 24973.594 , 24985.96 , 24992.674 , 24996.244 , 24998.11 , 25000. , 25000. ], dtype=float32)
- dzw(z_w)float32...
- long_name :
- midpoint of k to midpoint of k+1
- units :
- centimeters
array([ 500. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1000. , 1009.8404, 1038.0646, 1081.2218, 1136.901 , 1205.1101, 1286.6906, 1383.0544, 1496.1334, 1628.4027, 1782.946 , 1963.5574, 2174.8772, 2422.5496, 2713.4111, 3055.706 , 3459.305 , 3935.9016, 4499.1274, 5164.489 , 5948.973 , 6870.0684, 7943.9185, 9182.275 , 10588.062 , 12149.832 , 13836.271 , 15593.026 , 17344.922 , 19005.686 , 20494.096 , 21751.234 , 22751.502 , 23502.97 , 24038.332 , 24401.998 , 24638.896 , 24787.672 , 24878.15 , 24931.633 , 24962.443 , 24979.777 , 24989.316 , 24994.459 , 24997.176 , 24999.055 , 25000. ], dtype=float32)
- hflux_factor()float64...
- long_name :
- Convert Heat and Solar Flux to Temperature Flux
array(2.439086e-05)
- KMT(nlat, nlon)float64...
- long_name :
- k Index of Deepest Grid Cell on T Grid
[396805 values with dtype=float64]
- KMU(nlat, nlon)float64...
- long_name :
- k Index of Deepest Grid Cell on U Grid
[396805 values with dtype=float64]
- DIA_IMPVF_TEMP(time, z_w_bot, nlat, nlon)float32...
- long_name :
- TEMP Flux Across Bottom Face from Diabatic Implicit Vertical Mixing
- units :
- degC cm/s
- grid_loc :
- 3113
- cell_methods :
- time: mean
[24601910 values with dtype=float32]
- SHF(time, nlat, nlon)float32...
- long_name :
- Total Surface Heat Flux, Including SW
- units :
- watt/m^2
- grid_loc :
- 2110
- cell_methods :
- time: mean
[396805 values with dtype=float32]
- SHF_QSW(time, nlat, nlon)float32...
- long_name :
- Solar Short-Wave Heat Flux
- units :
- watt/m^2
- grid_loc :
- 2110
- cell_methods :
- time: mean
[396805 values with dtype=float32]
- title :
- g.e20.G.TL319_t13.control.001_hfreq
- history :
- none
- Conventions :
- CF-1.0; http://www.cgd.ucar.edu/cms/eaton/netcdf/CF-current.htm
- time_period_freq :
- day_5
- model_doi_url :
- https://doi.org/10.5065/D67H1H0V
- contents :
- Diagnostic and Prognostic Variables
- source :
- CCSM POP2, the CCSM Ocean Component
- revision :
- $Id: tavg.F90 89091 2018-04-30 15:58:32Z altuntas@ucar.edu $
- calendar :
- All years have exactly 365 days.
- start_time :
- This dataset was created on 2018-12-14 at 16:05:58.8
- cell_methods :
- cell_methods = time: mean ==> the variable values are averaged over the time interval between the previous time coordinate and the current one. cell_methods absent ==> the variable values are at the time given by the current time coordinate.
# get lola inds from somewhere for indexing later on
lola_inds = {}
inds_lat = [-8, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 8]
for j in inds_lat:
if j < 0:
lola_inds["j_" + str(j)[1:] + "s"] = pop_find_lat_ind(j, ds_g.TLAT)
else:
lola_inds["j_" + str(j) + "n"] = pop_find_lat_ind(j, ds_g.TLAT)
inds_lon = range(95, 185, 5)
for i in inds_lon:
lola_inds["i_" + str(i) + "_w"] = pop_find_lon_ind(i, ds_g.TLONG)
# just making sure everything works
ds.TEMP.isel(z_t=0).mean(dim="time").plot(levels=np.arange(20, 30.5, 0.5), cmap="RdYlBu_r")
plt.scatter(lola_inds["i_140_w"], lola_inds["j_0n"], marker="*");

Set up vertical thickness and volume for scaling#
ds["DZT"] = ds_g.DZT
ds["DZU"] = ds_g.DZU
ds.DZT.attrs["long_name"] = "Thickness of T cells"
ds.DZT.attrs["units"] = "centimeter"
ds.DZT.attrs["grid_loc"] = "3111"
ds.DZU.attrs["long_name"] = "Thickness of U cells"
ds.DZU.attrs["units"] = "centimeter"
ds.DZU.attrs["grid_loc"] = "3221"
# make sure we have the cell volumne for calculations
VOL = (ds.DZT * ds.DXT * ds.DYT).compute()
KMT = ds_g.KMT.compute()
for j in tqdm(range(len(KMT.nlat))):
for i in range(len(KMT.nlon)):
k = KMT.values[j, i].astype(int)
VOL.values[k:, j, i] = 0.0
ds["VOL"] = VOL
ds.VOL.attrs["long_name"] = "volume of T cells"
ds.VOL.attrs["units"] = "centimeter^3"
ds.VOL.attrs["grid_loc"] = "3111"
100%|██████████| 305/305 [00:01<00:00, 179.47it/s]
Set up dataset to gather budget terms#
budget = xr.Dataset()
Set grid and data set for xgcm with metrics#
metrics = {
("X",): ["DXU", "DXT"], # X distances
("Y",): ["DYU", "DYT"], # Y distances
("Z",): ["DZU", "DZT"], # Z distances
("X", "Y"): ["UAREA", "TAREA"],
}
# here we get the xgcm compatible dataset
gridxgcm, dsxgcm = pop_tools.to_xgcm_grid_dataset(
ds,
periodic=False,
metrics=metrics,
boundary={"X": "extend", "Y": "extend", "Z": "extend"},
)
for coord in ["nlat", "nlon"]:
if coord in dsxgcm.coords:
dsxgcm = dsxgcm.drop_vars(coord)
0) Tendency#
budget['TEND_TEMP'] = dsxgcm.TEND_TEMP
i) Total heat advection#
We use grid.diff and multiply and divide by the volumes ourselves is the way POP outputs the fluxes. It performs a division by the cell area before saving the terms, which would not be accounted for if we used grid.derivative. Note that we also multiply by dsxgcm.VOL.values and then divide by dsxgcm.VOL. This is due to the same issue, there is a mis-alignment in the grid in this output term that xgcm would not like, and we are getting around it this way. This might be specific to POP and it should likely be possible to use grid.derivative for other models. The variables used here are online accumulated transports as output by the model. If you calculate according to udT/dx etc you will have the transport from the mean fields, and miss the eddy contribution below the timescale of your averaging operator (e.g. monthly for monthly output).
budget["UET"] = -(gridxgcm.diff(dsxgcm.UET * dsxgcm.VOL.values, axis="X") / dsxgcm.VOL)
budget["VNT"] = -(gridxgcm.diff(dsxgcm.VNT * dsxgcm.VOL.values, axis="Y") / dsxgcm.VOL)
budget["WTT"] = (
gridxgcm.diff(dsxgcm.WTT.fillna(0) * (dsxgcm.dz * dsxgcm.DXT * dsxgcm.DYT).values, axis="Z")
/ dsxgcm.VOL
)
budget["TOT_ADV"] = budget["UET"] + budget["VNT"] + budget["WTT"]
ii) Heat flux due to vertical mixing:#
includes surface flux at the 0th layer
budget["DIA_IMPVF_TEMP"] = -(
gridxgcm.diff(dsxgcm.DIA_IMPVF_TEMP * dsxgcm.TAREA, axis="Z") / dsxgcm.VOL
)
# set surface flux at 0th layer
SRF_TEMP_FLUX = (dsxgcm.SHF - dsxgcm.SHF_QSW) * dsxgcm.hflux_factor
budget["DIA_IMPVF_TEMP"][:, 0, :, :] = (
SRF_TEMP_FLUX * dsxgcm.TAREA - dsxgcm.DIA_IMPVF_TEMP.isel(z_w_bot=0) * dsxgcm.TAREA
) / dsxgcm.VOL.values[0, :, :]
budget["KPP_SRC_TMP"] = dsxgcm.KPP_SRC_TEMP
budget["VDIF"] = budget["DIA_IMPVF_TEMP"] + budget["KPP_SRC_TMP"]
iii) Heat flux due to horizontal diffusion#
budget["HDIFE_TEMP"] = gridxgcm.diff(dsxgcm.HDIFE_TEMP * dsxgcm.VOL.values, axis="X") / dsxgcm.VOL
budget["HDIFN_TEMP"] = gridxgcm.diff(dsxgcm.HDIFN_TEMP * dsxgcm.VOL.values, axis="Y") / dsxgcm.VOL
budget["HDIF"] = budget["HDIFE_TEMP"] + budget["HDIFN_TEMP"]
iv) Solar Penetration#
budget["QSW_3D"] = -gridxgcm.diff((dsxgcm.QSW_3D * dsxgcm.hflux_factor), axis="Z") / dsxgcm.DZT
plot to make sure it closes at (0N, 140W)#
subset = budget.isel(nlon_t=lola_inds["i_140_w"], nlat_t=lola_inds["j_0n"], time=0)
%%time
fig, ax = plt.subplots(1, 2, figsize=(10, 3), sharey=True)
# plot individual components
subset.VDIF.plot(y="z_t", ylim=(300e2, 0), label="VDIF", ax=ax[0])
subset.HDIF.plot(y="z_t", ylim=(300e2, 0), label="HDIF", ax=ax[0])
subset.QSW_3D.plot(y="z_t", ylim=(300e2, 0), label="QSW_3D", ax=ax[0])
subset.TOT_ADV.plot(y="z_t", ylim=(300e2, 0), label="DIV", ax=ax[0])
# plot sum
(subset.QSW_3D + subset.HDIF + subset.VDIF + subset.TOT_ADV).plot(
y="z_t", ylim=(300e2, 0), label="SUM", ls="--", ax=ax[0]
)
# plot tendency
subset.TEND_TEMP.plot(y="z_t", ylim=(300e2, 0), label="TEND_TEMP", ax=ax[0])
ax[0].legend()
# plot sum
(subset.QSW_3D + subset.HDIF + subset.VDIF + subset.TOT_ADV).plot(
y="z_t", ylim=(300e2, 0), label="SUM", ls="--", ax=ax[1]
)
# plot tendency
subset.TEND_TEMP.plot(y="z_t", ylim=(300e2, 0), label="TEND_TEMP", ax=ax[1])
ax[1].legend();
CPU times: user 126 ms, sys: 4.58 ms, total: 130 ms
Wall time: 289 ms
<matplotlib.legend.Legend at 0x7fe4be71b310>

# You may need to install watermark (conda install -c conda-forge watermark)
import xgcm # just to display the version
%load_ext watermark
%watermark -d -iv -m -g
print("Above are the versions of the packages this works with.")
Compiler : Clang 11.0.1
OS : Darwin
Release : 19.6.0
Machine : x86_64
Processor : i386
CPU cores : 8
Architecture: 64bit
Git hash: c1b73199f169ce40d3a7ded125ec0fbbac90b285
matplotlib: 3.3.4
xarray : 0.17.0
xgcm : 0.5.1
numpy : 1.20.1
pop_tools : 2020.12.15.post10+dirty
Above are the versions of the packages this works with.