Add Rotational Methods function

demos/reanalysis-forced.ipynb

@@ -49,7 +49,7 @@
    "outputs": [],
    "source": [
     "import warnings\n",
-    "warnings.filterwarnings(\"ignore\", module=\"esmpy\")\n",
+    "warnings.filterwarnings(\"ignore\")\n",
     "import regional_mom6 as rmom6\n",
     "\n",
     "import os\n",
@@ -136,21 +136,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "NameError",
-     "evalue": "name 'longitude_extent' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Input \u001b[0;32mIn [2]\u001b[0m, in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m expt \u001b[38;5;241m=\u001b[39m rmom6\u001b[38;5;241m.\u001b[39mexperiment(\n\u001b[0;32m----> 2\u001b[0m     longitude_extent \u001b[38;5;241m=\u001b[39m \u001b[43mlongitude_extent\u001b[49m,\n\u001b[1;32m      3\u001b[0m     latitude_extent \u001b[38;5;241m=\u001b[39m latitude_extent,\n\u001b[1;32m      4\u001b[0m     date_range \u001b[38;5;241m=\u001b[39m date_range,\n\u001b[1;32m      5\u001b[0m     resolution \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.05\u001b[39m,\n\u001b[1;32m      6\u001b[0m     number_vertical_layers \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m75\u001b[39m,\n\u001b[1;32m      7\u001b[0m     layer_thickness_ratio \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m10\u001b[39m,\n\u001b[1;32m      8\u001b[0m     depth \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m4500\u001b[39m,\n\u001b[1;32m      9\u001b[0m     minimum_depth \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m5\u001b[39m,\n\u001b[1;32m     10\u001b[0m     mom_run_dir \u001b[38;5;241m=\u001b[39m run_dir,\n\u001b[1;32m     11\u001b[0m     mom_input_dir \u001b[38;5;241m=\u001b[39m input_dir,\n\u001b[1;32m     12\u001b[0m     toolpath_dir \u001b[38;5;241m=\u001b[39m toolpath_dir,\n\u001b[1;32m     13\u001b[0m     boundaries\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnorth\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msouth\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124meast\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwest\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m     14\u001b[0m )\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'longitude_extent' is not defined"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "expt = rmom6.experiment(\n",
     "    longitude_extent = longitude_extent,\n",

regional_mom6/regional_mom6.py

@@ -1357,23 +1357,14 @@ def setup_initial_condition(
         print("Regridding Velocities... ", end="")
         regridded_u = regridder_u(ic_raw_u)
         regridded_v = regridder_v(ic_raw_v)
-        if rotational_method == rot.RotationMethod.GIVEN_ANGLE:
-            rotated_u, rotated_v = rotate(
-                regridded_u,
-                regridded_v,
-                radian_angle=np.radians(self.hgrid.angle_dx.values),
-            )
-        elif rotational_method == rot.RotationMethod.EXPAND_GRID:
-            self.hgrid["angle_dx_rm6"] = (
-                rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid)
-            )
-            rotated_u, rotated_v = rotate(
-                regridded_u,
-                regridded_v,
-                radian_angle=np.radians(self.hgrid.angle_dx_rm6.values),
-            )
-        elif rotational_method == rot.RotationMethod.NO_ROTATION:
-            rotated_u, rotated_v = regridded_u, regridded_v
+        rotated_u, rotated_v = rotate(
+            regridded_u,
+            regridded_v,
+            radian_angle=np.radians(
+                rot.get_rotation_angle(rotational_method, self.hgrid).values
+            ),
+        )
+
         # Slice the velocites to the u and v grid.
         u_points = rgd.get_hgrid_arakawa_c_points(self.hgrid, "u")
         v_points = rgd.get_hgrid_arakawa_c_points(self.hgrid, "v")
@@ -3017,9 +3008,7 @@ def regrid_velocity_tracers(self, rotational_method=rot.RotationMethod.GIVEN_ANG
         Args:
         rotational_method (rot.RotationMethod): The method to use for rotation of the velocities. Currently, the default method, GIVEN_ANGLE, works even with non-rotated grids
         """
-        if rotational_method == rot.RotationMethod.NO_ROTATION:
-            if not is_rectilinear_hgrid(self.hgrid):
-                raise ValueError("NO_ROTATION method only works with rectilinear grids")
+
         rawseg = xr.open_dataset(self.infile, decode_times=False, engine="netcdf4")
 
         coords = rgd.coords(self.hgrid, self.orientation, self.segment_name)
@@ -3043,36 +3032,15 @@ def regrid_velocity_tracers(self, rotational_method=rot.RotationMethod.GIVEN_ANG
             )
 
             ## Angle Calculation & Rotation
-            if rotational_method == rot.RotationMethod.GIVEN_ANGLE:
-                rotated_u, rotated_v = rotate(
-                    regridded[self.u],
-                    regridded[self.v],
-                    radian_angle=np.radians(coords.angle.values),
-                )
-            elif rotational_method == rot.RotationMethod.EXPAND_GRID:
-
-                # Recalculate entire hgrid angles
-                self.hgrid["angle_dx_rm6"] = (
-                    rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid)
-                )
-
-                # Get just the boundary
-                degree_angle = rgd.coords(
-                    self.hgrid,
-                    self.orientation,
-                    self.segment_name,
-                    angle_variable_name="angle_dx_rm6",
-                )["angle"]
-
-                # Rotate
-                rotated_u, rotated_v = rotate(
-                    regridded[self.u],
-                    regridded[self.v],
-                    radian_angle=np.radians(degree_angle.values),
-                )
-            elif rotational_method == rot.RotationMethod.NO_ROTATION:
-                # Just transfer values
-                rotated_u, rotated_v = regridded[self.u], regridded[self.v]
+            rotated_u, rotated_v = rotate(
+                regridded[self.u],
+                regridded[self.v],
+                radian_angle=np.radians(
+                    rot.get_rotation_angle(
+                        rotational_method, self.hgrid, orientation=self.orientation
+                    ).values
+                ),
+            )
 
             rotated_ds = xr.Dataset(
                 {
@@ -3102,32 +3070,15 @@ def regrid_velocity_tracers(self, rotational_method=rot.RotationMethod.GIVEN_ANG
             )
 
             # See explanation of the rotational methods in the A grid section
-            if rotational_method == rot.RotationMethod.GIVEN_ANGLE:
-                velocities_out["u"], velocities_out["v"] = rotate(
-                    velocities_out["u"],
-                    velocities_out["v"],
-                    radian_angle=np.radians(coords.angle.values),
-                )
-            elif rotational_method == rot.RotationMethod.EXPAND_GRID:
-                self.hgrid["angle_dx_rm6"] = (
-                    rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid)
-                )
-                degree_angle = rgd.coords(
-                    self.hgrid,
-                    self.orientation,
-                    self.segment_name,
-                    angle_variable_name="angle_dx_rm6",
-                )["angle"]
-                velocities_out["u"], velocities_out["v"] = rotate(
-                    velocities_out["u"],
-                    velocities_out["v"],
-                    radian_angle=np.radians(degree_angle.values),
-                )
-            elif rotational_method == rot.RotationMethod.NO_ROTATION:
-                velocities_out["u"], velocities_out["v"] = (
-                    velocities_out["u"],
-                    velocities_out["v"],
-                )
+            velocities_out["u"], velocities_out["v"] = rotate(
+                velocities_out["u"],
+                velocities_out["v"],
+                radian_angle=np.radians(
+                    rot.get_rotation_angle(
+                        rotational_method, self.hgrid, orientation=self.orientation
+                    ).values
+                ),
+            )
 
             segment_out = xr.merge(
                 [
@@ -3167,29 +3118,16 @@ def regrid_velocity_tracers(self, rotational_method=rot.RotationMethod.GIVEN_ANG
             regridded_v = regridder_vvelocity(rawseg[[self.v]])
 
             # See explanation of the rotational methods in the A grid section
-            if rotational_method == rot.RotationMethod.GIVEN_ANGLE:
-                rotated_u, rotated_v = rotate(
-                    regridded_u,
-                    regridded_v,
-                    radian_angle=np.radians(coords.angle.values),
-                )
-            elif rotational_method == rot.RotationMethod.EXPAND_GRID:
-                self.hgrid["angle_dx_rm6"] = (
-                    rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid)
-                )
-                degree_angle = rgd.coords(
-                    self.hgrid,
-                    self.orientation,
-                    self.segment_name,
-                    angle_variable_name="angle_dx_rm6",
-                )["angle"]
-                rotated_u, rotated_v = rotate(
-                    regridded_u,
-                    regridded_v,
-                    radian_angle=np.radians(degree_angle.values),
-                )
-            elif rotational_method == rot.RotationMethod.NO_ROTATION:
-                rotated_u, rotated_v = regridded_u, regridded_v
+            rotated_u, rotated_v = rotate(
+                regridded_u,
+                regridded_v,
+                radian_angle=np.radians(
+                    rot.get_rotation_angle(
+                        rotational_method, self.hgrid, orientation=self.orientation
+                    ).values
+                ),
+            )
+
             rotated_ds = xr.Dataset(
                 {
                     self.u: rotated_u,
@@ -3357,9 +3295,6 @@ def regrid_tides(
         Type: Python Functions, Source Code
         Web Address: https://github.com/jsimkins2/nwa25
         """
-        if rotational_method == rot.RotationMethod.NO_ROTATION:
-            if not is_rectilinear_hgrid(self.hgrid):
-                raise ValueError("NO_ROTATION method only works with rectilinear grids")
 
         # Establish Coords
         coords = rgd.coords(self.hgrid, self.orientation, self.segment_name)
@@ -3450,31 +3385,13 @@ def regrid_tides(
         # and convert ellipse back to amplitude and phase.
         SEMA, ECC, INC, PHA = ap2ep(ucplex, vcplex)
 
-        if rotational_method == rot.RotationMethod.GIVEN_ANGLE:
-
-            # Get user-provided angle
-            angle = coords["angle"]
-
-            # Rotate
-            INC -= np.radians(angle.data[np.newaxis, :])
-
-        elif rotational_method == rot.RotationMethod.EXPAND_GRID:
-
-            # Generate entire hgrid angles using pseudo_hgrid
-            self.hgrid["angle_dx_rm6"] = (
-                rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid)
-            )
-
-            # Get just boundary angles
-            degree_angle = rgd.coords(
-                self.hgrid,
-                self.orientation,
-                self.segment_name,
-                angle_variable_name="angle_dx_rm6",
-            )["angle"]
+        # Rotate
+        INC -= np.radians(
+            rot.get_rotation_angle(
+                rotational_method, self.hgrid, orientation=self.orientation
+            ).data[np.newaxis, :]
+        )
 
-            # Rotate
-            INC -= np.radians(degree_angle.data[np.newaxis, :])
         ua, va, up, vp = ep2ap(SEMA, ECC, INC, PHA)
         # Convert to real amplitude and phase.
 

regional_mom6/regridding.py

@@ -32,7 +32,7 @@
 import dask.array as da
 import numpy as np
 import netCDF4
-from .utils import setup_logger
+from regional_mom6.utils import setup_logger
 
 
 regridding_logger = setup_logger(__name__, set_handler=False)

regional_mom6/rotation.py

@@ -1,12 +1,11 @@
-from .utils import setup_logger
-import logging
+from regional_mom6 import utils
+from regional_mom6.regridding import get_hgrid_arakawa_c_points, coords
 
-rotation_logger = setup_logger(__name__, set_handler=False)
+rotation_logger = utils.setup_logger(__name__, set_handler=False)
 # An Enum is like a dropdown selection for a menu, it essentially limits the type of input parameters. It comes with additional complexity, which of course is always a challenge.
 from enum import Enum
 import xarray as xr
 import numpy as np
-from .regridding import get_hgrid_arakawa_c_points
 
 
 class RotationMethod(Enum):
@@ -249,3 +248,79 @@ def create_expanded_hgrid(hgrid: xr.Dataset, expansion_width=1) -> xr.Dataset:
         }
     )
     return pseudo_hgrid
