Pint Addition (reopened with respect to main) (COSIMA#261)

* Retry

* black formatting

* more black formatting

* new black formatting

* ForOM2

* Black

* New black grr

* add a comment to explain why we keep manual conversion

---------

Co-authored-by: Ashley Barnes <53282288+ashjbarnes@users.noreply.github.com>
Co-authored-by: ashjbarnes <ashjbarnes97@gmail.com>

Author: manishvenu

pyproject.toml

@@ -17,7 +17,8 @@ dependencies = [
   "xarray",
   "xesmf >= 0.8.4",
   "f90nml >= 1.4.1",
-  "copernicusmarine >= 2.0.0,<2.1.0"
+  "copernicusmarine >= 2.0.0,<2.1.0",
+  "pint_xarray"
 ]
 
 [build-system]
@@ -31,6 +32,7 @@ packages = ["regional_mom6", "regional_mom6.demos.premade_run_directories"]
 "regional_mom6.demos" = "demos"
 
 [tool.setuptools.package-data]
+"regional_mom6" = ["*.txt"]  # include rm6_unit_defs.txt here
 "regional_mom6.demos.premade_run_directories" = ["**/*"]
 
 [tool.setuptools_scm]

regional_mom6/MOM_parameter_tools.py

@@ -35,7 +35,6 @@ def change_MOM_parameter(
     MOM_override_dict = read_MOM_file_as_dict("MOM_override", directory)
     original_val = "No original val"
     if not delete:
-
         if param_name in MOM_override_dict:
             original_val = MOM_override_dict[param_name]["value"]
             print(

regional_mom6/regional_mom6.py

@@ -25,9 +25,9 @@
     ep2ap,
     rotate,
     find_files_by_pattern,
+    try_pint_convert,
 )
 
-
 warnings.filterwarnings("ignore")
 
 __all__ = [
@@ -39,6 +39,14 @@
     "get_glorys_data",
 ]
 
+# If the array is pint possible, ensure we have the right units for main fields (eta, u, v, temp),
+# salinity and bgc tracers are a bit more abstract and should be already in the correct units, a TODO: would be to add functionality to convert these tracers
+main_field_target_units = {
+    "eta": "m",
+    "u": "m/s",
+    "v": "m/s",
+    "temp": "degC",
+}
 
 ## Mapping Functions
 
@@ -135,7 +143,6 @@ def longitude_slicer(data, longitude_extent, longitude_coords):
 
         for i in range(-1, 2, 1):
             if data[lon][0] <= central_longitude + 360 * i <= data[lon][-1]:
-
                 ## Shifted version of target midpoint; e.g., could be -90 vs 270
                 ## integer i keeps track of what how many multiples of 360 we need to shift entire
                 ## grid by to match central_longitude
@@ -231,11 +238,9 @@ def get_glorys_data(
 
     file = open(Path(path / "get_glorys_data.sh"), "w")
 
-    lines.append(
-        f"""
+    lines.append(f"""
 copernicusmarine subset --dataset-id cmems_mod_glo_phy_my_0.083deg_P1D-m --variable so --variable thetao --variable uo --variable vo --variable zos --start-datetime {str(timerange[0]).replace(" ","T")} --end-datetime {str(timerange[1]).replace(" ","T")} --minimum-longitude {longitude_extent[0] - buffer} --maximum-longitude {longitude_extent[1] + buffer} --minimum-latitude {latitude_extent[0] - buffer} --maximum-latitude {latitude_extent[1] + buffer} --minimum-depth 0 --maximum-depth 6000 -o {str(path)} -f {segment_name}.nc\n
-"""
-    )
+""")
     file.writelines(lines)
     file.close()
     return Path(path / "get_glorys_data.sh")
@@ -607,7 +612,6 @@ def __init__(
         regridding_method="bilinear",
         fill_method=rgd.fill_missing_data,
     ):
-
         # Creates an empty experiment object for testing and experienced user manipulation.
         if create_empty:
             return
@@ -830,7 +834,6 @@ def bathymetry_path(self):
             return "Not Found"
 
     def __getattr__(self, name):
-
         ## First, check whether the attribute is an input file
         if "segment" in name:
             try:
@@ -904,7 +907,6 @@ def _make_hgrid(self):
         ), "only even_spacing grid type is implemented"
 
         if self.hgrid_type == "even_spacing":
-
             # longitudes are evenly spaced based on resolution and bounds
             nx = int(
                 (self.longitude_extent[1] - self.longitude_extent[0])
@@ -1023,6 +1025,22 @@ def setup_initial_condition(
         if type(reprocessed_var_map["depth_coord"]) == list:
             reprocessed_var_map["depth_coord"] = reprocessed_var_map["depth_coord"][0]
 
+        # Convert zdim if possible & needed
+        ic_raw[reprocessed_var_map["depth_coord"]] = try_pint_convert(
+            ic_raw[reprocessed_var_map["depth_coord"]],
+            "m",
+            reprocessed_var_map["depth_coord"],
+        )
+
+        # Convert values
+        for var in main_field_target_units:
+            if var == "temp" or var == "salt":
+                value_name = reprocessed_var_map["tracer_var_names"][var]
+            else:
+                value_name = reprocessed_var_map[var + "_var_name"]
+            ic_raw[value_name] = try_pint_convert(
+                ic_raw[value_name], main_field_target_units[var], var
+            )
         # Remove time dimension if present in the IC.
         # Assume that the first time dim is the intended one if more than one is present
 
@@ -1060,6 +1078,8 @@ def setup_initial_condition(
 
         ## if min(temperature) > 100 then assume that units must be degrees K
         ## (otherwise we can't be on Earth) and convert to degrees C
+        ## Although we now attempt a pint convert, we're leaving this manual conversion in for now
+        ## just in case, as K->C is absolutely necessary, and for some inputs pint may fail where this won't.
         if np.nanmin(ic_raw[reprocessed_var_map["tracer_var_names"]["temp"]]) > 100:
             ic_raw[reprocessed_var_map["tracer_var_names"]["temp"]] -= 273.15
             ic_raw[reprocessed_var_map["tracer_var_names"]["temp"]].attrs[
@@ -2377,7 +2397,6 @@ def setup_run_directory(
             )
         # Tides OBC adjustments
         if with_tides:
-
             # Include internal tide forcing
             MOM_override_dict["TIDES"]["value"] = "True"
 
@@ -2655,6 +2674,14 @@ def regrid_velocity_tracers(
 
         coords = rgd.coords(self.hgrid, self.orientation, self.segment_name)
 
+        # Convert z coordinates to meters if pint-enabled
+        if type(reprocessed_var_map["depth_coord"]) != list:
+            dc_list = [reprocessed_var_map["depth_coord"]]
+        else:
+            dc_list = reprocessed_var_map["depth_coord"]
+        for dc in dc_list:
+            rawseg[dc] = try_pint_convert(rawseg[dc], "m", dc)
+
         regridders = create_vt_regridders(
             reprocessed_var_map,
             rawseg,
@@ -2789,6 +2816,17 @@ def regrid_velocity_tracers(
             ## Rename each variable in dataset
             segment_out = segment_out.rename({allfields[var]: v})
 
+            # Try Pint Conversion
+            if var in main_field_target_units:
+                # Apply raw data units if they exist
+                units = rawseg[allfields[var]].attrs.get("units")
+                if units is not None:
+                    segment_out[v].attrs["units"] = units
+
+                segment_out[v] = try_pint_convert(
+                    segment_out[v], main_field_target_units[var], var
+                )
+
             # Find out if the tracer has depth, and if so, what is it's z dimension (z dimension being a list is an edge case for MARBL BGC)
             variable_has_depth = False
             depth_coord = None
@@ -3093,7 +3131,6 @@ def encode_tidal_files_and_output(self, ds, filename):
         ## Expand Tidal Dimensions ##
 
         for var in ds:
-
             ds = rgd.add_secondary_dimension(ds, str(var), coords, self.segment_name)
 
         ## Rename Tidal Dimensions ##

regional_mom6/regridding.py

@@ -381,7 +381,6 @@ def add_secondary_dimension(
     )
     insert_behind_by = 0
     if not to_beginning:
-
         if any(
             coord.startswith("nz") or coord == "constituent" for coord in ds[var].dims
         ):
@@ -579,7 +578,6 @@ def mask_dataset(
             mask = mask[np.newaxis, :]
 
         for var in ds.data_vars.keys():
-
             # Drop to just the Boundary Dim
             da = ds[var].isel({dim: 0 for dim in list(ds[var].dims)[:-2]}).squeeze()
 

regional_mom6/rm6_unit_defs.txt

@@ -0,0 +1,6 @@
+# Insert Custom NetCDF Units Here
+
+# Units defs for CESM output (degrees are dimensionless, so 1)
+degrees_east = 1 
+degrees_north = 1
+degrees_N = 1

regional_mom6/rotation.py

@@ -121,7 +121,6 @@ def modulo_around_point(x, x0, L):
     if L <= 0:
         return x
     else:
-
         # Find that boundary point x0 + L/2
         edge_indexes = np.where((x == x0 + L / 2))
 

regional_mom6/utils.py

@@ -4,6 +4,97 @@
 import xarray as xr
 from regional_mom6 import regridding as rgd
 from pathlib import Path
+import pint
+import pint_xarray
+import importlib.resources
+
+# from pint_xarray.errors import PintExceptionGroup # This is only supported when pint_xarray is 0.6.0, which is not currently supported in the CI
+
+
+# Handle Unit Registry (only done once)
+ureg = pint.UnitRegistry(
+    force_ndarray_like=True
+)  # The force option is required for pint_xarray
+
+unit_path = Path(importlib.resources.files("regional_mom6") / "rm6_unit_defs.txt")
+ureg.load_definitions(unit_path)
+
+
+def try_pint_convert(da, target_units, var_name=None, debug=False):
+    """
+    Attempt to quantify and convert an xarray DataArray using Pint.
+
+    Steps:
+    1. Check if the DataArray has a 'units' attribute.
+       - If not, Pint cannot quantify it, so we raise a ValueError.
+    2. Quantify the DataArray using Pint (attach units).
+       - This converts the DataArray into a pint-aware object.
+       - If already a Pint Quantity, skip quantification.
+    3. Convert the DataArray to the target units if necessary.
+       - Uses Pint's `.to()` to perform unit conversion.
+       - Dequantify afterwards to return a plain xarray DataArray.
+    4. If any step fails (missing units, invalid units, etc.),
+       - Log a warning and return the original DataArray unchanged.
+
+    Parameters
+    ----------
+    da : xarray.DataArray
+        The DataArray to quantify and/or convert.
+    target_units : str
+        Units to convert the DataArray to.
+    var_name : str, optional
+        Name of the variable (used for logging messages).
+    debug : bool, optional
+        If True, print debug information about the error with subexceptions.
+
+    Returns
+    -------
+    xarray.DataArray
+        A DataArray with units converted if successful; otherwise the original.
+    """
+    try:
+        # Get the units string from the DataArray attributes
+        source_units = da.attrs.get("units", None)
+        if not source_units:
+            raise ValueError(f"DataArray '{var_name}' has no units; cannot quantify.")
+
+        # Only quantify if not already a Pint Quantity
+        if not isinstance(da.data, pint.Quantity):
+            da_quantified = da.pint.quantify(unit_registry=ureg)
+
+            # code for PintExceptionGroup (not supported in current CI until we can use pint 0.6.0)
+            # Allows catching multiple quantification errors at once
+            # except PintExceptionGroup as ex_group:
+            #     print(f"PintExceptionGroup: could not quantify some elements of {var_name}")
+            #     for idx, exc in enumerate(ex_group.exceptions):
+            #         print(f"  Sub-exception {idx+1}: {exc}")
+            #     raise ex_group
+
+        # Convert to the target units if they differ
+        if source_units != target_units:
+            da_converted = da_quantified.pint.to(target_units).pint.dequantify()
+            utils_logger.warning(
+                f"Converted {var_name} from {source_units} to {target_units}"
+            )
+            return da_converted
+        else:
+            utils_logger.info(f"Units for {var_name} did not need to be converted")
+
+    except Exception as e:
+        # If any error occurs (bad units, missing Pint, etc.), fall back gracefully
+        utils_logger.warning(
+            f"regional_mom6 could not use pint for data array {var_name}, assuming it's in the correct units"
+        )
+        if debug:
+            if hasattr(e, "exceptions"):
+                for i, exc in enumerate(e.exceptions):
+                    print(f"\n--- Sub-exception {i} ---")
+                    print(type(exc).__name__, exc)
+            else:
+                print(e)
+
+    # Return the original DataArray if quantification or conversion failed
+    return da
 
 
 def vecdot(v1, v2):
@@ -432,3 +523,6 @@ def get_edge(ds, edge, x_name=None, y_name=None):
         return ds.isel({x_name: -1})
     if edge == "west":
         return ds.isel({x_name: 0})
+
+
+utils_logger = setup_logger(__name__, set_handler=False)

tests/test_config_class.py

@@ -101,7 +101,6 @@ def test_write_config(create_expt, tmp_path):
 
 
 def test_read_config(create_expt, tmp_path):
-
     expt = create_expt
     path = tmp_path / "testing_config.json"
     Config.save_to_json(expt, path)

tests/test_expt_class.py

@@ -249,7 +249,6 @@ def test_rectangular_boundaries(
     hgrid_type,
     tmp_path,
 ):
-
     eastern_boundary = xr.Dataset(
         {
             "temp": xr.DataArray(