New VariableObj for remap

fmsgridtools/__init__.py

@@ -5,6 +5,7 @@
 from .make_mosaic import coupler_mosaic
 from .make_topog import make_topog
 from .remap import remap
+from .remap.variableobj import VariableObj, SrcFileObj, TgtFileObj, DimObj, FileDimsObj
 from .shared.gridtools_utils import check_file_is_there, get_provenance_attrs
 from .shared.gridobj import GridObj
 from .shared.mosaicobj import MosaicObj

fmsgridtools/remap/conservative.py

@@ -1,31 +1,65 @@
 import numpy as np
 
+from fmsgridtools.remap.variableobj import SrcFileObj, TgtFileObj, VariableObj
 from fmsgridtools.shared.xgridobj import XGridObj
 
-def remap(src_mosaic: str,
-          input_dir: str = "./",
-          output_dir: str = None,
-          input_file: str = None,
-          output_file: str = None,
-          tgt_mosaic: str = None,
-          tgt_nlon: int = None,
-          tgt_nlat: int =  None,
-          lon_bounds: list = None,
-          lat_bounds: list = None,
-          kbounds: list = None,
-          tbounds: list = None,
-          order: int = 1,
-          static_file: str = None,
-          check_conserve: bool = False):
+import pyfms
 
-    #create an xgrid object
-    xgrid = XGridObj(input_dir, src_mosaic_file=src_mosaic, tgt_mosaic_file=tgt_mosaic)
 
-    #create xgrid
-    xgrid.create_xgrid()
+def remap(input_dir: str = "./",
+      src_mosaicfile: str = None,
+      tgt_mosaicfile: str = None,
+      input_file: str = None,
+      output_dir: str = "./",
+      output_file: str = None,
+      scalar_variables: list[str] = None,
+      lon_bounds: list = None,
+      lat_bounds: list = None,
+      kbounds: list = None,
+      tbounds: list = None,
+      order: int = 1,
+      check_conserve: bool = False,
+      gpu: bool = False):
 
-    #write
-    xgrid.write()
+      xgrid = XGridObj(input_dir=input_dir, src_mosaicfile=src_mosaicfile, tgt_mosaicfile=tgt_mosaicfile)
+      tgt_tiles = list(xgrid.tgt.grid.keys())
+      src_tiles = list(xgrid.src.grid.keys())
 
+      pyfms.fms.init(ndomain=len(tgt_tiles))
 
+      for tgt_tile in tgt_tiles:
+
+            pyfms.horiz_interp.init(len(src_tiles))
+
+            nx_tgt, ny_tgt = xgrid.tgt.grid[tgt_tile].nx, xgrid.tgt.grid[tgt_tile].ny
+
+            #get xgrid, make sure mpi works cause i don't think it will
+            xgrid.domain = pyfms.mpp_domains.define_domains([0, nx_tgt-1, 0, ny_tgt-1])
+            xgrid.set_target_tile(tgt_tile)
+            xgrid.get_interp()
+
+            if input_file is None: return
+            if output_file is None: output_file = input_file
+            if output_dir is None: output_dir = input_dir
+
+            src_fileobj = SrcFileObj(datafile=input_file, input_dir=input_dir, tiles=src_tiles)
+            tgt_fileobj = TgtFileObj(datafile=output_file, output_dir=output_dir, nx=nx_tgt, ny=ny_tgt)
+
+            for var in src_fileobj.variables:
+
+                variable = VariableObj(var, src_fileobj)
+                variable.init()
+
+                for itime in variable.timeslist:
+                    for klevel in variable.zlist:
+
+                        input_data = variable.slice(tile=src_tiles[0], timepoint=itime, klevel=klevel, prepare_data=True)
+                        data_slice = pyfms.horiz_interp.interp(xgrid.interps[src_tile].interp_id, input_data, convert_cf_order=False)
+
+                        for src_tile in src_tiles[1:]:
+                            input_data = variable.slice(tile=src_tile, timepoint=itime, klevel=klevel, prepare_data=True)
+                            data_slice += pyfms.horiz_interp.interp(xgrid.interps[src_tile].interp_id, input_data, convert_cf_order=False)
+
+                            #gather if parallel, for now, no
+                        variable.set_tgt_data(data_slice, timepoint=itime, klevel=klevel)
 

fmsgridtools/remap/remap.py

@@ -17,30 +17,42 @@
               Only 1 or 2 order is supported for conservative interpolation
               """
 )
-@click.option("--static_file",
-              type = click.Path(exists=True),
+@click.option("--check_conserve",
+              type = bool,
               help =
               """
-              To remap data where cell_methods = CELL_METHODS_MEAN, the static_file
-              will src grid cell areas should be provided
+              If true, output grid area conservative will be checked
               """
 )
-@click.option("--check_conserve",
+@click.option("--gpu",
               type = bool,
               help =
               """
-              If true, output grid area conservative will be checked
+              If true, the exchange grid will be created on the GPU
               """
 )
-def conservative_method(input_dir, output_dir, input_file, output_file, #common_options
-                        src_mosaic, tgt_mosaic, tgt_nlon, tgt_nlat,     #common_options
-                        lon_bounds, lat_bounds, kbounds, tbounds,       #common_options
-                        debug, order, static_file, check_conserve):
-
+def conservative_method(input_dir, output_dir, input_mosaic_dir, output_mosaic_dir,
+                        input_file, output_file, #common_options
+                        src_mosaic, tgt_mosaic, tgt_nlon, tgt_nlat, #common_options
+                        scalar_variables, lon_bounds, lat_bounds,
+                        kbounds, tbounds, debug, order, check_conserve, gpu):
+
     setlogger.setconfig("remap.log", debug)
     logger.info("Starting conservative remapping")
-
-    conservative.remap(src_mosaic, input_dir, output_dir, input_file,
-                       output_file, tgt_mosaic, tgt_nlon, tgt_nlat,
-                       lon_bounds, lat_bounds, kbounds, tbounds,
-                       order, static_file, check_conserve)    
+
+    xgrid = conservative.remap(input_dir=input_dir,
+                               output_dir=output_dir,
+                               input_mosaic_dir=input_mosaic_dir,
+                               output_mosaic_dir=output_mosaic_dir,
+                               input_file=input_file,
+                               src_mosaic=src_mosaic,
+                               tgt_mosaic=tgt_mosaic,
+                               output_file=output_file,
+                               scalar_variables=scalar_variables,
+                               lon_bounds=lon_bounds,
+                               lat_bounds=lat_bounds,
+                               kbounds=kbounds,
+                               tbounds=tbounds,
+                               order=order,
+                               check_conserve=check_conserve,
+                               gpu=gpu)