Merge pull request #3264 from samsrabin/plumber2_surf_wrapper-fixes

Fix plumber2_surf_wrapper

Author: samsrabin

.git-blame-ignore-revs

@@ -68,3 +68,4 @@ cdf40d265cc82775607a1bf25f5f527bacc97405
 3dd489af7ebe06566e2c6a1c7ade18550f1eb4ba
 742cfa606039ab89602fde5fef46458516f56fd4
 4ad46f46de7dde753b4653c15f05326f55116b73
+75db098206b064b8b7b2a0604d3f0bf8fdb950cc

python/ctsm/pft_utils.py

@@ -0,0 +1,21 @@
+"""
+Constants and functions relating to PFTs
+"""
+
+MIN_PFT = 0  # bare ground
+MIN_NAT_PFT = 1  # minimum natural pft (not including bare ground)
+MAX_NAT_PFT = 14  # maximum natural pft
+MAX_PFT_GENERICCROPS = 16  # for runs with generic crops
+MAX_PFT_MANAGEDCROPS = 78  # for runs with explicit crops
+
+
+def is_valid_pft(pft_num, managed_crops):
+    """
+    Given a number, check whether it represents a valid PFT (bare ground OK)
+    """
+    if managed_crops:
+        max_allowed_pft = MAX_PFT_MANAGEDCROPS
+    else:
+        max_allowed_pft = MAX_PFT_GENERICCROPS
+
+    return MIN_PFT <= pft_num <= max_allowed_pft

python/ctsm/site_and_regional/plumber2_shared.py

@@ -0,0 +1,21 @@
+"""
+Things shared between plumber2 scripts
+"""
+
+import os
+import pandas as pd
+from ctsm.path_utils import path_to_ctsm_root
+
+PLUMBER2_SITES_CSV = os.path.join(
+    path_to_ctsm_root(),
+    "tools",
+    "site_and_regional",
+    "PLUMBER2_sites.csv",
+)
+
+
+def read_plumber2_sites_csv(file=PLUMBER2_SITES_CSV):
+    """
+    Read PLUMBER2_sites.csv using pandas
+    """
+    return pd.read_csv(file, skiprows=5)

python/ctsm/site_and_regional/plumber2_surf_wrapper.py

@@ -22,16 +22,18 @@
 
 import argparse
 import logging
-import os
-import subprocess
+import sys
 import tqdm
 
-import pandas as pd
+# pylint:disable=wrong-import-position
+from ctsm.site_and_regional.plumber2_shared import PLUMBER2_SITES_CSV, read_plumber2_sites_csv
+from ctsm import subset_data
+from ctsm.pft_utils import MAX_PFT_MANAGEDCROPS, is_valid_pft
 
 
-def get_parser():
+def get_args():
     """
-    Get parser object for this script.
+    Get arguments for this script.
     """
     parser = argparse.ArgumentParser(
         description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
@@ -45,137 +47,148 @@ def get_parser():
         help="Verbose mode will print more information. ",
         action="store_true",
         dest="verbose",
-        default=False,
     )
 
     parser.add_argument(
-        "--16pft",
-        help="Create and/or modify 16-PFT surface datasets (e.g. for a FATES run) ",
+        "--crop",
+        help=f"Create and/or modify {MAX_PFT_MANAGEDCROPS}-PFT "
+        "surface datasets (e.g. for a non-FATES run)",
         action="store_true",
-        dest="pft_16",
-        default=True,
+        dest="use_managed_crops",
     )
 
-    return parser
+    parser.add_argument(
+        "--overwrite",
+        help="Overwrite any existing files",
+        action="store_true",
+    )
+
+    parser.add_argument(
+        "--plumber2-sites-csv",
+        help=f"Comma-separated value (CSV) file with Plumber2 sites. Default: {PLUMBER2_SITES_CSV}",
+        default=PLUMBER2_SITES_CSV,
+    )
+
+    return parser.parse_args()
 
 
 def execute(command):
     """
-    Function for running a command on shell.
+    Runs subset_data with given arguments.
     Args:
-        command (str):
-            command that we want to run.
+        command (list):
+            list of args for command that we want to run.
     Raises:
-        Error with the return code from shell.
+        Whatever error subset_data gives, if any.
     """
     print("\n", " >>  ", *command, "\n")
 
-    try:
-        subprocess.check_call(command, stdout=open(os.devnull, "w"), stderr=subprocess.STDOUT)
-
-    except subprocess.CalledProcessError as err:
-        # raise RuntimeError("command '{}' return with error
-        # (code {}): {}".format(e.cmd, e.returncode, e.output))
-        # print (e.ouput)
-        print(err)
+    sys.argv = command
+    subset_data.main()
 
 
 def main():
     """
     Read plumber2_sites from csv, iterate through sites, and add dominant PFT
     """
 
-    args = get_parser().parse_args()
+    args = get_args()
 
     if args.verbose:
         logging.basicConfig(level=logging.DEBUG)
 
-    plumber2_sites = pd.read_csv("PLUMBER2_sites.csv", skiprows=4)
+    plumber2_sites = read_plumber2_sites_csv(args.plumber2_sites_csv)
 
     for _, row in tqdm.tqdm(plumber2_sites.iterrows()):
         lat = row["Lat"]
         lon = row["Lon"]
         site = row["Site"]
