fwd.py

import copy
import os
import pandas as pd
import pytest
import shutil
import xarray as xr
from pycif.utils.yml import ordered_dump


@pytest.fixture(
    params=
    [
        {"perturb_CH4": True, "datef": "2010-01-02", "no_MCF": True},
        {"perturb_CH4": False, "datef": "2010-01-02", "no_MCF": True},
        {"perturb_CH4": True, "datef": "2010-01-05", "no_MCF": True},
        {"perturb_CH4": False, "datef": "2010-01-05", "no_MCF": True},
        {"perturb_CH4": True, "datef": "2010-01-02"},
        {"perturb_CH4": False, "datef": "2010-01-02"},
        pytest.param({"perturb_CH4": True, "datef": "2010-01-05"},
                     marks=pytest.mark.ref_config),
        pytest.param({"perturb_CH4": False, "datef": "2010-01-05"},
                     marks=pytest.mark.inversion)
    ])
def dummy_config_fwd(ref_dummy, request):
    """
    Fixture that returns a temporary folder and a configuration file for the
    dummy_forward model.
    """
    tmpdir_str, config = ref_dummy

    # Adding perturbation to reference fluxes
    if request.param.get("perturb_CH4", False):
        params = config["datavect"]["components"]["fluxes"]["parameters"]
        params["CH4_ref"] = copy.copy(params["CH4"])
        params["CH4_perturb"] = copy.copy(params["CH4"])
        params["CH4_perturb"].update(
            {
                "xb_scale": 0.2,
                "flx_formula":
                    [{"sum":
                        [
                            {"cos": None,
                             "variable": "zlat",
                             "period": 200},
                            {"sin": None,
                             "variable": "zlon",
                             "period": 300}]}]
            }
        )
        del params["CH4"]
    
        # Adding families in controlvect
        config["controlvect"] = {
            "plugin": {"name": "standard", "version": "std"},
            "save_out_netcdf": True,
            "transform_pipe": {
                "trans1": {
                    "plugin": {"name": "families",
                               "version": "std",
                               "type": "transform"},
                    "component": "fluxes",
                    "parameters_in": ["CH4_ref", "CH4_perturb"],
                    "parameter_out": "CH4"
                }
            }
        }
    
    # Removing MCF if requires
    if request.param.get("no_MCF", False):
        del config["model"]["chemistry"]["acspecies"]["MCF"]
        del config["datavect"]["components"]["fluxes"]["parameters"]["MCF"]
        del config["datavect"]["components"]["concs"]["parameters"]["MCF"]
    
    # Changing end date
    if "datef" in request.param:
        config["datef"] = pd.to_datetime(request.param["datef"]).to_pydatetime()
    
    # Tag for config file to dump in examples
    tag = ("long" if "datef" in request.param else "short") \
          + ("_families" if request.param.get("perturb_CH4", False) else "") \
          + ("_noMCF" if request.param.get("no_MCF", False) else "")
    
    yield tmpdir_str, config, tag
    
    # Change config to be compatible in the docker
    config["workdir"] = "/workdir/"
    config["model"]["file_pg"] = \
        "/tmp/CIF/model_sources/dummy_gauss/Pasquill-Gifford.txt"
    
    # Dump configuration into CIF examples
    current_dir = os.path.abspath(os.path.dirname(os.path.realpath(__file__)))
    example_dir = \
        os.path.abspath(os.path.join(current_dir, "../../../examples/dummy/"))
    
    dummy_config_file = \
        os.path.join(example_dir, "config_fwd_{}.yml".format(tag))
    
    # Write only if mode is forward
    if config["mode"]["plugin"]["name"] == "forward":
        with open(dummy_config_file, "w") as outfile:
            ordered_dump(outfile, config)
    
        # Copy simulated observation vector for reading by inversions
        pytest_dir = tmpdir_str + "/../"
        target_dir = pytest_dir + "ref_obsvect"
        if os.path.isdir(target_dir):
            shutil.rmtree(target_dir)
    
        shutil.copytree(
            "{}/obsvect".format(tmpdir_str),
            target_dir
        )
        
        # Copy the H matrix to avoid re-compiling it every time
        Hfile_ref = "{}/model/H_matrix.pickle".format(tmpdir_str)
        Hfile_target = "{}/../H_matrix.pickle".format(tmpdir_str)
        
        if os.path.isfile(Hfile_target):
            shutil.rmtree(Hfile_target)
            
        shutil.copy(Hfile_ref, Hfile_target)