+
+
+def get_rotation_angle(
+    rotational_method: RotationMethod, hgrid: xr.Dataset, orientation=None
+):
+    """
+    This function returns the rotation angle - THIS IS ALWAYS BASED ON THE ASSUMPTION OF DEGREES - based on the rotational method and provided hgrid, if orientation & coords are provided, it will assume the boundary is requested
+    Parameters
+    ----------
+    rotational_method: RotationMethod
+        The rotational method to use
+    hgrid: xr.Dataset
+        The hgrid dataset
+    orientation: xr.Dataset
+        The orientation, which also lets us now that we are on a boundary
+    Returns
+    -------
+    xr.DataArray
+        angle in degrees
+    """
+    rotation_logger.info("Getting rotation angle")
+    boundary = False
+    if orientation != None:
+        rotation_logger.debug(
+            "The rotational angle is requested for the boundary: {}".format(orientation)
+        )
+        boundary = True
+
+    if rotational_method == RotationMethod.NO_ROTATION:
+        rotation_logger.debug("Using NO_ROTATION method")
+        if not utils.is_rectilinear_hgrid(hgrid):
+            raise ValueError("NO_ROTATION method only works with rectilinear grids")
+        angles = xr.zeros_like(hgrid.x)
+
+        if boundary:
+            # Subset to just boundary
+            # Add zeroes to hgrid
+            hgrid["zero_angle"] = angles
+
+            # Cut to boundary
+            zero_angle = coords(
+                hgrid,
+                orientation,
+                "doesnt_matter",
+                angle_variable_name="zero_angle",
+            )["angle"]
+
+            return zero_angle
+        else:
+            return angles
+    elif rotational_method == RotationMethod.GIVEN_ANGLE:
+        rotation_logger.debug("Using GIVEN_ANGLE method")
+        if boundary:
+            return coords(
+                hgrid, orientation, "doesnt_matter", angle_variable_name="angle_dx"
+            )["angle"]
+        else:
+            return hgrid["angle_dx"]
+    elif rotational_method == RotationMethod.EXPAND_GRID:
+        rotation_logger.debug("Using EXPAND_GRID method")
+        hgrid["angle_dx_rm6"] = initialize_grid_rotation_angles_using_expanded_hgrid(
+            hgrid
+        )
+
+        if boundary:
+            degree_angle = coords(
+                hgrid,
+                orientation,
+                "doesnt_matter",
+                angle_variable_name="angle_dx_rm6",
+            )["angle"]
+            return degree_angle
+        else:
+            return hgrid["angle_dx_rm6"]
+    else:
+        raise ValueError("Invalid rotational method")

