Bill Ladwig 6 years ago
parent
commit
3c03cd2916
  1. 138
      test/comp_utest.py

138
test/comp_utest.py

@ -4,7 +4,8 @@ import unittest as ut
import numpy.testing as nt import numpy.testing as nt
import numpy as np import numpy as np
import numpy.ma as ma import numpy.ma as ma
import os, sys import os
import sys
import subprocess import subprocess
from netCDF4 import Dataset as nc from netCDF4 import Dataset as nc
@ -21,36 +22,36 @@ NCGROUP = [NCFILE, NCFILE, NCFILE]
if sys.version_info > (3,): if sys.version_info > (3,):
xrange = range xrange = range
ROUTINE_MAP = {"avo": avo,
"eth": eth,
"cape_2d": cape_2d,
"cape_3d": cape_3d,
"ctt": ctt,
"dbz": dbz,
"helicity": srhel,
"omg": omega,
"pvo": pvo,
"pw": pw,
"rh": rh,
"slp": slp,
"td": td,
"tk": tk,
"tv": tvirtual,
"twb": wetbulb,
"updraft_helicity": udhel,
"uvmet": uvmet,
"cloudfrac": cloudfrac}
ROUTINE_MAP = {"avo" : avo,
"eth" : eth,
"cape_2d" : cape_2d,
"cape_3d" : cape_3d,
"ctt" : ctt,
"dbz" : dbz,
"helicity" : srhel,
"omg" : omega,
"pvo" : pvo,
"pw" : pw,
"rh" : rh,
"slp" : slp,
"td" : td,
"tk" : tk,
"tv" : tvirtual,
"twb" : wetbulb,
"updraft_helicity" : udhel,
"uvmet" : uvmet,
"cloudfrac" : cloudfrac}
class ProjectionError(RuntimeError): class ProjectionError(RuntimeError):
pass pass
def get_args(varname, wrfnc, timeidx, method, squeeze): def get_args(varname, wrfnc, timeidx, method, squeeze):
if varname == "avo": if varname == "avo":
ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "MAPFAC_U", varnames = ("U", "V", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", "F")
"MAPFAC_V", "MAPFAC_M", ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
"F"), cache=None, meta=True)
method, squeeze, cache=None, meta=True)
attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
u = ncvars["U"] u = ncvars["U"]
@ -66,10 +67,10 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (u, v, msfu, msfv, msfm, cor, dx, dy) return (u, v, msfu, msfv, msfm, cor, dx, dy)
if varname == "pvo": if varname == "pvo":
ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "T", "P", varnames = ("U", "V", "T", "P", "PB", "MAPFAC_U", "MAPFAC_V",
"PB", "MAPFAC_U", "MAPFAC_M", "F")
"MAPFAC_V", "MAPFAC_M", ncvars = extract_vars(wrfnc, timeidx,
"F"), varnames,
method, squeeze, cache=None, meta=True) method, squeeze, cache=None, meta=True)
attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
@ -92,7 +93,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (u, v, full_t, full_p, msfu, msfv, msfm, cor, dx, dy) return (u, v, full_t, full_p, msfu, msfv, msfm, cor, dx, dy)
if varname == "eth": if varname == "eth":
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
t = ncvars["T"] t = ncvars["T"]
@ -107,7 +108,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (qv, tkel, full_p) return (qv, tkel, full_p)
if varname == "cape_2d": if varname == "cape_2d":
varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC") varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
@ -187,7 +188,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
qice = np.zeros(qv.shape, qv.dtype) qice = np.zeros(qv.shape, qv.dtype)
haveqci = 0 haveqci = 0
else: else:
qice = icevars["QICE"] * 1000.0 #g/kg qice = icevars["QICE"] * 1000.0 # g/kg
try: try:
cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", cldvars = extract_vars(wrfnc, timeidx, "QCLOUD",
@ -195,7 +196,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
except KeyError: except KeyError:
raise RuntimeError("'QCLOUD' not found in NetCDF file") raise RuntimeError("'QCLOUD' not found in NetCDF file")
else: else:
qcld = cldvars["QCLOUD"] * 1000.0 #g/kg qcld = cldvars["QCLOUD"] * 1000.0 # g/kg
full_p = p + pb full_p = p + pb
p_hpa = full_p * ConversionFactors.PA_TO_HPA p_hpa = full_p * ConversionFactors.PA_TO_HPA
@ -269,8 +270,9 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (u, v, z, ter) return (u, v, z, ter)
if varname == "updraft_helicity": if varname == "updraft_helicity":
ncvars = extract_vars(wrfnc, timeidx, ("W", "PH", "PHB", "MAPFAC_M"), varnames = ("W", "PH", "PHB", "MAPFAC_M")
method, squeeze, cache=None, meta=True) ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True)
wstag = ncvars["W"] wstag = ncvars["W"]
ph = ncvars["PH"] ph = ncvars["PH"]
@ -297,7 +299,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (zstag, mapfct, u, v, wstag, dx, dy) return (zstag, mapfct, u, v, wstag, dx, dy)
if varname == "omg": if varname == "omg":
varnames=("T", "P", "W", "PB", "QVAPOR") varnames = ("T", "P", "W", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
t = ncvars["T"] t = ncvars["T"]
@ -314,7 +316,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (qv, tkel, wa, full_p) return (qv, tkel, wa, full_p)
if varname == "pw": if varname == "pw":
varnames=("T", "P", "PB", "PH", "PHB", "QVAPOR") varnames = ("T", "P", "PB", "PH", "PHB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
@ -335,7 +337,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (full_p, tkel, qv, ht) return (full_p, tkel, qv, ht)
if varname == "rh": if varname == "rh":
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
t = ncvars["T"] t = ncvars["T"]
@ -351,7 +353,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (qvapor, full_p, tkel) return (qvapor, full_p, tkel)
if varname == "slp": if varname == "slp":
varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB") varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
@ -375,7 +377,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (destag_ph, tkel, full_p, qvapor) return (destag_ph, tkel, full_p, qvapor)
if varname == "td": if varname == "td":
varnames=("P", "PB", "QVAPOR") varnames = ("P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
@ -390,7 +392,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (full_p, qvapor) return (full_p, qvapor)
if varname == "tk": if varname == "tk":
varnames=("T", "P", "PB") varnames = ("T", "P", "PB")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
t = ncvars["T"] t = ncvars["T"]
@ -403,7 +405,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (full_p, full_t) return (full_p, full_t)
if varname == "tv": if varname == "tv":
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
@ -419,7 +421,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
return (tkel, qv) return (tkel, qv)
if varname == "twb": if varname == "twb":
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True) cache=None, meta=True)
t = ncvars["T"] t = ncvars["T"]
@ -449,9 +451,9 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ") map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ")
map_proj = map_proj_attrs["MAP_PROJ"] map_proj = map_proj_attrs["MAP_PROJ"]
if map_proj in (0,3,6): if map_proj in (0, 3, 6):
raise ProjectionError("Map projection does not need rotation") raise ProjectionError("Map projection does not need rotation")
elif map_proj in (1,2): elif map_proj in (1, 2):
lat_attrs = extract_global_attrs(wrfnc, attrs=("TRUELAT1", lat_attrs = extract_global_attrs(wrfnc, attrs=("TRUELAT1",
"TRUELAT2")) "TRUELAT2"))
radians_per_degree = Constants.PI/180.0 radians_per_degree = Constants.PI/180.0
@ -463,15 +465,16 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
lon_attrs = extract_global_attrs(wrfnc, attrs="STAND_LON") lon_attrs = extract_global_attrs(wrfnc, attrs="STAND_LON")
except AttributeError: except AttributeError:
try: try:
cen_lon_attrs = extract_global_attrs(wrfnc, attrs="CEN_LON") cen_lon_attrs = extract_global_attrs(wrfnc,
attrs="CEN_LON")
except AttributeError: except AttributeError:
raise RuntimeError("longitude attributes not found in NetCDF") raise RuntimeError("longitude attributes not found in "
"NetCDF")
else: else:
cen_lon = cen_lon_attrs["CEN_LON"] cen_lon = cen_lon_attrs["CEN_LON"]
else: else:
cen_lon = lon_attrs["STAND_LON"] cen_lon = lon_attrs["STAND_LON"]
varname = "XLAT" varname = "XLAT"
xlat_var = extract_vars(wrfnc, timeidx, varname, xlat_var = extract_vars(wrfnc, timeidx, varname,
method, squeeze, cache=None, meta=True) method, squeeze, cache=None, meta=True)
@ -485,13 +488,15 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
if map_proj == 1: if map_proj == 1:
if((fabs(true_lat1 - true_lat2) > 0.1) and if((fabs(true_lat1 - true_lat2) > 0.1) and
(fabs(true_lat2 - 90.) > 0.1)): (fabs(true_lat2 - 90.) > 0.1)):
cone = (log(cos(true_lat1*radians_per_degree)) cone = (log(cos(true_lat1 * radians_per_degree))
- log(cos(true_lat2*radians_per_degree))) - log(cos(true_lat2 * radians_per_degree)))
cone = (cone / cone = (cone /
(log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) (log(tan((45.-fabs(true_lat1/2.)) *
- log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) radians_per_degree))
- log(tan((45.-fabs(true_lat2/2.)) *
radians_per_degree))))
else: else:
cone = sin(fabs(true_lat1)*radians_per_degree) cone = sin(fabs(true_lat1) * radians_per_degree)
else: else:
cone = 1 cone = 1
@ -499,8 +504,10 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
if varname == "cloudfrac": if varname == "cloudfrac":
from wrf.g_geoht import get_height from wrf.g_geoht import get_height
vars = extract_vars(wrfnc, timeidx, ("P", "PB", "QVAPOR", "T"),
method, squeeze, cache=None, meta=True) varnames = ("P", "PB", "QVAPOR", "T")
vars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
cache=None, meta=True)
p = vars["P"] p = vars["P"]
pb = vars["PB"] pb = vars["PB"]
@ -522,6 +529,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
class WRFVarsTest(ut.TestCase): class WRFVarsTest(ut.TestCase):
longMessage = True longMessage = True
def make_func(varname, wrfnc, timeidx, method, squeeze, meta): def make_func(varname, wrfnc, timeidx, method, squeeze, meta):
def func(self): def func(self):
@ -532,7 +540,7 @@ def make_func(varname, wrfnc, timeidx, method, squeeze, meta):
routine = ROUTINE_MAP[varname] routine = ROUTINE_MAP[varname]
kwargs = {"meta" : meta} kwargs = {"meta": meta}
result = routine(*args, **kwargs) result = routine(*args, **kwargs)
ref = getvar(wrfnc, varname, timeidx, method, squeeze, cache=None, ref = getvar(wrfnc, varname, timeidx, method, squeeze, cache=None,
@ -651,6 +659,7 @@ def test_cape2d_1d(wrfnc):
return func return func
if __name__ == "__main__": if __name__ == "__main__":
from wrf import (omp_set_num_threads, omp_set_schedule, omp_get_schedule, from wrf import (omp_set_num_threads, omp_set_schedule, omp_get_schedule,
omp_set_dynamic, omp_get_num_procs, OMP_SCHED_STATIC) omp_set_dynamic, omp_get_num_procs, OMP_SCHED_STATIC)
@ -664,11 +673,6 @@ if __name__ == "__main__":
"theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va",
"wa", "uvmet10", "uvmet", "z", "cloudfrac"] "wa", "uvmet10", "uvmet", "z", "cloudfrac"]
#varnames = ["helicity"]
varnames=["avo", "pvo", "eth", "dbz", "helicity", "updraft_helicity",
"omg", "pw", "rh", "slp", "td", "tk", "tv", "twb", "uvmet",
"cloudfrac", "ctt"]
omp_set_num_threads(omp_get_num_procs()-1) omp_set_num_threads(omp_get_num_procs()-1)
omp_set_schedule(OMP_SCHED_STATIC, 0) omp_set_schedule(OMP_SCHED_STATIC, 0)
omp_set_dynamic(False) omp_set_dynamic(False)
@ -677,8 +681,8 @@ if __name__ == "__main__":
varnames += ["cape_2d", "cape_3d"] varnames += ["cape_2d", "cape_3d"]
for varname in varnames: for varname in varnames:
for i,wrfnc in enumerate((NCFILE, NCGROUP)): for i, wrfnc in enumerate((NCFILE, NCGROUP)):
for j,timeidx in enumerate((0, ALL_TIMES)): for j, timeidx in enumerate((0, ALL_TIMES)):
for method in ("cat", "join"): for method in ("cat", "join"):
for squeeze in (True, False): for squeeze in (True, False):
for meta in (True, False): for meta in (True, False):
@ -686,10 +690,11 @@ if __name__ == "__main__":
squeeze, meta) squeeze, meta)
ncname = "single" if i == 0 else "multi" ncname = "single" if i == 0 else "multi"
timename = "t0" if j == 0 else "all" timename = "t0" if j == 0 else "all"
squeeze_name = "squeeze" if squeeze else "nosqueeze" squeeze_name = ("squeeze" if squeeze
else "nosqueeze")
meta_name = "meta" if meta else "nometa" meta_name = "meta" if meta else "nometa"
test_name = "test_{}_{}_{}_{}_{}_{}".format(varname, test_name = "test_{}_{}_{}_{}_{}_{}".format(
ncname, timename, method, varname, ncname, timename, method,
squeeze_name, meta_name) squeeze_name, meta_name)
setattr(WRFVarsTest, test_name, func) setattr(WRFVarsTest, test_name, func)
@ -700,7 +705,4 @@ if __name__ == "__main__":
func = test_cape2d_1d(wrfnc) func = test_cape2d_1d(wrfnc)
setattr(WRFVarsTest, "test_cape2d_1d", func) setattr(WRFVarsTest, "test_cape2d_1d", func)
ut.main() ut.main()
Loading…
Cancel
Save