vortex-nwp 2.0.0__py3-none-any.whl

vortex/nwp/tools/drhook.py

@@ -0,0 +1,66 @@
+"""
+Common interest classes to help setup the DrHook library environment.
+"""
+
+import footprints
+from bronx.fancies import loggers
+
+from vortex.algo.components import (
+    AlgoComponentDecoMixin,
+    Parallel,
+    algo_component_deco_mixin_autodoc,
+)
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+@algo_component_deco_mixin_autodoc
+class DrHookDecoMixin(AlgoComponentDecoMixin):
+    """Handle DrHook settings in AlgoComponents.
+
+    This mixin class is intended to be used with AlgoComponent classes. It will
+    automatically add footprints' arguments related to DrHook (namely the
+    drhookprof boolean attribute that is optional and False by default),
+    and set up DrHook environment variables (:meth:`_drhook_varexport`) depending
+    on the context (MPI run or not).
+    """
+
+    _MIXIN_EXTRA_FOOTPRINTS = [
+        footprints.Footprint(
+            attr=dict(
+                drhookprof=dict(
+                    info="Activate the DrHook profiling.",
+                    optional=True,
+                    type=bool,
+                    default=False,
+                    doc_zorder=-50,
+                ),
+            ),
+        )
+    ]
+
+    def _drhook_varexport(self, rh, opts):  # @UnusedVariable
+        """Export proper DrHook variables"""
+        drhook_vars = (
+            [
+                ("DR_HOOK", "1"),
+                ("DR_HOOK_OPT", "prof"),
+                ("DR_HOOK_IGNORE_SIGNALS", "-1"),
+            ]
+            if self.drhookprof
+            else [("DR_HOOK", "0"), ("DR_HOOK_IGNORE_SIGNALS", "-1")]
+        )
+        if not isinstance(self, Parallel):
+            drhook_vars += [
+                ("DR_HOOK_SILENT", "1"),
+                ("DR_HOOK_NOT_MPI", "1"),
+                ("DR_HOOK_ASSERT_MPI_INITIALIZED", "0"),
+            ]
+        for k, v in drhook_vars:
+            logger.info("Setting DRHOOK env %s = %s", k, v)
+            self.env[k] = v
+
+    _MIXIN_PREPARE_HOOKS = (_drhook_varexport,)

vortex/nwp/tools/grib.py

