Merge pull request #26 from FESOM/smm_interpolation

Smm interpolation

environment.yml

@@ -13,3 +13,8 @@ dependencies:
   - pip
   - cartopy
   - shapely
+  - cdo
+  - pip:
+    - smmregrid
+
+

src/fint/fint.py

@@ -7,6 +7,7 @@
 import pandas as pd
 import xarray as xr
 import subprocess
+from smmregrid import Regridder
 from scipy.interpolate import (
     CloughTocher2DInterpolator,
     LinearNDInterpolator,
@@ -282,23 +283,50 @@ def interpolate_cdo(target_grid,gridfile,original_file,output_file,variable_name
     """
 
     command = [
-            "cdo",
-            f"-{interpolation.split('_')[1]},{target_grid}",
-            f"-setgrid,{gridfile}",
-            f"{original_file}",
-            f"{output_file}"
-        ]
+        "cdo",
+        f"-{interpolation.split('_')[1]},{target_grid}",
+        f"-setgrid,{gridfile}",
+        f"{original_file}",
+        f"{output_file}"
+    ]
     if mask_zero:
         command.insert(1, "-setctomiss,0")
 
-
     # Execute the command
     subprocess.run(command)
 
     interpolated = xr.open_dataset(output_file)[variable_name].values
     os.remove(output_file)
     return interpolated
 
+def generate_cdo_weights(target_grid,gridfile,original_file,output_file,interpolation):
+    """
+    Generate CDO weights for interpolation using ssmregrid.
+
+    Args:
+        target_grid (str): Path to the target grid file.
+        gridfile (str): Path to the grid file associated with the original file.
+        original_file (str): Path to the original file containing the data to be remapped.
+        output_file (str): Path to the output file where the weights will be saved.
+
+    Returns:
+        xr.Dataset: Generated weights as an xarray Dataset.
+    """
+    command = [
+        "cdo",
+        f"-gen{interpolation.split('_')[-1]},{target_grid}",
+        f"-setgrid,{gridfile}",
+        f"{original_file}",
+        f"{output_file}"
+    ]
+
+    # Execute the command
+    subprocess.run(command)
+
+    weights = xr.open_dataset(output_file)
+    os.remove(output_file)
+    return weights
+
 def parse_depths(depths, depths_from_file):
     """
     Parses the selected depths from the available depth values and returns the corresponding depth indices and values.
@@ -584,7 +612,9 @@ def fint(args=None):
     )
     parser.add_argument(
         "--interp",
-        choices=["nn", "mtri_linear", "linear_scipy", "cdo_remapcon","cdo_remaplaf","cdo_remapnn", "cdo_remapdis"],  # "idist", "linear", "cubic"],
+        choices=["nn", "mtri_linear", "linear_scipy",
+                 "cdo_remapcon","cdo_remaplaf","cdo_remapnn", "cdo_remapdis",
+                 "smm_con","smm_laf","smm_nn","smm_dis"],  # "idist", "linear", "cubic"],
         default="nn",
         help="Interpolation method. Options are \
             nn - nearest neighbor (KDTree implementation, fast), \
@@ -767,7 +797,7 @@ def fint(args=None):
         trifinder = triang2.get_trifinder()
     elif interpolation == "nn":
         distances, inds = create_indexes_and_distances(x2, y2, lon, lat, k=1, workers=4)
-    elif interpolation in ["cdo_remapcon", "cdo_remaplaf", "cdo_remapnn", "cdo_remapdis"]:
+    elif interpolation in ["cdo_remapcon", "cdo_remaplaf", "cdo_remapnn", "cdo_remapdis", "smm_con", "smm_laf", "smm_nn", "smm_dis"]:
         gridtype = 'latlon'
         gridsize = x.size*y.size
         xsize = x.size
@@ -926,7 +956,6 @@ def fint(args=None):
                 input_data = xr.Dataset({variable_name: (["nod2"], data_in)})
                 if args.rotate:
                     input_data = xr.Dataset({variable_name: (["nod2"], data_in2)})
-
                 output_file_path = out_path.replace(".nc", "output_cdo_file.nc")
                 input_file_path = args.data.replace(".nc","cdo_interpolation.nc")
                 input_data.to_netcdf(input_file_path,encoding={
@@ -943,6 +972,29 @@ def fint(args=None):
                                                 )
                 os.remove(input_file_path)
 
+            elif interpolation in ["smm_con", "smm_laf", "smm_nn", "smm_dis"]:
+                input_data = xr.Dataset({variable_name: (["nod2"], data_in)})
+                if args.rotate:
+                    input_data = xr.Dataset({variable_name: (["nod2"], data_in2)})
+                input_file_path = args.data.replace(".nc","cdo_interpolation.nc")
+                input_data.to_netcdf(input_file_path,encoding={
+                            variable_name: {"dtype": np.dtype("double")},
+                        },
+                    )
+                output_file_path = out_path.replace(".nc", "weighs_cdo.nc")
+                weights = generate_cdo_weights(target_grid_path,
+                                            gridfile,
+                                            input_file_path,
+                                            output_file_path,
+                                            interpolation)
+                os.remove(input_file_path)
+                interpolator = Regridder(weights=weights)
+                interpolated = interpolator.regrid(input_data)
+                interpolated = interpolated[variable_name].values
+                mask_zero=args.no_mask_zero
+                if mask_zero:
+                    interpolated[interpolated == 0] = np.nan
+
             # masking of the data
             if mask_file is not None:
                 mask_level = mask_data[0, dind, :, :].values
@@ -996,7 +1048,7 @@ def fint(args=None):
                     lat,
                     out_path_one2,
                 )
-    if interpolation in ["cdo_remapcon", "cdo_remaplaf", "cdo_remapnn", "cdo_remapdis"]:      
+    if interpolation in ["cdo_remapcon","cdo_remaplaf","cdo_remapnn","cdo_remapdis","smm_con","smm_nn","smm_laf","smm_dis"]:    
         os.remove(target_grid_path)
 
     # save data (always 4D array)

test/tests.sh

@@ -47,6 +47,12 @@ fint ${FILE} ${MESH} ${INFL} -t 0 -d 0 --interp cdo_remaplaf
 fint ${FILE} ${MESH} ${INFL} -t 0 -d 0 -b "-150, 150, -50, 70" --interp cdo_remaplaf
 fint ${FILE} ${MESH} ${INFL} -t 0 -d 0 -b arctic --interp cdo_remapcon
 
+#smm_regrid
+fint ${FILE} ${MESH} --influence 500000 -t 0 -d -1  --interp smm_con --no_shape_mask
+fint ${FILE} ${MESH} ${INFL} -t 0 -d 0 -b "-150, 150, -50, 70" --interp smm_laf
+fint ${FILE} ${MESH} ${INFL} -t 0 -d 0 -b arctic --interp smm_nn
+fint ${FILE} ${MESH} ${INFL} -t 0 -d 0 --interp smm_dis
+
 
 # create mask
 fint ${FILE} ${MESH} ${INFL} -t 0 -d -1  --interp mtri_linear -o mask.nc