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




def get_domain(ref_dir, ref_file, input_interval, target_dir, tracer=None):

    # 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("{}/*nc".format(ref_dir))

    domain_file = None
    # Either a file is specified in the Yaml
    if ref_file in list_file:
        domain_file = "{}/{}".format(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 Exception(
            "VPRM domain could not be initialized as no file was found"
        )

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

    # Read lon/lat in
    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 = "{}/{}".format(ref_dir, cornersf)
    info('corner_file: {}'.format(corner_file))

    # Read lon/lat in corner_file
    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