import xarray as xr
import numpy as np
import pandas as pd
import os
import multiprocessing as mp
from itertools import repeat
from logging import info, warning
from .tracers import change_tracers_xml_lbc
from ..utils import LBC_DIRNAME
from ......utils import path
from ......utils.hdf5 import _hdf5_lock
[docs]
def make_lbc(self, datastore, ddi, ddf, runsubdir, mode):
"""Fill the meteorological lbc files with lbc for tracers.
"""
# Inputs and outputs
trids_in = list(datastore.keys())
trids_out = [("lbc", s) for s in self.chemistry.acspecies.attributes]
# Flags to detect the ensemble
is_ensemble = False
is_perturbed_comp = False
# Create the dictionary to store the input data
if not hasattr(self, "dict_lbc_dataout"):
self.dict_lbc_dataout = {}
if ddi not in self.dict_lbc_dataout:
self.dict_lbc_dataout[ddi] = {}
# Create the OEM dir
oem_dir = os.path.join(runsubdir, 'OEM')
path.init_dir(oem_dir)
# Loop over the species that must be transported
for trid_out in trids_out:
trid_in = trid_out
spec_ref = spec = trid_out[1]
# If ensemble, check what the datastore contains
if "__sample#" in spec:
is_ensemble = True
is_perturbed_comp = True
spec_ref = spec.split("__sample#")[0]
sample_id = spec.split("__sample#")[1]
if int(sample_id) > 0:
continue
if trid_in not in datastore:
trid_in = ("lbc", spec_ref)
is_perturbed_comp = False
if trid_in not in datastore:
continue
tracer = datastore[trid_in]
tracer_data = tracer["data"][ddi]
# Emitted species corresponding to reference species
emi_specs = self.chemistry.mapping_active2emi_ref[spec_ref]
# --------------------------------------------------------------------------
# -- Get the prior data
# --------------------------------------------------------------------------
# If lbc are not loaded, read the input files
if "spec" not in tracer_data:
varname = datastore[trid_in]['varname']
input_dates = [dd[0] for dd in datastore[trid_in]['input_dates'][ddi]]
input_files = datastore[trid_in]['input_files'][ddi]
func_arguments = zip(repeat(varname), input_files)
with mp.Pool(min(len(input_dates), 36)) as pool:
list_data2dump = pool.starmap(mp_read_data, func_arguments)
# If lbc are loaded, read the datastore
else:
da_lbc_prior = tracer_data["spec"].copy()
input_dates = pd.to_datetime(da_lbc_prior.time)
# TODO: problem with last timestep == 0 sometimes...
if da_lbc_prior[-1].mean() == 0:
da_lbc_prior[-1] = da_lbc_prior[-2]
list_data2dump = [da_lbc_prior.sel(time=dd, lat=0).values[np.newaxis]
for dd in input_dates]
list_data2dump_post = list_data2dump
# --------------------------------------------------------------------------
# -- Ensemble processing
# --------------------------------------------------------------------------
if is_ensemble:
if is_perturbed_comp:
info(f"The {spec_ref} lbc is perturbed.")
info(f"Calculating lbc scaling factors for {spec_ref}...")
list_da = [datastore[t]["data"][ddi]["spec"].isel(lev=-1)
for t in trids_in]
lbc_lambdas = xr.concat(list_da, dim="ens")
lbc_lambdas = lbc_lambdas / lbc_lambdas[0]
lbc_lambdas = lbc_lambdas.squeeze()
lbc_lambdas = lbc_lambdas.isel(time=slice(0, -1)).mean(dim='time')
lbc_lambdas = lbc_lambdas.fillna(1)
warning("At present, note that the optimization of LBCs with ICON-ART "
"can only be performed with hpixels=global or regions "
"and vpixels=columns." )
else:
info(f"The {spec_ref} lbc is NOT perturbed.")
info(f"Creating fake emission scaling factors (1.0) for {spec_ref}...")
nsamples = getattr(self.chemistry, "nsamples", 1)
lbc_lambdas = xr.DataArray(
np.ones((nsamples, 1)),
dims=("ens", "reg")
)
# Get the posterior background
lbc_lambdas_post = lbc_lambdas[2]
da_lbc_post = da_lbc_prior * lbc_lambdas_post.values
list_data2dump_post = [da_lbc_post.sel(time=dd, lat=0).values[np.newaxis]
for dd in input_dates]
# Define the regions
if 'tracer' in tracer and hasattr(tracer['tracer'], 'regions'):
data_regions = getattr(tracer['tracer'], 'regions')[0]
da_regions = xr.DataArray(data_regions, dims=('lon'))
else:
da_regions = xr.DataArray(np.zeros((self.domain.nlon)), dims=('lon'))
# Calculate the scaling factors for each region
regions_id = xr.DataArray(np.unique(da_regions), dims=('reg'))
lbc_lambdas = lbc_lambdas.where(da_regions == regions_id).mean(dim='lon')
# Store the scaling factors
self.lbc_lambdas[ddi][spec_ref] = lbc_lambdas
# Create and dump the regions file
lbc_ens_regions_filepath = os.path.join(oem_dir, 'boundary_regions.nc')
if not os.path.exists(lbc_ens_regions_filepath):
info(f"Creating the ensemble regions file for the lbc...")
da_lbc_reg = xr.DataArray((da_regions == regions_id).astype(int),
dims=("cell", "reg"))
ds_lbc_reg = da_lbc_reg.to_dataset(name='boundaryregion')
ds_lbc_reg['global_cell_idx'] = xr.DataArray(range(1, self.domain.nlon + 1),
dims=("cell"))
lbc_ens_regions_filepath = os.path.join(oem_dir, 'boundary_regions.nc')
with _hdf5_lock:
ds_lbc_reg.to_netcdf(lbc_ens_regions_filepath)
# --------------------------------------------------------------------------
# -- Merge the lbc data to meteo files
# --------------------------------------------------------------------------
for date in input_dates:
if date not in self.dict_lbc_dataout[ddi]:
self.dict_lbc_dataout[ddi][date] = None
func_arguments = zip(
list(self.dict_lbc_dataout[ddi].values()),
input_dates,
list_data2dump,
list_data2dump_post,
repeat(self.meteo_lbc_dir),
repeat(self.meteo_lbc_file),
repeat(spec_ref),
repeat(is_ensemble),
repeat(is_perturbed_comp),
repeat(self.lbc_dry2moist),
repeat(self.domain.extpar_file),
repeat(emi_specs)
)
with mp.Pool(min(len(input_dates), 36)) as pool:
list_ds_lbc = pool.starmap(mp_merge_data_with_meteo_lbc_file, func_arguments)
for date, ds_lbc in zip(input_dates, list_ds_lbc):
self.dict_lbc_dataout[ddi][date] = ds_lbc
info(f"Added {spec_ref} to the lbc dictionary.")
# --------------------------------------------------------------------------
# -- Adapt tracers.xml
# --------------------------------------------------------------------------
info(f"Modifying tracers.xml for {spec_ref} lbc ({is_ensemble=})...")
for t in trids_out:
change_tracers_xml_lbc(self, t[1], is_ensemble=is_ensemble)
return
# --------------------------------------------------------------------
# -- OTHERS FUNCTIONS FOR MULTIPROCESSING
# --------------------------------------------------------------------
[docs]
def dump_lbc_files(self, ddi, runsubdir):
"""Write accumulated lateral boundary condition datasets to disk.
Iterates over all LBC dates stored in ``self.dict_lbc_dataout[ddi]``
and writes each as a separate ``ifs_YYYYMMDDHH_lbc.nc`` file in the
``LBC/`` sub-directory of *runsubdir*, using parallel processes.
Args:
self: ICON-ART model plugin instance.
ddi (datetime): period start date.
runsubdir (str): path to the period run directory.
"""
lbc_dir = os.path.join(runsubdir, '..', LBC_DIRNAME)
path.init_dir(lbc_dir)
list_dates = list(self.dict_lbc_dataout[ddi].keys())
list_ds_lbc = list(self.dict_lbc_dataout[ddi].values())
func_arguments = zip(list_dates, list_ds_lbc, repeat(lbc_dir))
with mp.Pool(min(len(list_ds_lbc), 36)) as pool:
pool.starmap(mp_dump_lbc_files, func_arguments)
del self.dict_lbc_dataout[ddi]
return
[docs]
def mp_read_data(varname, file):
"""Read a single variable from a NetCDF file (multiprocessing-safe helper).
Args:
varname (str): variable name to extract.
file (str): path to the NetCDF file.
Returns:
np.ndarray: the variable's values array.
"""
with _hdf5_lock:
return xr.open_dataset(file)[varname].values
[docs]
def mp_merge_data_with_meteo_lbc_file(ds_lbc,
date,
data,
data_post,
meteo_lbc_dir,
meteo_lbc_file,
spec_ref,
is_ensemble,
is_perturbed_comp,
lbc_dry2moist,
extpar_file,
emi_specs):
"""Merge CIF tracer LBC data with the IFS meteorological LBC file for one date.
Reads the IFS LBC NetCDF for *date*, inserts CIF tracer fields (*data*
/ *data_post*), applies an optional dry-to-moist VMR conversion, and
returns the merged dataset.
Args:
ds_lbc: in-progress LBC xr.Dataset accumulator.
date (datetime): LBC date.
data: CIF tracer concentration array.
data_post: post-processed (interpolated) tracer array.
meteo_lbc_dir (str): directory of IFS LBC files.
meteo_lbc_file (str): IFS LBC filename pattern.
spec_ref (str): reference species name.
is_ensemble (bool): whether running in ensemble mode.
is_perturbed_comp (bool): whether the component is a perturbed ensemble member.
lbc_dry2moist (bool): convert dry VMR to moist VMR.
extpar_file (str): external parameter file path.
emi_specs (list): emitted species names.
Returns:
xr.Dataset: merged LBC dataset for *date*.
"""
# Fetch the meteo file
if ds_lbc is None:
meteo_file = os.path.join(meteo_lbc_dir,
date.strftime(meteo_lbc_file))
with _hdf5_lock:
ds_lbc = xr.open_dataset(meteo_file)
# Create a DataArray with the correct data
da = ds_lbc['T'].copy(data=data)
da_post = ds_lbc['T'].copy(data=data_post)
# Copy topography_c from Extpar in GEOP_SFC
with _hdf5_lock:
ds_extpar = xr.open_dataset(extpar_file)
ds_lbc['GEOP_SFC'][:] = ds_extpar["topography_c"].values * 9.80665
ds_lbc['GEOSP'][:] = ds_extpar["topography_c"].values * 9.80665
# If needed, convert from dry vmr/mmr to moist vmr/mmr
# In ICON-ART, lbc data must be in moist air mmr
if lbc_dry2moist:
info(f"LBC for {spec_ref} at {date} converted from dry to moist air.")
qv = ds_lbc['QV']
da = da.copy() * (1 - qv.data)
da_post = da_post.copy() * (1 - qv.data)
# Change the attributes
da.attrs['standard_name'] = spec_ref
da.attrs['long_name'] = f'LBC for {spec_ref} generated by cif '
da.attrs['units'] = 'mol mol-1'
da_post.attrs['standard_name'] = spec_ref
da_post.attrs['long_name'] = f'LBC posterior for {spec_ref} generated by cif '
da_post.attrs['units'] = 'mol mol-1'
# Change the time dimension if needed
ds_lbc['time'] = np.array([date])
# Change the coords
da['time'] = ds_lbc['time']
da['lev'] = ds_lbc['lev']
da['ncells'] = ds_lbc['ncells']
da_post['time'] = ds_lbc['time']
da_post['lev'] = ds_lbc['lev']
da_post['ncells'] = ds_lbc['ncells']
# Fill the prior and posterior background variable with the input data
ds_lbc[spec_ref + '_BG'] = da
if is_ensemble:
ds_lbc[spec_ref + '_BG_POST'] = da_post
# Fill the reference and emitted variables with the input data
if not is_ensemble:
ds_lbc[spec_ref] = da
for emi_spec in emi_specs:
ds_lbc[emi_spec] = da
return ds_lbc
[docs]
def mp_dump_lbc_files(date, ds_lbc, lbc_dir):
"""Write a single LBC dataset to disk (multiprocessing-safe helper).
Writes *ds_lbc* as ``{lbc_dir}/ifs_YYYYMMDDHH_lbc.nc``, overwriting
any existing file at the same path.
Args:
date (datetime): LBC date (used to format the filename).
ds_lbc (xr.Dataset): LBC dataset to write.
lbc_dir (str): destination directory.
"""
lbc_file = date.strftime(
f"{lbc_dir}/ifs_%Y%m%d%H_lbc.nc"
)
if os.path.exists(lbc_file):
os.remove(lbc_file)
with _hdf5_lock:
ds_lbc.to_netcdf(lbc_file)
info(f"Dumped lbc file for {date}.")
