import xarray as xr
import numpy as np
import glob
import os
import multiprocessing as mp
import pandas as pd
import time as tm
from itertools import repeat
from logging import info, debug
import math
from ..utils import OUTPUT_DIRNAME, OUTPUT_PATTERN
from ......utils.hdf5 import _hdf5_lock
#TODO: remove occurences of "height"", not supposed to know the name
# If all observations are processed
[docs]
def apply_interpolation_by_chunk_full(segment_idx_obs, ds_icon, ds_interp, all_trcrs):
"""Apply full 3-D interpolation to a chunk of observations.
Uses pre-computed adjacent-cell indices and vertical level indices
(``ilev_below``, ``ilev_above``) from *ds_interp* to perform bilinear
horizontal + linear vertical interpolation from the ICON icosahedral
grid to observation coordinates.
Args:
segment_idx_obs (tuple[int, int]): (start, end) slice of the
observation index in *ds_interp*.
ds_icon (xr.Dataset): ICON output dataset on the icosahedral grid.
ds_interp (xr.Dataset): pre-computed interpolation metadata.
all_trcrs (list[str]): tracer variable names to interpolate.
Returns:
xr.Dataset: interpolated concentrations for the observation chunk.
"""
debug(segment_idx_obs)
idx_obs_start, idx_obs_end = segment_idx_obs
ds_interp_red = ds_interp.isel(index=slice(idx_obs_start, idx_obs_end))
ilev_below = ds_interp_red['ilev_below'].astype(int)
ilev_above = ds_interp_red['ilev_above'].astype(int)
iadjcells = ds_interp_red['iadjcells']
vert_scaling_fact = ds_interp_red['vert_scale_fact']
inv_dist_adjcells = ds_interp_red['inv_dist_adjcells']
vals_lev_above = ilev_above.values + 1
vals_lev_above = xr.DataArray(vals_lev_above, dims=("obs", "adjcell"))
vals_lev_below = ilev_below.values + 1
vals_lev_below = xr.DataArray(vals_lev_below, dims=("obs", "adjcell"))
vals_cell = iadjcells.values
vals_cell = xr.DataArray(vals_cell, dims=("obs", "adjcell"))
vals_vert_scal = vert_scaling_fact
vals_inv_dist = inv_dist_adjcells
ds_above_adj = ds_icon[all_trcrs].sel(height=vals_lev_above, ncells=vals_cell)
ds_below_adj = ds_icon[all_trcrs].sel(height=vals_lev_below, ncells=vals_cell)
ds_adj = ds_below_adj + vals_vert_scal.values * (ds_above_adj - ds_below_adj)
ds_out = (ds_adj * vals_inv_dist.values).sum(dim='adjcell') / vals_inv_dist.sum(dim='adjcell').values
return ds_out
# If observations are grouped by index
[docs]
def apply_interpolation_by_chunk_reduced(segment_idx_obs, ds_icon, ds_interp, df_metadata, all_trcrs):
"""Apply reduced (level-fixed) interpolation to a chunk of observations.
Like :func:`apply_interpolation_by_chunk_full` but does not perform
vertical interpolation: uses the observation's prescribed level index
directly. Used when ``full_interpolation=False`` or when observations
specify only a pressure level without altitude.
Args:
segment_idx_obs (tuple[int, int]): (start, end) observation slice.
ds_icon (xr.Dataset): ICON output on icosahedral grid.
ds_interp (xr.Dataset): pre-computed interpolation metadata.
df_metadata: observation metadata DataFrame (carries level column).
all_trcrs (list[str]): tracer variable names to interpolate.
Returns:
xr.Dataset: interpolated concentrations for the observation chunk.
"""
debug(segment_idx_obs)
idx_obs_start, idx_obs_end = segment_idx_obs
ds_interp_red = ds_interp.isel(index=slice(idx_obs_start, idx_obs_end))
# Data that do not need to be interpolated vertically
mask_alt_nan = np.isnan(ds_interp_red['ilev_below'][:, 0].values)
index_without_alt = ds_interp_red.index[mask_alt_nan].values
iadjcells = ds_interp_red['iadjcells'][mask_alt_nan]
inv_dist_adjcells = ds_interp_red['inv_dist_adjcells'][mask_alt_nan]
vals_cell = iadjcells.values
vals_cell = xr.DataArray(vals_cell, dims=("obs_red", "adjcell"))
vals_inv_dist = inv_dist_adjcells
min_ilev = 1 #TODO: Find a better solution
max_ilev = 60 #TODO: Find a better solution
ds_adj = ds_icon[all_trcrs].sel(ncells=vals_cell, height=slice(min_ilev, max_ilev + 1))
ds_out = (ds_adj * vals_inv_dist.values).sum(dim='adjcell') / vals_inv_dist.sum(dim='adjcell').values
ds_out = ds_out.reindex(height=list(reversed(ds_out.height)))
ds_out = ds_out.stack(obs=("obs_red", "height"), create_index=False)
ds_out = ds_out.drop_vars('height')
ds_out['obs'] = df_metadata.loc[index_without_alt, 'iobs'].values
# Data that need to be interpolated vertically
if np.sum(~mask_alt_nan):
index_with_alt = ds_interp_red.index[~mask_alt_nan].values
ilev_below = ds_interp_red['ilev_below'][~mask_alt_nan]
ilev_above = ds_interp_red['ilev_above'][~mask_alt_nan]
iadjcells = ds_interp_red['iadjcells'][~mask_alt_nan]
vert_scaling_fact = ds_interp_red['vert_scale_fact'][~mask_alt_nan]
inv_dist_adjcells = ds_interp_red['inv_dist_adjcells'][~mask_alt_nan]
vals_lev_above = ilev_above.values + 1
vals_lev_above = xr.DataArray(vals_lev_above, dims=("obs", "adjcell"))
vals_lev_below = ilev_below.values + 1
vals_lev_below = xr.DataArray(vals_lev_below, dims=("obs", "adjcell"))
vals_cell = xr.DataArray(iadjcells.values, dims=("obs", "adjcell"))
vals_vert_scal = vert_scaling_fact
vals_inv_dist = inv_dist_adjcells
ds_above_adj = ds_icon[all_trcrs].sel(height=vals_lev_above, ncells=vals_cell)
ds_below_adj = ds_icon[all_trcrs].sel(height=vals_lev_below, ncells=vals_cell)
ds_adj = ds_below_adj + vals_vert_scal.values * (ds_above_adj - ds_below_adj)
ds_out2 = (ds_adj * vals_inv_dist.values).sum(dim='adjcell') / vals_inv_dist.sum(dim='adjcell').values
ds_out2['obs'] = df_metadata.loc[index_with_alt, 'iobs']
ds_out = xr.concat([ds_out, ds_out2], dim='obs')
ds_out = ds_out.sortby('obs')
return ds_out
[docs]
def apply_interpolation(self, runsubdir, data2dump):
"""Interpolate ICON-ART output fields to all observation locations.
Reads NetCDF output files from ``{runsubdir}/OUTPUT/``, computes
distance-weighted horizontal interpolation to observation station
coordinates using pre-computed adjacent-cell metadata, applies vertical
interpolation (full or reduced), and stores results in
``self.sim_data``.
Uses :func:`apply_interpolation_by_chunk_full` or
:func:`apply_interpolation_by_chunk_reduced` depending on whether
altitude is available in the observation metadata.
Args:
self: ICON-ART model plugin instance.
runsubdir (str): path to the period run directory.
data2dump (dict): tracer-ID-keyed data-store (provides the
observation metadata for interpolation targets).
"""
output_path = os.path.join(runsubdir, OUTPUT_DIRNAME)
# ---------------------------------------------------------
# -- Fetch data, interpolate and convert to DataFrame
# ---------------------------------------------------------
info('Concatenating the files...')
file_path_day = os.path.join(output_path,
"concatenated_byday",
OUTPUT_PATTERN)
list_files_day = sorted(glob.glob(str(file_path_day)))
with _hdf5_lock:
ds_icon = xr.open_mfdataset(list_files_day,
concat_dim='time',
combine='nested',
decode_cf=False)
# Load the interpolation data
interpolation_filename = os.path.join(
runsubdir, OUTPUT_DIRNAME, "interpolation/ds_data_interpolation.nc")
with _hdf5_lock:
ds_interp = xr.open_dataset(interpolation_filename)
# Create dataout
dataout = data2dump.copy()
dataout = pd.concat([dataout['metadata'], dataout['maindata']], axis=1)
dataout = dataout[[
'date', 'station', 'network', 'parameter', 'lon', 'lat',
'alt', 'i', 'j', 'level', 'tstep', 'tstep_glo', 'dtstep',
'duration', 'is_obsvect', 'obs', 'obserror'
]]
# Extract station information
nlev = self.domain.nlev
# TODO: get tstep from ds_interp, dont use anymore data2dump, must be species-agnostic
df_metadata = data2dump[('metadata')].copy()
df_metadata['iobs'] = range(len(df_metadata))
obs_tsteps = df_metadata['tstep'].values.astype(int)
vals_time = xr.DataArray(obs_tsteps, dims=("obs"))
# Extract variable information
dim_height = ds_icon['z_mc'].dims[1]
all_trcrs = [tr for tr in list(ds_icon.data_vars) \
if ds_icon.data_vars[tr].dims == ('time', dim_height, 'ncells')]
# Create all columns for output data
dataout_all_trcrs = pd.DataFrame(index=dataout.index, columns=all_trcrs)
dataout = pd.concat([dataout, dataout_all_trcrs], axis=1)
# Prepare the chunks to apply the interpolation over all obs in ds_interp
# Number of obs in ds_interp might be different from number
# of obs in df_metadata
nchunks = 36
nobs = len(ds_interp.index)
nobs_per_chunk = int(np.ceil(nobs / nchunks))
bounds_idx_obs = list(range(0, nobs, nobs_per_chunk)) + [nobs]
list_segments_idx_obs = [[idx_start, idx_end] for (idx_start, idx_end)
in zip(bounds_idx_obs[:-1], bounds_idx_obs[1:])]
# Apply the interpolation for every chunk
info(f'Applying inverse-distance weighting interpolation with {nchunks} nchunks...')
if self.full_interpolation:
func_arguments = zip(list_segments_idx_obs, repeat(ds_icon),
repeat(ds_interp), repeat(all_trcrs))
with mp.Pool(min(nchunks, 36)) as pool:
list_ds_out = pool.starmap(apply_interpolation_by_chunk_full, func_arguments)
# Fill dataout with interpolated data
ds_out_concat = xr.concat(list_ds_out, dim='obs')
ds_out_concat = ds_out_concat.isel(time=vals_time)
dataout.loc[:, all_trcrs] = ds_out_concat.to_array().values.T
else:
func_arguments = zip(list_segments_idx_obs, repeat(ds_icon),
repeat(ds_interp), repeat(df_metadata),
repeat(all_trcrs))
with mp.Pool(min(nchunks, 36)) as pool:
list_ds_out = pool.starmap(apply_interpolation_by_chunk_reduced, func_arguments)
# Fill dataout with interpolated data
ds_out_concat = xr.concat(list_ds_out, dim='obs')
ds_out_concat = ds_out_concat.isel(time=vals_time)
dataout.loc[:, all_trcrs] = ds_out_concat.to_array().values.T
return dataout
