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

import numpy as np
from logging import debug
import os
import xarray as xr
import itertools

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




def get_domain(ref_dir, ref_file, input_interval, target_dir, tracer=None):
    """Read information from the reference file
    to define the data horizontal and, if relevant, vertical domain.

    Args:
        ref_dir (str): the path to the input files
        ref_file (str): format of the input files
        input_interval (list): simulation interval (start and end dates)
        target_dir (str): where to copy
        tracer: the tracer Plugin, corresponding to the paragraph
            :bash:`datavect/components/fluxes/parameters/my_species` in the
            configuration yaml; can be needed to fetch extra information
            given by the user

    Return:
        domain (Domain): a domain class object, with the definition of the center grid
            cells coordinates, as well as corners

    """

    list_files = list(itertools.chain.from_iterable(tracer.input_files.values()))
    no_file = True if len(list_files) == 0 else not os.path.isfile(list_files[0])
    if no_file:
        raise Exception(
            f"Could not initialize the domain as no reference file is available. "
            f"Expecting the following file: {ref_dir}/{ref_file}")

    # Find lat/lon variables
    ref_file = list_files[0]
    ds = xr.open_dataset(ref_file)
    if "lat" in ds.variables:
        lat_name = "lat"
        lon_name = "lon"

    elif "latitude" in ds.variables:
        lat_name = "latitude"
        lon_name = "longitude"

    else:
        raise Exception("Could not find lat/lon coordinate variables "
                        "in reference file: {}".format(ref_file))

    # Grid cell centers
    coords = xr.open_dataset(ref_file)[[lat_name, lon_name]]
    lat = coords[lat_name].values
    lon = coords[lon_name].values

    nlat = len(lat)
    nlon = len(lon)

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

    # Grid cell corners
    dlat = np.unique(np.diff(lat))[0]
    dlon = np.unique(np.diff(lon))[0]

    latc = np.append(lat - dlat / 2, lat[-1] + dlat / 2)
    lonc = np.append(lon - dlon / 2, lon[-1] + dlon / 2)
    zlonc, zlatc = np.meshgrid(lonc, latc)

    # 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