Source code for pycif.plugins.transforms.complex.diagmet.propagate_incompatible

import numpy as np
from .....utils.classes.setup import Setup
from .....utils.check.errclass import CifError




def surface_domain(domain_3d):
    """Build a single-level (surface) domain from a 3-D domain.

    Same construction as the 2-D meteo domain of the CHIMERE mapper.

    Args:
        domain_3d (Plugin): 3-D domain.

    Returns:
        Plugin: copy of ``domain_3d`` with ``nlev = 1`` and truncated
        vertical coordinates.
    """
    domain_2d = Setup.load_registered(
        domain_3d.plugin.name, domain_3d.plugin.version,
        "domain", plg_orig=domain_3d
    )

    for attr in ["nlon", "nlat", "zlon", "zlat", "zlonc", "zlatc",
                 "nlon_side", "nlat_side", "zlonc_side", "zlatc_side",
                 "zlon_side", "zlat_side", "pressure_unit"]:
        if hasattr(domain_3d, attr):
            setattr(domain_2d, attr, getattr(domain_3d, attr))

    domain_2d.nlev = 1
    domain_2d.sigma_a = domain_3d.sigma_a[:2]
    domain_2d.sigma_b = domain_3d.sigma_b[:2]
    domain_2d.sigma_a_mid = domain_3d.sigma_a_mid[:1]
    domain_2d.sigma_b_mid = domain_3d.sigma_b_mid[:1]

    return domain_2d





def propagate_incompatible_domain(self, trans_mapper, trid_in, mode="backward"):
    """Assign the correct domain to a diagmet input tracer.

    2-D (surface) and 3-D (column) meteorological fields use different
    domain objects.  This function inspects the output domains already
    present in the mapper, selects the appropriate 2-D or 3-D domain for
    *trid_in*, and assigns it.

    Args:
        self (Plugin): diagmet plugin instance (carries
            ``meteo_parameters_2d_in`` listing 2-D field names).
        trans_mapper (dict): transform mapper; ``'inputs'[trid_in]["domain"]``
            is set in-place.
        trid_in (tuple): ``('meteo', parameter_name)`` key of the input
            tracer whose domain should be resolved.
        mode (str): propagation direction (unused; kept for API consistency).

    Raises:
        Exception: if multiple distinct 3-D or 2-D domains are found in the
            output mapper (should not occur in a well-formed configuration).
    """

    meteo_parameters_3d = self.model.meteo_parameters_3d
    meteo_parameters_2d = self.model.meteo_parameters_2d

    domains = [trans_mapper["outputs"][trid_out]["domain"]
               for trid_out in trans_mapper["outputs"]
               ]

    # Get 3D domains and check if incompatible
    domains_3d = [
        d for d in domains if d.nlev > 1
    ]
    domain_3d = domains_3d[0]
    for dom3d in domains_3d:
        if dom3d != domain_3d:
            raise CifError(
                "More than two output 3D domains in diagmet. This should not happen"
            )

    # Get 3D domains and check if incompatible
    domains_2d = [
        d for d in domains if d.nlev == 1
    ]

    # No 2D outputs (e.g., only_convection): build a 2D domain from the 3D one
    if domains_2d == []:
        if not hasattr(self, "domain_2d"):
            self.domain_2d = surface_domain(domain_3d)
        domains_2d = [self.domain_2d]

    domain_2d = domains_2d[0]
    for dom2d in domains_2d:
        if dom2d != domain_2d:
            raise CifError(
                "More than two output 2D domains in diagmet. This should not happen"
            )

    # Now fetch the correct domain
    trans_mapper["inputs"][trid_in]["domain"] = \
        domain_2d if trid_in in self.meteo_parameters_2d_in else domain_3d