Fix Longitude comparison error for regional subset_data

.git-blame-ignore-revs

@@ -64,3 +64,4 @@ e8fc526e0d7818d45f171488c78392c4ff63902a
 cdf40d265cc82775607a1bf25f5f527bacc97405
 251e389b361ba673b508e07d04ddcc06b2681989
 8ec50135eca1b99c8b903ecdaa1bd436644688bd
+3b7a2876933263f8986e4069f5d23bd45635756f

python/ctsm/longitude.py

@@ -3,6 +3,8 @@
 """
 
 import logging
+from argparse import ArgumentTypeError
+import numpy as np
 
 logger = logging.getLogger(__name__)
 
@@ -11,16 +13,22 @@ def _check_lon_type_180(lon_in):
     """
     Checks value range of longitude with type 180
     """
-    if not -180 <= lon_in <= 180:
-        raise ValueError(f"lon_in needs to be in the range [-180, 180]: {lon_in}")
+    lon_min = np.min(lon_in)
+    lon_max = np.max(lon_in)
+    for lon in [lon_min, lon_max]:
+        if not -180 <= lon <= 180:
+            raise ValueError("(All values of) lon_in must be in the range [-180, 180]")
 
 
 def _check_lon_type_360(lon_in):
     """
     Checks value range of longitude with type 360
     """
-    if not 0 <= lon_in <= 360:
-        raise ValueError(f"lon_in needs to be in the range [0, 360]: {lon_in}")
+    lon_min = np.min(lon_in)
+    lon_max = np.max(lon_in)
+    for lon in [lon_min, lon_max]:
+        if not 0 <= lon <= 360:
+            raise ValueError("(All values of) lon_in must be in the range [0, 360]")
 
 
 def _check_lon_value_given_type(lon_in, lon_type_in):
@@ -50,6 +58,28 @@ def _convert_lon_type_180_to_360(lon_in):
     return lon_out
 
 
+def _detect_lon_type(lon_in):
+    """
+    Detect longitude type of a given numeric. If lon_in contains more than one number (as in a list
+    or Numpy array), this function will assume all members are of the same type if (a) there is at
+    least one unambiguous member and (b) all unambiguous members are of the same type.
+    """
+    lon_min = np.min(lon_in)
+    lon_max = np.max(lon_in)
+    if lon_min < -180:
+        raise ValueError(f"(Minimum) longitude < -180: {lon_min}")
+    if lon_max > 360:
+        raise ValueError(f"(Maximum) longitude > 360: {lon_max}")
+    min_type_180 = lon_min < 0
+    max_type_360 = lon_max > 180
+    if min_type_180 and max_type_360:
+        raise RuntimeError("Longitude array contains values of both types 180 and 360")
+    if not min_type_180 and not max_type_360:
+        raise ArgumentTypeError("Longitude(s) ambiguous; could be type 180 or 360")
+    lon_type = 180 if min_type_180 else 360
+    return lon_type
+
+
 def _convert_lon_type_360_to_180(lon_in):
     """
     Convert a longitude from type 360 to type 180
@@ -70,6 +100,18 @@ def _convert_lon_type_360_to_180(lon_in):
     return lon_out
 
 
+def check_other_is_lontype(other):
+    """
+    We could try to coerce non-Longitude `other` to Longitude, but that might result in
+    situations where tests think everything works but code will fail if `other` is
+    ambiguous.
+    """
+    if not isinstance(other, Longitude):
+        raise TypeError(
+            f"Comparison not supported between instances of 'Longitude' and '{type(other)}'"
+        )
+
+
 class Longitude:
     """
     A class to keep track of a longitude and its type
@@ -93,6 +135,41 @@ def __init__(self, lon, lon_type):
         self._lon = lon
         self._lon_type = lon_type
 
+    # __eq__ makes it so that == and != both work.
+    def __eq__(self, other):
+        check_other_is_lontype(other)
+        return self._lon == other.get(self._lon_type)
+
+    def _check_lons_same_type(self, other):
+        """
+        If you're comparing two Longitudes of different types in different hemispheres, then
+        `lon1 > lon2` and `lon2 < lon1` will give different answers! We could make it so that
+        this doesn't fail as long as symmetricality isn't violated, but that might lead to
+        unexpected failures in practice.
+        """
+        if self.lon_type() != other.lon_type():
+            raise TypeError("Comparison not supported between Longitudes of different types")
+
+    def __lt__(self, other):
+        check_other_is_lontype(other)
+        self._check_lons_same_type(other)
+        return self._lon < other._lon
+
+    def __gt__(self, other):
+        check_other_is_lontype(other)
+        self._check_lons_same_type(other)
+        return self._lon > other._lon
+
+    def __le__(self, other):
+        check_other_is_lontype(other)
+        self._check_lons_same_type(other)
+        return self._lon <= other._lon
+
+    def __ge__(self, other):
+        check_other_is_lontype(other)
+        self._check_lons_same_type(other)
+        return self._lon >= other._lon
+
     def get(self, lon_type_out):
         """
         Get the longitude value, converting longitude type if needed
