Trying to use icon-remap-helper.py

templates/event-generator/adapt_member.sh

@@ -57,4 +57,12 @@ integrate_sst_to_extpar "${INTERPOLATED_SST}" "${EXTERNAL_PARAMETERS_FILE}"
 
 #old_integrate_sst_to_analysis "${INTERPOLATED_SST}" "${ANALYSIS_FILE}"
 # Integrate sst to analysis
-integrate_sst_to_analysis "${DESTINATION_GRID}" "${ANALYSIS_FILE}"
+#integrate_sst_to_analysis "${DESTINATION_GRID}" "${ANALYSIS_FILE}"
+
+
+python ${PROJ_FOLDER}/templates/event-generator/icon-remap-helper.py --src-grid ${DESTINATION_GRID} \
+                                                                --dst-grid ${DESTINATION_GRID} \
+                                                                --source ${ANALYSIS_FILE} \
+                                                                --dest . \
+                                                                --output-format nc \
+                                                                --rename "analysis.nc"
\ No newline at end of file

templates/event-generator/incon-remap-helper.py

+#!/usr/bin/env python3
+"""
+automatically create namelist files for icon grid to icon grid remapping and run iconremap
+"""
+import argparse
+from enstools.io import read, write
+from enstools.interpolation import nearest_neighbour
+from subprocess import run
+import logging
+import os
+import numpy as np
+import xarray
+
+
+def load_vgrid(grid_file):
+    """
+    read HHL from an input file and calculate also the full level heights
+    """
+    logging.info(f"Reading input file with vertical grid information {grid_file}...")
+    data = read(grid_file)
+    if not "HHL" in data:
+        logging.error(f"HHL not found in {grid_file}")
+        exit(-1)
+
+    # store both result arrays in one dataset without time dimension
+    result = xarray.Dataset()
+    result["HHL"] = data["HHL"][0, ...].compute()
+
+    FHL = xarray.DataArray(np.empty((result["HHL"].shape[0] - 1, result["HHL"].shape[1])), name="FHL",
+                           dims=("generalVertical2", "cell"))
+    for layer in range(FHL.shape[0]):
+        FHL[layer, ...] = (result["HHL"][layer, ...] + result["HHL"][layer + 1, ...]) / 2
+    result["FHL"] = FHL
+    return result
+
+
+def vertical_interpolation_one_variable(src_hl, dst_hl, values):
+    """
+    perform the interpolation using numpy.interp on one variable
+    """
+    # perform the interpolation gridpointwise
+    result = np.empty((values.shape[0], dst_hl.shape[0], values.shape[2]))
+    for time in range(values.shape[0]):
+        for cell in range(values.shape[2]):
+            # all the flipping is neccessary as the function interp expects increasing values
+            result[time, :, cell] = np.flip(np.interp(np.flip(dst_hl[:, cell], 0), np.flip(src_hl[:, cell], 0),
+                                                      np.flip(values.values[time, :, cell], 0)), 0)
+
+    # create the new xarray DataArray
+    new_array = xarray.DataArray(result, dims=values.dims, name=values.name, attrs=values.attrs)
+    return new_array
+
+
+def vertical_interpolation(src_vgrid, dst_vgrid, input_name, output_name):
+    """
+    perform vertical interpolation
+    """
+    logging.info("starting vertical interpolation...")
+    # read source and destination grids
+    src_vgrid_hl = load_vgrid(src_vgrid)
+    dst_vgrid_hl = load_vgrid(dst_vgrid)
+    src_hhl_dim = src_vgrid_hl["HHL"].shape[0]
+    src_fhl_dim = src_vgrid_hl["FHL"].shape[0]
+    dst_hhl_dim = dst_vgrid_hl["HHL"].shape[0]
+    dst_fhl_dim = dst_vgrid_hl["FHL"].shape[0]
+
+    # read input file
+    infile = read(input_name).compute()
+
+    # create output file
+    outfile = xarray.Dataset()
+
+    # loop over all variables of the input file
+    for var in infile.variables:
+        # VN is special, it is defined on the edges of the grid. find nearest hgith coordinates
+        if var == "VN" and infile[var].shape[1] == src_fhl_dim:
+            logging.info(f"    -> interpolating {var} onto FHL")
+            logging.info("        -> interpolating of height array to the edges")
+            fint = nearest_neighbour(infile["clon"], infile["clat"], infile["elon"], infile["elat"],
+                                     src_grid="unstructured", dst_grid="unstructured", npoints=2, method="mean")
+            src_vgrid_fhl_vn = fint(src_vgrid_hl["FHL"])
+            dst_vgrid_fhl_vn = fint(dst_vgrid_hl["FHL"])
+            outfile[var] = vertical_interpolation_one_variable(src_vgrid_fhl_vn.values, dst_vgrid_fhl_vn.values,
+                                                               infile[var])
+        elif not var.startswith("height") and len(infile[var].shape) > 1 and infile[var].shape[1] == src_hhl_dim:
+            logging.info(f"    -> interpolating {var} onto HHL")
+            outfile[var] = vertical_interpolation_one_variable(src_vgrid_hl["HHL"].values, dst_vgrid_hl["HHL"].values,
+                                                               infile[var])
+            continue
+        elif not var.startswith("height") and len(infile[var].shape) > 1 and infile[var].shape[1] == src_fhl_dim:
+            logging.info(f"    -> interpolating {var} onto FHL")
+            outfile[var] = vertical_interpolation_one_variable(src_vgrid_hl["FHL"].values, dst_vgrid_hl["FHL"].values,
+                                                               infile[var])
+            continue
+        else:
+            if var.startswith("height") and infile[var].shape[0] == src_hhl_dim:
+                if len(infile[var].shape) == 2:
+                    continue
+                logging.info(f"    -> replacing old height coordinate '{var}'")
+                outfile[var] = xarray.DataArray(np.arange(1, dst_hhl_dim + 1, 1) + 0.5, name=var, dims=infile[var].dims,
+                                                attrs=infile[var].attrs)
+                if var + "_bnd" in infile:
+                    bnds = xarray.DataArray(np.empty((dst_hhl_dim, 2)), name=var + "_bnds",
+                                            dims=infile[var + "_bnds"].dims, attrs=infile[var + "_bnds"].attrs)
+                    bnds[dst_hhl_dim, 0] = outfile[var].values - 0.5
+                    bnds[dst_hhl_dim, 1] = outfile[var].values + 0.5
+                    outfile[var + "_bnds"] = bnds
+            elif var.startswith("height") and infile[var].shape[0] == src_fhl_dim:
+                if len(infile[var].shape) == 2:
+                    continue
+                logging.info(f"    -> replacing old height coordinate '{var}'")
+                outfile[var] = xarray.DataArray(np.arange(1, dst_fhl_dim + 1, 1) + 0.5, name=var, dims=infile[var].dims,
+                                                attrs=infile[var].attrs)
+                if var + "_bnd" in infile:
+                    bnds = xarray.DataArray(np.empty((dst_fhl_dim, 2)), name=var + "_bnds",
+                                            dims=infile[var + "_bnds"].dims, attrs=infile[var + "_bnds"].attrs)
+                    bnds[:, 0] = outfile[var].values - 0.5
+                    bnds[:, 1] = outfile[var].values + 0.5
+                    outfile[var + "_bnds"] = bnds
+            else:
+                logging.info(f"    -> storing {var} without interpolation")
+                if var in infile.coords:
+                    outfile.coords[var] = infile[var]
+                else:
+                    outfile[var] = infile[var]
+
+    # store the result
+    logging.info(f"writing file {output_name}")
+    outfile.attrs = infile.attrs
+    outfile.to_netcdf(output_name, engine="scipy")
+
+
+def remap_one_file(in_grid, out_grid, one_file, dst_fodler, rename=None, src_vgrid=None, dst_vgrid=None):
+    """
+    write the remapping namelist and run iconremap
+
+    Parameters
+    ----------
+    in_grid
+    out_grid
+    one_file
+    dst_fodler
+    """
+    # read the file content to get a list of all variables
+    content = read(one_file)
+    all_vars = list(content.data_vars)
+    remap_vars = []
+    for var in all_vars:
+        if not "bnds" in var and not '_vertices' in var and not 'lat' in var and not 'lon' in var:
+            remap_vars.append(var)
+
+    # make sure that destination folder exists
+    if not os.path.exists(dst_fodler):
+        os.makedirs(dst_fodler)
+
+    # is vertical remapping requested?
+    if src_vgrid is not None and dst_vgrid is not None:
+        vinp = True
+        if args.output_format != "nc":
+            logging.error("vertical regridding is only supported for netcdf output!")
+            exit(-1)
+    else:
+        vinp = False
+
+    # rename the file if requested
+    if rename is not None:
+        # read the time stamp
+        if content["time"].size != 1:
+            logging.error("more then one timestep, unable to rename the file!")
+            exit(-1)
+        if content["time"].attrs["units"] == "day as %Y%m%d.%f":
+            date_part = int(content["time"][0])
+            time_part = float(content["time"][0]) - date_part
+            year = str(date_part)[0:4]
+            month = str(date_part)[4:6]
+            day = str(date_part)[6:8]
+            hour = "%02d" % round(time_part * 24)
+        else:
+            logging.error("unsupported timeformat!")
+            exit(-1)
+        # replace ICON-style placeholders:
+        rename = rename.replace("<y>", year)
+        rename = rename.replace("<m>", month)
+        rename = rename.replace("<d>", day)
+        rename = rename.replace("<h>", hour)
+        rename = os.path.join(dst_fodler, rename)
+    else:
+        rename = os.path.join(dst_fodler, (os.path.basename(one_file)))
+
+    # output grib or netcdf
+    filename, ext = os.path.splitext(rename)
+    if args.output_format == "grb":
+        ext = ".grb"
+    elif args.output_format == "nc":
+        ext = ".nc"
+    if ext in ["grib", "grb", "grib2", "grb2"]:
+        out_filetype = 2
+    else:
+        out_filetype = 4
+    rename = filename + ext
+
+    # create namelist for the input file
+    namelist = f"""
+    &remap_nml
+        in_grid_filename   = '{os.path.abspath(in_grid)}'
+        in_filename        = '{os.path.abspath(one_file)}'
+        in_type            = 2
+        out_grid_filename  = '{os.path.abspath(out_grid)}'
+        out_filename       = '{os.path.abspath(rename)}'
+        out_type           = 2
+        out_filetype       = {out_filetype}
+    /
+    """
+    # add all variables
+    for var in remap_vars:
+        # skip VN when vertical interpolation is done. it has a different number of points
+        if vinp and var == "VN":
+            logging.warning("skipping VN due to requested vertical interpolation. Make sure you have U and V!")
+            # continue
+        if "soiltyp" in var.lower() or content[var].dtype in [np.int32, np.int64]:
+            intp_method = 4
+        else:
+            intp_method = 3
+        namelist += f"""
+        &input_field_nml
+            inputname      = "{var}"
+            outputname     = "{var}"
+            intp_method    = {intp_method}
+        /
+        """
+    nml_file = os.path.abspath(os.path.join(dst_fodler, (os.path.basename(one_file))) + ".namelist")
+    with open(nml_file, "w") as nml:
+        nml.write(namelist + "\n")
+
+    # run the remapping tool
+    p = run(["iconremap", "-vvv", "--remap_nml", nml_file], cwd=dst_fodler)
+    if p.returncode != 0:
+        logging.error(f"remapping of {one_file} failed")
+        exit(-1)
+
+    # perform vertical regridding using numpy and enstools
+    if vinp:
+        # move horizontally remapped file to temporal name
+        tmpname = filename + ".no-vinp" + ext
+        vinp_filename = filename + ".vinp" + ext
+        os.rename(rename, tmpname)
+
+        # perform the actual vertical interpolation
+        vertical_interpolation(src_vgrid, dst_vgrid, tmpname, vinp_filename)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src-grid", required=True, help="source grid file")
+    parser.add_argument("--src-vgrid", required=False,
+                        help="optional: HHL from the source grid. Used for vertical interpolation.")
+    parser.add_argument("--dst-grid", required=True, help="destination grid file")
+    parser.add_argument("--dst-vgrid", required=False,
+                        help="optional: HHL for the destination grid. Used for vertical interpolation.")
+    parser.add_argument("--source", nargs="+", required=True, help="source data file(s)")
+    parser.add_argument("--dest", required=True, help="destination folder")
+    parser.add_argument("--output-format", choices=["input", "nc", "grb"], default="input",
+                        help="select type of output: input=same as input; nc=netcdf, grb=grib")
+    parser.add_argument("--rename", required=False,
+                        help="change the filename. Example: 'latbc_DOM01_ML_<y>-<m>-<d>T<h>.nc'")
+    args = parser.parse_args()
+
+    # if a vertical source grid is given, interpolate it at first horizontally onto the destination grid
+    if args.src_vgrid is not None:
+        src_vgrid_name, ext = os.path.splitext(os.path.basename(args.src_vgrid))
+        if ext != ".grb":
+            logging.error("--src-vgrid is expected to be a grib file with extension .grb!")
+            exit(-1)
+        remap_one_file(args.src_grid, args.dst_grid, args.src_vgrid, args.dest)
+        src_vgrid_name1 = os.path.join(args.dest, src_vgrid_name + ".nc")
+        src_vgrid_name2 = os.path.join(args.dest, src_vgrid_name + ".hinp.nc")
+        os.rename(src_vgrid_name1, src_vgrid_name2)
+        args.src_vgrid = src_vgrid_name2
+
+    # loop over all source files
+    for one_file in args.source:
+        remap_one_file(args.src_grid, args.dst_grid, one_file, args.dest, args.rename, args.src_vgrid, args.dst_vgrid)