Source code for pycif.plugins.models.TM5.io.inputs.make_inicond

import filecmp
import os
import shutil

import numpy as np
from netCDF4 import Dataset

from ......utils import path
from ......utils.hdf5 import _hdf5_lock
from logging import debug
from ......utils.check.errclass import CifError




def make_inicond(self, datastore, runsubdir, mode, datei, datef):
    """Write or symlink the TM5 initial-condition restart file.

    For the first sub-simulation period only.  Filters *datastore* to
    ``'inicond'`` entries, then for each active species either symlinks the
    original restart file or copies and overwrites the species variable.

    Args:
        self: TM5 model plugin instance.
        datastore (dict): tracer-ID-keyed CIF data-store entries.
        runsubdir (str): path to the period run directory.
        mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
        datei (datetime): period start date.
        datef (datetime): period end date.
    """
    
    datastore = {
        trid: datastore[trid]
        for trid in datastore
        if trid[0] == "inicond"
    }

    # Fixed name for initial condition files
    fileout = f"{runsubdir}/ini.nc"
    fileoutincr = f"{runsubdir}/ini.increment.nc"

    # Loop on all active species
    # If in datastore, take data, otherwise, link to original INI_CONCS
    for spec in self.chemistry.acspecies.attributes:
        trid = ("inicond", spec)
        if trid in datastore:
            pass
        # If spec not explicitly defined in datastore,
        # fetch general component information if available
        elif trid not in datastore and ("inicond", "") in datastore:
            trid = ("inicond", "")
        else:
            continue

        tracer = datastore[trid]
        trcr_data = tracer["data"][datei]
        dirorig = tracer["dirorig"]
        fileorig = tracer["fileorig"]
        fileini = f"{dirorig}/{fileorig}"

        # If no data is provided, just copy from original file
        if "spec" not in trcr_data:
            # If does not exist, just link
            linked = False
            if not os.path.isfile(fileout):
                path.link(fileini, fileout)
                linked = True

            # Otherwise, check for difference
            if not linked:
                with _hdf5_lock:
                    with Dataset(fileini, "r") as ds:
                        if spec not in ds.variables:
                            debug(f"{spec} is not accounted for in initial conditions")
                            continue

                    with Dataset(fileout, "r") as ds:
                        if spec not in ds.variables:
                            raise CifError(
                                f"I need info for {spec} in initial conditions but could not find any"
                            )

            # Repeat operations for tangent linear
            if mode != "tl":
                continue

            # If does not exist, just link
            if not os.path.isfile(fileoutincr):
                shutil.copy(fileini, fileoutincr)

            with _hdf5_lock:
                with Dataset(fileoutincr, "a") as fout:
                    if spec in fout.variables:
                        fout.variables[spec][:] = 0.0

        else:
            # Replace existing link by copy
            # of original file to modify it
            path.copyfromlink(fileout)

            # Write initial conditions
            ini_fwd = trcr_data["spec"]
            self.inicond.write(spec, fileout, ini_fwd, comp_type="inicond")

            if mode == "tl":
                path.copyfromlink(fileoutincr)
                ini_tl = trcr_data.get("incr", 0.0 * ini_fwd)
                self.inicond.write(
                    spec, fileoutincr, ini_tl, comp_type="inicond"
                )