Xgridobj rewrite

FREnctools_lib/pyfrenctools/shared/create_xgrid.py

@@ -108,11 +108,7 @@ def transfer_data_gpu(nxcells: int, src_nlon: int, tgt_nlon: int):
                 xarea=xarea[:nxcells])
 
 
-def get_2dx2d_order1_gpu(src_nlon: int,
-                         src_nlat: int,
-                         tgt_nlon: int,
-                         tgt_nlat: int,
-                         src_lon: npt.NDArray,
+def get_2dx2d_order1_gpu(src_lon: npt.NDArray,
                          src_lat: npt.NDArray,
                          tgt_lon: npt.NDArray,
                          tgt_lat: npt.NDArray,
@@ -121,11 +117,18 @@ def get_2dx2d_order1_gpu(src_nlon: int,
 
     create_xgrid_order1_gpu_wrapper = _lib.create_xgrid_order1_gpu_wrapper
 
+    nyp, nxp = src_lon.shape
+    src_nlon = nxp - 1
+    src_nlat = nyp - 1
+
+    nyp, nxp = tgt_lon.shape
+    tgt_nlon = nxp - 1
+    tgt_nlat = nyp - 1
+
     if src_mask is None: src_mask = np.ones((src_nlon*src_nlat), dtype=np.float64)
     if tgt_mask is None: tgt_mask = np.ones((tgt_nlon*tgt_nlat), dtype=np.float64)
 
     arrayptr_double = np.ctypeslib.ndpointer(dtype=np.float64, flags="C_CONTIGUOUS")
-
     create_xgrid_order1_gpu_wrapper.restype = np.int32
     create_xgrid_order1_gpu_wrapper.argtypes = [c_int, #src_nlon
                                                 c_int, #src_nlat

fmsgridtools/shared/gridobj.py

@@ -88,23 +88,16 @@ def __init__(self, name: str = None, data=None):
 
 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
+    def __init__(self, nx: int = None, ny: int = None, nxp: int = None, nyp: int = None,
+                    x: npt.NDArray = None, y: npt.NDArray = None):
+
+        self.nx = nx #number of cells
+        self.ny = ny #number of cells
+        self.nxp = nxp #number of gridpoints
+        self.nyp = nyp #number of gridpoints
         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:
     """
     Class for grid information
@@ -171,7 +164,7 @@ def to_domain(self, domain: dict = None):
         else:
             if self.domain is not None:
                 logger.warning("Overwriting %s with %s", self.domain, domain)
-                self.domain = domain
+            self.domain = domain
 
         if not pyfms.fms.module_is_initialized():
             logger.error("Please initialize pyfms first")
@@ -218,7 +211,7 @@ def to_radians(self):
                 logger.info("Converting %s to radians", {obj.name})
                 obj.data = np.radians(obj.data, dtype=np.float64)
 
-    def get_fms_area(self):
+    def get_fms_area(self, gridc: bool = False):
         """
         Compute grid cell areas
         """
@@ -228,8 +221,13 @@ def get_fms_area(self):
         if not pyfms.fms.module_is_initialized():
             logger.error("Please initialize pyfms first")
 
-        x = np.ascontiguousarray(self.x, dtype=np.float64)
-        y = np.ascontiguousarray(self.y, dtype=np.float64)
+        if gridc:
+            x = self.gridc.x
+            y = self.gridc.y
+        else:
+            x = np.ascontiguousarray(self.x, dtype=np.float64)
+            y = np.ascontiguousarray(self.y, dtype=np.float64)
+
         self.area = pyfms.grid_utils.get_grid_area(lon=x, lat=y, convert_cf_order=False)
         return self.area
 
@@ -244,7 +242,10 @@ def get_gridc(self):
 
         self.gridc.x = np.ascontiguousarray(self.x[::2, ::2])
         self.gridc.y = np.ascontiguousarray(self.y[::2, ::2])
-        self.gridc.set_size()
+
+        self.gridc.nyp, self.gridc.nxp = self.gridc.x.shape
+        self.gridc.ny = self.gridc.nyp - 1
+        self.gridc.nx = self.gridc.nxp - 1
 
         return self.gridc
 
@@ -259,7 +260,7 @@ def get_gridt(self):
 
         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()
+        self.gridt.nyp, self.gridt.nxp = self.gridt.x.shape
 
         return self.gridt