Clean up Python scripts (#232)

* Python: apply PEP8 formatting * validation-plots.py: fix wrong argument in a wrong place * validation-plots.py: remove unused argument * validation-plots.py: add missing argument * validation-plots.py: fix variable shadowing * stage_files.py: remove broken branch * compare-to-reference.py: optimize environment queries
earth-system-radiation · Sep 30, 2023 · c45e7bc · c45e7bc
1 parent 809428a
commit c45e7bc
Show file tree

Hide file tree

Showing 6 changed files with 112 additions and 97 deletions.
diff --git a/examples/all-sky/make_problem_size_loop.py b/examples/all-sky/make_problem_size_loop.py
@@ -3,60 +3,62 @@
 # This script...
 #
 import argparse
-import csv 
-import os
+import csv
 
-# template: exe ncol nlay nloops output_file k_dist clouds aeorsols 
+# template: exe ncol nlay nloops output_file k_dist clouds aeorsols
 
-# specify: kdist, optional clouds, aerosols. specify nloops 
-# Toggle clouds and aerosols?  
-# Loop over sets of ncol, nlay, 
-# output name 
+# specify: kdist, optional clouds, aerosols. specify nloops
+# Toggle clouds and aerosols?
+# Loop over sets of ncol, nlay,
+# output name
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(
         description="Description here ")
     # Argument parseing described at 
     # https://stackoverflow.com/questions/15753701/how-can-i-pass-a-list-as-a-command-line-argument-with-argparse
-    parser.add_argument("-x", "--executable", 
-                        type=str, 
+    parser.add_argument("-x", "--executable",
+                        type=str,
                         default="./rrtmgp_allsky",
                         help="Path to exectuable")
-    parser.add_argument("-c", "--ncol", 
-                        type=lambda items: [int(i) for i in list(csv.reader([items]))[0]], 
+    parser.add_argument("-c", "--ncol",
+                        type=lambda items: [int(i) for i in
+                                            list(csv.reader([items]))[0]],
                         default="2,4,8,16",
                         help="Number of columns  (multiple e.g. 2,4,8,16)")
-    parser.add_argument("-l", "--nlay", 
-                        type=lambda items: [int(i) for i in list(csv.reader([items]))[0]], 
+    parser.add_argument("-l", "--nlay",
+                        type=lambda items: [int(i) for i in
+                                            list(csv.reader([items]))[0]],
                         default="32, 64, 96",
                         help="Number of layers (multiple e.g. 32,64.96)")
-    parser.add_argument("-i", "--nloops", 
+    parser.add_argument("-i", "--nloops",
                         type=int, default=1,
                         help="Number of loops (same for all ncol)")
-    parser.add_argument("-o", "--output_file", 
-                        type=str, 
+    parser.add_argument("-o", "--output_file",
+                        type=str,
                         default="rrtmgp-allsky-output.nc",
                         help="Path to output file")
-    parser.add_argument("-k", "--k_distribution", 
-                        type=str, 
-                        required = True, 
+    parser.add_argument("-k", "--k_distribution",
+                        type=str,
+                        required=True,
                         help="Path to gas optics file [required]")
-    parser.add_argument("-cl", "--cloud_optics", 
+    parser.add_argument("-cl", "--cloud_optics",
                         type=str, default="",
                         help="Path to cloud optics file")
-    parser.add_argument("-a", "--aerosol_optics", 
+    parser.add_argument("-a", "--aerosol_optics",
                         type=str, default="",
                         help="Path to aerosol optics file")
     args = parser.parse_args()
 
-    # Can't supply aerosols without clouds 
-    if(args.cloud_optics == "" and args.aerosol_optics != ""): 
-        raise AssertionError("Need to supply cloud optics if providing aerosol optics")
+    # Can't supply aerosols without clouds
+    if args.cloud_optics == "" and args.aerosol_optics != "":
+        raise AssertionError(
+            "Need to supply cloud optics if providing aerosol optics")
 
-    # Every combo of ncol, nlay 
-    for l in args.nlay: 
-        for i in args.ncol: 
-          print(f"{args.executable} {i:6d} {l:4d} {args.nloops:3d} " + \
-                f"{args.output_file} {args.k_distribution}"         + \
-                f"{args.cloud_optics} {args.aerosol_optics}") 
+    # Every combo of ncol, nlay
+    for l in args.nlay:
+        for i in args.ncol:
+            print(f"{args.executable} {i:6d} {l:4d} {args.nloops:3d} " +
+                  f"{args.output_file} {args.k_distribution}" +
+                  f"{args.cloud_optics} {args.aerosol_optics}")
diff --git a/examples/all-sky/run-allsky-example.py b/examples/all-sky/run-allsky-example.py
@@ -3,9 +3,7 @@
 # This script runs runs the RTE+RRTMGP all-sky examples
 #
 import argparse
-import glob
 import os
-import shutil
 import subprocess
 
 rte_rrtmgp_dir = os.environ["RRTMGP_DATA"]
@@ -38,17 +36,21 @@
                              "cloudy columns")
 
     args = parser.parse_args()
-    ncol_str   = '{0:5d}'.format(args.ncol)
-    nlay_str   = '{0:5d}'.format(args.nlay)
+    ncol_str = '{0:5d}'.format(args.ncol)
+    nlay_str = '{0:5d}'.format(args.nlay)
     nloops_str = '{0:5d}'.format(args.nloops)
     if args.run_command:
         print("using the run command")
         all_sky_exe_name = args.run_command + " " + all_sky_exe_name
     os.chdir(all_sky_dir)
     # Remove cloudy-sky fluxes from the file containing the atmospheric profiles
     subprocess.run(
-        [all_sky_exe_name, ncol_str, nlay_str, nloops_str, "rrtmgp-allsky-lw-no-aerosols.nc", lw_gas_coeffs_file, lw_clouds_coeff_file])
+        [all_sky_exe_name, ncol_str, nlay_str, nloops_str,
+         "rrtmgp-allsky-lw-no-aerosols.nc", lw_gas_coeffs_file,
+         lw_clouds_coeff_file])
     subprocess.run(
-        [all_sky_exe_name, ncol_str, nlay_str, nloops_str, "rrtmgp-allsky-sw-no-aerosols.nc", sw_gas_coeffs_file, sw_clouds_coeff_file])
+        [all_sky_exe_name, ncol_str, nlay_str, nloops_str,
+         "rrtmgp-allsky-sw-no-aerosols.nc", sw_gas_coeffs_file,
+         sw_clouds_coeff_file])
 
 # end main
diff --git a/examples/compare-to-reference.py b/examples/compare-to-reference.py
@@ -9,43 +9,45 @@
 # Thresholds come from environement variables when set? 
 # 
 #
-import argparse
-import os
 import sys
-import warnings 
 
+import argparse
 import numpy as np
+import os
+import warnings
 import xarray as xr
 
-REPORT  = os.getenv("REPORTING_THRESHOLD") if os.getenv("REPORTING_THRESHOLD") is not None else 0.
-FAILURE = os.getenv("FAILURE_THRESHOLD")   if os.getenv("FAILURE_THRESHOLD")   is not None else 1.e-5
 #
 # Comparing reference and test results
 #
 if __name__ == '__main__':
     warnings.simplefilter("ignore", xr.SerializationWarning)
-    
+
     parser = argparse.ArgumentParser(
         description="Compares example output to file in reference directory")
-    parser.add_argument("--ref_dir", type=str, 
+    parser.add_argument("--ref_dir", type=str,
                         help="Directory containing reference files")
-    parser.add_argument("--tst_dir", type=str, default = ".",  
+    parser.add_argument("--tst_dir", type=str, default=".",
                         help="Directory contining test values")
-    parser.add_argument("--file_names", type=str, nargs='+', default = [],  
+    parser.add_argument("--file_names", type=str, nargs='+', default=[],
                         help="Name[s] of files to compare")
-    parser.add_argument("--variables", type=str, nargs='+', default = None,  
+    parser.add_argument("--variables", type=str, nargs='+', default=None,
                         help="Name[s] of files to compare")
-    parser.add_argument("--report_threshold", type=float, 
-                        default=REPORT,
+    parser.add_argument("--report_threshold", type=float,
+                        default=os.getenv("REPORTING_THRESHOLD", 0.),
                         help="Threshold for reporting differences")
-    parser.add_argument("--failure_threshold", type=float, 
-                        default=FAILURE,
+    parser.add_argument("--failure_threshold", type=float,
+                        default=os.getenv("FAILURE_THRESHOLD", 1.e-5),
                         help="Threshold at which differences cause failure "
                              "(for continuous integration)")
     args = parser.parse_args()
 
-    tst = xr.open_mfdataset([os.path.join(args.tst_dir, f) for f in args.file_names],combine='by_coords')
-    ref = xr.open_mfdataset([os.path.join(args.ref_dir, f) for f in args.file_names],combine='by_coords')
+    tst = xr.open_mfdataset(
+        [os.path.join(args.tst_dir, f) for f in args.file_names],
+        combine='by_coords')
+    ref = xr.open_mfdataset(
+        [os.path.join(args.ref_dir, f) for f in args.file_names],
+        combine='by_coords')
     variables = args.variables if args.variables is not None else tst.variables
 
     failed = False
@@ -62,23 +64,25 @@
         #
         # Reporting 
         #
-        if not np.allclose(tst[v], ref[v], atol=args.report_threshold, rtol=0): 
-            diff = abs( (tst - ref)[v].values)
+        if not np.allclose(tst[v], ref[v], atol=args.report_threshold, rtol=0):
+            diff = abs((tst - ref)[v].values)
             avg = 0.5 * (tst + ref)[v].values
-            # Division raises a runtime warning when we divide by zero even if the
-            # values in those locations will be ignored.
+            # Division raises a runtime warning when we divide by zero even if
+            # the values in those locations will be ignored.
             with np.errstate(divide='ignore', invalid='ignore'):
                 frac_diff = np.where(
                     (avg > 2. * np.finfo(float).eps), diff / avg, 0)
                 print('Variable %s differs (max abs difference: %e; '
-                      'max percent difference: %e%%)' % (v, diff.max(), 100.0 * frac_diff.max()))
+                      'max percent difference: %e%%)' % (
+                          v, diff.max(), 100.0 * frac_diff.max()))
         else:
             print('Variable %s: No diffs' % v)
         #
         # Failure
         #
-        if not np.allclose(tst[v], ref[v], atol=args.failure_threshold, rtol=0): failed = True
+        if not np.allclose(tst[v], ref[v], atol=args.failure_threshold,
+                           rtol=0): failed = True
 
-    if failed: print ("Tests failed")
+    if failed:
+        print("Tests failed")
     sys.exit(1) if failed else sys.exit(0)
-
diff --git a/examples/rfmip-clear-sky/run-rfmip-examples.py b/examples/rfmip-clear-sky/run-rfmip-examples.py
@@ -3,7 +3,6 @@
 # This script runs runs RRTMGP on the RFMIP off-line test cases
 #
 import argparse
-import glob
 import os
 import subprocess
 
@@ -15,11 +14,13 @@
 rfmip_dir = "."
 # Code should be run in the rfmip_dir directory
 
-conds_file         = os.path.join(os.environ["RRTMGP_DATA"],
-                   "examples", "rfmip-clear-sky", "inputs",
-                   "multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc")
-lw_gas_coeffs_file = os.path.join(os.environ["RRTMGP_DATA"], "rrtmgp-gas-lw-g256.nc")
-sw_gas_coeffs_file = os.path.join(os.environ["RRTMGP_DATA"], "rrtmgp-gas-sw-g224.nc")
+conds_file = os.path.join(
+    os.environ["RRTMGP_DATA"], "examples", "rfmip-clear-sky", "inputs",
+    "multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc")
+lw_gas_coeffs_file = os.path.join(
+    os.environ["RRTMGP_DATA"], "rrtmgp-gas-lw-g256.nc")
+sw_gas_coeffs_file = os.path.join(
+    os.environ["RRTMGP_DATA"], "rrtmgp-gas-sw-g224.nc")
 
 rfmip_lw_exe_name = "./rrtmgp_rfmip_lw"
 rfmip_sw_exe_name = "./rrtmgp_rfmip_sw"

diff --git a/examples/rfmip-clear-sky/stage_files.py b/examples/rfmip-clear-sky/stage_files.py
@@ -3,8 +3,9 @@
 # This script downloads and/or creates files needed for the RFMIP off-line test
 # cases
 #
-import subprocess
 import sys
+
+import subprocess
 import urllib.request
 from pathlib import Path
 
@@ -39,13 +40,6 @@
 print("Downloading RFMIP input files")
 urllib.request.urlretrieve(conds_url, conds_file)
 
-generate_templates = False
-if generate_templates:
-    print("Downloading scripts for generating output templates")
-    urllib.request.urlretrieve(templ_scr_url, templ_scr)
-    subprocess.run(
-        [sys.executable, templ_scr, "--source_id", "RTE-RRTMGP-181204"])
-else:
-    print("Downloading output templates")
-    for f in output_files:
-        urllib.request.urlretrieve(conds_url.replace(conds_file, f), f)
+print("Downloading output templates")
+for f in output_files:
+    urllib.request.urlretrieve(conds_url.replace(conds_file, f), f)