Source code for pycif.plugins.modes.response_functions.base_function.yaml_config

import os
from typing import Any, Union

from .....utils.yml import ordered_dump
from ..ref_forward import get_inicond_from_ref_forward
from .base_function import BaseFunction
from .base_function_batch import BaseFunctionSamplingBatch

# Aliases for type hinting
Mode = Any
ControlVect = Any




def dump_yaml_config(
    self: Mode,
    base_function: Union[BaseFunction, BaseFunctionSamplingBatch],
) -> None:
    """Generates and dump a YAML configuration file for a base function

    Args:
        self (Mode): the mode plugin
        base_function (BaseFunction or BaseFunctionSamplingBatch):
            the base function or sampling batch of base functions 
    """
    # Creating directory
    os.makedirs(base_function.rundir, exist_ok=True)

    # Getting base configuration
    yaml_file = self.reference_instances["reference_setup"].def_file
    yaml_dict = self.from_yaml(yaml_file)

    # Updating configuration
    yaml_dict.update({
        'workdir': base_function.rundir,
        'datei': base_function.date_start,
        'datef': base_function.date_end,
        'mode': {
            'plugin': {'name': "forward", 'version': "std"},
            'run_mode': self.run_mode,
            'use_xb': False
        }
    })

    if self.use_batch_sampling:
        assert isinstance(base_function, BaseFunctionSamplingBatch)

        # Updating observation operator configuration
        yaml_dict['obsoperator'].update({
            'plugin': {'name': "standard", 'version': "std"},
            'batch_computation': {
                'nsamples': base_function.n_samples,
                'dir_samples': base_function.rundir,
                'file_samples': base_function.controlvect_filename,
                'dont_propagate': base_function.dont_propagate,
            }
        })

        if self.independant_parameters:
            yaml_dict['model']['dont_perturb_species'] = \
                [param for _, param in base_function.independant_obs]

            yaml_dict['obsoperator']['batch_computation'][
                'dont_propagate_obsvect'] = base_function.independant_obs

            # Removing independant obs species from approximated operator check
            if 'approx_operator' in yaml_dict['model']:
                for _, spec in base_function.independant_obs:
                    if spec in yaml_dict['model']['approx_operator']['species_threshold']:
                        yaml_dict['model']['approx_operator'][
                            'species_threshold'].pop(spec)

    if self.use_model_approximation:
        yaml_dict['obsoperator'].update({'ref_fwd_dir': self.ref_fwd_dir})

    # Updating observation vector configuration
    yaml_dict['obsvect'].update({
        'plugin': {'name': "standard", 'version': "std"},
        'dir_obsvect': self.dir_obsvect,
        'dump_type': "nc"
    })

    # Updating control vector configuration
    yaml_dict['controlvect'].update({
        'plugin': {'name': "standard", 'version': "std"},
        'reload_xb': True,
        'reload_file': base_function.controlvect_path
    })

    if self.run_mode == "tl" and base_function.date_start != self.datei:
        # Initial datetime is not the main initial datetime, using initial
        # condition file from the reference forward simulation
        component, data = get_inicond_from_ref_forward(
            self, base_function.date_start)
        yaml_dict['datavect']['components'][component] = data

    # Writing configuration file
    with open(base_function.yaml_config_path, "w") as f:
        ordered_dump(f, yaml_dict)