Source code for pycif.plugins.modes.co2mvs.execute

import numpy as np
import copy
from logging import debug, info
import pandas as pd
from ....utils.datastores.dump import dump_datastore

from ....utils.classes.setup import Setup

from .utils.forward import run_forward
from .utils.adjoint import run_adjoint




def execute(self, **kwargs):
    """Runs the CO2MVS mode at a given outer loop

    Args:
        setup (Plugin): definition of the full set-up

    """

    # Working directory
    workdir = self.workdir

    # Observation operator
    obsoper = self.obsoperator

    # Control vector
    controlvect = obsoper.controlvect

    # Simulation window
    datei = self.datei
    datef = self.datef
    
    # Dump the full prior control vector
    controlvect.dump(
            f"{workdir}/controlvect_prior.pickle",
            to_netcdf=controlvect.save_out_netcdf or self.save_out_netcdf,
            dir_netcdf=f"{workdir}/controlvect_prior/",
        )
    
    # Loop over the windows
    for window_datei in self.list_windows_datei:
        window_datef = window_datei + pd.to_timedelta(self.window_length)
        
        info(f"Running CO2MVS window from {window_datei} to {window_datef}")
        
    
            
    print(__file__)
    import code
    code.interact(local=dict(locals(), **globals()))
    
    
    # Some verbose
    info("Running an outer loop of the CO2MVS system")
    
    # Save sensitivities for later use
    ref_dy = copy.deepcopy(self.obsvect.dy)
    
    # First run a reference forward
    obsvect = obsoper.obsoper(
        controlvect, obsoper.obsvect,
        "fwd", datei=datei, datef=datef, workdir=workdir,
        run_id="ref_forward",
        reload_results=self.reload_results,
        **kwargs
    )
    
    # Put dy back for backward propagation
    obsvect.dy[:] = ref_dy[:]
    
    # Now runs the backward propagation
    controlvect = obsoper.obsoper(
        controlvect,
        obsvect,
        "adj",
        datei=datei,
        datef=datef,
        workdir=self.workdir,
        run_id="backward_propagation",
        reload_results=self.reload_results,
        **kwargs,
    )
    
    # Updating the control vector by applying the B matrix
    controlvect.x += controlvect.sqrtbprod_ad(
        controlvect.dx, compute_sqrt=False, **kwargs)
    controlvect.dump(
            f"{workdir}/controlvect_final.pickle",
            to_netcdf=controlvect.save_out_netcdf or self.save_out_netcdf,
            dir_netcdf=f"{workdir}/controlvect/",
        )
    
    # Running updated forward model
    obsvect = obsoper.obsoper(
        controlvect, obsoper.obsvect,
        "fwd", datei=datei, datef=datef, workdir=workdir,
        run_id="posterior_trajectory",
        reload_results=self.reload_results,
        **kwargs
    )

    return obsvect