Source code for pycif.plugins.datastreams.fields.lmdz_outfields_nc.get_domain

import numpy as np
import xarray as xr
import itertools
import os
from .....utils.classes.domains import Domain
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 pyCIF domain from an LMDz ``trajq`` reference file.

    Opens the first fetched ``trajq`` input file, builds a cyclic regular
    lon-lat grid (with corner cells extended/wrapped at the poles and at
    +/-180 degrees) from its ``lon``/``lat`` variables, and reads hybrid
    vertical coefficients ``ap``/``bp`` from
    ``tracer.dir_vcoord/tracer.file_vcoord``.

    Args:
        ref_dir: Unused directly (kept for interface consistency).
        ref_file: Used only in the error message if no reference file is
            available.
        input_interval: Unused directly (kept for interface consistency).
        target_dir: Unused directly (kept for interface consistency).
        tracer: Tracer/component configuration; ``tracer.input_files``,
            ``tracer.dir_vcoord`` and ``tracer.file_vcoord`` are used.

    Returns:
        pycif.utils.classes.domains.Domain: The domain built from the
        reference file's horizontal grid and vertical coefficients, with
        ``lon_cyclic=True``.

    Raises:
        CifError: If no reference ``trajq`` file is available.
    """

    domain_file = list(itertools.chain.from_iterable(
        tracer.input_files.values()))[0]

    if not os.path.isfile(domain_file):
        raise CifError(f"Could not initialize the domain as no reference file is available. Expecting the following file: {ref_file}")

    # Read the file
    with _hdf5_lock:
        ds = xr.open_dataset(domain_file)

        lon = ds["lon"].values
        lat = ds["lat"].values
    zlon, zlat = np.meshgrid(lon, lat)
    nlat, nlon = zlon.shape

    # Corner coordinates
    zlonc = np.concatenate((zlon, zlon[:, np.newaxis, 1]), axis=1)
    zlonc = np.concatenate((zlonc, zlonc[-1, np.newaxis, :]), axis=0)
    zlonc -= 360.0 / (nlon - 1) / 2.0
    zlatc = zlat + 180.0 / (nlat - 1) / 2.0
    zlatc = np.concatenate((zlatc, -90.0 * np.ones((1, nlon))), axis=0)
    zlatc[0, :] = 90
    zlatc = np.concatenate((zlatc, zlatc[:, -1, np.newaxis]), axis=1)

    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.size
    nlat = lat.size

    # Reconstruct alpha and beta
    with _hdf5_lock:
        ds = xr.open_dataset(f"{tracer.dir_vcoord}/{tracer.file_vcoord}")
        sigma_a = ds["ap"].values
        sigma_b = ds["bp"].values

    # Take middle of levels
    nlevs = sigma_a.size - 1

    # Initializes domain
    domain = Domain(nlon=nlon, nlat=nlat,
                    zlon=zlon, zlat=zlat,
                    zlonc=zlonc, zlatc=zlatc,
                    nlev=nlevs, pressure_unit="Pa",
                    sigma_b=sigma_b, sigma_a=sigma_a,
                    lon_cyclic=True)

    # Cyclic domain in longitude
    domain.lon_cyclic = True

    return domain