vortex-nwp 2.0.0b1__py3-none-any.whl

vortex/nwp/tools/drhook.py

@@ -0,0 +1,62 @@
+"""
+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,268 @@
+"""
+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,99 @@
+"""
+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,163 @@
+"""
+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)