Add MiniGridObj gridc and gridt (#106)

Co-authored-by: mlee03 <Mikyung.Lee@lscamd50-d.gfdl.noaa.gov>

Author: mlee03

.github/workflows/test.yaml

@@ -29,7 +29,7 @@ jobs:
           pytest tests/shared/test_libs.py
           pytest tests/mosaic/test_mosaic.py
           pytest tests/shared/test_gridobj.py::test_read_write
-          pytest tests/shared/test_gridobj.py::test_center_option
+          pytest tests/shared/test_gridobj.py::test_minigrids_option
           pytest tests/shared/test_gridobj.py::test_to_domain
           #pytest tests/shared/test_xgridobj.py
           #pytest tests/hgrid/test_hgrid.py

fmsgridtools/shared/gridobj.py

@@ -86,6 +86,24 @@ def __init__(self, name: str = None, data=None):
         self.name = name
         self.data = data
 
+class MiniGridObj:
+
+    def __init__(self, xsize: int = None, ysize: int = None, x: npt.NDArray = None, y: npt.NDArray = None):
+        self.xsize = xsize
+        self.ysize = ysize
+        self.x = x
+        self.y = y
+
+    def set_size(self):
+
+        xsize, ysize = self.x.shape
+
+        if self.x.shape != self.y.shape:
+            logger.error("MiniGrid, x and y differ in dimensions")
+
+        self.xsize = xsize
+        self.ysize = ysize
+
 
 class GridObj:
     """
@@ -136,6 +154,9 @@ def __init__(
 
         self._set_dims()
 
+        self.gridc = MiniGridObj()
+        self.gridt = MiniGridObj()
+
         logger.info("Created new GridObj named:\n %s", self.__repr__())
 
     def to_domain(self, domain: dict = None):
@@ -212,12 +233,54 @@ def get_fms_area(self):
         self.area = pyfms.grid_utils.get_grid_area(lon=x, lat=y, convert_cf_order=False)
         return self.area
 
+    def get_gridc(self):
+
+        """
+        Save the edge points
+        """
+
+        if self.x is None or self.y is None:
+            logger.error("Cannot set gridc.  Please set x and y values first")
+
+        self.gridc.x = np.ascontiguousarray(self.x[::2, ::2])
+        self.gridc.y = np.ascontiguousarray(self.y[::2, ::2])
+        self.gridc.set_size()
+
+        return self.gridc
+
+    def get_gridt(self):
+
+        """
+        Save the center points
+        """
+
+        if self.x is None or self.y is None:
+            logger.error("Cannot set gridt.  Please set and y values first")
+
+        self.gridt.x = np.ascontiguousarray(self.x[1::2, 1::2])
+        self.gridt.y = np.ascontiguousarray(self.x[1::2, 1::2])
+        self.gridt.set_size()
+
+        return self.gridt
+
+    def free_supergrid(self):
+
+        self.x = None
+        self.y = None
+
+    def free_gridc(self):
+
+        self.gridc = None
+
+    def free_gridt(self):
+
+        self.gridt = None
+
     def read(
         self,
         gridfile: str = None,
         domain: dict = None,
         radians: bool = False,
-        center: bool = False,
         on_domain: bool = False,
         xy_only: bool = False,
     ):
@@ -252,9 +315,6 @@ def read(
             for obj in objlist:
                 if obj.name in ds:
                     obj.data = ds[obj.name].data
-                    if center:
-                        logger.info("saving center points for %s", {obj.name})
-                        obj.data = np.ascontiguousarray(obj.data[::2, ::2])
                 else:
                     logger.warning("could not %s in %s", {obj.name}, {self.gridfile})
 

fmsgridtools/shared/mosaicobj.py

@@ -159,7 +159,6 @@ def get_grid(
         self,
         input_dir: str | Path = "./",
         radians: bool = False,
-        center: bool = False,
         domain: pyfms.Domain = None,
     ) -> dict:
         """
@@ -183,7 +182,6 @@ def get_grid(
             readfile = Path(input_dir) / gridfile
             grid[gridtile] = GridObj(gridfile=readfile).read(
                 radians=radians,
-                center=center,
                 domain=domain,
                 on_domain=False if domain is None else True,
                 xy_only=True,

pyFMS

@@ -155,7 +155,7 @@ def test_read_write(set_fms_files):
     pyfms.fms.end()
 
 
-def test_center_option(set_fms_files):
+def test_minigrids_option(set_fms_files):
 
     pyfms.fms.init()
 
@@ -171,22 +171,40 @@ def test_center_option(set_fms_files):
 
     GridObj(gridfile=gridfile, x=x, y=y).write()
 
-    grid = GridObj(gridfile=gridfile).read(center=True, radians=True, xy_only=True)
+    #test gridc
+    grid = GridObj(gridfile=gridfile).read(radians=True, xy_only=True)
+    gridc = grid.get_gridc()
 
-    assert grid.nx == nx2
-    assert grid.ny == ny2
-    assert grid.nxp == nx2 + 1
-    assert grid.nyp == ny2 + 1
+    assert gridc.xsize == nx2p
+    assert gridc.ysize == ny2p
 
     answer = np.radians(np.ones((ny2p, nx2p), dtype=np.float64))
+    np.testing.assert_array_equal(gridc.x, answer)
+    np.testing.assert_array_equal(gridc.y, answer)
 
-    np.testing.assert_array_equal(grid.x, answer)
-    np.testing.assert_array_equal(grid.y, answer)
+    grid.free_gridc()
+    assert grid.gridc is None
+
+    #test gridt
+    gridt = grid.get_gridt()
+
+    assert gridt.xsize == nx2
+    assert gridt.ysize == ny2
+
+    answer = np.zeros((ny2, nx2), dtype=np.float64)
+    np.testing.assert_array_equal(gridt.x, answer)
+    np.testing.assert_array_equal(gridt.y, answer)
+
+    grid.free_gridt()
+    assert grid.gridt is None
+
+    grid.free_supergrid()
+    assert grid.x is grid.y is None
 
     gridfile.unlink()
     pyfms.fms.end()
 
-
+@pytest.mark.skip("error in pyFMS")
 def test_to_domain(set_fms_files):
 
     nx, ny = 8, 8

tests/shared/test_gridobj.py

@@ -19,7 +19,7 @@ def generate_mosaic(nx: int = 90, ny: int = 45, refine: int = 2):
     x_tgt = np.linspace(xstart, xend, nx*refine+1)
     y_tgt = np.linspace(ystart, yend, ny*refine+1)
     x_tgt, y_tgt = np.meshgrid(x_tgt, y_tgt)
-    
+
     area_src = np.ones((ny, nx), dtype=np.float64)
     area_tgt = np.ones((ny*refine, nx*refine), dtype=np.float64)
 
@@ -64,9 +64,9 @@ def test_create_xgrid(on_gpu):
     xgrid.create_xgrid()
     xgrid.to_dataset()
     xgrid.dataset["tile1"]["tile1"].to_netcdf("remap.nc")
-    
+
     del xgrid
-    
+
     xgrid = fmsgridtools.XGridObj(restart_remap_file="remap.nc")
 
     #check nxcells
@@ -79,7 +79,7 @@ def test_create_xgrid(on_gpu):
 
     src_i = [xgrid.src_cell[i][0] for i in range(nxcells)]
     src_j = [xgrid.src_cell[i][1] for i in range(nxcells)]
-    
+
     assert src_i == answer_i
     assert src_j == answer_j
 
@@ -94,7 +94,7 @@ def test_create_xgrid(on_gpu):
         for i in range(nx):
             for ixcell in range(refine):
                 answer_j += [j*refine + ixcell + 1]*refine
-                
+
     tgt_i = [xgrid.tgt_cell[i][0] for i in range(nxcells)]
     tgt_j = [xgrid.tgt_cell[i][1] for i in range(nxcells)]