forward.py

import numpy as np
import pandas as pd
import xarray as xr
import os
import itertools

from logging import info, debug, warning
from .....utils.path import init_dir
from .....utils.datastores.dump import dump_datastore, read_datastore
from .apply_AK import apply_ak
from .apply_AK import apply_ak_tl
from .vinterp import vertical_interp


def forward(
        transf,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        save_debug=False,
        **kwargs
):
    """Aggregate simulations at the grid scale to total columns.
    Re-interpolate the model pressure levels to the satellite averaging kernel
    levels. Average using the averaging kernel formula

    """

    ddi = min(di, df)

    if onlyinit:
        return

    ref_parameter = transf.parameter[0]
    ref_component = transf.component[0]

    ref_datastore = inout_datastore["outputs"][(ref_component, ref_parameter)]
    ref_ds = ref_datastore[ddi]
    target_datastore = \
        inout_datastore["inputs"][("concs", ref_parameter)][ddi]

    # Put auxiliary values into target datastore
    target_datastore.loc[:, ("metadata", "pressure")] = \
        inout_datastore["inputs"][
            ("pressure", ref_parameter)][ddi][("maindata", "spec")]
    target_datastore.loc[:, ("metadata", "dp")] = \
        inout_datastore["inputs"][
            ("dpressure", ref_parameter)][ddi][("maindata", "spec")]
    if transf.product == "column":
        target_datastore.loc[:, ("metadata", "airm")] = \
            inout_datastore["inputs"][
                ("airm", ref_parameter)][ddi][("maindata", "spec")]
        target_datastore.loc[:, ("metadata", "hlay")] = \
            inout_datastore["inputs"][(
                "hlay", ref_parameter)][ddi][("maindata", "spec")]

    ddi = min(di, df)
    ref_indexes = ~target_datastore.duplicated(subset=[("metadata", "indorig")])
    y0 = target_datastore.loc[ref_indexes]
    
    # Pass infor from previous transforms if debug
    if save_debug:
        # df_debug = {}
        # for trid in inout_datastore["inputs"]:
        #     if trid[0] == "stratosphere":
        #         continue
        #     print(trid, inout_datastore["inputs"][trid][ddi].columns)
        #     for c in inout_datastore["inputs"][trid][ddi].columns:
        #         if c[0] not in y0.columns.get_level_values(level=0):
        #             df_debug[c] = inout_datastore["inputs"][trid][ddi][c].loc[ref_indexes]
        
        df_debug = pd.DataFrame({
            c: inout_datastore["inputs"][trid][ddi][c].loc[ref_indexes]
            for trid in inout_datastore["inputs"]
            if trid[0] != "stratosphere"
            for c in inout_datastore["inputs"][trid][ddi].columns
            if c[0] not in y0.columns.get_level_values(level=0)
        })
        
        y0 = pd.concat([y0, df_debug], axis=1)

    # Exit if empty observation datastore
    if target_datastore.size == 0:
        return

    # Init firectory for saving forward datastore for later use by adjoint
    dir_fwd = ddi.strftime("{}/chain/satellites/{}".format(
        transf.model.adj_refdir, transf.transform_id))
    if not os.path.isdir(dir_fwd):
        init_dir(dir_fwd)

    # Building the extended dataframe
    iq1 = y0["metadata"]["station"]
    list_satIDs = iq1.unique()
    ds_p = target_datastore["metadata"].set_index("indorig")[
        ["pressure", "dp"] + (transf.product == "column") * ["airm", "hlay"]]
    for satID in list_satIDs:
        info("Processing satellite: {}".format(satID))
        satmask = iq1 == satID
        nobs = np.sum(satmask)

        # Stacking output datastore into levels * nobs
        native_ind_stack = (
                np.flatnonzero(ref_indexes)[satmask]
                + np.arange(transf.model.domain.nlev)[:, np.newaxis]
        )

        # If all nans in datasim, meaning that the species was not simulated
        sim = target_datastore.loc[:, ("maindata", "spec")].values[native_ind_stack]
        if not np.any(~np.isnan(sim)):
            continue

        # Grouping all data from this satellite
        datasim = xr.Dataset(
            {
                "pressure": (
                    ["level", "index"],
                    np.log(ds_p["pressure"].values[native_ind_stack]),
                ),
                "dp": (
                    ["level", "index"],
                    ds_p["dp"].values[native_ind_stack],
                ),
                "sim": (["level", "index"], sim),
            },
            coords={
                "index": np.arange(nobs),
                "level": np.arange(transf.model.domain.nlev),
            },
        )

        if transf.product == "column":
            datasim = datasim.assign(
                {
                    "airm": (
                        ["level", "index"],
                        ds_p["airm"].values[native_ind_stack],
                    ),
                    "hlay": (
                        ["level", "index"],
                        ds_p["hlay"].values[native_ind_stack],
                    ),
                }
            )

        if mode == "tl":
            datasim["sim_tl"] = (
                ["level", "index"],
                target_datastore.loc[:, ("maindata", "incr")].values[native_ind_stack],
            )

        # convert CHIMERE fields to the correct unit
        # from ppb to molec.cm-2 if the satellite product is a column
        if transf.product == "column":
            keys = ["sim"] + (["sim_tl"] if mode == "tl" else [])
            for k in keys:
                datasim[k] *= datasim["hlay"] / (1e9 / datasim["airm"])

        # Check whether there is some ak
        ref_mapper = mapper["outputs"][(ref_component, ref_parameter)]
        files_aks = \
            list(set(list(itertools.chain(
                *ref_mapper["tracer"].input_files.values()))))
        files_aks.sort()
        info("Fetching satellite infos from files: {}".format(files_aks))

        try:
            colsat = ["qa0", "ak", "pavg0", "pwf", "date", "index", "station"]
            if transf.formula == 5:
                colsat = colsat + ["dryair"]
            coord2dump = ["index"]
            all_sat_aks = []
            for ind_file, file_aks in enumerate(files_aks):
                sat_aks = \
                    read_datastore(file_aks,
                                   col2dump=colsat,
                                   coord2dump=coord2dump,
                                   keep_default=False,
                                   to_pandas=False).reset_index("index")

                # Keep in memory the file name
                sat_aks = sat_aks.assign(
                    file_aks=("index", sat_aks.dims["index"] * [ind_file]))

                if len(sat_aks) == 0:
                    warning("WARNING: There is no data in file {}".format(file_aks))
                    continue

                if len(all_sat_aks) == 0:
                    all_sat_aks = sat_aks
                else:
                    all_sat_aks = xr.concat([all_sat_aks, sat_aks], "index")

            # Selecting only lines used in simulation
            mask = all_sat_aks["date"].isin(ref_ds["metadata"]["date"]) \
                   & all_sat_aks.index.isin(ref_ds.index) \
                   & (all_sat_aks["station"] == satID)
            sat_aks = all_sat_aks.loc[{"index": mask}]

        except IOError:
            # Assumes total columns?
            raise IOError("Could not fetch "
                          "satellite info from {}".format(file_aks))
            # datastore['qdp'] = datastore['sim'] * datastore['dp']
            # groups = datastore.groupby(['indorig'])
            # y0.loc[:, 'sim'] += \
            #     groups['qdp'].sum().values / groups['dp'].sum().values
            #
            # if 'sim_tl' in datastore:
            #     datastore['qdp'] = datastore['sim_tl'] * datastore['dp']
            #     groups = datastore.groupby(['indorig'])
            #     y0.loc[:, 'sim_tl'] += \
            #         groups['qdp'].sum().values / groups['dp'].sum().values
            #    continue

        if transf.pressure == "Pa":
            pavgs = sat_aks["pavg0"][:, ::-1].T
        else:
            pavgs = 100 * sat_aks["pavg0"][:, ::-1].T

        # Coordinates of the layer boundaries
        coords_pb = {"index": np.arange(nobs),
                     "level": np.arange(sat_aks.level_pressure.size)}
        # Coordinates of the retrieval. Can be pressure layer centers
        # ("level") or layer boundaries ("level_pressure", e.g. OCO-2).
        # Read which it is from dimensions of averaging kernel.
        ret_level_dim = sat_aks["ak"].dims[1]
        coords_ret = {"index": np.arange(nobs),
                     "level": np.arange(sat_aks[ret_level_dim].size)}
        dims = ("level", "index")

        # Define pressure axis of the retrieval
        pavgs_pb = xr.DataArray(pavgs, coords_pb, dims).bfill("level")
        # Case: retrieval on layer midpoints
        if ret_level_dim=="level":
            pavgs_ret = xr.DataArray(
                np.log(0.5 * (pavgs_pb[:-1].values + pavgs_pb[1:].values)),
                coords_ret, dims)
        # Case: retrieval on layer boundaries (e.g. OCO-2)
        elif ret_level_dim=="level_pressure":
            pavgs_ret = np.log(pavgs_pb)
        else:
            raise ValueError("Unknown dimension: '{}'".format(ret_level_dim))

        # If present, read pressure weights from file
        if "pwf" in sat_aks.keys():
            dpavgs = xr.DataArray(sat_aks["pwf"][:, ::-1].T, coords_ret, dims)
        # Else, construct pressure weights
        else:
            # Case: retrieval on layer midpoints
            if ret_level_dim=="level":
                dpavgs = xr.DataArray(np.diff(-pavgs, axis=0), coords_ret, dims)
            # Case: retrieval on layer boundaries
            elif ret_level_dim=="level_pressure":
                raise NotImplementedError(
                    "Constructing pressure weights is not " + \
                    "for retrieval on layer boundaries. " + \
                    "Provide 'pwf' in observation file.")

        aks = xr.DataArray(sat_aks["ak"][:, ::-1].T, coords_ret, dims)
        qa0 = xr.DataArray(sat_aks["qa0"][:, ::-1].T, coords_ret, dims)

        # Exclude observations where there are all zeros in the simulation
        exclude_zeros = np.where(sim.sum(axis=0) != 0)[0]
        datasim = datasim.isel(index=exclude_zeros)
        sat_aks = sat_aks.isel(index=exclude_zeros)
        aks = aks.isel(index=exclude_zeros)
        pavgs_ret = pavgs_ret.isel(index=exclude_zeros)
        dpavgs = dpavgs.isel(index=exclude_zeros)
        qa0 = qa0.isel(index=exclude_zeros)

        # Adding dry air mole fraction if formula 5
        if transf.formula == 5:
            drycols = sat_aks["dryair"][:, ::-1].T
        else:
            drycols = qa0 * 0.0

        # Fetch values if fill strato
        if transf.fill_strato:
            strato = inout_datastore["inputs"][
                ("stratosphere", ref_parameter)][ddi].loc[exclude_zeros]
            strato_mapper = mapper["inputs"][("stratosphere",
                                              ref_parameter)]
            strato_sigma_a = strato_mapper["domain"].sigma_a
            strato_sigma_b = strato_mapper["domain"].sigma_b
            strato_nlev = len(strato_sigma_a)
            
            # Expand debug_datastore with info from previous transforms in strato
            if save_debug:
                trid_strato = ("stratosphere", ref_parameter)
                df_debug = pd.DataFrame({
                    c: inout_datastore["inputs"][trid_strato][ddi][c].iloc[::strato_nlev]
                    for c in inout_datastore["inputs"][trid_strato][ddi].columns
                    if c[0] not in y0.columns.get_level_values(level=0)
                })
                y0 = pd.concat([y0, df_debug], axis=1)
            
        # Interpolating simulated values to averaging kernel pressures
        # Doing it by chunk to fasten the process
        # A single chunk overloads the memory,
        # while too many chunks do not take advantage
        # of scipy automatic parallelisation
        # 50 chunks seems to be fairly efficient
        sim_ak = 0.0 * pavgs_ret
        sim_ak_tl = 0.0 * pavgs_ret
        
        if transf.split_tropo_strato:
            sim_ak_tropo = 0.0 * pavgs_ret
            sim_ak_tl_tropo = 0.0 * pavgs_ret
            sim_ak_strato = 0.0 * pavgs_ret
            sim_ak_tl_strato = 0.0 * pavgs_ret
        
        nchunks = transf.nchunks
        chunks = np.linspace(0, len(datasim.index), num=nchunks + 1, dtype=int)
        cropstrato = transf.cropstrato
        for k1, k2 in zip(chunks[:-1], chunks[1:]):
            # Skip chunks that are too short
            if k1 == k2:
                continue

            debug("Compute chunk for satellite {}: {}-{}".format(satID, k1, k2))

            sim_pressure = datasim["pressure"][:, k1:k2].values
            sim = datasim["sim"][:, k1:k2].values
            if mode == "tl":
                sim_tl = datasim["sim_tl"][:, k1:k2].values
            
            # Fetch missing values from stratosphere
            if transf.fill_strato:
                psurf_strato = np.exp(sim_pressure[0])
                pstrato = np.log(
                    strato_sigma_b * psurf_strato[:, np.newaxis]
                    + strato_sigma_a).T

                # Here merge sim_pressure
                # and pstrato properly, then apply vertical_interp
                missing_levels = np.argmax(pstrato < sim_pressure[-1],
                                           axis=0)[0]
                sim_pressure = \
                    np.concatenate([sim_pressure, pstrato[missing_levels:]],
                                   axis=0)

                # Stacking strato dataframe to column shape
                sim_strato_ref = \
                    strato["maindata"]["spec"].values.reshape(
                        (strato_nlev, -1), order="F")[:, k1:k2]

                incr_strato = 0. * sim_strato_ref
                if "incr" in strato:
                    incr_strato = \
                        strato["maindata"]["incr"].values.reshape(
                            (strato_nlev, -1), order="F")[:, k1:k2]

                # ppb to molec/cm2 if column product
                if transf.product == "column":
                    dpstrato = \
                        np.diff(np.concatenate(
                            [psurf_strato[np.newaxis, :], np.exp(pstrato)],
                            axis=0), axis=0) / 100  # hPa
                    G = 9.88
                    dmass = np.abs(dpstrato / G)  # kg/m2
                    column = dmass * 1e3  # g/m2
                    column /= transf.molmass * 1e4  # mol/cm2
                    column *= 6.02214076 * 10 ** 23  # molec / cm2

                    sim_strato_ref *= column
                    incr_strato *= column
                    
                if transf.split_tropo_strato:
                    sim_tropo = np.concatenate(
                        [sim, 0 * sim_strato_ref[missing_levels:]], axis=0)
                    sim_strato = np.concatenate(
                        [0 * sim, sim_strato_ref[missing_levels:]], axis=0)
                    
                    if mode == "tl":
                        sim_tropo_tl = np.concatenate(
                            [sim_tl, 0 * incr_strato[missing_levels:]], axis=0)
                        sim_strato_tl = np.concatenate(
                            [0. * sim_tl, incr_strato[missing_levels:]], axis=0)
                
                sim = np.concatenate([sim, sim_strato_ref[missing_levels:]],
                                     axis=0)
                if mode == "tl":
                    sim_tl = np.concatenate([sim_tl,
                                             incr_strato[missing_levels:]],
                                            axis=0)

            # Vertical interpolation
            xlow, xhigh, alphalow, alphahigh = vertical_interp(
                sim_pressure,
                pavgs_ret[:, k1:k2].values,
                cropstrato,
            )

            # Applying coefficients
            meshout = np.array(pavgs_ret.shape[0] * [list(range(k2 - k1))])
            sim_ak[:, k1:k2] = (
                alphalow * sim[xlow, meshout] + alphahigh * sim[xhigh, meshout]
            )

            if mode == "tl":
                sim_ak_tl[:, k1:k2] = (
                    alphalow * sim_tl[xlow, meshout]
                    + alphahigh * sim_tl[xhigh, meshout]
                )
            
            # Apply coefficient to "partial" columns
            if transf.split_tropo_strato:
                sim_ak_tropo[:, k1:k2] = (
                    alphalow * sim_tropo[xlow, meshout]
                    + alphahigh * sim_tropo[xhigh, meshout]
                )
                sim_ak_strato[:, k1:k2] = (
                    alphalow * sim_strato[xlow, meshout]
                    + alphahigh * sim_strato[xhigh, meshout]
                )
    
                if mode == "tl":
                    sim_ak_tl_tropo[:, k1:k2] = (
                        alphalow * sim_tropo_tl[xlow, meshout]
                        + alphahigh * sim_tropo_tl[xhigh, meshout]
                    )
                    sim_ak_tl_strato[:, k1:k2] = (
                        alphalow * sim_strato_tl[xlow, meshout]
                        + alphahigh * sim_strato_tl[xhigh, meshout]
                    )

        # Correction with the pressure thickness
        target_datastore[("metadata", "pthick")] = np.nan
        if transf.correct_pthick:
            scale_pthick = (datasim['dp'] * datasim['sim']).sum(axis=0).values \
                           / (dpavgs * sim_ak).sum(axis=0).values \
                           * dpavgs.sum(axis=0).values / datasim['dp'].sum(
                axis=0).values
            raise Exception("debug here")
            sim_ak *= scale_pthick
            if 'sim_tl' in datasim:
                sim_ak_tl *= scale_pthick

            target_datastore[("metadata", "pthick")].iloc[
                np.flatnonzero(ref_indexes)[satmask][exclude_zeros]] = scale_pthick
            
            if transf.split_tropo_strato:
                sim_ak_tropo *= scale_pthick
                sim_ak_strato *= scale_pthick
                if 'sim_tl' in datasim:
                    sim_ak_tl_tropo *= scale_pthick
                    sim_ak_tl_strato *= scale_pthick

        # Applying aks
        nbformula = transf.formula
        chosenlevel = transf.chosenlev

        debug("product: {}".format(transf.product))
        debug("nbformula: {}".format(nbformula))
        debug("chosenlev: {}".format(chosenlevel))

        y0.loc[satmask.iloc[exclude_zeros].index, ("maindata", "spec")] = \
            apply_ak(
                sim_ak.values, dpavgs.values, aks.values,
                nbformula, qa0.values, chosenlevel,
                drycols.values
            )

        if mode == "tl":
            y0.loc[satmask.iloc[exclude_zeros].index, ("maindata", "incr")] = \
                apply_ak_tl(
                    sim_ak_tl.values,
                    dpavgs.values,
                    aks.values,
                    nbformula,
                    qa0.values,
                    chosenlevel,
                    sim_ak.values,
                    drycols.values
                )

        if transf.split_tropo_strato:
            y0.loc[satmask.iloc[exclude_zeros].index,
                   (transf.transform_id, "tropo")] = \
                apply_ak(
                    sim_ak_tropo.values, dpavgs.values, aks.values,
                    nbformula, qa0.values, chosenlevel,
                    drycols.values
                )
            y0.loc[satmask.iloc[exclude_zeros].index,
                   (transf.transform_id, "strato")] = \
                apply_ak(
                    sim_ak_strato.values, dpavgs.values, aks.values,
                    nbformula, qa0.values, chosenlevel,
                    drycols.values
                )
    
            if mode == "tl":
                y0.loc[satmask.iloc[exclude_zeros].index,
                   (transf.transform_id, "tropo_tl")] = \
                    apply_ak_tl(
                        sim_ak_tl_tropo.values,
                        dpavgs.values,
                        aks.values,
                        nbformula,
                        qa0.values,
                        chosenlevel,
                        sim_ak_tropo.values,
                        drycols.values
                    )
                y0.loc[satmask.iloc[exclude_zeros].index,
                   (transf.transform_id, "strato_tl")] = \
                    apply_ak_tl(
                        sim_ak_tl_strato.values,
                        dpavgs.values,
                        aks.values,
                        nbformula,
                        qa0.values,
                        chosenlevel,
                        sim_ak_strato.values,
                        drycols.values
                    )

        # Saves sim_ak for later use by adjoint
        if nbformula == 3:
            file_dump = ddi.strftime(
                "{}/sim_ak_{}_%Y%m%d%H%M.nc".format(dir_fwd, satID))
            sim_ak.to_netcdf(file_dump)

    inout_datastore["outputs"][(ref_component, ref_parameter)][ddi] = y0

    # Dump information about what monitor file was used with which observation
    infos = sat_aks["file_aks"].to_dataframe().rename(
        columns={"index_": "orig_index", "file_aks": "file_id"})
    infos.loc[:, "file_aks"] = pd.DataFrame(
        {"files_aks": files_aks}, index=range(len(files_aks))
    ).loc[infos["file_id"].values].values
    del infos["file_id"]

    todump = pd.DataFrame(index=y0.index)
    todump.loc[satmask.iloc[exclude_zeros].index,
               [("metadata", "orig_index"), ("metadata", "file_aks")]] = infos.values

    file_infos = ddi.strftime(
        "{}/infos_%Y%m%d%H%M.nc".format(dir_fwd)
    )
    dump_datastore(
        todump,
        file_monit=file_infos,
        dump_default=False,
        col2dump=["orig_index", "file_aks"],
        mode="w",
    )

    # Save exclude_zeros in forward datastore
    target_datastore[("metadata", "exclude_zeros")] = True
    target_datastore.loc[exclude_zeros, ("metadata", "exclude_zeros")] = False

    # Save forward datastore for later use by adjoint
    file_monit = ddi.strftime(
        "{}/monit_%Y%m%d%H%M.nc"
            .format(dir_fwd)
    )
    col2dump = ["pressure", "dp", "indorig",
                "hlay", "airm", "sim", "pthick", "exclude_zeros"]
    dump_datastore(
        target_datastore,
        file_monit=file_monit,
        dump_default=False,
        col2dump=col2dump,
        mode="w",
    )