Source code for pycif.plugins.models.lmdz_ico.io.inputs.inicond

from __future__ import annotations

import datetime
from logging import warning
from os import PathLike
from pathlib import Path
from typing import Any, Literal

import numpy as np

from ......utils import path
from .ensemble import ensemble_trid




def make_inicond(
    self,
    datastore: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
    input_type: Literal["inicond", "restart_inicond"],
    datei: datetime.datetime,
    runsubdir: str | PathLike,
    mode: Literal["fwd", "tl", "adj"],
    onlyinit: bool,
) -> None:
    """Write or symlink the initial-condition restart file"""
    target_path = Path(runsubdir, "start.nc")

    # In adjoint mode (not onlyinit), do nothing as no concentration is needed
    if mode == "adj" and not onlyinit:
        return

    # Generating reference initial conditions if first sub-simulation
    for spec in self.chemistry.active_species:
        trid = ensemble_trid(self, (input_type, spec), datastore)

        if trid not in datastore:
            continue

        data = datastore[trid]["data"][datei]

        if mode in ("fwd", "tl") and "spec" not in data:
            # If restart_inicond restart file is already formatted and contains
            # all species variables -> simply link the restart file
            if input_type == "restart_inicond":
                ref_dir = Path(datastore[trid]["dirorig"])
                ref_file = datastore[trid]["fileorig"]
                source_path = ref_dir / datei.strftime(ref_file)

                path.link(source_path, target_path)
                continue

            # Reading input file with datavect plugin 'read' method
            dates2read = [
                [x.to_pydatetime() for x in row]
                for row in datastore[trid]["input_dates"][datei].itertuples(
                    index=False, name=None
                )
            ]
            _, tracer_name = trid
            input_tracer = datastore[trid]["tracer"]
            da = input_tracer.read(
                tracer_name,
                input_tracer.varname,
                dates=dates2read,
                files=datastore[trid]["input_files"][datei],
                tracer=input_tracer,
            )

            # Writing data
            self.inicond.write(spec, target_path, da)

        # Update data
        elif mode in ("fwd", "tl") and datei == self.datei:
            self.inicond.write(spec, target_path, data["spec"])

            if mode == "tl":
                if "incr" in data:
                    self.inicond.write(f"{spec}_tl", target_path, data["incr"])

                else:
                    warning(
                        "Setting LMDZ tangent-linear initial conditions "
                        f"to zeros for species '{spec}'."
                    )

                    da = data["spec"]
                    da_tl = da.copy(data=np.zeros(da.shape))
                    self.inicond.write(f"{spec}_tl", target_path, da_tl)