From 7ca7ba79fa9641486202d798cb7ccca7d64e29ab Mon Sep 17 00:00:00 2001 From: John Krasting Date: Sat, 10 Feb 2024 16:53:05 -0500 Subject: [PATCH] Option to keep N^2 on the tracer grid for PV calc - While vorticity should be calculated on the cell corners for a C-grid models, this is not always practical for other grids that do not provide the cell corners or other cases where vorticity is desired at the cell centers - Fix introduces an optional flag to retain the calculation on the cell centers --- src/momlevel/derived.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/momlevel/derived.py b/src/momlevel/derived.py index a06366d..67b61b9 100644 --- a/src/momlevel/derived.py +++ b/src/momlevel/derived.py @@ -487,7 +487,14 @@ def calc_pdens(thetao, so, level=0.0, patm=101325, eos="Wright"): def calc_pv( - zeta, coriolis, n2, gravity=9.8, coord_dict=None, symmetric=False, units="m" + zeta, + coriolis, + n2, + gravity=9.8, + coord_dict=None, + symmetric=False, + units="m", + interp_n2=True, ): """Function to calculate ocean potential vorticity @@ -507,6 +514,9 @@ def calc_pv( coord_dict : :obj:`dict`, optional Dictionary of xgcm coordinate name mappings, if different from the MOM6 default values, by default None + interp_n2 : bool, optional + Interpolate N^2 to the corner points. Set to `False` to calculate + potential vorticity on the cell centers, by default True symmetric : bool Flag denoting symmetric grid, by default False units : str, optional @@ -525,12 +535,13 @@ def calc_pv( calc_rel_vort : Calculate relative vorticity (zeta) """ - # create an internal dataset for xgcm purposes - _dset = xr.Dataset({"zeta": zeta, "coriolis": coriolis, "n2": n2}) - grid = util.get_xgcm_grid(_dset, coord_dict=coord_dict, symmetric=symmetric) + if interp_n2 is True: + # create an internal dataset for xgcm purposes + _dset = xr.Dataset({"zeta": zeta, "coriolis": coriolis, "n2": n2}) + grid = util.get_xgcm_grid(_dset, coord_dict=coord_dict, symmetric=symmetric) - # interpolate N2 to the corner points - n2 = grid.interp(n2, axis=["X", "Y"], boundary="fill") + # interpolate N2 to the corner points + n2 = grid.interp(n2, axis=["X", "Y"], boundary="fill") # calculate potential vorticity swpotvort = (zeta + coriolis) * (n2 / gravity)