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

import pandas as pd
import datetime
import os
import glob
from netCDF4 import Dataset
import xarray as xr
import numpy as np

from .....utils.hdf5 import _hdf5_lock
from .outputs.fetch_end import fetch_end




def native2inputs_adj(
    self, datastore, input_type, datei, datef, runsubdir, mode="fwd",
    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 'flux'
        datastore: data to convert
            if input_type == 'flux',
        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:
        - CHIMERE expects hourly inputs;

    """

    if datastore == {}:
        return datastore

    ddi = min(datei, datef)

    # List of CHIMERE dates
    dref = datetime.datetime.strptime(
        os.path.basename(os.path.normpath(runsubdir)), "%Y-%m-%d_%H-%M"
    )
    list_dates = self.input_dates[ddi]

    # Reading only output files related to given input_type
    ref_names = {
        "inicond": "ini",
        "flux": "aemis",
        "bioflux": "bemis",
        "latcond": "bc",
        "topcond": "bc",
        "meteo": "met"
    }

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

    # Read sensitivity for other types of adjoint outputs
    if input_type not in ref_names:
        return datastore

    for trid in datastore:
        file_list = glob.glob(
            f"{runsubdir}/aout.*{ref_names[trid[0]]}*.nc"
        )
        if len(file_list) == 0:
            continue

        sensit_file = file_list[0]
        sensit_basename = os.path.basename(sensit_file)
        with _hdf5_lock:
            with Dataset(sensit_file, "r") as f:
                # Load list of species and reformat it
                if input_type != "meteo":
                    list_species = [
                        b"".join(s).strip().decode("ascii")
                        for s in f.variables["species"][:]
                    ]
                else:
                    list_species = ['kzzz']

                if trid[1] not in list_species:
                    continue

                # Different output structure between LBC and others
                if "ini" in sensit_basename \
                        or "emis" in sensit_basename \
                        or "met" in sensit_basename:
                    data = f.variables[trid[1]][:]

                elif "bc" in sensit_basename:
                    k = list_species.index(trid[1])
                    data_lat = f.variables["lat_conc"][..., k]
                    data_top = f.variables["top_conc"][..., k]

        out_data = datastore[trid]["data"][ddi]
        if "ini" in sensit_basename:
            out_data["adj_out"] = xr.DataArray(
                data[np.newaxis, ...],
                coords={"time": np.array([dref])},
                dims=("time", "lev", "lat", "lon"),
            )

        elif "bc" in sensit_basename:
            if input_type == "latcond":
                out_data["adj_out"] = xr.DataArray(
                    data_lat[..., np.newaxis, :],
                    coords={"time": list_dates},
                    dims=("time", "lev", "lat", "lon"),
                )
            if input_type == "topcond":
                out_data["adj_out"] = xr.DataArray(
                    data_top[:, np.newaxis, ...],
                    coords={"time": list_dates},
                    dims=("time", "lev", "lat", "lon"),
                )

        elif "aemis" in sensit_basename or "bemis" in sensit_basename:
            if "aemis" in sensit_basename:
                emis_type = "flux"
            else:
                emis_type = "bioflux"

            out_data["adj_out"] = xr.DataArray(
                data,
                coords={"time": list_dates},
                dims=("time", "lev", "lat", "lon"),
            )

        elif 'met' in sensit_basename:
            for spec in data:
                out_data['adj_out'] = xr.DataArray(
                    data,
                    coords={"time": list_dates},
                    dims=("time", "lev", "lat", "lon"),
                )

    return datastore