PEP 8

test/comp_utest.py

@@ -4,7 +4,8 @@ import unittest as ut
 import numpy.testing as nt
 import numpy as np
 import numpy.ma as ma
-import os, sys
+import os
+import sys
 import subprocess
 
 from netCDF4 import Dataset as nc
@@ -21,36 +22,36 @@ NCGROUP = [NCFILE, NCFILE, NCFILE]
 if sys.version_info > (3,):
     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):
     pass
 
+
 def get_args(varname, wrfnc, timeidx, method, squeeze):
     if varname == "avo":
-        ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "MAPFAC_U",
-                                           "MAPFAC_V", "MAPFAC_M",
-                                           "F"),
-                          method, squeeze, cache=None, meta=True)
+        varnames = ("U", "V", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", "F")
+        ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
+                              cache=None, meta=True)
 
         attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
         u = ncvars["U"]
@@ -66,11 +67,11 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (u, v, msfu, msfv, msfm, cor, dx, dy)
 
     if varname == "pvo":
-        ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "T", "P",
-                                           "PB", "MAPFAC_U",
-                                           "MAPFAC_V", "MAPFAC_M",
-                                           "F"),
-                          method, squeeze, cache=None, meta=True)
+        varnames = ("U", "V", "T", "P", "PB", "MAPFAC_U", "MAPFAC_V",
+                    "MAPFAC_M", "F")
+        ncvars = extract_vars(wrfnc, timeidx,
+                              varnames,
+                              method, squeeze, cache=None, meta=True)
         attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
 
         u = ncvars["U"]
@@ -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)
 
     if varname == "eth":
-        varnames=("T", "P", "PB", "QVAPOR")
+        varnames = ("T", "P", "PB", "QVAPOR")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
         t = ncvars["T"]
@@ -107,7 +108,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (qv, tkel, full_p)
 
     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,
                               cache=None, meta=True)
 
@@ -177,7 +178,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         ph = ncvars["PH"]
         phb = ncvars["PHB"]
         ter = ncvars["HGT"]
-        qv = ncvars["QVAPOR"] * 1000.0 # g/kg
+        qv = ncvars["QVAPOR"] * 1000.0  # g/kg
 
         haveqci = 1
         try:
@@ -187,7 +188,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
             qice = np.zeros(qv.shape, qv.dtype)
             haveqci = 0
         else:
-            qice = icevars["QICE"] * 1000.0 #g/kg
+            qice = icevars["QICE"] * 1000.0  # g/kg
 
         try:
             cldvars = extract_vars(wrfnc, timeidx, "QCLOUD",
@@ -195,7 +196,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         except KeyError:
             raise RuntimeError("'QCLOUD' not found in NetCDF file")
         else:
-            qcld = cldvars["QCLOUD"] * 1000.0 #g/kg
+            qcld = cldvars["QCLOUD"] * 1000.0  # g/kg
 
         full_p = p + pb
         p_hpa = full_p * ConversionFactors.PA_TO_HPA
@@ -269,15 +270,16 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (u, v, z, ter)
 
     if varname == "updraft_helicity":
-        ncvars = extract_vars(wrfnc, timeidx, ("W", "PH", "PHB", "MAPFAC_M"),
-                          method, squeeze, cache=None, meta=True)
+        varnames = ("W", "PH", "PHB", "MAPFAC_M")
+        ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
+                              cache=None, meta=True)
 
         wstag = ncvars["W"]
         ph = ncvars["PH"]
         phb = ncvars["PHB"]
         mapfct = ncvars["MAPFAC_M"]
 
-        attrs  = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
+        attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
         dx = attrs["DX"]
         dy = attrs["DY"]
 
@@ -297,7 +299,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (zstag, mapfct, u, v, wstag, dx, dy)
 
     if varname == "omg":
-        varnames=("T", "P", "W", "PB", "QVAPOR")
+        varnames = ("T", "P", "W", "PB", "QVAPOR")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
         t = ncvars["T"]
@@ -314,7 +316,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (qv, tkel, wa, full_p)
 
     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,
                               cache=None, meta=True)
 
@@ -328,14 +330,14 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         # Change this to use real virtual temperature!
         full_p = p + pb
         ht = (ph + phb)/Constants.G
-        full_t =  t + Constants.T_BASE
+        full_t = t + Constants.T_BASE
 
         tkel = tk(full_p, full_t, meta=False)
 
         return (full_p, tkel, qv, ht)
 
     if varname == "rh":
-        varnames=("T", "P", "PB", "QVAPOR")
+        varnames = ("T", "P", "PB", "QVAPOR")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
         t = ncvars["T"]
@@ -351,7 +353,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (qvapor, full_p, tkel)
 
     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,
                               cache=None, meta=True)
 
@@ -375,7 +377,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (destag_ph, tkel, full_p, qvapor)
 
     if varname == "td":
-        varnames=("P", "PB", "QVAPOR")
+        varnames = ("P", "PB", "QVAPOR")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
 
@@ -390,7 +392,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (full_p, qvapor)
 
     if varname == "tk":
-        varnames=("T", "P", "PB")
+        varnames = ("T", "P", "PB")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
         t = ncvars["T"]
@@ -403,7 +405,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (full_p, full_t)
 
     if varname == "tv":
-        varnames=("T", "P", "PB", "QVAPOR")
+        varnames = ("T", "P", "PB", "QVAPOR")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
 
@@ -419,7 +421,7 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (tkel, qv)
 
     if varname == "twb":