+
+        clmsite = "1x1_PLUMBER2_" + site
+        print("Now processing site :", site)
+
+        # Set up part of subset_data command that is shared among all options
+        subset_command = [
+            "./subset_data",
+            "point",
+            "--lat",
+            str(lat),
+            "--lon",
+            str(lon),
+            "--site",
+            clmsite,
+            "--create-surface",
+            "--uniform-snowpack",
+            "--cap-saturation",
+            "--lon-type",
+            "180",
+        ]
+
+        # Read info for first PFT
         pft1 = row["pft1"]
+        if not is_valid_pft(pft1, args.use_managed_crops):
+            raise RuntimeError(f"pft1 must be a valid PFT; got {pft1}")
         pctpft1 = row["pft1-%"]
         cth1 = row["pft1-cth"]
         cbh1 = row["pft1-cbh"]
-        pft2 = row["pft2"]
-        pctpft2 = row["pft2-%"]
-        cth2 = row["pft2-cth"]
-        cbh2 = row["pft2-cbh"]
-        # overwrite missing values from .csv file
-        if pft1 == -999:
-            pft1 = 0
-            pctpft1 = 0
-            cth1 = 0
-            cbh1 = 0
-        if pft2 == -999:
-            pft2 = 0
-            pctpft2 = 0
-            cth2 = 0
-            cbh2 = 0
-        clmsite = "1x1_PLUMBER2_" + site
-        print("Now processing site :", site)
 
-        if args.pft_16:
-            # use surface dataset with 16 pfts, but overwrite to 100% 1 dominant PFT
-            # don't set crop flag
-            # set dominant pft
-            subset_command = [
-                "./subset_data",
-                "point",
-                "--lat",
-                str(lat),
-                "--lon",
-                str(lon),
-                "--site",
-                clmsite,
+        # Read info for second PFT, if a valid one is given in the .csv file
+        pft2 = row["pft2"]
+        if is_valid_pft(pft2, args.use_managed_crops):
+            pctpft2 = row["pft2-%"]
+            cth2 = row["pft2-cth"]
+            cbh2 = row["pft2-cbh"]
+
+        # Set dominant PFT(s)
+        if is_valid_pft(pft2, args.use_managed_crops):
+            subset_command += [
                 "--dompft",
                 str(pft1),
                 str(pft2),
                 "--pctpft",
                 str(pctpft1),
                 str(pctpft2),
-                "--cth",
-                str(cth1),
-                str(cth2),
-                "--cbh",
-                str(cbh1),
-                str(cbh2),
-                "--create-surface",
-                "--uniform-snowpack",
-                "--cap-saturation",
-                "--verbose",
-                "--overwrite",
             ]
         else:
-            # use surface dataset with 78 pfts, and overwrite to 100% 1 dominant PFT
-            # NOTE: FATES will currently not run with a 78-PFT surface dataset
-            # set crop flag
-            # set dominant pft
-            subset_command = [
-                "./subset_data",
-                "point",
-                "--lat",
-                str(lat),
-                "--lon",
-                str(lon),
-                "--site",
-                clmsite,
-                "--crop",
+            subset_command += [
                 "--dompft",
                 str(pft1),
-                str(pft2),
                 "--pctpft",
                 str(pctpft1),
-                str(pctpft2),
-                "--create-surface",
-                "--uniform-snowpack",
-                "--cap-saturation",
-                "--verbose",
-                "--overwrite",
             ]
+
+        if not args.use_managed_crops:
+            # use surface dataset with 78 pfts, but overwrite to 100% 1 dominant PFT
+            # don't set crop flag
+            # set canopy top and bottom heights
+            if is_valid_pft(pft2, args.use_managed_crops):
+                subset_command += [
+                    "--cth",
+                    str(cth1),
+                    str(cth2),
+                    "--cbh",
+                    str(cbh1),
+                    str(cbh2),
+                ]
+            else:
+                subset_command += [
+                    "--cth",
+                    str(cth1),
+                    "--cbh",
+                    str(cbh1),
+                ]
+        else:
+            # use surface dataset with 78 pfts, and overwrite to 100% 1 dominant PFT
+            # NOTE: FATES will currently not run with a 78-PFT surface dataset
+            # set crop flag
+            subset_command += ["--crop"]
+            # don't set canopy top and bottom heights
+
+        if args.verbose:
+            subset_command += ["--verbose"]
+        if args.overwrite:
+            subset_command += ["--overwrite"]
+
         execute(subset_command)
 
 

python/ctsm/site_and_regional/plumber2_usermods.py

@@ -13,7 +13,8 @@
 import os
 import tqdm
 
-import pandas as pd
+# pylint:disable=wrong-import-position
+from ctsm.site_and_regional.plumber2_shared import read_plumber2_sites_csv
 
 
 # Big ugly function to create usermod_dirs for each site
@@ -155,7 +156,7 @@ def main():
     """
 
     # For now we can just run the 'main' program as a loop
-    plumber2_sites = pd.read_csv("PLUMBER2_sites.csv", skiprows=4)
+    plumber2_sites = read_plumber2_sites_csv()
 
     for _, row in tqdm.tqdm(plumber2_sites.iterrows()):
         lat = row["Lat"]