@@ -0,0 +1,294 @@
+"""
+TODO: Module documentation.
+"""
+
+import json
+
+import footprints
+from bronx.fancies import loggers
+from vortex import sessions
+from vortex.algo.components import AlgoComponentError
+from vortex.layout.contexts import Context
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+class _GenericFilter:
+    """This class could be the start of filtering classes for different formats."""
+
+    def __init__(self):
+        """
+
+        No parameters.
+        """
+        self._filters = list()
+        self._sh = sessions.system()
+
+    def add_filters(self, *filters):
+        """Add one or more filters to the filters list.
+
+        :param filters: a list of filters
+
+        Filters are described using dictionaries. Here is an example of such a
+        dictionary::
+
+            {
+              "fields_include": [
+                {
+                  "indicatorOfTypeOfLevel": 100,
+                  "shortName": "t",
+                  "level": [
+                    850,
+                    500,
+                    300
+                  ]
+                }
+              ],
+              "fid_format": "GRIB1",
+              "filter_name": "toto"
+            }
+
+        **fields_include** or **fields_exclude** lists depends on the data
+        format specified with the **fid_format** key.
+
+        The filters argument of this function accepts dictionaries but also
+        strings or Context objects :
+
+            * If a string is provided, it will be converted to a dictionary
+              using json.loads
+            * If a Context object is provided, the Context's sequence will be
+              used to find available resources of filtering_request kind. The
+              content of such resources will be used as a filter.
+
+        """
+
+        for a_filter in filters:
+            if isinstance(a_filter, dict):
+                self._filters.append(a_filter)
+            elif isinstance(a_filter, str):
+                self._filters.append(json.loads(a_filter))
+            elif isinstance(a_filter, Context):
+                for a_request in a_filter.sequence.effective_inputs(
+                    kind="filtering_request"
+                ):
+                    self._filters.append(a_request.rh.contents.data)
+
+    def __len__(self):
+        """Returns the number of active filters."""
+        return len(self._filters)
+
+    @staticmethod
+    def _is_dict_superset(full, subset):
+        """Finds out if the full dictionary contains and matches subset."""
+        superset_ok = True
+        for k, v in subset.items():
+            # Ignore the comments...
+            if k.startswith("comment"):
+                continue
+            # Check for the key inside the full dictionary
+            try:
+                fullvalue = full[str(k)]
+            except KeyError:
+                superset_ok = False
+                break
+            # Does the key match ?
+            if isinstance(v, (list, tuple)):
+                if fullvalue not in v:
+                    superset_ok = False
+                    break
+            else:
+                if fullvalue != v:
+                    superset_ok = False
+                    break
+        return superset_ok
+
+    def _filter_process(self, fid, a_filter):
+        """Check if the data's fid complies with the filter."""
+        includes = a_filter.get("fields_include", [])
+        excludes = a_filter.get("fields_exclude", [])
+        try:
+            fid = fid[a_filter["fid_format"]]
+        except KeyboardInterrupt:
+            raise ValueError(
+                "Please specify a valid fid_format in the filter description"
+            )
+        # First process includes
+        includes_ok = True
+        for include in includes:
+            includes_ok = self._is_dict_superset(fid, include)
+            if includes_ok:
+                break
+        # Process the excludes if necessary
+        if includes_ok:
+            for exclude in excludes:
+                includes_ok = not self._is_dict_superset(fid, exclude)
+                if not includes_ok:
+                    break
+        return includes_ok
+
+    def __call__(self, inputfile, outfile_fmt):
+        """Apply the various filters on *inputfile*. Should be implemented..."""
+        raise NotImplementedError("This method have to be overwritten.")
+
+
+class GRIBFilter(_GenericFilter):
+    """Class in charge of filtering GRIB files."""
+
+    CONCATENATE_FILTER = "concatenate"
+
+    def __init__(self, concatenate=False):
+        """
+
+        :param bool concatenate: Wether to generate a concatenated GRIB file
+        """
+        super().__init__()
+        self.concatenate = concatenate
+        self._xgrib_support = "grib" in self._sh.loaded_addons()
+
+    def __len__(self):
+        """Returns the number of active filters (concatenate included)."""
+        return super().__len__() + (1 if self.concatenate else 0)
+
+    def _simple_cat(self, gribfile, outfile_fmt, intent):
+        """Just concatenate a multipart GRIB."""
+        if self._xgrib_support and self._sh.is_xgrib(gribfile):
+            self._sh.xgrib_pack(
+                gribfile,
+                outfile_fmt.format(filtername=self.CONCATENATE_FILTER),
+                intent=intent,
+            )
+        else:
+            # Just make a copy with the appropriate name...
+            self._sh.cp(
+                gribfile,
+                outfile_fmt.format(filtername=self.CONCATENATE_FILTER),
+                intent=intent,
+                fmt="grib",
+            )
+
+    def __call__(self, gribfile, outfile_fmt, intent="in"):
+        """Apply the various filters on *gribfile*.
+
+        :param gribfile: The path to the input GRIB file
+        :param outfile_fmt: The path of output files
+
+        The *outfile_fmt* must be a format string such as
+        **GRIBOUTPUT_{filtername:s}.grib** where **filtername** will be replaced
+        by the name of the filter.
+        """
+
+        if not self._sh.path.exists(gribfile):
+            raise OSError("{!s} doesn't exist".format(gribfile))
+
+        # We just want to concatenate files...
+        if not self._filters:
+            if self.concatenate:
+                self._simple_cat(gribfile, outfile_fmt, intent=intent)
+                return [
+                    outfile_fmt.format(filtername=self.CONCATENATE_FILTER),
+                ]
+            else:
+                raise ValueError("Set concatenate=True or provide a filter.")
+
+        # Open the input file using Epygram
+        from ..util import usepygram
+
+        if not usepygram.epygram_checker.is_available(version="1.0.0"):
+            raise AlgoComponentError("Epygram (v1.0.0) needs to be available")
+
+        if self._xgrib_support and self._sh.is_xgrib(gribfile):
+            idx = self._sh.xgrib_index_get(gribfile)
+            in_data = [
+                footprints.proxy.dataformat(
+                    filename=self._sh.path.realpath(a_gribfile),
+                    openmode="r",
+                    format="GRIB",
+                )
+                for a_gribfile in idx
+            ]
+        else:
+            in_data = [
+                footprints.proxy.dataformat(
+                    filename=self._sh.path.realpath(gribfile),
+                    openmode="r",
+                    format="GRIB",
+                ),
+            ]
+
+        # Open output files
+        out_data = list()
+        out_filelist = list()
+        for a_filter in self._filters:
+            f_name = outfile_fmt.format(filtername=a_filter["filter_name"])
+            out_filelist.append(f_name)
+            # It would be a lot better to use io.open but grib_api is very annoying !
+            out_data.append(open(f_name, "wb"))
+        if self.concatenate:
+            f_name = outfile_fmt.format(filtername=self.CONCATENATE_FILTER)
+            out_filelist.append(f_name)
+            # It would be a lot better to use io.open but grib_api is very annoying !
+            out_cat = open(f_name, "wb")
+
+        with usepygram.epy_env_prepare(sessions.current()):
+            for a_in_data in in_data:
+                msg = a_in_data.iter_messages(headers_only=False)
+                while msg is not None:
+                    for a_out_data, a_filter in zip(out_data, self._filters):
+                        thefid = msg.genfid()
+                        if self._filter_process(thefid, a_filter):
+                            logger.debug(
+                                "Select succeed for filter %s: %s",
+                                a_filter["filter_name"],
+                                thefid,
+                            )
+                            msg.write_to_file(a_out_data)
+                    if self.concatenate:
+                        msg.write_to_file(out_cat)
+                    msg = a_in_data.iter_messages(headers_only=False)
+
+        # Close outpout files
+        for a_in_data in in_data:
+            a_in_data.close()
+        for a_out_data in out_data:
+            a_out_data.close()
+        if self.concatenate:
+            out_cat.close()
+
+        return out_filelist
+
+
+def grib_inplace_cat(t, rh):
+    """Ensure that a GRIB file is a usual single file (if not, concatenate it).
+
+    This function is designed to be used as a hook function.
+
+    :param t: A :class:`vortex.sessions.Ticket` object
+    :param rh: A :class:`vortex.data.handlers.Handler` object
+    """
+    xgrib_support = "grib" in t.sh.loaded_addons()
+    if xgrib_support:
+        if t.sh.is_xgrib(rh.container.localpath()):
+            # Some cleanup...
+            rh.reset_contents()
+            rh.container.close()
+            # Move the index file prior to the concatenation
+            tmpfile = (
+                rh.container.localpath() + "_concat" + t.sh.safe_filesuffix()
+            )
+            t.sh.move(rh.container.localpath(), tmpfile)
+            # Concatenate
+            t.sh.xgrib_pack(tmpfile, rh.container.localpath())
+            # Remove the multipart file
+            t.sh.grib_remove(tmpfile)
+            logger.info("The multipart GRIB has been concatenated.")
+        else:
+            logger.info(
+                "The localpath is not a multipart GRIB: nothing to do."
+            )
+    else:
+        logger.info(
+            "Multipart GRIB support is not activated: nothing can be done."
+        )

