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

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

from .....utils.classes.setup import Setup
from .....utils.classes.domains import Domain

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"

    cornersf = tracer.corners_file
    corner_file = "{}/{}".format(ref_dir, cornersf)

    # Read lon/lat in corner_file
    nc = xr.open_dataset(corner_file, decode_times=False)
    debug('Domain file for VPRM fluxes: {}'.format(corner_file))

    if tracer.regular:
        # Regular lon lat VPRM 1km (1/60, 1/120)
        # Grid cell corners
        lonc = nc["lon"].values
        latc = nc["lat"].values

        lonc = np.sort(np.unique(lonc.flatten()))
        latc = np.sort(np.unique(latc.flatten()))

        dlat = np.diff(latc)
        dlon = np.diff(lonc)
        info('Resolution for VPRM fluxes: lat {}, lon {}'.format(dlat, dlon))

        lonc = np.append(lonc, lonc[-1]+dlon[-1])
        latc = np.append(latc, latc[-1]+dlat[-1])

        zlonc, zlatc = np.meshgrid(lonc, latc)

        # Grid cell centers
        dlat = np.diff(latc)
        dlon = np.diff(lonc)
        lat = latc[:-1] + dlat / 2
        lon = lonc[:-1] + dlon / 2

        zlon, zlat = np.meshgrid(lon, lat)

    else:
        # Irregular lon lat VPRM 5km
        # Grid cell corners
        debug('Irregular lon lat VPRM 5km')
        zlonc = nc["XLONG_C"].values[0, :, :]
        zlatc = nc["XLAT_C"].values[0, :, :]

        # Grid cell centers
        # 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)
        zlon = nc["lon"].values
        zlat = nc["lat"].values

    nlon = zlon.shape[1]
    nlat = zlat.shape[0]
    info('Numder of pixels for VPRM fluxes: lat {}, lon {}'.format(nlat, nlon))

    # Vertical definition (surface level)
    pressure_unit = "Pa"
    nlev = 1
    sigma_a_mid = np.array([0])
    sigma_b_mid = np.array([1])

    # Put it to a domain Plugin
    domain = Domain(nlon=nlon, nlat=nlat,
                    zlon=zlon, zlat=zlat,
                    zlonc=zlonc, zlatc=zlatc,
                    nlev=nlev, pressure_unit=pressure_unit,
                    sigma_b_mid=sigma_b_mid, sigma_a_mid=sigma_a_mid,
                    )

    return domain