Source code for pycif.plugins.datastreams.fluxes.VPRM_nc.get_domain

import numpy as np
import xarray as xr
import glob
import datetime
import os

from .....utils.classes.setup import Setup
from logging import info, debug
from .....utils.check.errclass import CifError
from .....utils.hdf5 import _hdf5_lock




def get_domain(ref_dir, ref_file, input_interval, target_dir, tracer=None):
    """Build the VPRM domain from a reference data file and the corners file.

    Locates a reference data file (by exact name match, or by
    pattern-matching `ref_file` against filenames in `ref_dir`) to read the
    ``lon``/``lat`` cell centers and their min/max extents. Reads the
    ``XLONG_C``/``XLAT_C`` corner arrays from the separate
    `tracer.corners_file`. Initializes a ``dummy``/``std`` domain via
    `Setup` using the centers' extent and shape, then overwrites its
    ``zlon``/``zlat``/``zlonc``/``zlatc`` attributes with the arrays read
    from file.

    Args:
        ref_dir (str): Directory containing the reference data and corners
            files.
        ref_file (str): Filename (or ``strftime`` pattern) of the
            reference data file.
        input_interval (list): Unused directly, kept for interface
            consistency.
        target_dir (str): Unused directly, kept for interface consistency.
        tracer: The flux tracer plugin, providing ``corners_file``.

    Returns:
        Domain: a surface-only `Domain` built from the reference and
        corners files.

    Raises:
        CifError: If no matching reference data file is found.
    """

    # Inputs:
    # ---------
    # ref_dir: directory where the original files are found
    # ref_file: (template) name of the original files
    # input_interval: list of the periods to simulate, each item is the list of the dates of the period
    # target_dir: directory where the links to the orginal files are created
    #
    # Ouputs:
    # ---------
    # setup of the domain in section "Initializes domain"

    # Looking for a reference file to read lon/lat in
    list_file = glob.glob(f"{ref_dir}/*nc")

    domain_file = None
    # Either a file is specified in the Yaml
    if ref_file in list_file:
        domain_file = f"{ref_dir}/{ref_file}"

    # Or loop over available file regarding file pattern
    else:
        for flx_file in list_file:
            try:
                date = datetime.datetime.strptime(
                    os.path.basename(flx_file), ref_file
                )
                domain_file = flx_file
                break
            except ValueError:
                continue

    if domain_file is None:
        raise CifError(
            "VPRM domain could not be initialized as no file was found"
        )

    debug(f'Domain file for VPRM fluxes: {domain_file}')

    # Read lon/lat in
    with _hdf5_lock:
        nc = xr.open_dataset(domain_file, decode_times=False)

        lon = nc["lon"].values
        lat = nc["lat"].values

    # print('Get the min and max latitude and longitude of centers
    #       + the number of longitudes and latitudes')
    lon_min = lon.min()  # - (lon[1] - lon[0]) / 2
    lon_max = lon.max()  # + (lon[-1] - lon[-2]) / 2
    lat_min = lat.min()  # - (lat[1] - lat[0]) / 2
    lat_max = lat.max()  # + (lat[-1] - lat[-2]) / 2

    nlon = lon.shape[1]
    nlat = lat.shape[0]
    cornersf = tracer.corners_file
    corner_file = f"{ref_dir}/{cornersf}"
    info(f'corner_file: {corner_file}')

    # Read lon/lat in corner_file
    with _hdf5_lock:
        nc = xr.open_dataset(corner_file, decode_times=False)
        lonc = nc["XLONG_C"].values[0, :, :]
        latc = nc["XLAT_C"].values[0, :, :]
    punit = "Pa"
    nlevs = 1
    sigma_a_mid = np.array([0])
    sigma_b_mid = np.array([1])

    # Initializes domain
    setup = Setup.from_dict(
        {
            "domain": {
                "plugin": {
                    "name": "dummy",
                    "version": "std",
                    "type": "domain",
                },
                "xmin": lon_min,    # minimum longitude for centers
                "xmax": lon_max,    # maximum longitude for centers
                "ymin": lat_min,    # minimum latitude for centers
                "ymax": lat_max,    # maximum latitude for centers
                "nlon": nlon,       # number of longitudinal cells
                "nlat": nlat,       # number of latitudinal cells
                "nlev": nlevs,      # number of vertical levels
                "sigma_a_mid": sigma_a_mid,
                "sigma_b_mid": sigma_b_mid,
                "pressure_unit": punit  # adapted to sigmas
            }
        }
    )

    Setup.load_setup(setup, level=1)
    # if lon and lat are vectors, convert into a grid with
    # zlon, zlat = np.meshgrid(lon, lat)
    setup.domain.zlon = lon   # longitudes of centers
    setup.domain.zlat = lat   # latitudes of centers
    setup.domain.zlonc = lonc  # longitudes of corners
    setup.domain.zlatc = latc  # latitudes of corners

    return setup.domain