adjoint.py

import numpy as np
import xarray as xr
import itertools

from logging import info, debug
from .....utils.datastores.dump import read_datastore
from .apply_AK import apply_ak_ad
from .vinterp import vertical_interp
import copy


def adjoint(
        transf,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        **kwargs
):
    """De-aggregate total columns to the model level."""
    ddi = min(di, df)

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

    ref_ds = inout_datastore["outputs"][(ref_component, ref_parameter)][ddi]
    out_datastore = inout_datastore["inputs"]

    # Copy data to output datastore
    input_components = ["concs", "pressure", "dpressure"] \
                       + (transf.product == "column") * ["airm", "hlay"]
    for incomp in input_components:
        out_datastore[(incomp, ref_parameter)][ddi] = \
            copy.deepcopy(ref_ds)
        if incomp != "concs":
            out_datastore[(incomp, ref_parameter)][ddi].loc[
                :, ("maindata", "adj_out")] = 0

    y0 = out_datastore[("concs", ref_parameter)][ddi]

    # Exit if empty observations
    if len(y0) == 0:
        return

    # Number of levels to extract for satellites
    dlev = np.ones(len(y0), dtype=int) * transf.model.domain.nlev

    # Index in the original data of the level-extended dataframe
    native_inds_main = np.append([0], dlev.cumsum())

    # Output index
    idx = np.zeros((native_inds_main[-1]), dtype=int)
    idx[native_inds_main[:-1]] = np.arange(len(y0))
    np.maximum.accumulate(idx, out=idx)
    native_inds_main = native_inds_main[:-1]

    # Output dataframe
    datacol = "adj_out" if mode == "adj" else "obs"
    col2process = [
        "tstep",
        "tstep_glo",
        "i",
        "j",
        "level",
        "dtstep",
        "parameter",
        "duration",
        datacol,
    ]
    df_main = copy.deepcopy(y0.iloc[idx])

    # Levels
    sublevels = np.meshgrid(
        list(range(transf.model.domain.nlev)),
        np.ones(len(y0)))[0].flatten()
    df_main[("metadata", "level")] = sublevels

    # Building the extended dataframe
    # iq1 = (np.abs(y0['level']) - np.abs((y0['level'] / 10.)
    #                                    .astype(int) * 10)) \
    #    .astype(int)
    iq1 = y0[("metadata", "station")]
    list_satIDs = iq1.unique()

    # Saving original values for later re-aggregation
    df_main.loc[:, ("metadata", "indorig")] = idx
    df_main.loc[:, ("metadata", "iq1")] = iq1.iloc[idx]
    for incomp in input_components:
        out_datastore[(incomp, ref_parameter)][ddi] = \
            copy.deepcopy(df_main)
        if incomp != "concs":
            out_datastore[(incomp, ref_parameter)][ddi].loc[
                :, ("maindata", "adj_out")] = 0

    # Un-stack columns into dataframe for stratosphere
    if transf.fill_strato:
        nlev_strato = \
            mapper["inputs"][("stratosphere", ref_parameter)]["domain"].nlev

        # Number of levels to extract for satellites
        dlev = np.ones(len(y0), dtype=int) * nlev_strato

        # Index in the original data of the level-extended dataframe
        native_inds_strato = np.append([0], dlev.cumsum())

        # Output index
        idx = np.zeros((native_inds_strato[-1]), dtype=int)
        idx[native_inds_strato[:-1]] = np.arange(len(y0))
        np.maximum.accumulate(idx, out=idx)
        native_inds_strato = native_inds_strato[:-1]

        # Output dataframe
        df_strato = copy.deepcopy(y0.iloc[idx])

        # Levels
        sublevels = np.meshgrid(list(range(nlev_strato)),
                                np.ones(len(y0)))[0].flatten()
        df_strato[("metadata", "level")] = sublevels

        # Building the extended dataframe
        iq1 = y0[("metadata", "station")]
        list_satIDs = iq1.unique()

        # Saving original values for later re-aggregation
        df_strato.loc[:, ("metadata", "indorig")] = idx
        df_strato.loc[:, ("metadata", "iq1")] = iq1.iloc[idx]
        
        out_datastore[("stratosphere", ref_parameter)][ddi] = df_strato
        
    # Stop here if no adjoint to be fully computed
    # Just forward datastore to precursor transforms
    if onlyinit:
        return

    # Load pressure coordinates from previous run
    file_monit = ddi.strftime(
        "{}/chain/satellites/{}/monit_%Y%m%d%H%M.nc".format(
            transf.model.adj_refdir, transf.transform_id)
    )
    fwd_pressure = read_datastore(
        file_monit,
        col2dump=["pressure", "dp", "indorig",
                  "hlay", "airm", "sim", "pthick", "exclude_zeros"]
    )
    ref_indexes = ~fwd_pressure["metadata"].duplicated(subset=["indorig"])

    for satID in list_satIDs:
        satmask = iq1 == satID
        nobs = np.sum(satmask)

        # Getting the vector of increments
        obs_incr = y0.loc[satmask, ("maindata", "adj_out")]

        # If all increments are NaNs, just pass to next satellite
        if not np.any(obs_incr != 0.0):
            continue

        # Get target pressure
        native_ind_stack = (
                native_inds_main[satmask]
                + np.arange(transf.model.domain.nlev)[:, np.newaxis]
        )
        datasim = xr.Dataset(
            {
                "pressure": (
                    ["level", "index"],
                    np.log(fwd_pressure["metadata"]["pressure"].values[native_ind_stack]),
                ),
                "dp": (
                    ["level", "index"],
                    fwd_pressure["metadata"]["dp"].values[native_ind_stack],
                ),
                "sim": (["level", "index"],
                        fwd_pressure["maindata"]["sim"].values[native_ind_stack]),
            },
            coords={
                "index": np.arange(len(y0)),
                "level": np.arange(transf.model.domain.nlev),
            },
        )
        if transf.product == "column":
            datasim = datasim.assign(
                {
                    "airm": (
                        ["level", "index"],
                        fwd_pressure["metadata"]["airm"].values[native_ind_stack],
                    ),
                    "hlay": (
                        ["level", "index"],
                        fwd_pressure["metadata"]["hlay"].values[native_ind_stack],
                    ),
                }
            )

        # Getting averaging kernels
        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"]
            if transf.formula == 5:
                colsat = colsat + ["dryair"]
            coord2dump = ["index"]
            all_sat_aks = []
            for file_aks in files_aks:
                sat_aks = \
                    read_datastore(file_aks,
                                   col2dump=colsat,
                                   coord2dump=coord2dump,
                                   keep_default=False,
                                   to_pandas=False).reset_index("index")

                # sat_aks["index"] = np.arange(sat_aks.dims["index"])
                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)
            sat_aks = all_sat_aks.loc[{"index": mask}]

        except IOError:
            # Assumes total columns?
            # groups = fwd_pressure.groupby(['indorig'])
            # df['obs_incr'] = y0.ix[idx, 'obs_incr'] * fwd_pressure['dp'] \
            #                  / groups['dp'].sum().values[idx]
            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.")

        # Defining ak info
        #        aks = sat_aks['ak'][nblloc, ::-1][:,1:].T
        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(
            ~fwd_pressure["metadata"].groupby(['indorig']).min()["exclude_zeros"])[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
            
        # Applying aks
        nbformula = transf.formula
        chosenlevel = getattr(transf, "chosenlev", 0)

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

        # If nbformula 3, load sim_ak from forward
        if nbformula == 3:
            file_dump = ddi.strftime(
                "{}/chain/satellites/{}/sim_ak_{}_%Y%m%d%H%M.nc".format(
                    transf.model.adj_refdir, transf.transform_id, satID)
            )
            sim_ak = xr.open_dataarray(file_dump).values
        else:
            sim_ak = 0

        obs_incr = apply_ak_ad(
            sim_ak, dpavgs.values, aks.values,
            nbformula, qa0.values, chosenlevel, obs_incr.values[exclude_zeros],
            drycols.values
        )
        obs_incr[np.isnan(obs_incr)] = 0.

        # Correction with the pressure thickness
        # WARNING: there is an inconsistency in the number of levels
        if transf.correct_pthick:
            scale_pthick = fwd_pressure["metadata"]["pthick"].iloc[
                np.flatnonzero(ref_indexes)[satmask]]
            obs_incr *= scale_pthick.values

        # Adjoint of the log-pressure interpolation
        obs_incr_interp = 0.0 * datasim["pressure"].values

        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))

            # Fetch missing values from stratosphere
            sim_pressure = datasim["pressure"].values[:, k1:k2]
            if transf.fill_strato:
                strato_mapper = mapper["inputs"][("stratosphere",
                                                  ref_parameter)]

                sigma_a = strato_mapper["domain"].sigma_a
                sigma_b = strato_mapper["domain"].sigma_b
                psurf_strato = np.exp(sim_pressure[0])
                pstrato = np.log(
                    sigma_b * psurf_strato[:, np.newaxis] + 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)

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

            # Applying coefficients
            # WARNING: There might be repeated indexes in a given column
            # To deal with repeated index, np.add.at is recommended
            levmeshout = np.array(
                1 * [list(range(pavgs_ret.shape[0]))]
            ).T
            meshout = np.array(pavgs_ret.shape[0] * [list(range(k2 - k1))])
            meshin = np.array(pavgs_ret.shape[0] * [list(range(k1, k2))])
            
            tmp_obs_incr = 0.0 * sim_pressure
            np.add.at(
                tmp_obs_incr,
                (xlow, meshout),
                obs_incr[levmeshout, meshin] * alphalow,
            )

            np.add.at(
                tmp_obs_incr,
                (xhigh, meshout),
                obs_incr[levmeshout, meshin] * alphahigh,
            )

            # Deal with the stratosphere
            if transf.fill_strato:
                nstrato = strato_mapper["domain"].nlev
                nlon = strato_mapper["domain"].nlon
                nlat = strato_mapper["domain"].nlat
                incr_strato = np.zeros((nstrato, k2 - k1))
                incr_strato[missing_levels:] += \
                    tmp_obs_incr[-nstrato + missing_levels:]

                # 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

                    incr_strato *= column

                # Fill adj_out
                native_ind_stack_strato = (
                        native_inds_strato[satmask]
                        + np.arange(nstrato)[:, np.newaxis]
                )
                out_strato = out_datastore[
                    ("stratosphere", transf.parameter[0])][ddi]
                out_strato[("maindata", "adj_out")].iloc[
                    native_ind_stack_strato[
                        :, exclude_zeros[k1:k2]].flatten()] = incr_strato.flatten()
                tmp_obs_incr = tmp_obs_incr[:-nstrato + missing_levels]
                
            obs_incr_interp[:, k1:k2] = tmp_obs_incr[:]

        # Convert CHIMERE fields to the correct unit
        # from ppb to molec.cm-2 if the satellite product is a column
        if transf.product == "column":
            obs_incr_interp *= datasim["hlay"].values / (
                    1e9 / datasim["airm"].values
            )

        # Applying increments to the flattened datastore
        df_main[("maindata", "adj_out")].iloc[
            native_ind_stack[:, exclude_zeros].flatten()] = \
            obs_incr_interp.flatten()

    out_datastore[("concs", transf.parameter[0])][ddi] = df_main