@@ -104,3 +181,9 @@ def get(self, lon_type_out):
         if lon_type_out == 360:
             return _convert_lon_type_180_to_360(self._lon)
         raise RuntimeError(f"Add handling for lon_type_out {lon_type_out}")
+
+    def lon_type(self):
+        """
+        Getter method for self._lon_type
+        """
+        return self._lon_type

python/ctsm/site_and_regional/regional_case.py

@@ -19,6 +19,7 @@
 from ctsm.utils import add_tag_to_filename
 from ctsm.utils import abort
 from ctsm.config_utils import check_lon1_lt_lon2
+from ctsm.longitude import Longitude, _detect_lon_type
 
 logger = logging.getLogger(__name__)
 
@@ -132,6 +133,31 @@ def __init__(
         self.ni = None
         self.nj = None
 
+    def _subset_lon_lat(self, x_dim, y_dim, f_in):
+        """
+        subset longitude and latitude arrays
+        """
+        lat = f_in["lat"]
+        lon = f_in["lon"]
+
+        # Handle type of input longitude
+        f_lon_type = _detect_lon_type(lon)
+        lon1_type = self.lon1.lon_type()
+        lon2_type = self.lon2.lon_type()
+        if lon1_type != lon2_type:
+            raise RuntimeError(f"lon1 type ({lon1_type}) doesn't match lon2 type ({lon2_type})")
+        if f_lon_type != lon1_type:
+            # This may be overly strict; we might want to allow conversion to lon1 type.
+            raise RuntimeError(
+                f"File lon type ({f_lon_type}) doesn't match boundary lon type ({lon1_type})"
+            )
+        lon = Longitude(lon, lon1_type)
+
+        xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0]
+        yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0]
+        f_out = f_in.isel({y_dim: yind, x_dim: xind})
+        return f_out
+
     def create_tag(self):
         """
         Create a tag for a region which is either the region name
@@ -192,14 +218,12 @@ def create_domain_at_reg(self, indir, file):
         logger.info("Creating domain file at region: %s", self.tag)
 
         # create 1d coordinate variables to enable sel() method
-        f_in = self.create_1d_coord(fdomain_in, "xc", "yc", "ni", "nj")
-        lat = f_in["lat"]
-        lon = f_in["lon"]
+        x_dim = "ni"
+        y_dim = "nj"
+        f_in = self.create_1d_coord(fdomain_in, "xc", "yc", x_dim, y_dim)
 
         # subset longitude and latitude arrays
-        xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0]
-        yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0]
-        f_out = f_in.isel(nj=yind, ni=xind)
+        f_out = self._subset_lon_lat(x_dim, y_dim, f_in)
 
         # update attributes
         self.update_metadata(f_out)
@@ -240,14 +264,12 @@ def create_surfdata_at_reg(self, indir, file, user_mods_dir, specify_fsurf_out):
         logger.info("fsurf_out: %s", os.path.join(self.out_dir, fsurf_out))
 
         # create 1d coordinate variables to enable sel() method
-        f_in = self.create_1d_coord(fsurf_in, "LONGXY", "LATIXY", "lsmlon", "lsmlat")
-        lat = f_in["lat"]
-        lon = f_in["lon"]
+        x_dim = "lsmlon"
+        y_dim = "lsmlat"
+        f_in = self.create_1d_coord(fsurf_in, "LONGXY", "LATIXY", x_dim, y_dim)
 
         # subset longitude and latitude arrays
-        xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0]
-        yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0]
-        f_out = f_in.isel(lsmlat=yind, lsmlon=xind)
+        f_out = self._subset_lon_lat(x_dim, y_dim, f_in)
 
         # update attributes
         self.update_metadata(f_out)
@@ -304,14 +326,12 @@ def create_landuse_at_reg(self, indir, file, user_mods_dir):
         logger.info("fluse_out: %s", os.path.join(self.out_dir, fluse_out))
 
         # create 1d coordinate variables to enable sel() method
-        f_in = self.create_1d_coord(fluse_in, "LONGXY", "LATIXY", "lsmlon", "lsmlat")
-        lat = f_in["lat"]
-        lon = f_in["lon"]
+        x_dim = "lsmlon"
+        y_dim = "lsmlat"
+        f_in = self.create_1d_coord(fluse_in, "LONGXY", "LATIXY", x_dim, y_dim)
 
         # subset longitude and latitude arrays
-        xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0]
-        yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0]
-        f_out = f_in.isel(lsmlat=yind, lsmlon=xind)
+        f_out = self._subset_lon_lat(x_dim, y_dim, f_in)
 
         # update attributes
         self.update_metadata(f_out)

python/ctsm/subset_data.py

@@ -69,7 +69,7 @@
 from ctsm.path_utils import path_to_ctsm_root
 from ctsm.utils import abort
 from ctsm.config_utils import check_lon1_lt_lon2
-from ctsm.longitude import Longitude
+from ctsm.longitude import Longitude, _detect_lon_type
 
 # -- import ctsm logging flags
 from ctsm.ctsm_logging import (
@@ -626,46 +626,23 @@ def setup_files(args, defaults, cesmroot):
         logger.info("clmforcingindir does not exist: %s", clmforcingindir)
         abort("inputdata directory does not exist")
 
+    file_dict = {"main_dir": clmforcingindir}
+
     # DATM data
     # TODO Issue #2960: Make datm_type a user option at the command
     # line. For reference, this option affects three .cfg files:
     #      tools/site_and_regional/default_data_1850.cfg
     #      tools/site_and_regional/default_data_2000.cfg
     #      python/ctsm/test/testinputs/default_data.cfg
-    datm_type = "datm_crujra"  # also available: datm_type = "datm_gswp3"
-    dir_output_datm = "datmdata"
-    dir_input_datm = os.path.join(clmforcingindir, defaults.get(datm_type, "dir"))
     if args.create_datm:
+        datm_type = "datm_crujra"  # also available: datm_type = "datm_gswp3"
+        dir_output_datm = "datmdata"
+        dir_input_datm = os.path.join(clmforcingindir, defaults.get(datm_type, "dir"))
         if not os.path.isdir(os.path.join(args.out_dir, dir_output_datm)):
             os.mkdir(os.path.join(args.out_dir, dir_output_datm))
         logger.info("dir_input_datm : %s", dir_input_datm)
         logger.info("dir_output_datm: %s", os.path.join(args.out_dir, dir_output_datm))
-
-    # if the crop flag is on - we need to use a different land use and surface data file
-    num_pft = determine_num_pft(args.crop_flag)
-
-    fsurf_in = defaults.get("surfdat", "surfdat_" + num_pft + "pft")
-    fluse_in = defaults.get("landuse", "landuse_" + num_pft + "pft")
-    if args.out_surface:
-        fsurf_out = args.out_surface
-    else:
-        fsurf_out = None
-
-    file_dict = {
-        "main_dir": clmforcingindir,
-        "fdomain_in": defaults.get("domain", "file"),
-        "fsurf_dir": os.path.join(
-            clmforcingindir,
-            os.path.join(defaults.get("surfdat", "dir")),
-        ),
-        "fluse_dir": os.path.join(
-            clmforcingindir,
-            os.path.join(defaults.get("landuse", "dir")),
-        ),
-        "fsurf_in": fsurf_in,
-        "fsurf_out": fsurf_out,
-        "fluse_in": fluse_in,
-        "datm_tuple": DatmFiles(
+        file_dict["datm_tuple"] = DatmFiles(
             dir_input_datm,
             dir_output_datm,
             defaults.get(datm_type, "domain"),
@@ -678,8 +655,30 @@ def setup_files(args, defaults, cesmroot):
             defaults.get(datm_type, "solarname"),
             defaults.get(datm_type, "precname"),
             defaults.get(datm_type, "tpqwname"),
-        ),
-    }
+        )
+
+    # if the crop flag is on - we need to use a different land use and surface data file
+    num_pft = determine_num_pft(args.crop_flag)
+
+    if args.create_domain:
+        file_dict["fdomain_in"] = defaults.get("domain", "file")
+    if args.create_surfdata:
+        file_dict["fsurf_dir"] = os.path.join(
+            clmforcingindir,
+            os.path.join(defaults.get("surfdat", "dir")),
+        )
+        file_dict["fsurf_in"] = defaults.get("surfdat", "surfdat_" + num_pft + "pft")
+        if args.out_surface:
+            fsurf_out = args.out_surface
+        else:
+            fsurf_out = None
+        file_dict["fsurf_out"] = fsurf_out
+    if args.create_landuse:
+        file_dict["fluse_in"] = defaults.get("landuse", "landuse_" + num_pft + "pft")
+        file_dict["fluse_dir"] = os.path.join(
+            clmforcingindir,
+            os.path.join(defaults.get("landuse", "dir")),
+        )
 
     return file_dict
 
@@ -821,17 +820,6 @@ def subset_region(args, file_dict: dict):
     logger.info("Successfully ran script for a regional case.")
 
 
-def _detect_lon_type(lon_in):
-    if lon_in < 0:
-        lon_type = 180
-    elif lon_in > 180:
-        lon_type = 360
-    else:
-        msg = "When providing an ambiguous longitude, you must specify --lon-type 180 or 360"
-        raise argparse.ArgumentTypeError(msg)
-    return lon_type
-
-
 def process_args(args):
     """
     Process arguments after parsing