vortex/nwp/tools/gribdiff.py

@@ -0,0 +1,104 @@
+"""
+TODO: Module documentation.
+"""
+
+import footprints
+
+from vortex.tools.grib import GRIBAPI_Tool
+
+#: No automatic export
+__all__ = []
+
+
+class _GRIBDIFF_Plus_St:
+    """Status of the GRIB comparison."""
+
+    def __init__(self, rc, result):
+        self.rc = rc
+        self._result = result
+
+    def __str__(self):
+        return "{:s} | rc={:d}>".format(repr(self).rstrip(">"), self.rc)
+
+    @property
+    def result(self):
+        """Indicates whether the diff succeeded or not."""
+        return self._result
+
+    def __bool__(self):
+        return self.rc
+
+
+class _GRIBDIFF_Plus_Res:
+    """Detailed result of the GRIB comparison."""
+
+    def __init__(self, gapi, epydiff, epydiff_res):
+        self._gapi = gapi
+        self._epydiff = epydiff
+        self._epydiff_res = epydiff_res
+
+    def __str__(self):
+        return "{0:s} | gribapi_rc={1:d} epydiff_done={2:d}>".format(
+            repr(self).rstrip(">"), self._gapi, self._epydiff
+        )
+
+    def differences(self):
+        """Print detailed informations about the diff."""
+        print(self._epydiff_res)
+
+
+class GRIBDIFF_Plus(GRIBAPI_Tool):
+    """
+    Interface to gribapi commands + epygram diff (designed as a shell Addon).
+    """
+
+    _footprint = dict(
+        info="Default GRIBAPI system interface",
+        attr=dict(
+            maxepydiff=dict(
+                info="Epygram diffs are costfull, they will run only maxepydiff times",
+                type=int,
+                default=2,
+                optional=True,
+            ),
+        ),
+        priority=dict(
+            level=footprints.priorities.top.TOOLBOX  # @UndefinedVariable
+        ),
+    )
+
+    def __init__(self, *kargs, **kwargs):
+        super().__init__(*kargs, **kwargs)
+        self._epycount = 0
+        self._epyavail = None
+
+    def _actual_diff(self, grib1, grib2, skipkeys, **kw):
+        rc = super()._actual_diff(grib1, grib2, skipkeys, **kw)
+        if not rc:
+            if self._epyavail is None:
+                from ..util.usepygram import epygram_checker
+
+                self._epyavail = epygram_checker.is_available(version="1.0.0")
+            if self._epyavail:
+                if self._epycount < self.maxepydiff:
+                    from ..util.diffpygram import EpyGribDiff
+
+                    gdiff = EpyGribDiff(grib2, grib1)  # Ref file is first...
+                    self._epycount += 1
+                    res = _GRIBDIFF_Plus_Res(rc, True, str(gdiff))
+                    # Save the detailed diff
+                    with open(grib1 + "_epygram_diffstats.log", "w") as outfh:
+                        outfh.write(gdiff.format_diff(detailed=True))
+                else:
+                    res = _GRIBDIFF_Plus_Res(
+                        rc,
+                        False,
+                        "grib_compare failed (but the Epygram diffs max number is exceeded...)",
+                    )
+            else:
+                res = _GRIBDIFF_Plus_Res(
+                    rc, False, "grib_compare failed (Epygram unavailable)"
+                )
+        else:
+            res = _GRIBDIFF_Plus_Res(rc, False, "")
+        return _GRIBDIFF_Plus_St(rc, res)

