def test_decomp_sumcomp(adv,
avg_dims_error=None,
thresh=0.9999999999,
loc=None,
iloc=None,
plot=True,
**plot_kws):
"""Test that the total advective tendency is indeed the sum of the mean and
resolved turbulent components in all three spatial directions.
The test fails if the Nash-Sutcliffe efficiency coefficient (NSE)
is below the given threshold. If avg_dims_error is given, the averaging in the
NSE calculation is only carried out over these dimensions. Afterwards the minimum NSE
value is taken over the remaining dimensions.
Parameters
----------
adv : xarray DataArray
Advective tendencies.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the NSE. The default is None.
thresh : float, optional
Threshold value for NSE below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic NSE
"""
failed = False
err = []
for ID in adv.ID.values:
ref = adv.sel(ID=ID, comp="adv_r")
dat = adv.sel(ID=ID, comp=["mean", "trb_r"]).sum("comp")
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
e = tools.nse(dat, ref, dim=avg_dims_error).min().values
err.append(e)
if e < thresh:
log = "decomp_sumcomp, {} (XYZ) for ID={}: min. NSE less than {}: {:.11f}".format(
dat.description, ID, thresh, e)
print(log)
if plot:
ref.name = "adv_r"
dat.name = "mean + trb_r"
plotting.scatter_hue(dat, ref, title=log, **plot_kws)
failed = True
return failed, min(err)
def test_2nd(adv,
avg_dims_error=None,
thresh=0.999,
loc=None,
iloc=None,
plot=True,
**plot_kws):
"""Test that the advective tendencies resulting from 2nd-order and
correct advection order are equal in all three directions and components
(usually carried out if correct order is equal to 2nd order).
The test fails if the Nash-Sutcliffe efficiency coefficient (NSE)
is below the given threshold. If avg_dims_error is given, the averaging in the
NSE calculation is only carried out over these dimensions. Afterwards the minimum NSE
value is taken over the remaining dimensions.
Parameters
----------
adv : xarray DataArray
Advective tendencies.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the NSE. The default is None.
thresh : float, optional
Threshold value for NSE below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic NSE
"""
failed = False
ref = adv.sel(ID="cartesian")
dat = adv.sel(ID="cartesian 2nd")
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
err = tools.nse(dat, ref, dim=avg_dims_error).min().values
if err < thresh:
log = "test_2nd, {} (XYZ): min. NSE less than {}: {:.5f}".format(
dat.description, thresh, err)
print(log)
if plot:
ref.name = "correct order"
dat.name = "2nd order"
plotting.scatter_hue(dat, ref, title=log, **plot_kws)
failed = True
return failed, err
def test_decomp_sumdir(adv, corr, avg_dims_error=None, thresh=0.99999,
loc=None, iloc=None, plot=True, **plot_kws):
"""
Test that budget methods "native" and "cartesian" give equal advective tendencies
in all components if the three spatial directions are summed up.
The test fails if the coefficient of determination
is below the given threshold. If avg_dims_error is given, the averaging in the
R2 calculation is only carried out over these dimensions. Afterwards the minimum R2
value is taken over the remaining dimensions.
Parameters
----------
adv : xarray DataArray
Advective tendencies.
corr : xarray DataArray
Cartesian corrections for advective tendencies.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the R2. The default is None.
thresh : float, optional
Threshold value for R2 below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic R2
"""
data = adv.sel(dir="sum", comp=corr.comp)
ref = data.sel(ID="native")
dat = data.sel(ID="cartesian") - corr.sel(ID="cartesian", dir="T")
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
err = R2(dat, ref, dim=avg_dims_error).min().values
failed = False
if err < thresh:
log = "test_decomp_sumdir, {}: min. R2 less than {}: {:.7f}".format(dat.description, thresh, err)
print(log)
if plot:
dat.name = "cartesian"
ref.name = "native"
plotting.scatter_hue(dat, ref, title=log, **plot_kws)
failed = True
return failed, err
def test_w(dat_inst,
avg_dims_error=None,
thresh=0.995,
loc=None,
iloc=None,
plot=True,
**plot_kws):
"""Test that the instantaneous vertical velocity is very similar to the
instantaneous diagnosed vertical velocity used in the tendency calculations.
The test fails if the Nash-Sutcliffe efficiency coefficient (NSE)
is below the given threshold. If avg_dims_error is given, the averaging in the
NSE calculation is only carried out over these dimensions. Afterwards the minimum NSE
value is taken over the remaining dimensions.
Parameters
----------
adv : xarray DataArray
Advective tendencies.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the NSE. The default is None.
thresh : float, optional
Threshold value for NSE below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic NSE
"""
dat_inst = tools.loc_data(dat_inst, loc=loc, iloc=iloc)
ref = dat_inst["W"]
dat = dat_inst["W_DIAG"]
err = tools.nse(dat, ref, dim=avg_dims_error).min().values
failed = False
if err < thresh:
log = "test_w: min. NSE less than {}: {:.5f}".format(thresh, err)
print(log)
if plot:
plotting.scatter_hue(dat, ref, title=log, **plot_kws)
failed = True
return failed, err
def test_dz_out(adv, avg_dims_error=None, thresh=0.95, loc=None, iloc=None, plot=True, **plot_kws):
"""Test that the Cartesian corrections imposed by the budget methods
"cartesian" and "cartesian dz_out_z" lead to
similar advective tendencies in all three directions and components.
The test fails if the coefficient of determination
is below the given threshold. If avg_dims_error is given, the averaging in the
R2 calculation is only carried out over these dimensions. Afterwards the minimum R2
value is taken over the remaining dimensions.
Parameters
----------
adv : xarray DataArray
Advective tendencies.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the R2. The default is None.
thresh : float, optional
Threshold value for R2 below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic R2
"""
failed = False
ref = adv.sel(ID="cartesian")
dat = adv.sel(ID="cartesian dz_out_z")
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
err = R2(dat, ref, dim=avg_dims_error).min().values
if err < thresh:
log = "test_dz_out, {} (XYZ): min. R2 less than {}: {:.5f}".format(dat.description, thresh, err)
print(log)
if plot:
dat.name = "dz_out_z"
ref.name = "reference corr."
plotting.scatter_hue(dat, ref, title=log, **plot_kws)
failed = True
return failed, err
def test_adv_form(dat_mean, datout, var, cyclic=None, hor_avg=False, avg_dims=None,
avg_dims_error=None, thresh=0.9995, loc=None, iloc=None, plot=True, **plot_kws):
"""Compare implicit and explicit advective form calculations
Explicitly calculate 2nd order mean advection in advective form and compare with
implicit calculation.
Parameters
----------
dat_mean : xarray Dataset
WRF time-averaged output.
datout : dict
Postprocessed output for variable var.
var : str
Variable to process.
cyclic : dict of booleans for xy or None, optional
Defines which dimensions have periodic boundary conditions.
Use periodic boundary conditions to fill lateral boundary points.
The default is None.
hor_avg : bool, optional
Horizontal averaging was used in postprocessing. The default is False.
avg_dims : str or list of str, optional
Averaging dimensions if hor_avg=True. The default is None.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the R2. The default is None.
thresh : float, optional
Threshold value for R2 below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic R2
"""
adv, flux, grid = datout["tend"]["adv"], datout["flux"], datout["grid"]
dat_mean["bottom_top"] = flux["bottom_top"]
dat_mean["bottom_top_stag"] = flux["bottom_top_stag"]
vmean = xr.Dataset({"X": dat_mean["U_MEAN"], "Y": dat_mean["V_MEAN"], "Z": dat_mean["WD_MEAN"]})
if var == "w":
v = "ZWIND"
else:
v = var.upper()
var_mean = dat_mean[v + "_MEAN"]
if hor_avg:
var_mean = tools.avg_xy(var_mean, avg_dims, cyclic=cyclic)
vmean_c = xr.Dataset()
dd = xr.Dataset()
grad = xr.Dataset()
tend = xr.Dataset()
for dim in tools.XYZ:
if hor_avg:
vmean[dim] = tools.avg_xy(vmean[dim], avg_dims, cyclic=cyclic, **grid[tools.stagger_const])
ds = dim.lower()
if dim == "Z":
ds = "bottom_top"
cyc = cyclic[ds]
d = ds
if ds in var_mean.dims:
ds = ds + "_stag"
else:
d = d + "_stag"
if dim == "Z":
dd[dim] = tools.diff(grid["Z_STAG"], d, new_coord=flux[ds], cyclic=cyc)
else:
dd[dim] = grid["D" + dim]
if d in adv.dims:
grad[dim] = tools.diff(var_mean, d, new_coord=flux[ds], cyclic=cyc) / dd[dim]
vmean_c[dim] = tools.stagger_like(vmean[dim], ref=grad[dim], cyclic=cyclic, **grid[tools.stagger_const])
for dim in tools.XYZ:
if dim in grad:
adv_s = - vmean_c[dim] * grad[dim]
tend[dim] = tools.stagger_like(adv_s, ref=adv, cyclic=cyclic, **grid[tools.stagger_const])
for dim in ["X", "Y"]:
if dim in grad:
corr = grid[f"dzdt_{dim.lower()}"]
corr = grad["Z"]*tools.stagger_like(corr, ref=grad["Z"], cyclic=cyclic, **grid[tools.stagger_const])
corr = tools.stagger_like(corr, ref=adv, cyclic=cyclic, **grid[tools.stagger_const])
tend[dim] = tend[dim] - corr
tend = tend.to_array("dir")
fname = None
if "fname" in plot_kws:
fname = plot_kws.pop("fname")
dat = tools.loc_data(adv.sel(ID="cartesian adv_form", dir=["X", "Y", "Z"], comp="mean"), loc=loc, iloc=iloc)
ref = tools.loc_data(tend, loc=loc, iloc=iloc)
dat = dat.sel(dir=ref.dir)
if var == "w":
dat = dat.isel(bottom_top_stag=slice(1, None))
ref = ref.isel(bottom_top_stag=slice(1, None))
err = R2(dat, ref, dim=avg_dims_error).min().values
failed = False
if err < thresh:
failed = True
log = "test_adv_form: mean advective component: min. R2 less than {}: {:.10f}".format(thresh, err)
print(log)
if plot:
dat.name = "Implicit calculation"
ref.name = "Explicit calculation"
if fname is not None:
log = fname + "\n" + log
plotting.scatter_hue(dat, ref, title=log, fname=fname, **plot_kws)
return failed, err
def test_mass(tend_mass, avg_dims_error=None, thresh=0.99999999,
loc=None, iloc=None, plot=True, **plot_kws):
"""Test closure of continuity equation.
In the tendency calculations the vertical component of the continuity equation
is calculated as residual to improve the budget closure which leads to automatic
closure of the continuity equation.
This test ensures that this residual calculation does not produce larger changes
in the vertical component by comparing the residual calculation with the
explicit calculation which uses the vertical velocity.
For the dz_out type formulations, the continuity equation cannot be well closed.
Therefore, we only compare the individual components with the standard Cartesian
formulation.
The test fails if the coefficient of determination
is below the given threshold. If avg_dims_error is given, the averaging in the
R2 calculation is only carried out over these dimensions. Afterwards the minimum R2
value is taken over the remaining dimensions.
Parameters
----------
tend_mass : xarray DataArray
Components of continuity equation.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the R2. The default is None.
thresh : float, optional
Threshold value for R2 below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic R2
"""
ref = tend_mass.sel(dir="Z")
dat = tend_mass.sel(dir="T") - tend_mass.sel(dir="X") - tend_mass.sel(dir="Y")
failed = False
err = []
fname = ""
if "fname" in plot_kws:
fname = plot_kws.pop("fname")
for ID in dat.ID.values:
dat_i = dat.sel(ID=ID)
ref_i = ref.sel(ID=ID)
dat_i = tools.loc_data(dat_i, loc=loc, iloc=iloc)
ref_i = tools.loc_data(ref_i, loc=loc, iloc=iloc)
e = R2(dat_i, ref_i, dim=avg_dims_error).min().values
err.append(e)
if e < thresh:
log = "test_mass: vertical component of continuity equation\n for ID={}: min. R2 less than {}: {:.10f}".format(ID, thresh, e)
print(log)
if plot:
dat_i.name = "Residual calculation"
ref_i.name = "Calculation with vertical velocity"
fname_i = fname
if fname is not None:
fname_i = "ID=" + ID + "_" + fname
log = fname_i + "\n" + log
plotting.scatter_hue(dat_i, ref_i, title=log, fname=fname_i, **plot_kws)
failed = True
return failed, min(err)
def test_budget(tend, forcing, avg_dims_error=None, thresh=0.9999, thresh_cartesian=None,
budget_methods=("native", "adv_form", "cartesian", "cartesian adv_form"),
loc=None, iloc=None, plot=True, **plot_kws):
"""
Test closure of budget: tend = forcing.
The test fails if the coefficient of determination
is below the given threshold. If avg_dims_error is given, the averaging in the
R2 calculation is only carried out over these dimensions. Afterwards the minimum R2
value is taken over the remaining dimensions.
Parameters
----------
tend : xarray DataArray
Total tendency.
forcing : xarray DataArray
Total forcing.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the R2. The default is None.
thresh : float, optional
Threshold value for R2 below which the test fails
thresh_cartesian : float, optional
Use different threshold value for Cartesian coordinate system.
The default is None, for which 'thresh' is used in both formulations.
budget_methods : list of str
Budget methods to consider. By default, only "native", "adv_form", "cartesian", and "cartesian adv_form" are tested.
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic R2
"""
failed = False
err = []
fname = ""
if "fname" in plot_kws:
fname = plot_kws.pop("fname")
for ID in budget_methods:
thresh_i = thresh
if (ID == "cartesian") and (thresh_cartesian is not None):
thresh_i = thresh_cartesian
if ID not in tend.ID:
continue
ref = tend.sel(ID=ID, drop=True)
dat = forcing.sel(ID=ID, drop=True)
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
e = R2(dat, ref, dim=avg_dims_error).min().values
err.append(e)
if e < thresh_i:
log = "test_budget for ID='{}': min. R2 less than {}: {:.10f}\n".format(ID, thresh_i, e)
print(log)
if plot:
dat.name = dat.description[:8] + "forcing"
ref.name = ref.description
fname_i = fname
if fname is not None:
fname_i = "ID=" + ID + "_" + fname
log = fname_i + "\n" + log
plotting.scatter_hue(dat, ref, title=log, fname=fname_i, **plot_kws)
failed = True
return failed, min(err)
def test_decomp_sumcomp(adv, avg_dims_error=None, thresh=0.999995,
loc=None, iloc=None, plot=True, **plot_kws):
"""Test that the total advective tendency is indeed the sum of the mean and
resolved turbulent components in all three spatial directions.
The test fails if the coefficient of determination
is below the given threshold. If avg_dims_error is given, the averaging in the
R2 calculation is only carried out over these dimensions. Afterwards the minimum R2
value is taken over the remaining dimensions.
Parameters
----------
adv : xarray DataArray
Advective tendencies.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the R2. The default is None.
thresh : float, optional
Threshold value for R2 below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic R2
"""
ref = adv.sel(comp="trb_r")
dat = adv.sel(comp="adv_r") - adv.sel(comp="mean")
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
failed = False
err = []
fname = ""
if "fname" in plot_kws:
fname = plot_kws.pop("fname")
for ID in dat.ID.values:
dat_i = dat.sel(ID=ID)
ref_i = ref.sel(ID=ID)
e = R2(dat_i, ref_i, dim=avg_dims_error).min().values
err.append(e)
if e < thresh:
log = "decomp_sumcomp, {} (XYZ) for ID={}: min. R2 less than {}: {:.8f}".format(
dat.description, ID, thresh, e)
print(log)
if plot:
ref_i.name = "trb_r"
dat_i.name = "adv_r - mean"
fname_i = fname
if fname is not None:
fname_i = "ID=" + ID + "_" + fname
log = fname_i + "\n" + log
plotting.scatter_hue(dat_i, ref_i, title=log, fname=fname_i, **plot_kws)
failed = True
return failed, min(err)
def test_budget(tend,
forcing,
avg_dims_error=None,
thresh=0.9993,
loc=None,
iloc=None,
plot=True,
**plot_kws):
"""
Test closure of budget: tend = forcing.
Only the budget methods "native" and "cartesian" are tested.
The test fails if the Nash-Sutcliffe efficiency coefficient (NSE)
is below the given threshold. If avg_dims_error is given, the averaging in the
NSE calculation is only carried out over these dimensions. Afterwards the minimum NSE
value is taken over the remaining dimensions.
Parameters
----------
tend : xarray DataArray
Total tendency.
forcing : xarray DataArray
Total forcing.
avg_dims_error : str or list of str, optional
Dimensions over which to calculate the NSE. The default is None.
thresh : float, optional
Threshold value for NSE below which the test fails
loc : dict, optional
Mapping for label based indexing before running the test. The default is None.
iloc : dict, optional
Mapping for integer-location based indexing before running the test. The default is None.
plot : bool, optional
Create scatter plot if test fails. The default is True.
**plot_kws :
keyword arguments passed to plotting.scatter_hue.
Returns
-------
failed : bool
Test failed.
err : float
Test statistic NSE
"""
failed = False
err = []
for ID in ["native", "cartesian"]:
if ID not in tend.ID:
continue
ref = tend.sel(ID=ID, drop=True)
dat = forcing.sel(ID=ID, drop=True)
dat = tools.loc_data(dat, loc=loc, iloc=iloc)
ref = tools.loc_data(ref, loc=loc, iloc=iloc)
e = tools.nse(dat, ref, dim=avg_dims_error).min().values
err.append(e)
if e < thresh:
log = "test_budget for ID='{}': min. NSE less than {}: {:.5f}\n".format(
ID, thresh, e)
print(log)
if plot:
dat.name = dat.description[:2] + "forcing"
ref.name = ref.description
plotting.scatter_hue(dat, ref, title=log, **plot_kws)
failed = True
return failed, min(err)