some final tweaks (make sin/cos a little neater)

NOAA-PSL · May 2, 2024 · e75b3da · e75b3da
1 parent 9a0fe2a
commit e75b3da
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 19 deletions.
diff --git a/examples/replay/config-1.00-degree.yaml b/examples/replay/config-1.00-degree.yaml
@@ -235,23 +235,23 @@ MOM6Dataset:
     time:
         start: "1993-12-31T18"
         end: "1999-06-13T03"
-        freq: "3h"
+        freq: "6h"
 
     chunks_in:
         time: 1
-        z_l: 1
-        z_i: 1
-        zl: 1
+        z_l: -1
+        z_i: -1
+        zl: -1
         yh: -1
         xh: -1
         yq: -1
         xq: -1
 
     chunks_out:
         time: 1
-        z_l: 1
-        z_i: 1
-        zl: 1
+        z_l: -1
+        z_i: -1
+        zl: -1
         lat: -1
         lon: -1
 

diff --git a/ufs2arco/regrid/mom6regridder.py b/ufs2arco/regrid/mom6regridder.py
@@ -70,6 +70,8 @@ def create_grid_in(
         grid_in["lat_u"] = mom6_grid["geolat_u"]
         grid_in["lon_v"] = mom6_grid["geolon_v"]
         grid_in["lat_v"] = mom6_grid["geolat_v"]
+        grid_in['cos_rot'] = mom6_grid["cos_rot"]
+        grid_in['sin_rot'] = mom6_grid["sin_rot"]
         ny, nx = grid_in["lon"].shape
         lon_b = np.empty((ny + 1, nx + 1))
         lat_b = np.empty((ny + 1, nx + 1))
@@ -96,8 +98,9 @@ def create_grid_in(
         ds_in_v = grid_in[["lon_v", "lat_v", "lat_b", "lon_b"]].rename(
             {"lat_v": "lat", "lon_v": "lon"}
         )
+        ds_rot =  grid_in[['cos_rot','sin_rot']]
 
-        return ds_in_t, ds_in_u, ds_in_v
+        return ds_in_t, ds_in_u, ds_in_v, ds_rot
 
     def create_grid_out(
         self,
@@ -138,20 +141,14 @@ def _create_regridder(self, ds_in: xr.Dataset) -> None:
         self.ds_rot = None
         if self.rotation_file is not None:
             mom6_grid = xr.open_dataset(self.rotation_file)
-            ds_rot = mom6_grid[["cos_rot", "sin_rot"]]
-            self.ds_rot = ds_rot.rename({"xh": "lon", "yh": "lat"})
-        elif "cos_rot" in ds_in and "sin_rot" in ds_in:
-            self.ds_rot = xr.Dataset()
-            self.ds_rot["cos_rot"] = ds_in["cos_rot"]
-            self.ds_rot["sin_rot"] = ds_in["sin_rot"]
         else:
             warnings.warn(
-                f"MOM6Regridder._create_regridder: Could not find 'rotation_file' in configuration yaml. "
-                f"Vector fields will be silently ignored."
+                "MOM6Regridder._create_regridder: Could not find 'rotation_file' in configuration yaml. "
+                "If using conservative regridding, it will fail. Otherwise, just vector fields will be ignored."
             )
 
         # create renamed datasets
-        ds_in_t, ds_in_u, ds_in_v = self.create_grid_in(mom6_grid=mom6_grid)
+        ds_in_t, ds_in_u, ds_in_v, self.ds_rot = self.create_grid_in(mom6_grid=mom6_grid)
 
         # create output dataset
         grid_out = self.create_grid_out(

diff --git a/ufs2arco/regrid/ufsregridder.py b/ufs2arco/regrid/ufsregridder.py
@@ -183,10 +183,10 @@ def regrid_tripolar(
 
                 # rotate to earth-relative
                 urot = (
-                    interp_u * ds_rot.cos_rot.values + interp_v * ds_rot.sin_rot.values
+                    interp_u * ds_rot.cos_rot + interp_v * ds_rot.sin_rot
                 )
                 vrot = (
-                    interp_v * ds_rot.cos_rot.values - interp_u * ds_rot.sin_rot.values
+                    interp_v * ds_rot.cos_rot - interp_u * ds_rot.sin_rot
                 )
 
                 # interoplate