Source code for pycif.plugins.transforms.complex.diagmet.utils.omega_top

import numpy as np




def omega_top(transf, inout_datastore, ddi, mapper):
    """Compute ECMWF omega at the pressure of the CHIMERE model top.

    The output ``wtop`` (Pa/s, positive downward) lets CHIMERE set the air
    flux through its model top to the ECMWF one. See
    :doc:`/documentation/doc-models/chimere/diagmet` (section 15).

    Args:
        transf (Plugin): diagmet transform instance.
        inout_datastore (dict): mutable datastore.
        ddi (datetime): current sub-simulation date.
        mapper (dict): transform mapper.
    """
    # Omega on the ECMWF levels (no vertical interpolation), lowest level first
    omega = inout_datastore["inputs"][("meteo", "w")][ddi]["spec"].values
    psurf = inout_datastore["inputs"][("meteo", "psurf")][ddi]["spec"].values

    # ECMWF hybrid coefficients (Pa), bottom interface first
    domain_ecm = transf.datavect.components.meteo.parameters.w.domain
    a_ecm = np.asarray(domain_ecm.sigma_a, dtype=float)[None, :, None, None]
    b_ecm = np.asarray(domain_ecm.sigma_b, dtype=float)[None, :, None, None]
    if a_ecm.shape[1] != omega.shape[1] + 1:
        raise ValueError(
            f"omega_top: {omega.shape[1]} ECMWF levels for w but "
            f"{a_ecm.shape[1]} interface coefficients")

    # CHIMERE top interface pressure in each column (Pa)
    domain = mapper["inputs"][("meteo", "temp")]["domain"]
    a_top = domain.sigma_a[-1] * (100. if domain.pressure_unit == "hPa" else 1.)
    b_top = domain.sigma_b[-1]
    ptop = a_top + b_top * psurf

    # Log-pressure interpolation between the two ECMWF levels around ptop,
    # using each column's level pressures (pressure decreases with index)
    p_half = a_ecm + b_ecm * psurf
    p_mid = 0.5 * (p_half[:, 1:] + p_half[:, :-1])
    nlev = p_mid.shape[1]
    ilow = np.clip((p_mid > ptop).sum(axis=1, keepdims=True) - 1, 0, nlev - 2)
    p_lo = np.take_along_axis(p_mid, ilow, axis=1)
    p_hi = np.take_along_axis(p_mid, ilow + 1, axis=1)
    w_lo = np.take_along_axis(omega, ilow, axis=1)
    w_hi = np.take_along_axis(omega, ilow + 1, axis=1)
    weight = np.clip(np.log(ptop / p_lo) / np.log(p_hi / p_lo), 0., 1.)
    wtop = w_lo + weight * (w_hi - w_lo)

    ref_dataarray = 0. * inout_datastore["outputs"][("meteo", "hght")][ddi]["spec"]
    inout_datastore["outputs"][("meteo", "wtop")][ddi]["spec"] = ref_dataarray + wtop