-        varnames=("T", "P", "PB", "QVAPOR")
+        varnames = ("T", "P", "PB", "QVAPOR")
         ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze,
                               cache=None, meta=True)
         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 = 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")
-        elif map_proj in (1,2):
+        elif map_proj in (1, 2):
             lat_attrs = extract_global_attrs(wrfnc, attrs=("TRUELAT1",
                                                            "TRUELAT2"))
             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")
             except AttributeError:
                 try:
-                    cen_lon_attrs = extract_global_attrs(wrfnc, attrs="CEN_LON")
+                    cen_lon_attrs = extract_global_attrs(wrfnc,
+                                                         attrs="CEN_LON")
                 except AttributeError:
-                    raise RuntimeError("longitude attributes not found in NetCDF")
+                    raise RuntimeError("longitude attributes not found in "
+                                       "NetCDF")
                 else:
                     cen_lon = cen_lon_attrs["CEN_LON"]
             else:
                 cen_lon = lon_attrs["STAND_LON"]
 
-            
             varname = "XLAT"
             xlat_var = extract_vars(wrfnc, timeidx, varname,
                                     method, squeeze, cache=None, meta=True)
@@ -485,13 +488,15 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
             if map_proj == 1:
                 if((fabs(true_lat1 - true_lat2) > 0.1) and
                         (fabs(true_lat2 - 90.) > 0.1)):
-                    cone = (log(cos(true_lat1*radians_per_degree)) 
-                        - log(cos(true_lat2*radians_per_degree)))
+                    cone = (log(cos(true_lat1 * radians_per_degree))
+                            - log(cos(true_lat2 * radians_per_degree)))
                     cone = (cone /
-                            (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) 
-                        - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) 
+                            (log(tan((45.-fabs(true_lat1/2.)) *
+                                     radians_per_degree))
+                                - log(tan((45.-fabs(true_lat2/2.)) *
+                                          radians_per_degree))))
                 else:
-                    cone = sin(fabs(true_lat1)*radians_per_degree)
+                    cone = sin(fabs(true_lat1) * radians_per_degree)
             else:
                 cone = 1
 
@@ -499,8 +504,10 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
 
     if varname == "cloudfrac":
         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"]
         pb = vars["PB"]
@@ -522,17 +529,18 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
 class WRFVarsTest(ut.TestCase):
     longMessage = True
 
+
 def make_func(varname, wrfnc, timeidx, method, squeeze, meta):
     def func(self):
 
         try:
             args = get_args(varname, wrfnc, timeidx, method, squeeze)
-        except ProjectionError: # Don't fail for this
+        except ProjectionError:  # Don't fail for this
             return
 
         routine = ROUTINE_MAP[varname]
 
-        kwargs = {"meta" : meta}
+        kwargs = {"meta": meta}
         result = routine(*args, **kwargs)
 
         ref = getvar(wrfnc, varname, timeidx, method, squeeze, cache=None,
@@ -651,6 +659,7 @@ def test_cape2d_1d(wrfnc):
 
     return func
 
+
 if __name__ == "__main__":
     from wrf import (omp_set_num_threads, omp_set_schedule, omp_get_schedule,
                      omp_set_dynamic, omp_get_num_procs, OMP_SCHED_STATIC)
@@ -659,15 +668,10 @@ if __name__ == "__main__":
     omp_set_dynamic(False)
 
     varnames = ["avo", "eth", "cape_2d", "cape_3d", "ctt", "dbz", "mdbz",
-            "geopt", "helicity", "lat", "lon", "omg", "p", "pressure", 
-            "pvo", "pw", "rh2", "rh", "slp", "ter", "td2", "td", "tc", 
-            "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", 
-            "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"]
+                "geopt", "helicity", "lat", "lon", "omg", "p", "pressure",
+                "pvo", "pw", "rh2", "rh", "slp", "ter", "td2", "td", "tc",
+                "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va",
+                "wa", "uvmet10", "uvmet", "z", "cloudfrac"]
 
     omp_set_num_threads(omp_get_num_procs()-1)
     omp_set_schedule(OMP_SCHED_STATIC, 0)
@@ -677,20 +681,21 @@ if __name__ == "__main__":
     varnames += ["cape_2d", "cape_3d"]
 
     for varname in varnames:
-        for i,wrfnc in enumerate((NCFILE, NCGROUP)):
-            for j,timeidx in enumerate((0, ALL_TIMES)):
+        for i, wrfnc in enumerate((NCFILE, NCGROUP)):
+            for j, timeidx in enumerate((0, ALL_TIMES)):
                 for method in ("cat", "join"):
                     for squeeze in (True, False):
                         for meta in (True, False):
                             func = make_func(varname, wrfnc, timeidx, method,
-                                                  squeeze, meta)
+                                             squeeze, meta)
                             ncname = "single" if i == 0 else "multi"
                             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"
-                            test_name = "test_{}_{}_{}_{}_{}_{}".format(varname, 
-                                                        ncname, timename, method,
-                                                        squeeze_name, meta_name)
+                            test_name = "test_{}_{}_{}_{}_{}_{}".format(
+                                varname, ncname, timename, method,
+                                squeeze_name, meta_name)
 
                             setattr(WRFVarsTest, test_name, func)
 
@@ -700,7 +705,4 @@ if __name__ == "__main__":
     func = test_cape2d_1d(wrfnc)
     setattr(WRFVarsTest, "test_cape2d_1d", func)
 
-    
     ut.main()
-    
-