add inteface to horiz_interp_read_weights_conserve

compile_c_libs.sh

@@ -4,10 +4,10 @@ set -ex
 #set -o posix
 
 #yaml includes
-YAML_FLAGS+=""
+YAML_FLAGS+="-I/opt/libyaml/0.2.5/GNU/14.2.0/include"
 
 #yaml libraries
-YAML_LDFLAGS+=""
+YAML_LDFLAGS+="-L/opt/libyaml/0.2.5/GNU/14.2.0/lib"
 
 #fortran netcdf includes
 NF_FLAGS+=$(nf-config --fflags)

pyfms/__init__.py

@@ -4,7 +4,7 @@
 from .py_field_manager.py_field_manager import FieldTable
 from .py_fms import fms
 from .py_horiz_interp import horiz_interp
-from .py_horiz_interp.interp import Interp
+from .py_horiz_interp.interp import ConserveInterp
 from .py_mpp import mpp, mpp_domains
 from .py_mpp.domain import Domain
 from .utils import constants, grid_utils

pyfms/py_horiz_interp/_functions.py

@@ -74,6 +74,7 @@ def define(lib):
             POINTER(c_bool),  # src_modulo
             POINTER(c_bool),  # is_latlon_in
             POINTER(c_bool),  # is_latlon_out
+            POINTER(c_bool), # save_weights_as_fregrid
             POINTER(c_bool),  # convert_cf_order
         ]
 
@@ -97,6 +98,21 @@ def define(lib):
             POINTER(c_bool),  # convert_cf_order
         ]
 
+    lib.cFMS_horiz_interp_read_weights_conserve.restype = c_int
+    lib.cFMS_horiz_interp_read_weights_conserve.argtypes = [        
+        POINTER(c_char),  # weight_filename
+        POINTER(c_char),  # weight_file_src
+        POINTER(c_int),  # nlon_src
+        POINTER(c_int),  # nlat_src
+        POINTER(c_int),  # nlon_dst
+        POINTER(c_int),  # nlat_dst
+        POINTER(c_int),  # isw
+        POINTER(c_int),  # iew
+        POINTER(c_int),  # jsw
+        POINTER(c_int),  # jew
+        POINTER(c_int)  # src_tile
+    ]
+
     # getter routines for individual fields
     lib.cFMS_get_wti_cfloat.restype = None
     lib.cFMS_get_wti_cfloat.argtypes = [
@@ -146,6 +162,13 @@ def define(lib):
         ndpointer(np.float64, ndim=1, flags=C),
     ]
 
+    # cFMS_get_area_frac_dst
+    lib.cFMS_get_xgrid_area.restype = None
+    lib.cFMS_get_xgrid_area.argtypes = [
+        POINTER(c_int),
+        ndpointer(np.float64, ndim=1, flags=C),
+    ]
+
     # cFMS_get_nlon_src
     lib.cFMS_get_nlon_src.restype = None
     lib.cFMS_get_nlon_src.argtypes = [POINTER(c_int), POINTER(c_int)]

pyfms/py_horiz_interp/horiz_interp.py

@@ -12,6 +12,7 @@
     set_c_int,
     set_c_str,
 )
+from pyfms.py_mpp.domain import Domain
 
 
 # enumerations used by horiz_interp_types.F90 (FMS)
@@ -29,6 +30,7 @@
 _cFMS_horiz_interp_2d_new_cfloat = None
 _cFMS_horiz_interp_2d_base_cdouble = None
 _cFMS_horiz_interp_2d_base_cfloat = None
+_cFMS_horiz_interp_read_weights_conserve = None
 _cFMS_get_i_src = None
 _cFMS_get_j_src = None
 _cFMS_get_i_dst = None
@@ -38,6 +40,7 @@
 _cFMS_get_nlon_dst = None
 _cFMS_get_nlat_dst = None
 _cFMS_get_interp_method = None
+_cFMS_get_xgrid_area = None
 _cFMS_get_area_frac_dst_double = None
 _cFMS_get_nxgrid = None
 _cFMS_horiz_interp_new_dict = {}
