vortex-nwp 2.0.0__py3-none-any.whl

vortex/nwp/util/ens.py

@@ -0,0 +1,279 @@
+"""
+A collection of utility functions used in the context of Ensemble forecasts.
+"""
+
+import io
+import json
+import re
+import time
+
+from bronx.compat import random
+from bronx.fancies import loggers
+from bronx.stdtypes.date import Period
+
+from vortex import sessions
+from vortex.data.stores import FunctionStoreCallbackError
+from vortex.util import helpers
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+def drawingfunction(options):
+    """Draw a random sample from a *set* of values.
+
+    This function is designed to be executed by a
+    :obj:`vortex.data.stores.FunctionStore` object.
+
+    The *set* of values is computed using the resource's argument:
+    *set = [resource.start, resource.start + resource.nbset - 1]*. If
+    *resource.start* does not exists, *resource.start=1* is assumed.
+
+    The size of the sample is given by the *nblot* argument of the resource
+
+    The random generator is initialised using the resource's date. Consequently,
+    for a given date, the drawing is reproducible.
+
+    :param dict options: All the options passed to the store plus anything from
+        the query part of the URI.
+
+    :return: Content of a :obj:`nwp.data.ens.Sample` resource
+
+    :rtype: A file like object
+    """
+    rhdict = options.get("rhandler", None)
+    if rhdict:
+        date = rhdict["resource"]["date"]
+        rgen = random.Random()
+        rgen.seed(int(date[:-2]))
+        nbsample = rhdict["resource"].get("nbsample", 0)
+        if not nbsample:
+            raise FunctionStoreCallbackError(
+                "The resource must hold a non-null nbsample attribute"
+            )
+        population = rhdict["resource"].get("population", [])
+        if not population:
+            raise FunctionStoreCallbackError(
+                "The resource must hold a non-empty population attribute"
+            )
+        nbset = len(population)
+
+        tirage = rgen.sample(
+            population * (nbsample // nbset), (nbsample // nbset) * nbset
+        ) + rgen.sample(population, nbsample % nbset)
+        logger.info(
+            "List of random elements: %s", ", ".join([str(x) for x in tirage])
+        )
+    else:
+        raise FunctionStoreCallbackError("no resource handler here :-(")
+    # NB: The result have to be a file like object !
+    outdict = dict(
+        vapp=rhdict["provider"].get("vapp", None),
+        vconf=rhdict["provider"].get("vconf", None),
+        cutoff=rhdict["resource"].get("cutoff", None),
+        date=rhdict["resource"].get("date", None),
+        resource_kind=rhdict["resource"].get("kind", None),
+        drawing=tirage,
+        population=population,
+    )
+    if rhdict["provider"].get("experiment", None) is not None:
+        outdict["experiment"] = rhdict["provider"]["experiment"]
+    return io.BytesIO(json.dumps(outdict, indent=4).encode(encoding="utf_8"))
+
+
+def _checkingfunction_dict(options):
+    """
+    Internal function that returns a dictionnary that describes the available
+    inputs.
+    """
+    rhdict = options.get("rhandler", None)
+    if rhdict:
+        # If no nbsample is provided, easy to achieve...
+        nbsample = rhdict["resource"].get("nbsample", None)
+        # ...and if no explicit minimum of resources, nbsample is the minimum
+        nbmin = int(
+            options.get(
+                "min",
+                [
+                    (0 if nbsample is None else nbsample),
+                ],
+            ).pop()
+        )
+        if nbsample is not None and nbsample < nbmin:
+            logger.warning(
+                "%d resources needed, %d required: sin of gluttony ?",
+                nbsample,
+                nbmin,
+            )
+        # What to look for ?
+        checkrole = rhdict["resource"].get("checkrole", None)
+        if not checkrole:
+            raise FunctionStoreCallbackError(
+                "The resource must hold a non-empty checkrole attribute"
+            )
+        rolematch = re.match(r"(\w+)(?:\+(\w+))?$", checkrole)
+        cur_t = sessions.current()
+        if rolematch:
+            ctx = cur_t.context
+            checklist = [
+                sec.rh
+                for sec in ctx.sequence.filtered_inputs(
+                    role=rolematch.group(1)
+                )
+            ]
+            mandatorylist = (
+                [
+                    sec.rh
+                    for sec in ctx.sequence.filtered_inputs(
+                        role=rolematch.group(2)
+                    )
+                ]
+                if rolematch.group(2)
+                else []
+            )
+        else:
+            raise FunctionStoreCallbackError(
+                "checkrole is not properly formatted"
+            )
+        # Other options
+        nretries = int(
+            options.get(
+                "nretries",
+                [
+                    0,
+                ],
+            ).pop()
+        )
+        retry_wait = Period(
+            options.get(
+                "retry_wait",
+                [
+                    "PT5M",
+                ],
+            ).pop()
+        )
+        comp_delay = Period(
+            options.get(
+                "comp_delay",
+                [
+                    0,
+                ],
+            ).pop()
+        )
+        fakecheck = options.get(
+            "fakecheck",
+            [
+                False,
+            ],
+        ).pop()
+
+        def _retry_cond(the_ntries, the_acceptable_time):
+            return (
+                the_acceptable_time is None and the_ntries <= nretries
+            ) or (
+                the_acceptable_time
+                and (time.time() - the_acceptable_time)
+                < comp_delay.total_seconds()
+            )
+
+        # Ok let's work...
+        ntries = 0
+        acceptable_time = None
+        found = []
+        while _retry_cond(ntries, acceptable_time):
+            if ntries:
+                logger.info(
+                    "Let's sleep %d sec. before the next check round...",
+                    retry_wait.total_seconds(),
+                )
+                cur_t.sh.sleep(retry_wait.total_seconds())
+            ntries += 1
+            try:
+                logger.info("Starting an input check...")
+                found, candidates = helpers.colorfull_input_checker(
+                    nbmin,
+                    checklist,
+                    mandatory=mandatorylist,
+                    fakecheck=fakecheck,
+                )
+                if acceptable_time is None and (found or nbmin == 0):
+                    acceptable_time = time.time()
+                    if comp_delay.total_seconds() and len(found) != len(
+                        candidates
+                    ):
+                        logger.info(
+                            "The minimum required size was reached (nbmin=%d). "
+                            + "That's great but we are waiting a little longer "
+                            + "(for at most %d sec.)",
+                            nbmin,
+                            comp_delay.total_seconds(),
+                        )
+
+                if len(found) == len(candidates):
+                    # No need to wait any longer...
+                    break
+            except helpers.InputCheckerError as e:
+                if not _retry_cond(ntries, acceptable_time):
+                    raise FunctionStoreCallbackError(
+                        "The input checher failed ({!s})".format(e)
+                    )
+        return found
+    else:
+        raise FunctionStoreCallbackError("no resource handler here :-(\n")
+
+
+def checkingfunction(options):
+    """Check what are the available resources and returns the list.
+
+    This function is designed to be executed by a
+    :obj:`vortex.data.stores.FunctionStore` object.
+
+    The *checkrole* resource attribute is used to look into the current context
+    in order to establish the list of resources that will checked.
+
+    :param dict options: All the options passed to the store plus anything from
+        the query part of the URI.
+
+    :return: Content of a :obj:`nwp.data.ens.PopulationList` resource
+
+    :rtype: A file like object
+    """
+    rhdict = options.get("rhandler", None)
+    avail_list = _checkingfunction_dict(options)
+    outdict = dict(
+        vapp=rhdict["provider"].get("vapp", None),
+        vconf=rhdict["provider"].get("vconf", None),
+        cutoff=rhdict["resource"].get("cutoff", None),
+        date=rhdict["resource"].get("date", None),
+        resource_kind=rhdict["resource"].get("kind", None),
+        population=avail_list,
+    )
+    if rhdict["provider"].get("experiment", None) is not None:
+        outdict["experiment"] = rhdict["provider"]["experiment"]
+    return io.BytesIO(json.dumps(outdict, indent=4).encode(encoding="utf_8"))
+
+
+def safedrawingfunction(options):
+    """Combined called to :func:`checkingfunction` and :func:`drawingfunction`.
+
+    See the documentation of these two functions for more details.
+    """
+    checkedlist = _checkingfunction_dict(options)
+    options["rhandler"]["resource"]["population"] = checkedlist
+    return drawingfunction(options)
+
+
+def unsafedrawingfunction(options):
+    """Combined called to :func:`checkingfunction` and :func:`drawingfunction`...
+    but with a big lie on the checking: no real check, all the resources are assumed ok.
+
+    See the documentation of these two functions for more details.
+    """
+    options["fakecheck"] = [
+        True,
+    ]
+    checkedlist = _checkingfunction_dict(options)
+    options["rhandler"]["resource"]["population"] = checkedlist
+    return drawingfunction(options)

vortex/nwp/util/hooks.py

@@ -0,0 +1,139 @@
+"""
+Some useful hooks.
+"""
+
+import collections.abc
+import functools
+
+from bronx.fancies import loggers
+from bronx.stdtypes.date import Date, Period, Time
+
+from ..data.query import StaticCutoffDispenser
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+def update_namelist(t, rh, *completive_rh):
+    """Update namelist with resource handler(s) given in **completive_rh**."""
+    touched = False
+    for crh in completive_rh:
+        if not isinstance(crh, (list, tuple)):
+            crh = [
+                crh,
+            ]
+        for arh in crh:
+            logger.info(
+                "Merging: {!r} :\n{:s}".format(
+                    arh.container, arh.contents.dumps()
+                )
+            )
+            rh.contents.merge(arh.contents)
+            touched = True
+    if touched:
+        rh.save()
+
+
+def concatenate(t, rh, *rhlist):
+    """Concatenate *rhlist* after *rh*."""
+    blocksize = 32 * 1024 * 1024  # 32Mb
+    rh.container.close()
+    with rh.container.iod_context():
+        myfh = rh.container.iodesc(mode="ab")
+        for crh in rhlist:
+            if not isinstance(crh, (list, tuple)):
+                crh = [
+                    crh,
+                ]
+            for arh in crh:
+                logger.info("Appending %s to self.", str(arh.container))
+                with arh.container.iod_context():
+                    afh = arh.container.iodesc(mode="rb")
+                    stuff = afh.read(blocksize)
+                    while stuff:
+                        myfh.write(stuff)
+                        stuff = afh.read(blocksize)
+
+
+def insert_cutoffs(t, rh, rh_cutoff_source, fuse_per_obstype=False):
+    """Read the cutoff from *rh_cutoff_source* and feed them into *rh*.
+
+    If *fuse_per_obstype* is ``True``, the latest cutoff of a given obstype
+    will be used for all the occurences of this obstype.
+    """
+    # rh_cutoff_source may be a list
+    if isinstance(rh_cutoff_source, list):
+        if rh_cutoff_source:
+            rh_cutoff_source = rh_cutoff_source[0]
+        else:
+            ValueError("The resource handler's list is empty.")
+    # Get the CutoffDispenser
+    import vortex.tools.listings
+
+    assert vortex.tools.listings
+    if rh_cutoff_source.container.actualfmt == "bdmbufr_listing":
+        c_disp_callback = functools.partial(
+            rh_cutoff_source.contents.data.cutoffs_dispenser,
+            fuse_per_obstype=fuse_per_obstype,
+        )
+    else:
+        raise RuntimeError(
+            "Incompatible < {!s} > ressource handler".format(rh_cutoff_source)
+        )
+    # Fill the gaps in the original request
+    rh.contents.add_cutoff_info(c_disp_callback())
+    # Actually save the result to file
+    rh.save()
+
+
+def _new_static_cutoff_dispencer(base_date, cutoffs_def):
+    def x_period(p):
+        try:
+            return Period(p)
+        except ValueError:
+            return Period(Time(p))
+
+    if not isinstance(base_date, Date):
+        base_date = Date(base_date)
+    if isinstance(cutoffs_def, collections.abc.Mapping):
+        cutoffs_def = {
+            (k if isinstance(k, Period) else x_period(k)): v
+            for k, v in cutoffs_def.items()
+        }
+        cutoffs = {base_date + k: v for k, v in cutoffs_def.items()}
+        c_disp = StaticCutoffDispenser(max(cutoffs.keys()), cutoffs)
+    else:
+        if not isinstance(cutoffs_def, Period):
+            cutoffs_def = x_period(cutoffs_def)
+        c_disp = StaticCutoffDispenser(base_date + cutoffs_def)
+    return c_disp
+
+
+def insert_static_cutoffs(t, rh, base_date, cutoffs_def):
+    """Compute the cutoff from *cutoffs_def* and feed them into *rh*.
+
+    :param base_date: The current analysis time
+    :param cutoffs_def: The cutoff time represented as time offset with respect
+                        to *base_date*. *cutoffs_defs* may be a single value or
+                        a dictionary. If *cutoffs_def* is a dictionary, it
+                        associates a cutoff with a list of `obstypes`.
+    """
+    # Fill the gaps in the original request
+    rh.contents.add_cutoff_info(
+        _new_static_cutoff_dispencer(base_date, cutoffs_def)
+    )
+    # Actually save the result to files
+    rh.save()
+
+
+def arpifs_obs_error_correl_legacy2oops(t, rh):
+    """Convert a constant file that contains observation errors correlations."""
+    if rh.resource.realkind != "correlations":
+        raise ValueError("Incompatible resource: {!s}".format(rh))
+    if rh.contents[0].startswith("SIGMAO"):
+        logger.warning("Non conversion is needed...")
+    else:
+        rh.contents[:0] = ["SIGMAO unused\n", "1 1.2\n", "CORRELATIONS\n"]
+        rh.save()

vortex/nwp/util/taskdeco.py

@@ -0,0 +1,85 @@
+"""
+A collection of Tasks decorators (to add usual inputs/outputs to existing classes).
+"""
+
+from bronx.fancies import loggers
+
+from vortex import toolbox
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+def process_needs_lfi_stuff(*kargs, **kwargs):
+    """
+    Decorator that update the tasks's ``process`` method in order to retrieve the
+    things needed with the FA/LFI file format.
+
+    Example (the self.conf.cycle Genv/Uenv cycle will be used)::
+
+        @process_needs_lfi_stuff
+        class MyTask(Task):
+
+            def process(self)
+                pass
+
+    Example (the self.conf.arpege_cycle Genv/Uenv cycle will be used)::
+
+        @process_needs_lfi_stuff(cyclekey='arpege_cycle')
+        class MyOtherTask(Task):
+
+            def process(self)
+                pass
+
+    """
+    cyclekey = kwargs.pop("cyclekey", "cycle")
+
+    def decorate_process(cls):
+        """Decorator for Task: get LFI stuff before calling process."""
+        original_process = getattr(cls, "process", None)
+        if original_process is not None:
+
+            def process(self, *args, **kwargs):
+                _get_lfi_stuff(self, cyclekey)
+                original_process(self, *args, **kwargs)
+
+            process.__doc__ = original_process.__doc__
+            cls.process = process
+        return cls
+
+    if kargs:
+        return decorate_process(kargs[0])
+    else:
+        return decorate_process
+
+
+def _get_lfi_stuff(self, cyclekey):
+    """Get LFI stuff method (called from process)."""
+    if "early-fetch" in self.steps or "fetch" in self.steps:
+        actualcycle = getattr(self.conf, cyclekey)
+
+        self.sh.title("Toolbox input tblfiscripts")
+        toolbox.input(
+            role="LFIScripts",
+            genv=actualcycle,
+            kind="lfiscripts",
+            local="usualtools/tools.lfi.tgz",
+        )
+        self.sh.title("Toolbox input tbiopoll")
+        toolbox.input(
+            role="IOPoll",
+            format="unknown",
+            genv=actualcycle,
+            kind="iopoll",
+            language="perl",
+            local="usualtools/io_poll",
+        )
+        self.sh.title("Toolbox input tblfitools")
+        toolbox.input(
+            role="LFITOOLS",
+            genv=actualcycle,
+            kind="lfitools",
+            local="usualtools/lfitools",
+        )