Source code for pycif.plugins.models.lmdz_old.io.inputs.chemfields

import os
import shutil

import xarray as xr

from ......utils import path
from ......utils.check.errclass import CifValueError




def make_chemfields(self, datastore, input_type, ddi, ddf, runsubdir, mode):
    """Write a chemical concentration field for LMDZ-old.

    Reads the CIF data for *input_type* and writes it as a NetCDF input
    file in *runsubdir*.  If no CIF modification, symlinks the original.

    Args:
        self: LMDZ-old model plugin instance.
        datastore (dict): tracer-ID-keyed CIF data-store entries.
        input_type (str): chemistry field type (e.g. ``'prescrconcs'``).
        ddi (datetime): period start date.
        ddf (datetime): period end date.
        runsubdir (str): path to the period run directory.
        mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
    """

    if not hasattr(self, 'chemistry'):
        return

    chemistry_input = getattr(self.chemistry, input_type, None)
    if chemistry_input is None:
        return

    input_plugin = getattr(self, input_type)

    if input_type == "prodloss3d":
        bin_name = "prodscale"
        nc_name = "prodloss"
    elif input_type == "prescrconcs":
        bin_name = "scale"
        nc_name = "prescr"
    else:
        raise CifValueError(f"Unknown input_type '{input_type}'")

    for spec in chemistry_input.attributes:
        trid = (input_type, spec)

        if trid not in datastore:
            continue

        data = datastore[trid]['data'][ddi]
        target_path = os.path.join(runsubdir, f"{nc_name}_{spec}.nc")

        # if species is in the control vector
        if 'spec' in data:
            _, var_name = trid
            ds = xr.Dataset({var_name: data['spec']})
            input_plugin.write(var_name, target_path, data['spec'])

            if datastore[trid]['tracer'].iscontrol:
                if 'incr' not in data or mode != "tl":
                    data['incr'] = 0.0 * data['spec']

                # Put in dataset for writing by 'write'
                ds = xr.Dataset({'fwd': data['spec'], 'tl': data['incr']})

                # Write to FORTRAN binary
                target_path = os.path.join(runsubdir, f"mod_{bin_name}_{spec}.bin")
                input_plugin.write("", target_path, ds)

        else:
            # Links reference NetCDF files that are needed anyway by LMDZ
            dirorig = datastore[trid]["dirorig"]
            fileorig = datastore[trid]["fileorig"]

            if fileorig is None or dirorig is None:
                tracer = getattr(chemistry_input, spec)
                dirorig = tracer.dir
                fileorig = tracer.file

            origin_path = ddi.strftime(os.path.join(dirorig, fileorig))

            if input_type in self.copy_inputs:
                # Follow the symlink for 'origin_path' and copy it to 'target_path'
                origin_path = os.path.realpath(origin_path)
                shutil.copy(origin_path, target_path)
            else:
                # Simply link the file
                path.link(origin_path, target_path)





def make_kinetic(self, datastore, input_type, ddi, ddf, runsubdir, mode):
    """Write the LMDZ-old kinetic (pressure/temperature) field for one period.

    Reads the CIF kinetic data-store and writes ``kinetic.nc`` in *runsubdir*.

    Args:
        self: LMDZ-old model plugin instance.
        datastore (dict): kinetic CIF data-store entries.
        input_type (str): unused; kept for API consistency.
        ddi (datetime): period start date.
        ddf (datetime): period end date.
        runsubdir (str): path to the period run directory.
        mode (str): ``'fwd'``, ``'tl'``, or ``'adj'``.
    """
    target_path = os.path.join(runsubdir, "kinetic.nc")
    for trid in datastore:
        data = datastore[trid]["data"][ddi]

        if "spec" in data:
            _, var_name = trid
            self.kinetic.write(var_name, target_path, data['spec'])

        else:
            input_file_list = list(set(datastore[trid]["input_files"][ddi]))
            if len(input_file_list) != 1:
                raise CifValueError(
                    "There is not a single file to link "
                    "for the kinetic file of LMDZ."
                )

            (input_file,) = input_file_list

            if 'kinetic' in self.copy_inputs:
                # Follow the symlink for 'input_file' and copy it to 'target_path'
                input_file = os.path.realpath(input_file)
                shutil.copy(input_file, target_path)
            else:
                # Simply link the file
                path.link(input_file, target_path)