Source code for pycif.plugins.models.iconart.io.inputs.inicond
import numpy as np
import xarray as xr
import os
import shutil
from logging import info
from ..utils import INICOND_FILENAME
from .tracers import change_tracers_xml_inicond
from ......utils.check.errclass import CifKeyError
from ......utils.hdf5 import _hdf5_lock
[docs]
def make_inicond(self, datastore, ddi, ddf, runsubdir, mode):
"""Prepare the ICON-ART initial-condition file for one sub-simulation period.
Fetches the meteorological initial-condition file (via
:func:`fetch_meteo_inicond_file`), inserts the CIF-modified tracer
concentrations from *datastore*, and writes ``meteo_inicond.nc`` into
*runsubdir*. Handles ensemble/perturbed species by mapping sample names
back to their reference species.
Args:
self: ICON-ART model plugin instance.
datastore (dict): tracer-ID-keyed CIF data-store entries.
ddi (datetime): period start date.
ddf (datetime): period end date (unused).
runsubdir (str): path to the period run directory.
mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
"""
# Inputs and outputs
trids_in = list(datastore.keys())
trids_out = [("inicond", 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 output inicond data
if not hasattr(self, "dict_inicond_dataout"):
self.dict_inicond_dataout = {}
if ddi not in self.dict_inicond_dataout:
ds_ini_meteo = fetch_meteo_inicond_file(self, runsubdir)
else:
ds_ini_meteo = self.dict_inicond_dataout[ddi]
# 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 = ("inicond", spec_ref)
is_perturbed_comp = False
if trid_in not in datastore:
continue
tracer = datastore[trid_in]
tracer_data = tracer["data"][ddi]
# --------------------------------------------------------------------------
# -- Ensemble processing
# --------------------------------------------------------------------------
if is_ensemble and is_perturbed_comp:
info(f"The {spec_ref} inicond is perturbed.")
info(f"Calculating inicond scaling factors for {spec_ref}...")
# Check if only partial vertical data have been loaded
same_nlevs = len(np.unique(
[len(datastore[t]["data"][ddi]["spec"]["lev"]) for t in trids_in]
)) == 1
if same_nlevs:
# Fetch all the samples
list_da = [datastore[t]["data"][ddi]["spec"] for t in trids_in]
# Concatenate the data
da_ini_ens = xr.concat(list_da, dim='ens')
else:
# Fetch all the samples at the surface
list_da_surf = [datastore[t]["data"][ddi]["spec"].isel(lev=-1)
for t in trids_in]
# Calculate the scaling factors
inicond_lambdas = xr.concat(list_da_surf, dim="ens")
inicond_lambdas = inicond_lambdas / inicond_lambdas[0]
inicond_lambdas = inicond_lambdas.fillna(1)
inicond_lambdas = inicond_lambdas.squeeze()
# Concatenate the data
da_ini_ens = datastore[trids_in[0]]["data"][ddi]["spec"] * inicond_lambdas
da_ini_ens = da_ini_ens.transpose('ens', ...)
# If needed, convert from dry mmr (vmr) to moist mmr (vmr)
# In ICON-ART, inicond data must be in moist air mmr
if self.inicond_dry2moist:
info(f"Converting {spec_ref} ensemble from dry to moist air...")
qv = ds_ini_meteo['QV'].values[..., np.newaxis, :]
da_ini_ens = da_ini_ens.copy() * (1 - qv)
# Convert the dataarray to a dataset with multiple variables
ds_ini_ens = da_ini_ens.to_dataset(dim="ens")
ds_ini_ens = ds_ini_ens.isel(lat=0)
ds_ini_ens = ds_ini_ens.rename_dims({"lon": "ncells"})
# Change the attributes
for it, t in enumerate(trids_in):
ds_ini_ens[it].attrs['standard_name'] = f'{t[1]}'
ds_ini_ens[it].attrs['long_name'] = f'IC for {t[1]} generated by the CIF'
ds_ini_ens[it].attrs['units'] = 'mol per mol of moist air'
# Change the coords
ds_ini_ens['time'] = ds_ini_meteo['time']
ds_ini_ens['lev'] = ds_ini_meteo['lev']
ds_ini_ens['ncells'] = ds_ini_meteo['ncells']
# Change the name of variables
rename_map = {it: t[1] for it, t in enumerate(trids_in)}
ds_ini_ens = ds_ini_ens.rename_vars(rename_map)
# Add prior and posterior background tracer
ds_ini_ens[spec_ref + '_BG'] = ds_ini_ens[trids_in[0][1]].copy()
ds_ini_ens[spec_ref + '_BG_POST'] = ds_ini_ens[trids_in[2][1]].copy()
# Merge the ensemble inicond data with the meteo file
ds_ini_meteo = xr.merge([ds_ini_meteo, ds_ini_ens])
# --------------------------------------------------------------------------
# -- No ensemble processing
# --------------------------------------------------------------------------
else:
# Get the data
if "spec" not in tracer_data:
varname = datastore[trid_in]['varname']
ref_dir = datastore[trid_in]["dirorig"]
ref_file = datastore[trid_in]["fileorig"]
with _hdf5_lock:
ds_init_prior = xr.open_dataset(os.path.join(ref_dir, ref_file))
da_ini_prior = ds_init_prior[varname]
else:
da_ini_prior = tracer_data["spec"][:, :, 0, :]
da_ini_prior = da_ini_prior.rename({"lon": "ncells"})
# If needed, convert from dry vmr/mmr to moist vmr/mmr
# In ICON-ART, inicond data must be in moist air mmr
if self.inicond_dry2moist:
info(f"Converting {spec_ref} from dry to moist air...")
qv = ds_ini_meteo['QV']
da_ini_prior = da_ini_prior.copy() * (1 - qv)
# Change the attributes
da_ini_prior.attrs['standard_name'] = f'{spec_ref}'
da_ini_prior.attrs['long_name'] = f'IC for {spec_ref} generated by the CIF'
da_ini_prior.attrs['units'] = 'mol per mol of moist air'
# Change the coords
da_ini_prior['time'] = ds_ini_meteo['time']
da_ini_prior['lev'] = ds_ini_meteo['lev']
da_ini_prior['ncells'] = ds_ini_meteo['ncells']
# Add the new variable
ds_ini_meteo[spec_ref] = da_ini_prior
# Add background tracer
ds_ini_meteo[spec_ref + '_BG'] = da_ini_prior.copy()
if is_ensemble:
ds_ini_meteo[spec_ref + '_BG_POST'] = da_ini_prior.copy()
# --------------------------------------------------------------------------
# -- Save the data in the dictionary
# --------------------------------------------------------------------------
self.dict_inicond_dataout[ddi] = ds_ini_meteo
info(f"Added {spec_ref} to the inicond dictionary.")
# --------------------------------------------------------------------------
# -- Adapt tracers.xml
# --------------------------------------------------------------------------
info(f"Modifying tracers.xml for {spec_ref} inicond ({is_ensemble=})...")
for t in trids_out:
change_tracers_xml_inicond(self, t[1], is_ensemble=is_ensemble, is_perturbed_comp=is_perturbed_comp)
return
# --------------------------------------------------------------------
# -- OTHERS FUNCTIONS
# --------------------------------------------------------------------
[docs]
def dump_inicond_file(self, ddi, runsubdir):
"""Write the accumulated initial-condition dataset to ``meteo_inicond.nc``.
Retrieves the in-memory inicond dataset from ``self.dict_inicond_dataout[ddi]``,
updates its time dimension to *ddi*, and writes it to the run directory.
Args:
self: ICON-ART model plugin instance.
ddi (datetime): period start date (used as time coordinate).
runsubdir (str): path to the period run directory.
"""
# Get the saved inicond file
ds_ini = self.dict_inicond_dataout[ddi]
# Change the time dimension
ds_ini['time'] = np.array([ddi])
# Dump the updated inicond dataset
info(f"Dumping the updated inicond dataset...")
init_file = os.path.join(runsubdir, INICOND_FILENAME)
if os.path.exists(init_file):
os.remove(init_file)
with _hdf5_lock:
ds_ini.to_netcdf(init_file)
del self.dict_inicond_dataout[ddi]
return
[docs]
def fetch_meteo_inicond_file(self, runsubdir):
"""Symlink or copy the IFS meteorological initial-condition file into the run directory.
Checks that ``self.meteo_inicond_file`` exists, links it as
``meteo_inicond.nc`` in *runsubdir*, and opens it as an xarray Dataset
for in-place species modification.
Args:
self: ICON-ART model plugin instance.
runsubdir (str): path to the period run directory.
Returns:
xr.Dataset: the opened meteorological initial-condition dataset.
Raises:
FileNotFoundError: if ``self.meteo_inicond_file`` does not exist.
"""
# Check the meteo inicond file exists
ifs_init_file = os.path.join(self.meteo_inicond_file)
# Copy the file, as we will need to modify it maybe
init_file = os.path.join(runsubdir, INICOND_FILENAME)
shutil.copyfile(ifs_init_file, init_file)
info(f'Copied meteorological inicond file from {ifs_init_file} to {init_file}.')
# Add the surface pressure required by ART
with _hdf5_lock:
ds = xr.open_dataset(init_file)
if 'PS' not in ds:
if 'LNPS' not in ds:
raise CifKeyError(
f"'LNPS' must be found in the initial conditions file {init_file}"
)
ds['PS'] = np.exp(ds['LNPS'])
ds['PS'] = ds['PS'].squeeze(dim=['lev_2']) # Remove the lev 2 dimension
ds['PS'].attrs['standard_name'] = "Surface pressure" # Just change the attributes
ds['PS'].attrs['long_name'] = f'exp(LNPS), generated by cif '
ds['PS'].attrs['units'] = 'Pa'
# Copy topography_c from Extpar in GEOP_SFC
ds_extpar = xr.open_dataset(self.domain.extpar_file)
# ds = ds.drop_vars(['GEOSP'])
ds['GEOP_SFC'][:] = ds_extpar["topography_c"].values * 9.80665
ds['GEOSP'][:] = ds_extpar["topography_c"].values * 9.80665
# Add a Q variable
if 'Q' not in ds:
ds['Q'] = ds['QV']
return ds
