Source code for pycif.plugins.modes.response_functions.ref_forward

import datetime
import os
from logging import info
from typing import Any, Dict, Tuple

import pandas as pd

# Aliases for type hinting
Mode = Any




def run_ref_forward(self: Mode) -> str:
    """Run the reference forward and dumps the observation vector

    Args:
        self (Mode): the mode plugin

    Returns:
        str: path to the reference forward directory
    """
    if not self.dryrun:
        info("Running reference forward simulation with "
             f"'{self.model.plugin.name}' model")

        # Running reference forward run (dumps the obsvect)
        self.obsvect = self.obsoperator.obsoper(
            self.controlvect, self.obsvect, "fwd",
            datei=self.datei, datef=self.datef,
            workdir=self.workdir, reload_results=self.reload_results
        )

        info("The forward simulation has been successfully executed")
        return self.obsoperator.ref_fwd_dir

    else:
        self.obsvect.dump(self.dir_obsvect)

        # Dummy value
        info("Not running reference forward mode simulation (dryrun)")
        return os.path.join(self.workdir, "obsoperator/fwd_0000")





def get_inicond_from_ref_forward(
    self: Mode,
    datei: datetime.datetime
) -> Tuple[str, Dict[str, Any]]:
    """Get the path to the reference forward run initial conditions (forward)
    file for running a response functions starting on datetime 'datei, and
    returns the model initial condition component in the YAML configuration file
    and the data to use the initial conditions file

    Used for response functions in tangent mode on a period that is simulated
    a sub-period of the full simulation period

    The model used needs to be implemented here, be carefull to keep all control
    vector arguments when implementing a new model here.

    Args:
        self (Mode)
        datei (datetime.datetime): start date of the response function simulation period

    Returns:
        str, dict: component of the model initial condition paragraph in the YAML configuration,
            and data containing the path to the reference forward run initial conditions file
    """
    if not hasattr(self, 'ref_fwd_dir'):
        raise RuntimeError("the reference forward simulation has not been run")

    model_name = self.model.plugin.name
    model_version = self.model.plugin.version

    yaml_file = self.reference_instances["reference_setup"].def_file
    yaml_dict = self.from_yaml(yaml_file)

    # Add cases to the following "if, elif, ..., else" statement to handle new models

    # LMDZ std and LMDZ acc
    if model_name == "LMDZ" and model_version in ("std", "acc"):
        # Getting initial condition file path
        datei = pd.to_datetime(datei) - pd.offsets.MonthBegin()
        dirname = os.path.join(self.ref_fwd_dir, "chain")
        filename = datei.strftime("restart_%Y%m%d0000.nc")

        if not self.dryrun and not os.path.isfile(os.path.join(dirname, filename)):
            raise FileNotFoundError(
                f"file '{filename}' not found in '{dirname}'")

        # Making YAML configuration
        orig_dict = yaml_dict['datavect']['components']['inicond']
        out_dict = {
            'plugin': {'type': "field", 'name': "LMDZ", 'version': "ic"},
            'parameters': {}
        }

        # Keeping all original YAML configuration except:
        # 'plugin', 'dir', 'file', 'varname', 'restart_id'
        # in order to keep all control vector arguments
        keys_to_drop = ['plugin', 'dir', 'file', 'varname', 'restart_id']

        for spec in self.model.chemistry.acspecies.attributes:
            spec_plg = getattr(self.model.chemistry.acspecies, spec)
            restart_id = spec_plg.restart_id

            out_dict['parameters'][spec] = {
                'dir': dirname,
                'file': filename,
                'restart_id': restart_id,
                **{key: val
                   for key, val in orig_dict['parameters'][spec].items()
                   if key not in keys_to_drop}
            }

        return 'inicond', out_dict

    else:
        raise NotImplementedError(
            f"model '{model_name} / {model_version}' is not implemented in "
            "function 'get_inicond_from_ref_fwd'"
        )