Skip to content
Snippets Groups Projects
Commit 6bcab6c8 authored by Oriol Tintó's avatar Oriol Tintó
Browse files

Remove icon-remap-helper.py

parent 1cdc5e98
No related branches found
No related tags found
No related merge requests found
#!/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)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment