Source code for pycif.plugins.models.iconart.io.inputs.restart_inicond

import os
from ......utils import path
from ......utils.hdf5 import _hdf5_lock
import xarray as xr
import re
        


def make_restart_inicond(self, datastore, ddi, ddf, runsubdir, mode):
    """Link or modify the ICON-ART EnSRF restart initial-condition file.

    Used in ensemble Kalman smoother (EnSRF) mode to initialise ICON-ART
    from a CIF-modified restart file rather than the standard
    ``meteo_inicond.nc``.  Reads the source restart dataset, applies the
    CIF tracer modifications, and writes the result to *runsubdir*.

    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.
        runsubdir (str): path to the period run directory.
        mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
    """
    
    trids = list(datastore.keys())
    fileorig = datastore[trids[0]]["fileorig"]
    dirorig = datastore[trids[0]]["dirorig"]
    trid = trids[0]
    
    # Get the restart dataset
    with _hdf5_lock:
        ds_restart = xr.open_dataset(os.path.join(dirorig, fileorig))
    ds_restart.attrs["cif_reset_mean"] = "False"

    # Get the last output file of the previous simulation that 
    # is to be used as first output of this simulation
    self.last_runsubdir = ds_restart.attrs["cif_runsubdir"]
    last_chaindir = os.path.join(self.last_runsubdir, '..', 'chain')
    last_output_filename = os.path.join(
        last_chaindir,
        ddi.strftime('OUTPUT__%Y%m%dT%H%M%SZ.nc')
    )
    with _hdf5_lock:
        ds_last_output = xr.open_dataset(last_output_filename)
    ds_last_output.attrs["cif_reset_mean"] = "False"

    # Open the restart file and the last output file
    # and set all tracer concentrations to those of the mean tracer
    if self.ensrf_same_restart:

        spec = trid[1]

        # Restart file
        pattern = f'^{spec}-\\d+.TL\\d$'
        specs_restart = [var for var in list(ds_restart.data_vars) 
                        if re.match(pattern, var)]
        specs_restart = sorted(specs_restart)
        for sr in specs_restart[3:]:
            ds_restart[sr][:] = ds_restart[specs_restart[2]].values
        ds_restart.attrs["cif_reset_mean"] = "True"

        # Last output file
        pattern = f'^{spec}-\\d+$'
        specs_last_output = [var for var in list(ds_last_output.data_vars) 
                        if re.match(pattern, var)]
        specs_last_output = sorted(specs_last_output)
        for sr in specs_last_output[3:]:
            ds_last_output[sr][:] = ds_last_output[specs_last_output[2]].values
        ds_last_output.attrs["cif_reset_mean"] = "True"

    # Dump the restart file and link it to the ICON restart file
    source_file = f"{runsubdir}/restart.nc"
    target_file = f"{runsubdir}/restart_ATMO_DOM01.nc"
    if os.path.exists(source_file):
        os.remove(source_file)
    with _hdf5_lock:
        ds_restart.to_netcdf(source_file)
    path.link(source_file, target_file)

    # Dump the last output file
    if os.path.exists(last_output_filename):
        os.remove(last_output_filename)
    with _hdf5_lock:
        ds_last_output.to_netcdf(last_output_filename)