vortex/nwp/tools/ifstools.py

@@ -0,0 +1,203 @@
+"""
+Various tools related to the IFS code.
+"""
+
+import re
+
+from bronx.fancies import loggers
+from bronx.stdtypes.date import Time
+
+import footprints
+
+from vortex.syntax.stdattrs import model
+from ..syntax.stdattrs import arpifs_cycle
+
+logger = loggers.getLogger(__name__)
+
+
+class _IfsOutputsTimesListDesc:
+    """Convert the input data to Time objects."""
+
+    def __init__(self, attr, doc):
+        self._attr = attr
+        self.__doc__ = doc
+
+    def __get__(self, instance, owner):
+        return instance._tlists_store.get(self._attr, None)
+
+    def __set__(self, instance, value):
+        if value is None:
+            instance._tlists_store.pop(self._attr, None)
+        else:
+            if not isinstance(value, list):
+                raise ValueError("**value** should be a list.")
+            instance._tlists_store[self._attr] = [Time(t) for t in value]
+
+    def __delete__(self, instance):
+        instance._tlists_store.pop(self._attr, None)
+
+
+class IfsOutputsAbstractConfigurator(footprints.FootprintBase):
+    """Abstract utility class to configure the IFS model regarding output data."""
+
+    _abstract = True
+    _collector = ("ifsoutputs_configurator",)
+    _footprint = [
+        model,
+        arpifs_cycle,
+        dict(
+            attr=dict(
+                fcterm_unit=dict(
+                    info="The unit used in the *fcterm* attribute.",
+                    values=["h", "t"],
+                ),
+            )
+        ),
+    ]
+
+    def __init__(self, *kargs, **kwargs):
+        super().__init__(*kargs, **kwargs)
+        self._tlists_store = dict()
+
+    modelstate = _IfsOutputsTimesListDesc(
+        "modelstate", "The list of terms for modelstate outputs."
+    )
+
+    surf_modelstate = _IfsOutputsTimesListDesc(
+        "surf_modelstate",
+        "The list of terms for surface scheme modelstate outputs.",
+    )
+
+    spectral_diag = _IfsOutputsTimesListDesc(
+        "spectral_diag",
+        "The list of terms for spectral space diagnostics outputs.",
+    )
+
+    post_processing = _IfsOutputsTimesListDesc(
+        "post_processing",
+        "The list of terms for inline post-processing outputs.",
+    )
+
+    def _setup_nam_obj(self, namelist_object, namelist_name):
+        """Actualy tweak the IFS namelist."""
+        raise NotImplementedError()
+
+    def __call__(self, namelist_object, namelist_name):
+        """Tweak **namelist_object** that was read from the **namelist_name** file."""
+        if self._tlists_store:
+            self._setup_nam_obj(namelist_object, namelist_name)
+            return True
+        else:
+            return False
+
+
+class IfsOutputsConfigurator(IfsOutputsAbstractConfigurator):
+    """Utility class to configure the IFS model regarding output data."""
+
+    @staticmethod
+    def _get_namblock(namelist_object, nam_block):
+        """Get or create a **nam_block** namelist."""
+        if nam_block in namelist_object:
+            return namelist_object[nam_block]
+        else:
+            return namelist_object.newblock(nam_block)
+
+    @staticmethod
+    def _set_namvar_value(namblock, var, value, namname):
+        """Set a value in a **namblock** namelist and log it."""
+        namblock[var] = value
+        logger.info(
+            "Setup &%s %s=%s / (file: %s)",
+            namblock.name,
+            var,
+            namblock.nice(value),
+            namname,
+        )
+
+    @staticmethod
+    def _clean_namvar(namblock, var, namname):
+        """Clean the **var** value from the **namblock** namelist."""
+        todo = {k for k in namblock.keys() if re.match(var + r"($|\(|%)", k)}
+        if todo:
+            for k in todo:
+                namblock.delvar(k)
+            logger.info(
+                "Cleaning %s variable in namelist &%s (file: %s)",
+                var,
+                namblock.name,
+                namname,
+            )
+
+    def _generic_terms_setup(self, namct0, namct1, what, terms, namname):
+        """Setup a given kind of output data (in a generic way)."""
+        if terms is not None:
+            sign = -1 if self.fcterm_unit == "h" else 1
+            with_minutes = any([t.minute > 0 for t in terms])
+            self._clean_namvar(namct0, "NFR{:s}".format(what), namname)
+            self._clean_namvar(namct0, "N{:s}TS".format(what), namname)
+            self._clean_namvar(namct0, "N{:s}TSMIN".format(what), namname)
+            self._set_namvar_value(
+                namct1, "N1{:s}".format(what), 1 if terms else 0, namname
+            )
+            if terms:
+                self._set_namvar_value(
+                    namct0,
+                    "N{:s}TS(0)".format(what),
+                    sign * len(terms),
+                    namname,
+                )
+                if with_minutes:
+                    if (
+                        "cy46" <= self.cycle < "cy47"
+                    ):  # Temporary fix for cy46 only
+                        self._set_namvar_value(
+                            namct0,
+                            "N{:s}TSMIN(0)".format(what),
+                            len(terms),
+                            namname,
+                        )
+                logger.info(
+                    "Setting up N%sTS and N%sTSMIN in &%s (file: %s)",
+                    what,
+                    what,
+                    namct0.name,
+                    namname,
+                )
+                for i, t in enumerate(terms):
+                    namct0["N{:s}TS({:d})".format(what, i + 1)] = sign * t.hour
+                if with_minutes:
+                    for i, t in enumerate(terms):
+                        namct0["N{:s}TSMIN({:d})".format(what, i + 1)] = (
+                            t.minute
+                        )
+
+    def _setup_nam_obj(self, namelist_object, namelist_name):
+        """Actualy tweak the IFS namelist."""
+        namoph = self._get_namblock(namelist_object, "NAMOPH")
+        namct0 = self._get_namblock(namelist_object, "NAMCT0")
+        namct1 = self._get_namblock(namelist_object, "NAMCT1")
+        # First take into account the **fcterm_unit**
+        self._set_namvar_value(
+            namoph, "LINC", self.fcterm_unit == "h", namelist_name
+        )
+        # Setup outputs
+        self._generic_terms_setup(
+            namct0, namct1, "HIS", self.modelstate, namelist_name
+        )
+        self._generic_terms_setup(
+            namct0, namct1, "SFXHIS", self.surf_modelstate, namelist_name
+        )
+        self._generic_terms_setup(
+            namct0, namct1, "SDI", self.spectral_diag, namelist_name
+        )
+        self._generic_terms_setup(
+            namct0, namct1, "POS", self.post_processing, namelist_name
+        )
+        # Extra fixup for fullpos
+        if self.post_processing is not None:
+            if not self.post_processing:
+                self._set_namvar_value(namct0, "NFPOS", 0, namelist_name)
+            else:
+                # Do not overwrite a pre-existing positive value:
+                if "NFPOS" not in namct0 or namct0["NFPOS"] == 0:
+                    self._set_namvar_value(namct0, "NFPOS", 1, namelist_name)