Source code for pycif.plugins.datastreams.fluxes.dummy_txt.make.fromformula

import numpy as np
from ......utils.check.errclass import CifError




def makefromformula(formula, zlon, zlat, formula_type="sum"):
    """Recursively evaluate a nested-dict formula into a 2D flux array.

    The formula is a nested dictionary tree: a node either declares an
    aggregation operator (``"sum"`` or ``"product"`` of a list of
    sub-formulas), or a unary operator (``"cos"``, ``"sin"``, ``"exp"``,
    ``"log"`` or ``"square"``) applied to ``zlon`` or ``zlat`` (selected via
    a ``"variable"`` key), optionally divided by a ``"period"`` value
    (default 1) before the operator is applied.

    Args:
        formula (dict): the (sub-)formula to evaluate; either a leaf node
            with a ``"variable"`` key (``"zlon"`` or ``"zlat"``) and one
            operator key set to ``None``, plus an optional ``"period"`` key,
            or an aggregation node with a single key (``"sum"`` or
            ``"product"``) mapping to a list of sub-formulas.
        zlon (np.ndarray): 2D array of longitudes (grid centers).
        zlat (np.ndarray): 2D array of latitudes (grid centers).
        formula_type (str): aggregation operator to use when combining a
            list of sub-formula results (``"sum"`` or ``"product"``); only
            relevant for aggregation nodes.

    Returns:
        np.ndarray: the 2D array resulting from evaluating the formula,
        same shape as ``zlon``/``zlat``.

    Raises:
        CifError: if an unrecognized ``"variable"`` or operator is
            encountered.
    """
    # If variable in formula, means applying an operator
    if "variable" in formula:
        if formula["variable"] == "zlat":
            var = zlat
        elif formula["variable"] == "zlon":
            var = zlon
        else:
            raise CifError("Not recognized variable type (only zlon and zlat):"
                            + formula["variable"])

        formula_type = [k for k in formula if formula[k] is None][0]

        period = formula.get("period", 1)

        if formula_type == "cos":
            return np.cos(var / period)
        elif formula_type == "sin":
            return np.sin(var / period)
        elif formula_type == "exp":
            return np.exp(var / period)
        elif formula_type == "log":
            return np.exp(var / period)
        elif formula_type == "square":
            return (var / period) ** 2
        else:
            raise CifError("Not recognized operation:" + formula_type)

    # Otherwise
    formula_type = list(formula.keys())[0]
    tmp_vars = [
        makefromformula(formula_tmp, zlon, zlat,
                        formula_type=formula_type)
        for formula_tmp in formula[formula_type]
    ]

    if formula_type == "sum":
        out = 0.
        for var in tmp_vars:
            out += var
        return out
    elif formula_type == "product":
        out = 1.
        for var in tmp_vars:
            out *= var
        return out
    else:
        raise CifError("Not recognized operation:" + formula_type)