@@ -138,6 +141,7 @@ def get_weights(
     src_modulo: bool = False,
     is_latlon_in: bool = False,
     is_latlon_out: bool = False,
+    save_weights_as_fregrid: bool = False,
     convert_cf_order: bool = True,
 ) -> int:
     """
@@ -178,6 +182,7 @@ def get_weights(
     set_c_bool(src_modulo, arglist)
     set_c_bool(is_latlon_in, arglist)
     set_c_bool(is_latlon_out, arglist)
+    set_c_bool(save_weights_as_fregrid, arglist)
     set_c_bool(convert_cf_order, arglist)
 
     return _cFMS_horiz_interp_new(*arglist)
@@ -294,6 +299,17 @@ def get_area_frac_dst(interp_id: int):
     return area_frac_dst
 
 
+def get_xgrid_area(interp_id: int):
+
+    nxgrid = get_nxgrid(interp_id)
+    arglist = []
+    set_c_int(interp_id, arglist)
+    xgrid_area = set_array(np.zeros(nxgrid, dtype=np.float64), arglist)
+
+    _cFMS_get_xgrid_area(*arglist)
+    return xgrid_area
+
+
 def get_interp_method(interp_id: int):
 
     interp_method_dict = {
@@ -358,6 +374,43 @@ def interp(
     return data_out
 
 
+def read_weights_conserve(weight_filename: str,
+                          weight_file_src: str,
+                          nlon_src: int,
+                          nlat_src: int,
+                          domain: Domain = None,
+                          nlon_tgt: int = None,
+                          nlat_tgt: int = None,
+                          src_tile: int = None):
+
+    if domain is None:
+        if nlon_tgt is None: cFMS_error(FATAL, "must provide nlon_tgt if Domain is not specified")
+        if nlat_tgt is None: cFMS_error(FATAL, "must provide nlon_tgt if Domain is not specified")
+        isc, iec, jsc, jec = 0, nlon_tgt-1, 0, nlat_tgt-1
+    else:
+        nlon_tgt = domain.xsize_c
+        nlat_tgt = domain.ysize_c
+        isc = domain.isc
+        iec = domain.iec
+        jsc = domain.jsc
+        jec = domain.jec
+
+    arglist = []
+    set_c_str(weight_filename, arglist)
+    set_c_str(weight_file_src, arglist)
+    set_c_int(nlon_src, arglist)
+    set_c_int(nlat_src, arglist)
+    set_c_int(nlon_tgt, arglist)
+    set_c_int(nlat_tgt, arglist)
+    set_c_int(isc, arglist)
+    set_c_int(iec, arglist)
+    set_c_int(jsc, arglist)
+    set_c_int(jec, arglist)
+    set_c_int(src_tile, arglist)
+
+    return _cFMS_horiz_interp_read_weights_conserve(*arglist)
+
+
 def _init_functions():
 
     global _cFMS_create_xgrid_2dx2d_order1
@@ -370,6 +423,7 @@ def _init_functions():
     global _cFMS_horiz_interp_base_dict
     global _cFMS_horiz_interp_base_2d_cdouble
     global _cFMS_horiz_interp_base_2d_cfloat
+    global _cFMS_horiz_interp_read_weights_conserve
     global _cFMS_get_wti_cfloat
     global _cFMS_get_wti_cdouble
     global _cFMS_get_wtj_cfloat
@@ -383,6 +437,7 @@ def _init_functions():
     global _cFMS_get_nlon_dst
     global _cFMS_get_nlat_dst
     global _cFMS_get_interp_method
+    global _cFMS_get_xgrid_area
     global _cFMS_get_area_frac_dst_double
     global _cFMS_get_nxgrid
 
@@ -397,6 +452,8 @@ def _init_functions():
     _cFMS_horiz_interp_base_2d_cdouble = _lib.cFMS_horiz_interp_base_2d_cdouble
     _cFMS_horiz_interp_base_2d_cfloat = _lib.cFMS_horiz_interp_base_2d_cfloat
 
+    _cFMS_horiz_interp_read_weights_conserve = _lib.cFMS_horiz_interp_read_weights_conserve
+
     _cFMS_get_wti_cfloat = _lib.cFMS_get_wti_cfloat
     _cFMS_get_wti_cdouble = _lib.cFMS_get_wti_cdouble
     _cFMS_get_wtj_cfloat = _lib.cFMS_get_wtj_cfloat
@@ -412,6 +469,7 @@ def _init_functions():
     _cFMS_get_nlat_dst = _lib.cFMS_get_nlat_dst
     _cFMS_get_nxgrid = _lib.cFMS_get_nxgrid
     _cFMS_get_interp_method = _lib.cFMS_get_interp_method
+    _cFMS_get_xgrid_area = _lib.cFMS_get_xgrid_area
     _cFMS_get_area_frac_dst_double = _lib.cFMS_get_area_frac_dst_cdouble
 
     _cFMS_horiz_interp_new_dict = {

pyfms/py_horiz_interp/interp.py

@@ -1,8 +1,8 @@
 from pyfms.py_horiz_interp import horiz_interp
 
 
-class Interp:
-    def __init__(self, interp_id: int):
+class ConserveInterp:
+    def __init__(self, interp_id: int, weights_as_fregrid: bool = False):
         # set only for conservative
         self.interp_id = interp_id
         self.nxgrid = horiz_interp.get_nxgrid(interp_id)
@@ -16,3 +16,7 @@ def __init__(self, interp_id: int):
         self.nlat_dst = horiz_interp.get_nlat_dst(interp_id)
         self.interp_method = horiz_interp.get_interp_method(interp_id)
         self.get_area_frac_dst = horiz_interp.get_area_frac_dst(interp_id)
+
+        self.xgrid_area = None
+        if weights_as_fregrid:
+            self.xgrid_area = horiz_interp.get_xgrid_area(interp_id)

pyfms/py_mpp/_mpp_functions.py

@@ -82,3 +82,6 @@ def define(lib):
         NDPOINTER(dtype=np.int32, ndim=(1), flags=C),  # pelist
         POINTER(c_bool),  # no_sync
     ]
+
+    lib.cFMS_sync.restype = None
+    lib.cFMS_sync.argtypes = None

pyfms/py_mpp/mpp.py

@@ -30,6 +30,7 @@
 _cFMS_pe = None
 _cFMS_root_pe = None
 _cFMS_set_current_pelist = None
+_cFMS_sync = None
 
 
 def gather(
@@ -205,6 +206,11 @@ def set_current_pelist(pelist: list[int] = None, no_sync: bool = None):
     _cFMS_set_current_pelist(*arglist)
 
 
+def sync():
+
+    _cFMS_sync()
+
+
 def _init_functions():
 
     global _cFMS_declare_pelist
@@ -218,6 +224,7 @@ def _init_functions():
     global _cFMS_pe
     global _cFMS_root_pe
     global _cFMS_set_current_pelist
+    global _cFMS_sync
 
     _mpp_functions.define(_lib)
 
@@ -231,6 +238,7 @@ def _init_functions():
     _cFMS_pe = _lib.cFMS_pe
     _cFMS_root_pe = _lib.cFMS_root_pe
     _cFMS_set_current_pelist = _lib.cFMS_set_current_pelist
+    _cFMS_sync = _lib.cFMS_sync
 
     _cFMS_gather_pelist_2ds = {
         "int32": _cFMS_gather_pelist_2d_cint,

pyfms/py_mpp/mpp_domains.py

@@ -3,6 +3,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
+import pyfms.py_mpp.mpp as mpp
 from pyfms.py_mpp import _mpp_domains_functions
 from pyfms.py_mpp.domain import Domain
 from pyfms.utils.ctypes_utils import (
@@ -117,7 +118,7 @@ def get_compute_domain(
     _cFMS_get_compute_domain(*arglist)
 
     return dict(
-        domain_id_c=domain_id,
+        domain_id=domain_id,
         isc=xbegin.value,
         jsc=ybegin.value,
         iec=xend.value,
@@ -182,7 +183,7 @@ def get_data_domain(
 
 def define_domains(
     global_indices: list[int],
-    layout: list[int],
+    layout: list[int] = None,
     pelist: list[int] = None,
     xflags: int = None,
     yflags: int = None,
@@ -213,6 +214,12 @@ def define_domains(
     corresponds to the saved FmsMppDomain2D derived type in cFMS
     """
 