regional_mom6/utils.py

@@ -2,6 +2,7 @@
 import logging
 import sys
 import xarray as xr
+from regional_mom6 import regridding as rgd
 
 
 def vecdot(v1, v2):
@@ -372,11 +373,12 @@ def is_rectilinear_hgrid(hgrid: xr.Dataset, rtol: float = 1e-3) -> bool:
         hgrid (xarray.Dataset): The horizontal grid dataset.
         rtol (float): Relative tolerance. Default is 1e-3.
     """
+    ds_t = rgd.get_hgrid_arakawa_c_points(hgrid)
     if (
-        np.allclose(hgrid.tlon[:, 0], hgrid.tlon[0, 0], rtol=rtol)
-        and np.allclose(hgrid.tlon[:, -1], hgrid.tlon[0, -1], rtol=rtol)
-        and np.allclose(hgrid.tlat[0, :], hgrid.tlat[0, 0], rtol=rtol)
-        and np.allclose(hgrid.tlat[-1, :], hgrid.tlat[-1, 0], rtol=rtol)
+        np.allclose(ds_t.tlon[:, 0], ds_t.tlon[0, 0], rtol=rtol)
+        and np.allclose(ds_t.tlon[:, -1], ds_t.tlon[0, -1], rtol=rtol)
+        and np.allclose(ds_t.tlat[0, :], ds_t.tlat[0, 0], rtol=rtol)
+        and np.allclose(ds_t.tlat[-1, :], ds_t.tlat[-1, 0], rtol=rtol)
     ):
         return True
     return False