diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..48e92b197efeeea3df0be2da6be17201ef402ff7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,42 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# Interpreter-specific settings +*.pyc +*.pyo +*.pyd + +# Environment-specific settings +.env +.env.* + +# IDE settings +.idea/ +.vscode/ +*.swp +*.swo +*.swn diff --git a/assets/icon_grid_0044_R19B07_L.nc b/assets/icon_grid_0044_R19B07_L.nc new file mode 100644 index 0000000000000000000000000000000000000000..9db3cd09c9e5fb4a322c3338d32e567cc6a8c0d8 Binary files /dev/null and b/assets/icon_grid_0044_R19B07_L.nc differ diff --git a/assets/icon_grid_0047_R19B07_L.nc b/assets/icon_grid_0047_R19B07_L.nc new file mode 100644 index 0000000000000000000000000000000000000000..eb2b476edb9ec40c6086ef74d95344c58aba2ed9 Binary files /dev/null and b/assets/icon_grid_0047_R19B07_L.nc differ diff --git a/doc/pai_example_trigrid/XaXb.png b/doc/pai_example_trigrid/XaXb.png new file mode 100644 index 0000000000000000000000000000000000000000..4ce82e3714864dcfeeed6613e52a81c0e408d856 Binary files /dev/null and b/doc/pai_example_trigrid/XaXb.png differ diff --git a/doc/pai_example_trigrid/compute_pai.py b/doc/pai_example_trigrid/compute_pai.py new file mode 100644 index 0000000000000000000000000000000000000000..44a18bc15ecae4fe26b34de322de8de0eca44785 --- /dev/null +++ b/doc/pai_example_trigrid/compute_pai.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python3 + +from pathlib import Path + +import numpy as np +import xarray as xr + +import proplot as pplt +import matplotlib.patches as patches + +from utils import plot_oi_utils as utils + +from utils.get_lat_lon_of_grib_data import read_grib, read_grib_mf + + +from pai.observation import Observation +from pai.localization import get_dist_from_obs, gaspari_cohn, vertical_dist + +# from cosmo_state import CosmoState + +if __name__ == "__main__": + + import argparse + + parser = utils.define_parser() + args = parser.parse_args() + + d = args.startdate + + path = ( + Path(args.basedir) + / "feedback" + / (d + args.time) + / f"ekf{args.obstype}.nc.{d}140000" + ) + print(path) + + obs = Observation(path, var=args.obs_var, report=311) + + ana_path = ( + Path(args.basedir) / "data" / (d + args.time) / f"an_R19B07.20210610140000.mean" + ) + + print(ana_path) + ana_mean = read_grib( + ana_path, args.model_var, grid_path="./assets/icon_grid_0047_R19B07_L.nc" + ) + + ana = read_grib_mf( + ana_path, args.model_var, grid_path="./assets/icon_grid_0047_R19B07_L.nc" + ) + + # get the hloc area, in triangular grid. + # find the grid point that is closest to the observation + + h_loc = 25.0 + dist = get_dist_from_obs( + ana_mean.clat, ana_mean.clon, obs.data.lat, obs.data.lon, h_loc + ) + + cutoff = 2.0 * h_loc * np.sqrt(10.0 / 3.0) + + ana_mean_hloc = ana_mean.where(dist.isel(x=0, y=0) <= cutoff, drop=True) + ana_hloc = ana.where(dist.isel(x=0, y=0) <= cutoff, drop=True) + + dist_hloc = dist.where(dist.isel(x=0, y=0) <= cutoff, drop=True) + + gc_hloc = gaspari_cohn(dist_hloc.values, h_loc) + + Xa = ana_hloc - ana_mean_hloc + + YaT = obs.ekf.anaens() + + obserr = obs.ekf.obs(param="e_o") + + print(Xa.transpose().shape, YaT.shape, obserr.shape) + + # vertical localization + # pressure... + # pressure obs. + ana_mean_pres = read_grib( + ana_path, "pres", grid_path="./assets/icon_grid_0047_R19B07_L.nc" + ) + pres_hloc = ana_mean_pres.where(dist.isel(x=0, y=0) <= cutoff, drop=True) + vdist = vertical_dist(pres_hloc.values, obs.data.pres.values) + + gc_vloc = gaspari_cohn(vdist, 0.3) # obs.ekf.obs(param="v_loc")[0]) + + # TODO additional damping term between 300 and 100 hPa is missing + + # compute localized kalman gain + kalman = ( + np.matmul(Xa.transpose().values, YaT) + * obserr ** (-2) + * gc_hloc + * gc_vloc.T[..., np.newaxis] + / (10.0 - 1.0) + ) + + print(kalman.shape) + + pai = np.matmul(kalman, obs.ekf.fgmeandep()) + + # get the closest model point + print(pai.shape, ana_mean_hloc) + + closest = np.argmin(dist_hloc.isel(x=0, y=0).values) + fig, ax = pplt.subplots() + + ax.plot(gc_vloc[:, closest], pres_hloc[:, closest] / 100.0, label="gc") + ax.plot(pai[closest, :], pres_hloc[:, closest] / 100.0, label="PAI") + ax.format(yreverse=True) + fig.save("vdist.png") + + fig, ax = pplt.subplots() + + ax.scatter(obs.data.lon, obs.data.lat, color="red", marker="x", s=50, zorder=3) + ax.scatter( + ana_mean_hloc.clon.isel(cell=closest), + ana_mean_hloc.clat.isel(cell=closest), + color="black", + marker="o", + s=50, + zorder=3, + ) + + circle = patches.Circle( + (obs.data.lon, obs.data.lat), + cutoff, + linewidth=2, + edgecolor="red", + facecolor="none", + alpha=1, + zorder=3, + ) + # ax.add_patch(circle) + ax.tripcolor( + dist_hloc.clon, + dist_hloc.clat, + pai[:, 50], # .isel(generalVerticalLayer=50, ens=0), + # dist.clon, # .isel(cell=sl), + # dist.clat, # .isel(cell=sl), + # dist.isel(x=0, y=0), # , values=sl), + # # data=data.t.isel(generalVerticalLayer=10), + # facecolors=z, + # edgecolors="k", + colorbar=True, + ) + + # Show the plot + fig.save(f"distance_test.png") + + # get the horizontal distance from the choosen observation. + # fg_path = ( + # Path(args.basedir) / "data" / (d + args.time) / f"fc_R19B07.20210610140000.mean" + # ) + + # print(ana_path) + # fg_mean = read_grib( + # fg_path, args.model_var, grid_path="./assets/icon_grid_0047_R19B07_L.nc" + # ) + + # fg = read_grib_mf( + # fg_path, args.model_var, grid_path="./assets/icon_grid_0047_R19B07_L.nc" + # ) + + # Xb = fg - fg_mean + + # subfig = pplt.figure() + # gridspec = [1, 1, 2, 2, 3, 3] + + # ax = subfig.subplots(gridspec) + + # pa = np.matmul( + # Xa.isel(generalVerticalLayer=40, cell=slice(18000, 20300)).values.transpose(), + # Xa.isel(generalVerticalLayer=40, cell=slice(18000, 20300)).values, + # ) + # pb = np.matmul( + # Xb.isel(generalVerticalLayer=40, cell=slice(18000, 20300)).values.transpose(), + # Xb.isel(generalVerticalLayer=40, cell=slice(18000, 20300)).values, + # ) + # m = ax[0].imshow(pa) + # n = ax[1].imshow(pb) + # o = ax[2].imshow(pa - pb) + + # ax[0].colorbar(m, label=r"$P_a$") + # ax[1].colorbar(n, label=r"$P_b$") + # ax[2].colorbar(o, label=r"$P_a - P_b$") + + # print(Xa.isel(generalVerticalLayer=30).shape) + # subfig.save(f"XaXb.png") diff --git a/doc/pai_example_trigrid/distance_test.png b/doc/pai_example_trigrid/distance_test.png new file mode 100644 index 0000000000000000000000000000000000000000..b7f60245b2869405c2aecdf4ceb8eb21e0deb658 Binary files /dev/null and b/doc/pai_example_trigrid/distance_test.png differ diff --git a/doc/pai_example_trigrid/pai/__init__.py b/doc/pai_example_trigrid/pai/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e5a0d9b4834ec8f46d6e0d1256c6dcaad2e460fe --- /dev/null +++ b/doc/pai_example_trigrid/pai/__init__.py @@ -0,0 +1 @@ +#!/usr/bin/env python3 diff --git a/doc/pai_example_trigrid/pai/localization.py b/doc/pai_example_trigrid/pai/localization.py new file mode 100644 index 0000000000000000000000000000000000000000..4ce6c02a0c0c6c1a0a65bdc9355e767fadd191b0 --- /dev/null +++ b/doc/pai_example_trigrid/pai/localization.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +import numpy as np +from numba import jit + +EARTH_RAD = 6371.0 # in km according to WGS84 + + +def get_dist_from_obs(lats, lons, olat, olon, h_loc): + """Get lower left corner and upper right corner lats and lons from the area + which is influenced by the observation, confined by the cutoff distance. + (Add +.1 degrees to each side) + + Args: + model_state (ModelState): see model_state.py + obs (Observation): see observation.py + h_loc (float): localization length scale in km + + Returns: + (llcrnrlat, llcrnrlon, urcrnrlat, urcrnrlon) (4 x float): lat lons of corners + """ + + dist = haversine_distance(lats, lons, olat, olon) + cutoff = 2.0 * h_loc * np.sqrt(10.0 / 3.0) + + return dist + + +def haversine_distance(lat, lon, lat_0, lon_0): + dlat_rad = np.radians(lat - lat_0) + dlon_rad = np.radians(lon - lon_0) + s2_dlat = np.sin(0.5 * dlat_rad) ** 2 + s2_dlon = np.sin(0.5 * dlon_rad) ** 2 + ccs2 = np.cos(np.radians(lat)) * np.cos(np.radians(lat_0)) * s2_dlon + return 2.0 * EARTH_RAD * np.arcsin(np.sqrt(s2_dlat + ccs2)) + + +@jit(nopython=True, fastmath=True, parallel=True) +def gaspari_cohn(dist, l): + """Computes the Gaspari Cohn function + + for a constant localization length scale + + Args: + dist (numpy.ndarray): distance field + l (float): localization length scale + Returns: + numpy.ndarray: Gaspari Cohn Factors, same dimensions as dist + """ + a = l * np.sqrt(10.0 / 3.0) + da = np.abs(dist.flatten()) / a + gp = np.zeros_like(da) + + i = np.where(dist.flatten() <= a) + + gp[i] = np.maximum( + 0.0, + -0.25 * da[i] ** 5 + + 0.5 * da[i] ** 4 + + 0.625 * da[i] ** 3 + - 5.0 / 3.0 * da[i] ** 2 + + 1.0, + ) + + i = np.where((dist.flatten() > a) * (dist.flatten() < 2 * a)) + gp[i] = np.maximum( + 0.0, + 1.0 / 12.0 * da[i] ** 5 + - 0.5 * da[i] ** 4 + + 0.625 * da[i] ** 3 + + 5.0 / 3.0 * da[i] ** 2 + - 5.0 * da[i] + + 4.0 + - 2.0 / 3.0 / da[i], + ) + return gp.reshape(dist.shape) + + +# @jit(nopython=True, fastmath=True, parallel=True) +def vertical_dist(pres, obs_pres): + + distances = np.zeros_like(pres) + + for ilev, lev in enumerate(pres): + distances[ilev, :] = np.abs(np.log(lev) - np.log(obs_pres)) + return distances + + +def additional_damping(pres_model): + pres_model = pres_model / 100.0 + p_one = 300.0 + p_zero = 100.0 + damping = np.ones_like(pres_model) + i = np.where(pres_model < 100.0) + damping[i] = 0.0 + + i = np.where((pres_model >= 100.0) & (pres_model <= 300.0)) + dist = np.abs(np.log(pres_model[i]) - np.log(p_zero)) / np.abs( + np.log(p_zero) - np.log(p_one) + ) + damping[i] = 0.5 * (1 - np.cos(np.pi * dist)) + return damping[..., np.newaxis] diff --git a/doc/pai_example_trigrid/pai/observation.py b/doc/pai_example_trigrid/pai/observation.py new file mode 100644 index 0000000000000000000000000000000000000000..7cc3fa659f47bce9e708c1d8a3cb37c7e932fdc1 --- /dev/null +++ b/doc/pai_example_trigrid/pai/observation.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +import numpy as np +import xarray as xr + +from kendapy.ekf import Ekf + + +class Observation: + def __init__(self, path, var=None, report=None, plot=False): + self.path = path + self.var = var + self.ekf = Ekf(str(path), filter=f"state=active") + + self.ekf.add_filter(filter=f"variable={var}") + self.ekf.add_filter(filter=f"report={report}") + + self.ekftoxarray() + + def ekftoxarray(self): + + rep = self.ekf.reports()[0] + lat = self.ekf.obs(param="lat")[0] + lon = self.ekf.obs(param="lon")[0] + + pres = self.ekf.obs(param="plevel")[np.newaxis, np.newaxis, :, np.newaxis] + obs = self.ekf.obs().filled()[np.newaxis, np.newaxis, :, np.newaxis] + print(lon, lat) + datavars = {} + datavars[self.var] = (["x", "y", "z", "rep"], obs) + datavars["pres"] = (["x", "y", "z", "rep"], pres) + + self.data = xr.Dataset( + data_vars=datavars, + coords=dict( + lon=(["x"], [lon]), + lat=(["y"], [lat]), + z=(["z"], np.arange(0, pres.shape[2])), + report=(["rep"], [rep]), + ), + ) diff --git a/doc/pai_example_trigrid/test_map.png b/doc/pai_example_trigrid/test_map.png new file mode 100644 index 0000000000000000000000000000000000000000..94f4eea1e3e49f9b9a31a80549aaaa6f24c6b129 Binary files /dev/null and b/doc/pai_example_trigrid/test_map.png differ diff --git a/doc/pai_example_trigrid/utils/__init__.py b/doc/pai_example_trigrid/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/doc/pai_example_trigrid/utils/get_lat_lon_of_grib_data.py b/doc/pai_example_trigrid/utils/get_lat_lon_of_grib_data.py new file mode 100644 index 0000000000000000000000000000000000000000..48684020dbac2e568f797e5c8592e86592fbc79b --- /dev/null +++ b/doc/pai_example_trigrid/utils/get_lat_lon_of_grib_data.py @@ -0,0 +1,48 @@ +import xarray as xr + +import numpy as np + +from pathlib import Path +import cfgrib + +grid_path = "../assets/icon_grid_0047_R19B07_L.nc" + + +def read_grib(path, var, grid_path="../assets/icon_grid_0047_R19B07_L.nc"): + + data = xr.open_dataset( + path, + engine="cfgrib", + ).rename(values="cell") + + data = data[var] + grid = xr.open_dataset(grid_path) + + data = data.assign_coords(clat=("cell", grid.clat.values * 180 / np.pi)) + + data = data.assign_coords(clon=("cell", grid.clon.values * 180 / np.pi)) + + return data + + +def read_grib_mf(path, var, grid_path="../assets/icon_grid_0047_R19B07_L.nc"): + + d = [] + ens = [] + for i, p in enumerate(path.parent.glob("an_R19B07.20210610140000.0??")): + ens.append(i) + xda = xr.open_dataset( + p, + engine="cfgrib", + ).rename(values="cell") + d.append(xda[var]) + + data = xr.concat(d, dim="ens") + data = data.assign_coords(ens=("ens", ens)) + grid = xr.open_dataset(grid_path) + + data = data.assign_coords(clat=("cell", grid.clat.values * 180 / np.pi)) + + data = data.assign_coords(clon=("cell", grid.clon.values * 180 / np.pi)) + + return data diff --git a/doc/pai_example_trigrid/utils/plot_icon.py b/doc/pai_example_trigrid/utils/plot_icon.py new file mode 100644 index 0000000000000000000000000000000000000000..cf1103c13e968f90ffb83577c557bfdfd715f806 --- /dev/null +++ b/doc/pai_example_trigrid/utils/plot_icon.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 + +import xarray as xr +import proplot as pplt +import numpy as np + + +def plot_icon(ax, var, grid_path="../assets/icon_..."): + + grid = xr.open_dataset(grid_path) + + cmap = "hot_r" + + edgecolor = (0.7, 0.7, 0.7) + + +if __name__ == "__main__": + + grid = xr.open_dataset("../assets/icon_grid_0047_R19B07_L.nc") + + data = xr.open_dataset( + "/hpc/uwork/exttdief/ilam_jun_jul21_new/data/20210610120000/an_R19B07.20210610140000.001", + engine="cfgrib", + ) + fig, ax = pplt.subplots() + + z = np.zeros(grid.vertex_of_cell.shape[1]) + print(z.shape) + sl = slice(900, 2000) + + ax.tripcolor( + grid.clon, # .isel(cell=sl), + grid.clat, # .isel(cell=sl), + data.t.isel(generalVerticalLayer=50), # , values=sl), + # data=data.t.isel(generalVerticalLayer=10), + # facecolors=z, + # edgecolors="k", + colorbar=True, + ) + + # Show the plot + fig.save(f"trigrid_test.png") diff --git a/doc/pai_example_trigrid/utils/plot_oi_utils.py b/doc/pai_example_trigrid/utils/plot_oi_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b8eabe3cf702be19cf5c07b80b99bd9b74ff9b11 --- /dev/null +++ b/doc/pai_example_trigrid/utils/plot_oi_utils.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 + +import os +import argparse +from datetime import datetime, timedelta +import cartopy.crs as ccrs +import cartopy.feature as cfeature +import numpy as np +from numba import jit + +############# +# Constants # +############# + +EARTH_RAD = 6378.0 # in km + + +def define_parser(): + parser = argparse.ArgumentParser( + description="Plotting Script for ICON KENDA experiments" + ) + + parser.add_argument( + "-b", "--basedir", dest="basedir", default="", help="path to files" + ) + + parser.add_argument( + "-s", + "--startdate", + dest="startdate", + default="", + help="YYYYMMDD of the experiment", + ) + parser.add_argument( + "-time", + "--time", + dest="time", + default="", + help="HHMMSS of the experiment", + ) + parser.add_argument( + "-m", + "--member", + dest="member", + default="", + help="member of the experiment, format example: 001", + ) + parser.add_argument( + "-obstype", + "--obstype", + dest="obstype", + default="", + help="obstype for ekf read", + ) + parser.add_argument( + "-t", + "--timeseries", + dest="timeseries", + help="get timeseries plots", + action="store_true", + ) + + parser.add_argument( + "-o", + "--outfile", + dest="outfile", + default="../figures/tmp.png", + help="Name of outfile", + ) + + parser.add_argument( + "-e", + "--experiments", + dest="exps", + default=None, + help="comma-seperated list of experiments", + ) + parser.add_argument( + "-movie", + "--movie", + dest="movie", + help="make vertical velocity movied", + action="store_true", + ) + + parser.add_argument( + "-model_var", + "--model_var", + dest="model_var", + help="store a snapshots of var", + ) + parser.add_argument( + "-obs_var", + "--obs_var", + dest="obs_var", + help="store a snapshots of var", + ) + parser.add_argument( + "-cmap", + "--cmap", + dest="cmap", + help="cmap", + ) + parser.add_argument( + "-clevs", + "--clevs", + dest="clevs", + help="clevs for var snapshot, Format: from,to,no.steps (used in np.linspace).", + ) + + parser.add_argument( + "-wtg", + "--wtg", + dest="wtg", + help="compute the weak temperature diagnostic", + action="store_true", + ) + + parser.add_argument( + "-l", + "--lev", + dest="lev", + help="model level to plot", + ) + parser.add_argument( + "-d", + "--diff", + dest="diff", + help="plot the differences in w", + action="store_true", + ) + parser.add_argument( + "-vmd", + "--vmd", + dest="vmd", + help="plot vertical velocity masking = vertical motion diagnostic", + action="store_true", + ) + parser.add_argument( + "-r", + "--region", + dest="region", + help="evaluate only for a region", + action="store_true", + ) + return parser + + +def get_timeseries(dir, startdate, template="data_ML_%H%M%S.001.nc"): + ts = sorted(os.listdir(dir)) + timesteps = [] + for t in (t for t in ts if "13" in t and ".001.nc" in t and "tmp" not in t): + date = datetime.strptime(t, template) + date = date.replace( + year=startdate.year, month=startdate.month, day=startdate.day + ) + timesteps.append(date) + return [t for t in ts if "13" in t and ".001.nc" in t], timesteps + + +@jit(nopython=True, fastmath=True, parallel=True) +def haversine_distance(lat, lon, lat_0, lon_0): + dlat_rad = np.radians(lat - lat_0) + dlon_rad = np.radians(lon - lon_0) + s2_dlat = np.sin(0.5 * dlat_rad) ** 2 + s2_dlon = np.sin(0.5 * dlon_rad) ** 2 + ccs2 = np.cos(np.radians(lat)) * np.cos(np.radians(lat_0)) * s2_dlon + return 2.0 * EARTH_RAD * np.arcsin(np.sqrt(s2_dlat + ccs2)) + + +def compute_potential_temperature(data): + CP = 1005 # specific heat capacity [J / (kg * K )] at constant pressure + R = 287.058 # gas constant dry air + return data.temp * (1000.0 / data.pres) ** (R / CP) diff --git a/doc/pai_example_trigrid/utils/trigrid_test.png b/doc/pai_example_trigrid/utils/trigrid_test.png new file mode 100644 index 0000000000000000000000000000000000000000..9ec5d9cc3d4732a50d0ff768a758449bf33bf2b1 Binary files /dev/null and b/doc/pai_example_trigrid/utils/trigrid_test.png differ diff --git a/doc/pai_example_trigrid/vdist.png b/doc/pai_example_trigrid/vdist.png new file mode 100644 index 0000000000000000000000000000000000000000..7e2cce8aa60799ec587a113169dda6b42d6da057 Binary files /dev/null and b/doc/pai_example_trigrid/vdist.png differ