test_integration_3inversions.py

import os
import time
import shutil
import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
from pathlib import Path

import pytest

from pycif.utils.classes.setup import Setup
from pycif.utils.datastores.dump import read_datastore
from pycif.utils.yml import ordered_dump
from pycif.utils.path import init_dir


@pytest.mark.dummy
@pytest.mark.article
@pytest.mark.parametrize(
    "settings", [
        {"mode": "4dvar", "minimizer": "M1QN3"},
        pytest.param({"mode": "4dvar", "minimizer": "M1QN3", "montecarlo": 10},
                     marks=pytest.mark.uncertainties),
        {"mode": "4dvar", "minimizer": "congrad"},
        {"mode": "ensrf"},
        pytest.param({"mode": "ensrf", "nsample": 5},
                     marks=pytest.mark.uncertainties),
        {"mode": "analytical"}
    ]
)
def test_integration_inversion(dummy_config_inversion, settings, pytestconfig):
    """
    Integration test that runs the dummy_forward model.
    """

    tmpdir, config, tag = dummy_config_inversion
    
    # Force reloading observation operator
    # from forward to make computation faster
    config["model"]["reload_H"] = "{}/../H_matrix.pickle".format(tmpdir)
    
    # Update number of simulations depending on pytests options
    marker = pytestconfig.getoption('-m')
    nsimmax = 10
    montecarlo = settings.get("montecarlo", 10)
    nsample = settings.get("nsample", 5)
    if "allsimulations" not in marker:
        nsimmax = 3
        montecarlo = 3
        nsample = 1
        
    # Changing mode
    if config["datavect"]["components"]["flux"]\
            ["parameters"]["CH4"]["hresol"] == "hpixels":
        nsimmax = int(2.5 * nsimmax)
    
    elif config["datavect"]["components"]["flux"]\
            ["parameters"]["CH4"]["hresol"] == "global" \
            and settings.get("minimizer", "") == "congrad":
        nsimmax = 1
        
    if settings["mode"] == "4dvar":
        mode = {
            "plugin": {"name": "4dvar", "version": "std"},
            "minimizer": {
                "plugin": {"name": settings["minimizer"], "version": "std"},
                "simulator": {
                    "plugin": {"name": "gausscost", "version": "std"},
                    "reload_from_previous": True
                },
                "maxiter": nsimmax,
                "nsim": nsimmax,
                "epsg": 0.0002,
                "df1": 0.5
            },
            "save_out_netcdf": True
        }
        if settings["minimizer"] == "congrad":
            mode["minimizer"]["save_uncertainties"] = True
        
        if "montecarlo" in settings:
            mode["montecarlo"] = {
                "nsample": settings["montecarlo"],
                "perturb_x": True,
                "perturb_y": False
            }
        
        prior_dir = os.path.join(tmpdir, "obsoperator/fwd_-001/obsvect")
        posterior_dir = os.path.join(tmpdir, "obsoperator/fwd_-002/obsvect")
    
    elif settings["mode"] == "analytical":
        mode = {"plugin": {"name": "analytic", "version": "std",
                           "dump_nc_base_control": True}}
        prior_dir = os.path.join(tmpdir, "obsvect_prior")
        posterior_dir = os.path.join(tmpdir, "obsvect_posterior")
    
    elif settings["mode"] == "ensrf":
        nsimmax = 2 * settings.get("nsample", 1) * nsimmax
        mode = {"plugin": {"name": "EnSRF", "version": "std"},
                "nsample": nsimmax}
        prior_dir = os.path.join(tmpdir, "obsvect_prior")
        posterior_dir = os.path.join(tmpdir, "obsvect_posterior")
      
    config["mode"] = mode
    config["platform"] = {
        "plugin":
            {"name": "LSCE", "version": "obelix"}
            if os.getenv("PYCIF_PLATFORM") == "LSCE"
            else {"name": "docker", "version": "cif"}
    }

    # Save ID for later plot of cost function
    resolution = \
        config["datavect"]["components"]["flux"]["parameters"]["CH4"][
            "hresol"]
    sigma = config["datavect"]["components"]\
                ["flux"]["parameters"]["CH4"]["hcorrelations"]["sigma"]
    correlations = "lowcorr" if sigma == 500 else "highcorr"
    test_id = (settings["mode"], settings.get("minimizer", ""),
               resolution, correlations)

    pytest.test_ids[test_id] = (tmpdir, nsimmax)
    
    # Removing MCF to make computation quicker
    del config["model"]["chemistry"]["acspecies"]["MCF"]
    del config["datavect"]["components"]["flux"]["parameters"]["MCF"]
    del config["datavect"]["components"]["concs"]["parameters"]["MCF"]
    
    # Dump yml config file
    dummy_config_file = os.path.join(tmpdir, "dummy_config.yml")
    with open(dummy_config_file, "w") as outfile:
        ordered_dump(outfile, config)
    
    # Run as job
    setup = Setup.yaml_to_setup(dummy_config_file)
    setup.todo_init = ["platform"]
    setup = Setup.load_config(setup)
    platform = setup.platform

    job_file = os.path.join(tmpdir,
                            "job_pycif_{}".format(tmpdir[-5:]))

    exe = "{} -m pycif {}".format(platform.python, dummy_config_file)

    job_id = platform.submit_job(
        exe,
        job_file
    )
    
    # Wait for the end of the execution
    while not platform.check_jobs([job_id]):
        time.sleep(platform.sleep_time)
    
    # Get results and check that posterior closer to prior
    components = os.listdir(prior_dir)
    for comp in components:
        comp_dir_prior = os.path.join(prior_dir, comp)
        comp_dir_post = os.path.join(posterior_dir, comp)
        parameters = os.listdir(comp_dir_prior)
        for param in parameters:
            param_dir = os.path.join(comp_dir_prior, param)
            monitor_prior = \
                read_datastore(os.path.join(param_dir, "monitor.nc"))["maindata"]

            param_dir = os.path.join(comp_dir_post, param)
            monitor_post = \
                read_datastore(os.path.join(param_dir, "monitor.nc"))["maindata"]
            assert (
                    (monitor_prior.loc[:, "obs"]
                     - monitor_prior.loc[:, "sim"]).pow(2).sum()
                    - (monitor_post.loc[:, "obs"]
                       - monitor_post.loc[:, "sim"]).pow(2).sum()
                    > 0
            )

    # Dump configuration into CIF examples
    tag += "_{}_{}".format(settings["mode"],
                           "" if settings["mode"] != "4dvar"
                           else settings["minimizer"])

    current_dir = os.path.abspath(os.path.dirname(os.path.realpath(__file__)))
    root_dir = os.path.abspath(os.path.join(current_dir, "../../"))
    pytest_dir = os.path.abspath(tmpdir + "/../")
    example_dir = \
        os.path.abspath(os.path.join(root_dir, "examples_artifact/dummy/"))
    Path(example_dir).mkdir(parents=True, exist_ok=True)
    
    config["workdir"] = "{}/inversion_{}/".format(pytest_dir, tag)

    dummy_config_file = \
        os.path.join(example_dir, "config_inversion_{}.yml".format(tag))
    with open(dummy_config_file, "w") as outfile:
        ordered_dump(outfile, config,
                     ref_directories={"outdir": pytest_dir,
                                      "rootdir": root_dir},
                     replace_values={"rootdir": "/tmp/CIF/"})
    

    # Stop here if not full article computation
    marker = pytestconfig.getoption('-m')
    if "allsimulations" not in marker:
        return
    
    # Loop with different number of simulations for non analytical inversions
    title = ""
    obs_root = "obsvect_posterior"
    control_root = "controlvect"
    if settings["mode"] != "analytical":
        for nsim in range(1, nsimmax + 1, 2):
            file_config = "{}/dummy_config.yml".format(tmpdir)
            inv_setup = Setup.from_yaml(file_config)
            
            if settings.get("minimizer", None) == "congrad":
                title = "CONGRAD"
                obs_root = "obsoperator/fwd_-002/obsvect"
                control_root = "obsoperator/fwd_-002/controlvect"
                inv_setup["mode"]["minimizer"]["maxiter"] = \
                    max(1, int(nsim))
                nsim *= 2
                if os.path.isdir("{}/obsoperator/fwd_-002".format(tmpdir)):
                    shutil.rmtree("{}/obsoperator/fwd_-002".format(tmpdir))
            elif settings.get("minimizer", None) == "M1QN3":
                title = "M1QN3"
                obs_root = "obsoperator/fwd_-002/obsvect"
                control_root = "obsoperator/fwd_-002/controlvect"
                inv_setup["mode"]["minimizer"]["maxiter"] = \
                    max(1, int(nsim))
                inv_setup["mode"]["minimizer"]["nsim"] = \
                    max(1, int(nsim))
                nsim *= 2
                if os.path.isdir("{}/obsoperator/fwd_-002".format(tmpdir)):
                    shutil.rmtree("{}/obsoperator/fwd_-002".format(tmpdir))
            elif settings["mode"] == "ensrf":
                title = "EnSRF"
                inv_setup["mode"]["nsample"] = nsim
            
            # Execute a partial inversion with less simulations
            inv_setup = Setup.from_dict(inv_setup, convert_none=True)
            inv_setup = inv_setup.load_config(inv_setup)
            controlvect, obsvect = inv_setup.mode.execute()
            
            # Compute the cost function
            departures = obsvect.ysim - obsvect.yobs
            j_o = 0.5 * (departures * obsvect.rinvprod(departures)).sum()
            
            dx = controlvect.x - controlvect.xb
            bfull = np.linalg.inv(controlvect.build_b(inv_setup.controlvect))
            j_b = dx[np.newaxis, :].dot(bfull.dot(dx[:, np.newaxis])).sum() / 2
            
            # Dump to text file for later plot
            with open("{}/varying_cost_function.txt".format(tmpdir), "a") as f:
                f.write("{} {} {}\n".format(nsim, j_o, j_b))
    
    else:
        title = "Analytical"
        
        # Load results
        file_config = "{}/dummy_config.yml".format(tmpdir)
        inv_setup = Setup.from_yaml(file_config)
        inv_setup = Setup.from_dict(inv_setup, convert_none=True)
        inv_setup = inv_setup.load_config(inv_setup)
        controlvect, obsvect = inv_setup.mode.execute()
        
        # Compute the cost function
        departures = obsvect.ysim - obsvect.yobs
        j_o = 0.5 * (departures * obsvect.rinvprod(departures)).sum()
        
        dx = controlvect.x - controlvect.xb
        bfull = np.linalg.inv(
            inv_setup.controlvect
                .build_b(inv_setup.controlvect)
        )
        j_b = dx[np.newaxis, :].dot(bfull.dot(dx[:, np.newaxis])).sum() / 2
        
        # Dump to text file for later plot
        with open("{}/varying_cost_function.txt".format(tmpdir), "a") as f:
            f.write("{} {} {}\n".format(0, j_o, j_b))
            f.write("{} {} {}\n".format(nsimmax, j_o, j_b))
    
    marker = pytestconfig.getoption('-m')
    if "article" in marker and "not article" not in marker:
        # Domain limits
        xmin = inv_setup.domain.xmin
        xmax = inv_setup.domain.xmax
        ymin = inv_setup.domain.ymin
        ymax = inv_setup.domain.ymax
    
        # Read observations
        file_obs = "{}/obsvect/concs/CH4/monitor.nc".format(tmpdir)
        monitor_ref = read_datastore(file_obs)["metadata"]
        coords = monitor_ref.loc[:, ["lon", "lat", "alt"]].drop_duplicates()
        
        # Compute fluxes from control vector
        file_flx = "{}/{}/flux/" \
                   "controlvect_flux_CH4.nc".format(tmpdir, control_root)
        ds = xr.open_dataset(file_flx)
        dflx = ds["x_phys"].mean(axis=(0, 1)) - ds["xb_phys"].mean(axis=(0, 1))
        dx = ds["x"].mean(axis=(0, 1)) - ds["xb"].mean(axis=(0, 1))
        
        # Fetch resolution for figure name
        resol = ""
        if config["datavect"]["components"][
                "flux"]["parameters"]["CH4"]["hresol"] == "hpixels":
            if config["datavect"]["components"][
                    "flux"]["parameters"]["CH4"][
                    "hcorrelations"]["sigma"] == 500:
                resol = "lowcorr"
            else:
                resol = "highcorr"
    
        elif config["datavect"]["components"]["flux"] \
                ["parameters"]["CH4"]["hresol"] == "ibands":
            resol = "bands"
    
        elif config["datavect"]["components"]["flux"] \
                ["parameters"]["CH4"]["hresol"] == "global":
            resol = "global"
        
        # Plot the figure of physical fluxes increments
        plt.figure(figsize=(21, 11))
        
        ax0 = plt.axes([0.05, 0.05, 0.73, 0.87])
        im = plt.imshow(dflx, extent=(xmin, xmax, ymin, ymax), cmap="YlOrRd",
                        vmin=-0.2, vmax=0.5)
        sc = plt.scatter(coords["lon"], coords["lat"], c=coords["alt"],
                         cmap="Blues", linewidths=1, edgecolors="k", s=600)
        plt.yticks(fontsize=25)
        plt.xticks(fontsize=25)
        
        ax1 = plt.axes([0.74, 0.05, 0.05, 0.87])
        cb1 = plt.colorbar(im, cax=ax1)
        plt.yticks(fontsize=25)
        plt.ylabel("Fluxes (a.u.)", fontsize=30)
        
        ax2 = plt.axes([0.86, 0.05, 0.05, 0.87])
        cb2 = plt.colorbar(sc, cax=ax2)
        plt.yticks(fontsize=25)
        plt.ylabel("Station altitude (m a.g.l)", fontsize=30)
        
        ax0.set_title(title, fontsize=45)
        
        current_dir = os.path.abspath(
            os.path.dirname(os.path.realpath(__file__)))
        figure_dir = \
            os.path.abspath(os.path.join(current_dir,
                                         "../../figures_artifact/"))
        Path(figure_dir).mkdir(parents=True, exist_ok=True)
        plt.savefig("{}/map_dx_{}_{}_{}.pdf".format(
            figure_dir, title, resol, nsimmax))
        plt.close()
        
        # Plot the figure of scalar fluxes increments
        plt.figure(figsize=(21, 11))
        
        ax0 = plt.axes([0.05, 0.05, 0.73, 0.87])
        im = plt.imshow(dx, extent=(xmin, xmax, ymin, ymax), cmap="YlOrRd")
        sc = plt.scatter(coords["lon"], coords["lat"], c=coords["alt"],
                         cmap="Blues", linewidths=1, edgecolors="k", s=600)
        plt.yticks(fontsize=25)
        plt.xticks(fontsize=25)
        
        ax1 = plt.axes([0.74, 0.05, 0.05, 0.87])
        cb1 = plt.colorbar(im, cax=ax1)
        plt.yticks(fontsize=25)
        plt.ylabel("Fluxes (a.u.)", fontsize=30)
        
        ax2 = plt.axes([0.86, 0.05, 0.05, 0.87])
        cb2 = plt.colorbar(sc, cax=ax2)
        plt.yticks(fontsize=25)
        plt.ylabel("Station altitude (m a.g.l)", fontsize=30)
        
        ax0.set_title(title, fontsize=45)
        
        current_dir = os.path.abspath(
            os.path.dirname(os.path.realpath(__file__)))
        figure_dir = \
            os.path.abspath(os.path.join(current_dir,
                                         "../../figures_artifact/"))
        Path(figure_dir).mkdir(parents=True, exist_ok=True)
        plt.savefig("{}/map_dx_scale_{}_{}_{}.pdf".format(
            figure_dir, title, resol, nsimmax))
        plt.close()
        
        # Plot uncertainty reduction
        if "pa_std" in ds:
            dstd = ds["b_std"].mean(axis=(0, 1)) \
                   - ds["pa_std"].mean(axis=(0, 1))
        
            # Plot the figure
            plt.figure(figsize=(21, 11))
            
            ax0 = plt.axes([0.05, 0.05, 0.73, 0.87])
            im = plt.imshow(dstd, extent=(xmin, xmax, ymin, ymax),
                            cmap="YlOrRd", vmin=0, vmax=1)
            sc = plt.scatter(coords["lon"], coords["lat"], c=coords["alt"],
                             cmap="Blues", linewidths=1, edgecolors="k", s=600)
            plt.yticks(fontsize=25)
            plt.xticks(fontsize=25)
            
            ax1 = plt.axes([0.74, 0.05, 0.05, 0.87])
            cb1 = plt.colorbar(im, cax=ax1)
            plt.yticks(fontsize=25)
            plt.ylabel("Uncertainty reduction (a.u.)", fontsize=30)
            
            ax2 = plt.axes([0.86, 0.05, 0.05, 0.87])
            cb2 = plt.colorbar(sc, cax=ax2)
            plt.yticks(fontsize=25)
            plt.ylabel("Station altitude (m a.g.l)", fontsize=30)
            
            ax0.set_title(title, fontsize=45)
            
            current_dir = os.path.abspath(
                os.path.dirname(os.path.realpath(__file__)))
            figure_dir = \
                os.path.abspath(os.path.join(current_dir,
                                             "../../figures_artifact/"))
            Path(figure_dir).mkdir(parents=True, exist_ok=True)
            plt.savefig("{}/map_dstd_{}_{}_{}.pdf"
                        .format(figure_dir, title, resol, nsimmax))
            plt.close()
        
        # Plot matrix of uncertainty reduction
        if hasattr(controlvect, "pa"):
            pa = controlvect.pa
            
            plt.figure(figsize=(21, 11))
            ax0 = plt.axes([0.05, 0.05, 0.73, 0.87])
            im = plt.imshow(pa[:int(controlvect.dim / 2),
                               :int(controlvect.dim / 2)],
                            vmin=-0.5, vmax=0.5, cmap="RdBu")
            plt.xticks(fontsize=20)
            plt.yticks(fontsize=20)
            
            ax1 = plt.axes([0.74, 0.05, 0.05, 0.87])
            cb1 = plt.colorbar(im, cax=ax1)
            plt.yticks(fontsize=25)
            plt.ylabel("Posterior uncertainties", fontsize=30)

            ax0.set_title(title, fontsize=45)
            
            current_dir = os.path.abspath(
                os.path.dirname(os.path.realpath(__file__)))
            figure_dir = \
                os.path.abspath(os.path.join(current_dir,
                                             "../../figures_artifact/"))
            Path(figure_dir).mkdir(parents=True, exist_ok=True)
            plt.savefig("{}/posterior_matrix_{}_{}_{}.pdf"
                        .format(figure_dir, title, resol, nsimmax),
                        bbox_inches='tight')
            plt.close()
            
            # Prior matrix
            bfull = controlvect.build_b(controlvect)
            plt.figure(figsize=(21, 11))
            ax0 = plt.axes([0.05, 0.05, 0.73, 0.87])
            im = plt.imshow(bfull[:int(controlvect.dim / 2),
                                  :int(controlvect.dim / 2)],
                            vmin=-0.5, vmax=0.5, cmap="RdBu")
            plt.xticks(fontsize=20)
            plt.yticks(fontsize=20)

            ax1 = plt.axes([0.74, 0.05, 0.05, 0.87])
            cb1 = plt.colorbar(im, cax=ax1)
            plt.yticks(fontsize=25)
            plt.ylabel("Prior uncertainties", fontsize=30)

            current_dir = os.path.abspath(
                os.path.dirname(os.path.realpath(__file__)))
            figure_dir = \
                os.path.abspath(os.path.join(current_dir,
                                             "../../figures_artifact/"))
            Path(figure_dir).mkdir(parents=True, exist_ok=True)
            plt.savefig("{}/prior_matrix_{}_{}_{}.pdf"
                        .format(figure_dir, title, resol, nsimmax),
                        bbox_inches='tight')
            plt.close()