+    if layout is None:
+        if pelist is None:
+            layout = define_layout(global_indices, mpp.npes())
+        else:
+            layout = define_layout(global_indices, len(pelist))
+
     arglist = []
     set_list(global_indices, np.int32, arglist)
     set_list(layout, np.int32, arglist)

pyfms/utils/grid_utils.py

@@ -23,14 +23,16 @@ def get_grid_area(
     on lon and lat
     """
 
-    nlon, nlat = lon.shape
+    nlat, nlon = lon.shape
+    nlat -= 1
+    nlon -= 1
 
     arglist = []
     set_c_int(nlon, arglist)
     set_c_int(nlat, arglist)
     set_array(lon, arglist)
     set_array(lat, arglist)
-    area = set_array(np.zeros((nlon, nlat), dtype=np.float64), arglist)
+    area = set_array(np.zeros((nlat, nlon), dtype=np.float64), arglist)
 
     _cFMS_get_grid_area(*arglist)
 

tests/py_horiz_interp/test_horiz_interp.py

@@ -75,15 +75,7 @@ def test_horiz_interp_conservative():
         pes = pyfms.mpp.npes()
 
         global_indices = [0, ni_src - 1, 0, nj_src - 1]
-        layout = pyfms.mpp_domains.define_layout(global_indices, ndivs=pes)
-        domain = pyfms.mpp_domains.define_domains(
-            global_indices=global_indices,
-            layout=layout,
-            pelist=pyfms.mpp.get_current_pelist(npes=pes),
-            name="horiz_interp_conservative_test",
-            xflags=pyfms.mpp_domains.CYCLIC_GLOBAL_DOMAIN,
-            yflags=pyfms.mpp_domains.CYCLIC_GLOBAL_DOMAIN,
-        )
+        domain = pyfms.mpp_domains.define_domains(global_indices=global_indices)
         # get compute domain indices
         isc = domain.isc
         iec = domain.iec + 1  # grid has one more point
@@ -117,12 +109,14 @@ def test_horiz_interp_conservative():
             nlat_in=nlat_in,
             nlon_out=nlon_out,
             nlat_out=nlat_out,
+            save_weights_as_fregrid=True,
             convert_cf_order=convert_cf_order,
         )
 
         # check weights
         nxgrid = (jec - jsc) * (iec - isc)
-        interp = pyfms.Interp(interp_id)
+        interp = pyfms.ConserveInterp(interp_id, weights_as_fregrid=True)
+        assert interp.xgrid_area is not None
 
         j_answers = np.array([j for j in range(jsc, jec) for ilon in range(iec - isc)])
         if convert_cf_order: