Source code for pycif.plugins.models.lmdz_acc.io.native2inputs_adj

import datetime
import numpy as np
import os
import xarray as xr
import glob
from .outputs.fetch_end import fetch_end




def native2inputs_adj(
        self, datastore, input_type, datei, datef, runsubdir, mode="fwd",
        onlyinit=False, do_simu=True, check_transforms=False, **kwargs
):
    """Converts data at the model data resolution to model compatible input
    files.

    Args:
        self: the model Plugin
        input_type (str): one of 'fluxes', 'obs'
        datastore: data to convert
            if input_type == 'fluxes', a dictionary with flux maps
            if input_type == 'obs', a pandas dataframe with the observations
        datei, datef: date interval of the sub-simulation
        mode (str): running mode: one of 'fwd', 'adj' and 'tl'
        runsubdir (str): sub-directory for the current simulation
        workdir (str): the directory of the whole pycif simulation

    Notes:
        - LMDZ expects daily inputs; if the periods in the control vector are
        longer than one day, period values are uniformly de-aggregated to the
        daily scale; this is done with pandas function 'asfreq' and the option
        'ffill' as 'forward-filling'
        See Pandas page for details:
        https://pandas.pydata.org/pandas-docs/stable/generated/pandas
        .DataFrame.asfreq.html

    """
    ddi = min(datei, datef)

    nlon = self.domain.nlon
    nlat = self.domain.nlat
    nlev = self.domain.nlev

    # Stores daily dates of the period for later aggregation
    dref = datetime.datetime.strptime(
        os.path.basename(os.path.normpath(runsubdir)), "%Y-%m-%d_%H-%M"
    )

    if input_type == "flux":
        list_dates = self.flx_input_dates[ddi]
    else:
        list_dates = self.input_dates[ddi]

    # Reading only output files related to given input_type
    ref_names = {
        "inicond": "init",
        "flux": "fluxes",
        "prescrconcs": "scale",
        "prodloss3d": "prodscale",
    }

    # Fetch end concentrations of adjoint for chain simulation
    if input_type == "endconcs":
        datastore = fetch_end(
            self, datastore, runsubdir, mode, datei, datef,
            check_transforms=check_transforms, onlyinit=onlyinit
        )

    if input_type not in ref_names:
        return datastore

    for trid in datastore:
        out_data = datastore[trid]["data"][ddi]
        spec = trid[1]
        file_list = glob.glob(
            f"{runsubdir}/mod_{ref_names[trid[0]]}_{spec}_out.bin"
        )
        if len(file_list) == 0:
            continue

        out_file = file_list[0]

        with open(out_file, "rb") as f:
            data = np.fromfile(f, dtype=float)

        if input_type == "inicond":
            data = data.reshape((nlon, nlat, -1),
                                order="F").transpose((2, 1, 0))
            out_data["adj_out"] = xr.DataArray(
                data[np.newaxis, ...],
                coords={"time": np.array([dref])},
                dims=("time", "lev", "lat", "lon"),
            )
            continue

        elif input_type == "prescrconcs":
            data = data.reshape((nlon, nlat, nlev, -1),
                                order="F").transpose((3, 2, 1, 0))
            data = np.concatenate((data, np.zeros((1, nlev, nlat, nlon))),
                                  axis=0)
            out_data["adj_out"] = xr.DataArray(
                data,
                coords={"time": list_dates},
                dims=("time", "lev", "lat", "lon"),
            )
            continue

        elif input_type == "prodloss3d":
            data = data.reshape((nlon, nlat, nlev, -1),
                                order="F").transpose((3, 2, 1, 0))
            data = np.concatenate((data, np.zeros((1, nlev, nlat, nlon))),
                                  axis=0)
            out_data["adj_out"] = xr.DataArray(
                data,
                coords={"time": list_dates},
                dims=("time", "lev", "lat", "lon"),
            )
            continue

        # Adding one time stamp to fit with input dates
        # including the first stamp of the next month
        data = data.reshape((nlon, nlat, -1), order="F").transpose((2, 1, 0))
        # data = np.concatenate((data, np.zeros((1, nlat, nlon))), axis=0)
        data = data[:, np.newaxis, ...]
        out_data["adj_out"] = xr.DataArray(
            data,
            coords={"time": list_dates[:-1]},
            dims=("time", "lev", "lat", "lon"),
        )

    return datastore