diff --git a/.gitignore b/.gitignore index ef35fc65..8979fceb 100644 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,22 @@ docs/_build/ .docker-current .eggs +# tests .pytest_cache tests/*.idx +tests/data +major_basins_of_the_world_0_0_0* +Major_Basins_of_the_World.* +zzz_for_spectra.grib +# diagnostics +diag/data +diag/mem_result.yaml +diag/.benchmarks + +# vscode +.vscode/ +.env + +# mac +.DS_Store diff --git a/.travis.yml b/.travis.yml index 07f0ee15..3bdb389f 100644 --- a/.travis.yml +++ b/.travis.yml @@ -38,12 +38,12 @@ before_install: # auto-yes for conda - conda config --set always_yes yes # update conda - #- conda update -n base -c conda-forge conda + - conda update -n base -c conda-forge conda install: # install deps - - pip install netcdf4 xarray cfgrib pytest pytest-cov black coveralls - - conda install -c conda-forge metview-batch + - pip install netcdf4 xarray cfgrib scipy pytest pytest-cov black coveralls PyYAML + - conda install -c conda-forge "metview-batch>=5.14.1" #---------------------------------# # build configuration # diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 51bde4b8..9268575e 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,7 +1,133 @@ Changelog for Metview's Python interface + ======================================== +1.16.1 +------------------- +- added support for numpy 2.0 + + +1.16.0 +------------------- +- added helper functions to aid automatic conversion from Macro to Python +- added support for reading PNG files +- added string() function from Metview +- fixed warning message in select() method with pandas 2.2 + + +1.15.0 +------------------ +- Added arguments() function to mimic the behaviour of the Macro function of the same name +- Fixed case of mean(dim=...) where the input fieldset is not a complete hypercube + + +1.14.0 +------------------ +- Fieldset function stdev() now accepts 'dim' argument to compute over given dimension +- added new function plot_xs_avg() to provide high-level plotting interface for average cross sections +- Fixed issue where functions that take numpy arrays as input did not work properly with multi-dimensional arrays + + +1.13.1 +------------------ +- fixed memory leaks when passing strings, numbers and dates from the binary layer to the Python layer + + +1.13.0 +------------------ +- Fieldset functions mean() and sum() now accept 'dim' argument to compute over given dimension +- function gallery.load_dataset() now downloads from a new server +- added new functions smooth_n_point() and smooth_gaussian() to perform spatial smoothing on fieldsets with lat-lon grids +- added new function convolve() to perform spatial 2D convolution on fieldsets with lat-lon grids + + +1.12.0 +------------------ +- add 'reverse' and/or operators between Fieldsets and bools (e.g. 'True & my_fieldset') +- fixed issue where automatic indexing of a dataset fails + + +1.11.0 +------------------ +- make fieldset parameter selector operator work in pure Python Fieldset +- make vector data extraction operator work for 100m and 200m wind +- select(): do not try to interpret shortName values as vector field names +- select(): do not apply sorting on results +- Fieldsets: add sort method based on indexed metadata +- select(): fix issue where cannot filter by date when data date has no year component (e.g. monthly climatologies) +- Requests: allow True/False as argument values when creating Requests + + +1.10.5 +------------------ +- fix for automatic indexing of dataset + + +1.10.0 +------------------ +- experimental pure Python implementation of Fieldset - for internal testing at the moment + + +1.9.0 +------------------ +- the plot functions now automatically plot inline if running inside a Jupyter notebook + - it is no longer necessary to call setoutput('jupyter') + - call setoutput('screen') to force the interactive plot window to appear +- inline plots in Jupyter notebooks will be automatically trimmed of surrounding whitespace if pillow is installed +- new functions to build popup dialogs and read in user input. Available via the newly added ui module. + - ui.dialog() + - ui.any() + - ui.colour() + - ui.icon() + - ui.option_menu() + - ui.slider() + - ui.toggle() +- added high-level plotting functions to be used with Datasets or in Jupyter notebooks + - plot_maps() + - plot_diff_maps() + - plot_xs() + - plot_rmse() +- new object Track to represent a storm track +- new function make_geoview() to generate a geoview object with predefined settings +- new interface for Datasets + - a Dataset represents a collection of data files (GRIB and CSV) and a set of predefined styles to visualise the data. Ideal for training courses or case studies. + - see Jupyter notebook example at https://metview.readthedocs.io/en/latest/notebook_gallery.html +- added new keyword argument called check_local to gallery.load_dataset(). If it is True and the data file exists locally it will not be downloaded. +- fixed issue when describe() crashed when called with a paramId + + +1.8.1 +------------------ +- fixed case where map_area_gallery() crashed +- fixed case where map_style_gallery() crashed +- fixed issue where plot_maps() could not plot wind data +- fixed issue where a style could not be updated when verb argument is specified + + +1.8.0 +------------------ +- new functions/methods on Fieldset to give an overview of contents: + - fs.describe() + - fs.describe("tp") + - fs.ls() + - see Jupyter notebook example at https://metview.readthedocs.io/en/latest/notebook_gallery.html +- new GRIB filtering function, select(), offers different filtering options from read() and is faster + - see Jupyter notebook example at https://metview.readthedocs.io/en/latest/notebook_gallery.html +- new shorthand way to select parameters from Fieldsets, e.g. + - g = fs["t"] + - g = fs["t500"] + - g = fs["t500hPa"] +- the Fieldset constructor can now take a list of paths to GRIB files or a wildcard: + - e.g. a = mv.Fieldset(path=["/path1/to/data1.grib", "relpath/data2.grib"]) + - e.g. a = mv.Fieldset(path="data/*.grib") +- the result of a call to mcont() etc can now be modified, e.g. + - c = mv.mcont() ; c["contour_line_colour"] = "green" ; mv.plot(data, c) + - gv.update({"MAP_COASTLINE_land_SHADE_COLOUR": "green"}, sub="COASTlines") +- improved the output of print(Fieldset): + - "Fieldset (6 fields)" + + 1.7.2 ------------------ - new argument to setoutput(plot_widget=) - default True, set False to allow images to be saved into the notebook diff --git a/MANIFEST.in b/MANIFEST.in index e78cc396..ac67a476 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -4,6 +4,7 @@ include LICENSE include Makefile include *.rst include tox.ini +graft metview/etc recursive-include docs *.gitkeep recursive-include docs *.py diff --git a/README.rst b/README.rst index c66dbc67..64bb5316 100644 --- a/README.rst +++ b/README.rst @@ -3,7 +3,7 @@ Metview Python bindings ======================= Python interface to Metview, a meteorological workstation and batch system for accessing, examining, manipulating and visualising meteorological data. -See documentation at https://confluence.ecmwf.int/metview/Metview's+Python+Interface +See documentation at https://metview.readthedocs.io/en/latest/index.html Try the example notebooks on Binder! @@ -18,8 +18,8 @@ Requirements ------------ - A working Metview 5 installation (at least version 5.0.3, ideally 5.3.0 or above), either from binaries or built from source. - Binary installation packages are available for many Linux distributions. - See https://confluence.ecmwf.int/metview/Releases + Conda packages are available for Linux, and native packages are available for many Linux distributions. + See https://metview.readthedocs.io/en/latest/install.html - An alternative is to build from the Metview Source Bundle. See https://confluence.ecmwf.int/metview/The+Metview+Source+Bundle @@ -38,6 +38,11 @@ The package is installed from PyPI with:: $ pip install metview +or from conda-forge with:: + + $ conda install metview-python -c conda-forge + + Test ---- @@ -81,7 +86,7 @@ Code quality .. image:: https://travis-ci.com/ecmwf/metview-python.svg?branch=m License ------- -Copyright 2017-2020 European Centre for Medium-Range Weather Forecasts (ECMWF). +Copyright 2017-2021 European Centre for Medium-Range Weather Forecasts (ECMWF). Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/docs/conf.py b/docs/conf.py index 645dc2ce..5cda8703 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -52,8 +52,8 @@ master_doc = "index" # General information about the project. -project = u"metview" -copyright = u"2017-2019, European Centre for Medium-Range Weather Forecasts (ECMWF)." +project = "metview" +copyright = "2017-2019, European Centre for Medium-Range Weather Forecasts (ECMWF)." # The version info for the project you're documenting, acts as replacement # for |version| and |release|, also used in various other places throughout @@ -203,7 +203,7 @@ # (source start file, target name, title, author, documentclass # [howto/manual]). latex_documents = [ - ("index", "metview.tex", u"metview Documentation", u"Metview python API", "manual"), + ("index", "metview.tex", "metview Documentation", "Metview python API", "manual"), ] # The name of an image file (relative to this directory) to place at @@ -231,7 +231,7 @@ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [("index", "metview", u"metview Documentation", [u"Metview python API"], 1)] +man_pages = [("index", "metview", "metview Documentation", ["Metview python API"], 1)] # If true, show URL addresses after external links. # man_show_urls = False @@ -246,8 +246,8 @@ ( "index", "metview", - u"metview Documentation", - u"Metview python API", + "metview Documentation", + "Metview python API", "metview", "One line description of project.", "Miscellaneous", diff --git a/environment.yml b/environment.yml index a097c38f..86271b55 100644 --- a/environment.yml +++ b/environment.yml @@ -1,7 +1,7 @@ channels: - conda-forge dependencies: - - metview-batch + - metview-batch >=5.16.0 - pip - pip: - metview >=1.4.0 diff --git a/metview/__init__.py b/metview/__init__.py index cc68a97f..4a02c67d 100644 --- a/metview/__init__.py +++ b/metview/__init__.py @@ -8,6 +8,7 @@ # nor does it submit to any jurisdiction. # requires a Python 3 interpreter +import os import sys if sys.version_info[0] < 3: # pragma: no cover @@ -22,13 +23,22 @@ # catch errors differently if len(sys.argv) != 2 or sys.argv[0] != "-m" or sys.argv[1] != "selfcheck": + try: + from . import bindings as _bindings - from . import bindings as _bindings + _bindings.bind_functions(globals(), module_name=__name__) - _bindings.bind_functions(globals(), module_name=__name__) + # Remove "_bindings" from the public API. + del _bindings - # Remove "_bindings" from the public API. - del _bindings + from . import gallery + from . import style + from metview.metviewpy import indexer + from . import title + from . import layout + from . import ui + from . import compat - -from . import gallery + except Exception as exp: + if "METVIEW_PYTHON_ONLY" not in os.environ: + raise exp diff --git a/metview/bindings.py b/metview/bindings.py index 3558a430..4f0c40af 100644 --- a/metview/bindings.py +++ b/metview/bindings.py @@ -7,21 +7,38 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. - +import builtins import datetime +from enum import Enum import keyword +import logging import os import pkgutil import signal import tempfile -import builtins -from enum import Enum import cffi import numpy as np +from metview.metviewpy.indexdb import FieldsetDb +from metview.dataset import Dataset +from metview.style import ( + GeoView, + Style, + Visdef, + map_area_gallery, + map_style_gallery, + make_geoview, +) +from metview import plotting +from metview.metviewpy.ipython import is_ipython_active, import_widgets +from metview.metviewpy import utils + +__version__ = "1.16.1" -__version__ = "1.7.2" + +# logging.basicConfig(level=logging.DEBUG, format="%(levelname)s - %(message)s") +LOG = logging.getLogger(__name__) def string_from_ffi(s): @@ -85,11 +102,14 @@ def __init__(self): try: subprocess.Popen(metview_flags) except Exception as exp: # pragma: no cover - print( - "Could not run the Metview executable ('" + metview_startup_cmd + "'); " - "check that the binaries for Metview (version 5 at least) are installed " - "and are in the PATH." - ) + if "METVIEW_PYTHON_ONLY" not in os.environ: + print( + "Could not run the Metview executable ('" + + metview_startup_cmd + + "'); " + "check that the binaries for Metview (version 5 at least) are installed " + "and are in the PATH." + ) raise exp # wait for Metview to respond... @@ -259,48 +279,55 @@ class Request(dict, Value): verb = "UNKNOWN" def __init__(self, req, myverb=None): - self.val_pointer = None - # initialise from Python object (dict/Request) + if isinstance(req, Request): # copy an existing Request + self.xx = _request(req.get_verb(), req) # to avoid deletion of Macro object + self.val_pointer = self.xx.val_pointer + self.verb = req.get_verb() + self.update(req) # copy into dict + return + + if myverb: + self.verb = myverb if isinstance(req, dict): - self.update(req) - self.to_metview_style() - if isinstance(req, Request): - self.verb = req.verb - self.val_pointer = req.val_pointer - else: - if myverb is not None: - self.set_verb(myverb) + self.create_new(self.verb, req) + return # initialise from a Macro pointer else: Value.__init__(self, req) self.verb = string_from_ffi(lib.p_get_req_verb(req)) - n = lib.p_get_req_num_params(req) - for i in range(0, n): - param = string_from_ffi(lib.p_get_req_param(req, i)) - self[param] = self[param] - # self['_MACRO'] = 'BLANK' - # self['_PATH'] = 'BLANK' + super().update(self.to_dict()) # update dict def __str__(self): return "VERB: " + self.verb + super().__str__() - # translate Python classes into Metview ones where needed - def to_metview_style(self): - for k in list(self): + def to_dict(self): + keys = _keywords(self) + d = {} + for k in keys: + d[k] = self[k] + return d + + # translate Python classes into Metview ones where needed - single parameter + def item_to_metview_style(self, key, value): + modified = False + delete_original_key = False - # bool -> on/off - v = self.get(k) # self[k] returns 1 for True - if isinstance(v, bool): - conversion_dict = {True: "on", False: "off"} - self[k] = conversion_dict[v] + # bool -> on/off + if isinstance(value, bool): + conversion_dict = {True: "on", False: "off"} + value = conversion_dict[value] + modified = True - # class_ -> class (because 'class' is a Python keyword and cannot be - # used as a named parameter) - elif k == "class_": - self["class"] = v - del self["class_"] + # class_ -> class (because 'class' is a Python keyword and cannot be + # used as a named parameter) + elif key == "class_": + key = "class" + delete_original_key = True + modified = True + + return (key, value, modified, delete_original_key) def set_verb(self, v): self.verb = v @@ -308,27 +335,52 @@ def set_verb(self, v): def get_verb(self): return self.verb + def create_new(self, rverb, rdict): + r = definition() # new, empty definition + self.val_pointer = r.val_pointer + self.xx = r # to avoid deletion of Macro object + self.update(rdict) # will set all parameters via __setitem__ + return + def push(self): # if we have a pointer to a Metview Value, then use that because it's more # complete than the dict if self.val_pointer: Value.push(self) else: - r = lib.p_new_request(self.verb.encode("utf-8")) - - # to populate a request on the Macro side, we push each - # value onto its stack, and then tell it to create a new - # parameter with that name for the request. This allows us to - # use Macro to handle the addition of complex data types to - # a request - for k, v in self.items(): - push_arg(v) - lib.p_set_request_value_from_pop(r, k.encode("utf-8")) - + self.create_new(self.verb, self) lib.p_push_request(r) + def update(self, items, sub=""): + if sub: + if not isinstance(sub, str): + raise IndexError("sub argument should be a string") + subreq = self[sub.upper()] + if subreq: + subreq.update(items) + self[sub] = subreq + else: + raise IndexError("'" + sub + "' not a valid subrequest in " + str(self)) + else: + for key in items: + self.__setitem__(key, items[key]) + def __getitem__(self, index): - return subset(self, index) + item = subset(self, index) + # subrequests can return '#' if not uppercase + if isinstance(item, str) and item == "#": + item = subset(self, index.upper()) + return item + + def __setitem__(self, index, value): + if (self.val_pointer) and (value is not None): + new_key, new_val, _, _ = self.item_to_metview_style(index, value) + push_arg(new_key) + push_arg(new_val) + lib.p_set_subvalue_from_arg_stack(self.val_pointer) + dict.__setitem__(self, new_key, new_val) + else: + dict.__setitem__(self, index, value) def push_bytes(b): @@ -374,18 +426,18 @@ def push_vector(npa): dtype = npa.dtype if dtype == np.float64: # can directly pass the data buffer cffi_buffer = ffi.cast("double*", npa.ctypes.data) - lib.p_push_vector_from_double_array(cffi_buffer, len(npa), np.nan) + lib.p_push_vector_from_double_array(cffi_buffer, npa.size, np.nan) elif dtype == np.float32: # can directly pass the data buffer cffi_buffer = ffi.cast("float*", npa.ctypes.data) - lib.p_push_vector_from_float32_array(cffi_buffer, len(npa), np.nan) + lib.p_push_vector_from_float32_array(cffi_buffer, npa.size, np.nan) elif dtype == bool: # convert first to float32 f32_array = npa.astype(np.float32) cffi_buffer = ffi.cast("float*", f32_array.ctypes.data) - lib.p_push_vector_from_float32_array(cffi_buffer, len(f32_array), np.nan) + lib.p_push_vector_from_float32_array(cffi_buffer, f32_array.size, np.nan) elif dtype == int: # convert first to float64 f64_array = npa.astype(np.float64) cffi_buffer = ffi.cast("double*", f64_array.ctypes.data) - lib.p_push_vector_from_double_array(cffi_buffer, len(f64_array), np.nan) + lib.p_push_vector_from_double_array(cffi_buffer, f64_array.size, np.nan) else: raise TypeError( "Only float32 and float64 numPy arrays can be passed to Metview, not ", @@ -393,6 +445,39 @@ def push_vector(npa): ) +def push_style_object(s): + r = s.to_request() + if isinstance(r, list): + push_list(r) + else: + r.push() + + +def valid_date(*args, base=None, step=None, step_units=None): + if len(args) != 0: + return call("valid_date", *args) + else: + assert isinstance(base, datetime.datetime) + step = [] if step is None else step + step_units = datetime.timedelta(hours=1) if step_units is None else step_units + if not isinstance(step, list): + step = [step] + return [base + step_units * int(x) for x in step] + + +def sort(*args, **kwargs): + if "_cpp_implementation" in kwargs: + return call("sort", *args) + elif len(args) != 0 and isinstance(args[0], Fieldset): + if len(args) == 1: + return args[0].sort(**kwargs) + else: + args = args[1:] + return args[0].sort(*args, **kwargs) + else: + return call("sort", *args) + + class File(Value): def __init__(self, val_pointer): Value.__init__(self, val_pointer) @@ -474,9 +559,15 @@ def __abs__(self): def __and__(self, other): return met_and(self, other) + def __rand__(self, other): + return met_and(other, self) + def __or__(self, other): return met_or(self, other) + def __ror__(self, other): + return met_or(other, self) + def __invert__(self): return met_not(self) @@ -575,13 +666,24 @@ def __init__(self, val_pointer=None, path=None, fields=None): element_types=Fieldset, support_slicing=True, ) + self._db = None + self._ds_param_info = None + self._label = "" if (path is not None) and (fields is not None): raise ValueError("Fieldset cannot take both path and fields") if path is not None: - temp = read(path) - self.steal_val_pointer(temp) + if isinstance(path, list): + v = [] + for p in path: + v.extend(utils.get_file_list(p)) + path = v + else: + path = utils.get_file_list(path) + + # fill the 'fields' var - it will be used a few lines down + fields = [read(p) for p in path] if fields is not None: for f in fields: @@ -592,6 +694,7 @@ def append(self, other): if self.val_pointer is not None: # we will overwrite ourselves, so delete lib.p_destroy_value(self.val_pointer) self.steal_val_pointer(temp) + self._db = None def to_dataset(self, **kwarg): # soft dependency on cfgrib @@ -603,6 +706,163 @@ def to_dataset(self, **kwarg): dataset = xr.open_dataset(self.url(), engine="cfgrib", backend_kwargs=kwarg) return dataset + def index(self, path=""): + pass + + def load_index(self, path): + pass + + def _scan(self): + if self._db is None: + self._db = FieldsetDb(fs=self) + self._db.scan() + + def _get_db(self): + if self._db is None: + self._db = FieldsetDb(fs=self) + assert self._db is not None + return self._db + + def select(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + return self._get_db().select(**args[0]) + else: + return self._get_db().select(**kwargs) + + def describe(self, *args, **kwargs): + return self._get_db().describe(*args, **kwargs) + + def ls(self, **kwargs): + return self._get_db().ls(**kwargs) + + def sort(self, *args, **kwargs): + return self._get_db().sort(*args, **kwargs) + + def apply_function_over_dim(self, dim, preserve_dims, func_to_run, name, **kwargs): + if len(self) == 0: + raise ValueError(f"Input to function {name} is an empty Fieldset") + + if dim is None: + return func_to_run(self, **kwargs) + + import itertools + + if preserve_dims: + _preserve_dims = preserve_dims + else: + _preserve_dims = ["shortName", "level", "step", "number", "date", "time"] + if dim in _preserve_dims: + _preserve_dims.remove(dim) + dim_combos = {k: unique(self.grib_get_string(k)) for k in _preserve_dims} + keys, values = zip(*dim_combos.items()) + perms = [dict(zip(keys, v)) for v in itertools.product(*values)] + # e.g. [{level=1000,shortName="t",date=20220101, time=6}, ...] + fieldsets_to_apply_function_to = [self.select(**x) for x in perms] + result = Fieldset( + fields=[ + func_to_run(x, **kwargs) + for x in fieldsets_to_apply_function_to + if len(x) != 0 + ] + ) + return result + + def mean(self, dim=None, preserve_dims=None, missing=False): + return self.apply_function_over_dim( + dim, preserve_dims, met_mean, "mean", missing=missing + ) + + def sum(self, dim=None, preserve_dims=None, missing=False): + return self.apply_function_over_dim( + dim, preserve_dims, met_sum, "sum", missing=missing + ) + + def stdev(self, dim=None, preserve_dims=None, missing=False): + return self.apply_function_over_dim(dim, preserve_dims, met_stdev, "stdev") + + @property + def ds_param_info(self): + if self._ds_param_info is None: + self._ds_param_info = FieldsetDb.make_param_info(self) + return self._ds_param_info + + @property + def label(self): + if self._label is not None and self._label: + return self._label + elif self._db is not None: + return self._db.label + else: + return str() + + @label.setter + def label(self, value): + if self._db is not None: + if self._db.label: + print( + f"Warning: cannot set label! It is already set on index database as {self._db.label}!" + ) + return + else: + self._db.label = value + + self._label = value + + def _unique_metadata(self, key): + return self._get_db().unique(key) + + def ds_style(self, plot_type="map"): + from metview import style + + return style.get_db().style(self, plot_type=plot_type) + + def ds_style_list(self, plot_type="map"): + from metview import style + + return style.get_db().style_list(self, plot_type=plot_type) + + def speed(self, *args): + if len(args) == 0: + u = self[0::2] + v = self[1::2] + if len(u) != len(v): + raise Exception( + f"Fieldsets must contain an even number of fields for this operation! len={len(self)} is not even!" + ) + sp = u.speed(v) + sp._init_db_from(self) + return sp + else: + return call("speed", self, *args) + + def deacc(self, **kwargs): + r = utils.deacc(self, **kwargs) + r._init_db_from(self) + return r + + def convolve(self, *args, **kwargs): + return utils.convolve(self, *args, **kwargs) + + def smooth_n_point(self, *args, **kwargs): + return utils.smooth_n_point(self, *args, **kwargs) + + def smooth_gaussian(self, *args, **kwargs): + return utils.smooth_gaussian(self, *args, **kwargs) + + def _init_db_from(self, other): + if self._db is None and other._db is not None: + self._db = FieldsetDb(self, other.label) + self._db.load() + + def __getitem__(self, key): + if isinstance(key, str): + self._scan() + if self._db is not None: + return self._db.select_with_name(key) + return None + else: + return super().__getitem__(key) + def __getstate__(self): # used for pickling # we cannot (and do not want to) directly pickle the Value pointer @@ -619,6 +879,13 @@ def __setstate__(self, state): self.__dict__.update(state) self.__init__(val_pointer=None, path=state["url_path"]) + def __str__(self): + n = int(self.count()) + s = "s" + if n == 1: + s = "" + return "Fieldset (" + str(n) + " field" + s + ")" + class Bufr(FileBackedValue): def __init__(self, val_pointer): @@ -723,7 +990,7 @@ def dataset_to_fieldset(val, **kwarg): # we try to import xarray as locally as possible to reduce startup time # try to write the xarray as a GRIB file, then read into a fieldset import xarray as xr - import cfgrib + from cfgrib.xarray_to_grib import to_grib if not isinstance(val, xr.core.dataset.Dataset): raise TypeError( @@ -737,7 +1004,7 @@ def dataset_to_fieldset(val, **kwarg): try: # could add keys, e.g. grib_keys={'centre': 'ecmf'}) - cfgrib.to_grib(val, tmp, **kwarg) + to_grib(val, tmp, **kwarg) except: print( "Error trying to write xarray dataset to GRIB for conversion to Metview Fieldset" @@ -789,6 +1056,9 @@ def __init__(self): (datetime.date, lambda n: push_datetime_date(n)), (np.ndarray, lambda n: push_vector(n)), (File, lambda n: n.push()), + (GeoView, lambda n: push_style_object(n)), + (Style, lambda n: push_style_object(n)), + (Visdef, lambda n: push_style_object(n)), ) def push_value(self, val): @@ -851,6 +1121,7 @@ def datestring_from_metview(val): mdate = string_from_ffi(lib.p_value_as_datestring(val)) dt = datetime.datetime.strptime(mdate, "%Y-%m-%dT%H:%M:%S") + lib.p_destroy_value(val) return dt @@ -888,7 +1159,15 @@ def handle_error(val): def string_from_metview(val): - return string_from_ffi(lib.p_value_as_string(val)) + s = string_from_ffi(lib.p_value_as_string(val)) + lib.p_destroy_value(val) + return s + + +def number_from_metview(val): + n = lib.p_value_as_number(val) + lib.p_destroy_value(val) + return n class MvRetVal(Enum): @@ -916,7 +1195,7 @@ class ValueReturner: def __init__(self): self.funcs = {} - self.funcs[MvRetVal.tnumber.value] = lambda val: lib.p_value_as_number(val) + self.funcs[MvRetVal.tnumber.value] = lambda val: number_from_metview(val) self.funcs[MvRetVal.tstring.value] = lambda val: string_from_metview(val) self.funcs[MvRetVal.tgrib.value] = lambda val: Fieldset(val) self.funcs[MvRetVal.trequest.value] = lambda val: Request(val) @@ -938,9 +1217,17 @@ def translate_return_val(self, val): try: return self.funcs[rt](val) except Exception: - raise Exception( - "value_from_metview got an unhandled return type: " + str(rt) - ) + # if the type is unknown, it might be a type that is actually stored + # as a request rather than as a MARS type, e.g. PNG + try: + if rt == 99: + rt = MvRetVal.trequest.value + return self.funcs[rt](val) + except Exception: + raise Exception( + "value_from_metview got an unhandled return type and could not convert to a Request: " + + str(rt) + ) vr = ValueReturner() @@ -953,6 +1240,19 @@ def value_from_metview(val): return retval +# ----------------------------------------------------------------------------- +# +# ----------------------------------------------------------------------------- + + +def to_dataset(fs, *args, **kwargs): + return fs.to_dataset(args, kwargs) + + +def to_dataset(fs, *args, **kwargs): + return fs.to_dataset(args, kwargs) + + # ----------------------------------------------------------------------------- # Creating and calling Macro functions # ----------------------------------------------------------------------------- @@ -979,7 +1279,7 @@ def make(mfname): def wrapped(*args, **kwargs): err = _call_function(mfname, *args, **kwargs) if err: - pass # throw Exceception + pass # throw Exception val = lib.p_result_as_value() return value_from_metview(val) @@ -987,6 +1287,19 @@ def wrapped(*args, **kwargs): return wrapped +def _make_function_for_object(name): + """ + Creates a function to invoke the method called name on obj. This will make it + possible to call some object methods as global functions. E.g.: if name="ls" and + f is a Fieldset we could invoke the ls() method as mv.ls(f) on top of f.ls() + """ + + def fn(obj, *args, **kwargs): + return getattr(obj, name)(*args, **kwargs) + + return fn + + def bind_functions(namespace, module_name=None): """Add to the module globals all metview functions except operators like: +, &, etc.""" for metview_name in make("dictionary")(): @@ -1005,11 +1318,13 @@ def bind_functions(namespace, module_name=None): # else: # print('metview function %r not bound to python' % metview_name) - # HACK: some fuctions are missing from the 'dictionary' call. + # HACK: some functions are missing from the 'dictionary' call. namespace["neg"] = make("neg") namespace["nil"] = make("nil") + namespace["dialog"] = make("dialog") namespace["div"] = div namespace["mod"] = mod + namespace["string"] = make("string") # override some functions that need special treatment # FIXME: this needs to be more structured namespace["plot"] = plot @@ -1018,23 +1333,57 @@ def bind_functions(namespace, module_name=None): namespace["version_info"] = version_info namespace["merge"] = merge namespace["dataset_to_fieldset"] = dataset_to_fieldset - + namespace["valid_date"] = valid_date + namespace["sort"] = sort + namespace["load_dataset"] = Dataset.load_dataset + namespace["plot_maps"] = plotting.plot_maps + namespace["plot_diff_maps"] = plotting.plot_diff_maps + namespace["plot_xs"] = plotting.plot_xs + namespace["plot_xs_avg"] = plotting.plot_xs_avg + namespace["plot_stamp"] = plotting.plot_stamp + namespace["plot_rmse"] = plotting.plot_rmse + namespace["plot_cdf"] = plotting.plot_cdf + namespace["map_style_gallery"] = map_style_gallery + namespace["map_area_gallery"] = map_area_gallery + namespace["make_geoview"] = make_geoview + namespace["arguments"] = met_arguments namespace["Fieldset"] = Fieldset namespace["Request"] = Request + # some ui specific functions are prefixed with _. They will be exposed via the ui module! + for name in ["dialog", "any", "colour", "icon", "option_menu", "slider", "toggle"]: + namespace["_" + name] = namespace[name] + namespace.pop(name) + + # add some object methods the to global namespace + for name in [ + "to_dataset", + "to_dataframe", + "ls", + "describe", + "select", + "convolve", + "smooth_n_point", + "smooth_gaussian", + ]: + namespace[name] = _make_function_for_object(name) + # some explicit bindings are used here add = make("+") call = make("call") count = make("count") +definition = make("definition") div = make("/") download = make("download") equal = make("=") filter = make("filter") greater_equal_than = make(">=") greater_than = make(">") +_keywords = make("keywords") lower_equal_than = make("<=") lower_than = make("<") +met_mean = make("mean") met_merge = make("&") met_not_eq = make("<>") met_plot = make("plot") @@ -1049,10 +1398,14 @@ def bind_functions(namespace, module_name=None): metzoom = make("metzoom") sub = make("-") subset = make("[]") +met_stdev = make("stdev") +met_sum = make("sum") met_and = make("and") met_or = make("or") met_not = make("not") met_version_info = make("version_info") +_request = make("request") +unique = make("unique") write = make("write") @@ -1078,12 +1431,24 @@ def merge(*args): class Plot: + has_pillow = None + padding = np.array([x * 40 for x in [-1, -1, 1, 1]]) + def __init__(self): self.plot_to_jupyter = False self.plot_widget = True self.jupyter_args = {} + self.called_once = False + self.setoutput_called_once = False def __call__(self, *args, **kwargs): + # first time called? If we are in Jupyter and user did not specify, + # then plot inline to Jupyter by default + if not (self.called_once or self.setoutput_called_once): + if is_ipython_active(): + setoutput("jupyter") + self.called_once = True + if self.plot_to_jupyter: # pragma: no cover if self.plot_widget: return plot_to_notebook(args, **kwargs) @@ -1104,6 +1469,42 @@ def __call__(self, *args, **kwargs): # None is better for Python return None + def crop_image(self, path): + if Plot.has_pillow is None: + try: + import PIL + + Plot.has_pillow = True + except ImportError as e: + Plot.has_pillow = False + if Plot.has_pillow: + try: + from PIL import Image + from PIL import ImageOps + + im = Image.open(path) + im.load() + + # find inner part + im_invert = im.convert("RGB") + box = ImageOps.invert(im_invert).getbbox() + + # crop to box + if box[2] - box[0] > 100 or box[3] - box[1] > 100: + box = list(np.asarray(box) + Plot.padding) + box = ( + max(0, box[0]), + max(0, box[1]), + min(im.size[0], box[2]), + min(im.size[1], box[3]), + ) + im_crop = im.crop(box) + im.close() + im_crop.save(path) + except Exception as e: + # print(f"ERROR={e}") + pass + plot = Plot() @@ -1141,6 +1542,9 @@ def plot_to_notebook(*args, **kwargs): # pragma: no cover files = [os.path.join(tempdirpath, f) for f in sorted(filenames)] + for f in files: + plot.crop_image(f) + if (animation_mode == True) or (animation_mode == "auto" and len(filenames) > 1): frame_widget = widgets.IntSlider( value=1, @@ -1218,6 +1622,8 @@ def plot_to_notebook_return_image(*args, **kwargs): # pragma: no cover plot.jupyter_args.update(output_name=base, output_name_first_page_number="off") met_setoutput(png_output(plot.jupyter_args)) met_plot(*args) + plot.crop_image(tmp) + image = Image(tmp) os.unlink(tmp) return image @@ -1228,19 +1634,16 @@ def plot_to_notebook_return_image(*args, **kwargs): # pragma: no cover # functionality. Since this occurs within a function, we need a little trickery to # get the IPython functions into the global namespace so that the plot object can use them def setoutput(*args, **kwargs): + plot.setoutput_called_once = True if "jupyter" in args: # pragma: no cover - try: - import IPython - - get_ipython = IPython.get_ipython - except ImportError as imperr: - print("Could not import IPython module - plotting to Jupyter will not work") - raise imperr - - # test whether we're in the Jupyter environment - if get_ipython() is not None: + if is_ipython_active(): + global widgets plot.plot_to_jupyter = True plot.plot_widget = kwargs.get("plot_widget", True) + if plot.plot_widget: + widgets = import_widgets() + if not widgets: + plot.plot_widget = False if "plot_widget" in kwargs: del kwargs["plot_widget"] plot.jupyter_args = kwargs @@ -1250,15 +1653,20 @@ def setoutput(*args, **kwargs): ) raise (Exception("Could not set output to jupyter")) - try: - global widgets - widgets = __import__("ipywidgets", globals(), locals()) - except ImportError as imperr: - print( - "Could not import ipywidgets module - plotting to Jupyter will not work" - ) - raise imperr - else: plot.plot_to_jupyter = False met_setoutput(*args) + + +def met_arguments(): + """Emulate the Macro arguments() function""" + import sys + + args = sys.argv[1:] + # these will all come in as strings; but Macro does a little processing on them + # in order to intelligently decide whether their types or string or number; + # creating a Request with these values will simulate this behaviour + args_dict = {i: i for i in args} + modified_args_dict = Request(args_dict).to_dict() + modified_args = [modified_args_dict[j] for j in modified_args_dict] + return modified_args diff --git a/metview/compat.py b/metview/compat.py new file mode 100644 index 00000000..64ccd1c4 --- /dev/null +++ b/metview/compat.py @@ -0,0 +1,86 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +import numpy as np + +import metview as mv + + +def concat(*args): + if len(args) <= 1: + raise ValueError(f"concat requires at least 2 arguments, got only {len(args)}") + + vals = list(args) + if vals[0] is None: + vals.pop(0) + if len(vals) == 1: + return vals[0] + + # concatenate strings + if isinstance(vals[0], str): + return "".join([str(v) for v in vals]) + elif isinstance(vals[0], np.ndarray): + return np.concatenate(vals) + elif isinstance( + vals[0], + (mv.Fieldset, mv.bindings.Geopoints, mv.bindings.GeopointSet, mv.bindings.Bufr), + ): + return mv.merge(*vals) + elif isinstance(vals[0], list): + first = [] + for v in vals: + if isinstance(v, list): + first.extend(v) + else: + first.append(v) + return first + + raise ValueError(f"concat failed to handle the specified arguments={args}") + + +def index4(vals, start, stop, step, num): + """ + Return a boolean index ndarray to select a subset of elements from ``vals``. + + Parameters + ---------- + vals: ndarray + Input array. + start: int + Start index. + stop: int + Stop index. + step: int + Step. + num: int + Number of elements to be selected for each step. + + Returns + ------- + ndarray + Boolean index array + + Examples + -------- + + >>> import numpy as np + >>> import metview + >>> vals = np.array(list(range(12))) + >>> r = index4(vals, 0, 11, 4, 2) + [ True True True True True True False False False False False False] + a[] + >>> vals[r] + array([0, 1, 4, 5, 8, 9]) + + """ + num = min(num, step) + m = np.zeros(len(vals), dtype=bool) + for i in range(start, stop, step): + m[i : i + num] = 1 + return m diff --git a/metview/dataset.py b/metview/dataset.py new file mode 100644 index 00000000..0f905e87 --- /dev/null +++ b/metview/dataset.py @@ -0,0 +1,461 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import logging +import os +from pathlib import Path +import shutil + +import pandas as pd +import requests +import yaml + +import metview as mv +from metview.metviewpy.indexer import ExperimentIndexer +from metview.metviewpy.indexdb import IndexDb +from metview.track import Track +from metview.metviewpy.param import init_pandas_options +from metview.metviewpy import utils + + +ETC_PATH = os.path.join(os.path.dirname(__file__), "etc") + +# logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") +# logging.basicConfig(level=logging.DEBUG, format="%(levelname)s - %(message)s") +LOG = logging.getLogger(__name__) + + +class ExperimentDb(IndexDb): + def __init__(self, name, **kwargs): + super().__init__(name, **kwargs) + self.fs = {} + self.vector_loaded = True + self._indexer = None + self.fieldset_class = mv.Fieldset + LOG.debug(f"rootdir_placeholder_value={self.rootdir_placeholder_value}") + + @staticmethod + def make_from_conf(name, conf, root_dir, db_root_dir, regrid_conf, dataset): + # LOG.debug(f"conf={conf}") + db = ExperimentDb( + name, + label=conf.get("label", ""), + desc=conf.get("desc", ""), + path=conf.get("dir", "").replace( + IndexDb.ROOTDIR_PLACEHOLDER_TOKEN, root_dir + ), + rootdir_placeholder_value=root_dir + if IndexDb.ROOTDIR_PLACEHOLDER_TOKEN in conf.get("dir", "") + or "merge" in conf + else "", + file_name_pattern=conf.get("fname", ""), + db_dir=os.path.join(db_root_dir, name), + merge_conf=conf.get("merge", []), + mapped_params={v: k for k, v in conf.get("mapped_params", {}).items()}, + regrid_from=regrid_conf.get(name, []), + blocks={}, + dataset=dataset, + ) + return db + + def _clone(self): + return ExperimentDb( + self.name, + label=self.label, + db_dir=self.db_dir, + mapped_params=self.mapped_params, + regrid_from=self.regrid_from, + dataset=self.dataset, + data_files=self.data_files, + rootdir_placeholder_value=self.rootdir_placeholder_value, + ) + + @property + def indexer(self): + if self._indexer is None: + self._indexer = ExperimentIndexer(self) + return self._indexer + + def scan(self, vector=True): + print(f"Generate index for dataset component={self.name} ...") + self.data_files = [] + # self.blocks = {} + self.indexer.scan() + + def load(self, keys=None, vector=True): + keys = [] if keys is None else keys + ivk = [x for x in keys if x not in self.indexer.allowed_keys()] + if ivk: + raise Exception( + f"{self} keys={ivk} cannot be used! The allowed set of keys={self.indexer.allowed_keys()}" + ) + + self.load_data_file_list() + if len(self.data_files) == 0: + self.scan(vector=True) + self.load_data_file_list() + + if len(self.blocks) == 0: + for key in ExperimentIndexer.get_storage_key_list(self.db_dir): + self.blocks[key] = ExperimentIndexer.read_dataframe(key, self.db_dir) + + def load_data_file_list(self): + if len(self.data_files) == 0: + try: + file_path = os.path.join(self.db_dir, "datafiles.yaml") + with open(file_path, "rt") as f: + self.data_files = yaml.safe_load(f) + if self.rootdir_placeholder_value: + self.data_files = [ + x.replace( + IndexDb.ROOTDIR_PLACEHOLDER_TOKEN, + self.rootdir_placeholder_value, + ) + for x in self.data_files + ] + # LOG.debug(f"data_files={self.data_files}") + for f in self.data_files: + assert os.path.isfile(f) + except: + pass + + def __getitem__(self, key): + if isinstance(key, str): + self.load() + return self.select_with_name(key) + return None + + def _filter_blocks(self, dims): + self.load() + dfs = {} + # LOG.debug(f"data_files={self.data_files}") + # LOG.debug(f"dims={dims}") + cnt = 0 + for key, df in self.blocks.items(): + # LOG.debug(f"key={key}") + # df = self._load_block(key) + # LOG.debug(f"df={df}") + f_df = self._filter_df(df=df, dims=dims) + # LOG.debug(f"df={df}" + if f_df is not None and not f_df.empty: + cnt += len(f_df) + # LOG.debug(f" matching rows={len(df)}") + dfs[key] = f_df + + # LOG.debug(f"total matching rows={cnt}") + return dfs + + def _extract_fields(self, df, fs, max_count): + if df.empty: + return None + + if "_fileIndex3" in df.columns: + comp_num = 3 + elif "_fileIndex2" in df.columns: + comp_num = 2 + elif "_fileIndex1" in df.columns: + comp_num = 1 + else: + return None + + idx = [[] for k in range(comp_num)] + comp_lst = list(range(comp_num)) + for row in df.itertuples(): + for comp in comp_lst: + idx_file = row[-1 - (comp_num - comp - 1) * 2] + idx_msg = row[-2 - (comp_num - comp - 1) * 2] + if not idx_file in self.fs: + self.fs[idx_file] = mv.read(self.data_files[idx_file]) + fs.append(self.fs[idx_file][idx_msg]) + idx[comp].append(len(fs) - 1) + # generate a new dataframe + df = df.copy() + for k, v in enumerate(idx): + df[f"_msgIndex{k+1}"] = v + df.drop([f"_fileIndex{x+1}" for x in range(comp_num)], axis=1, inplace=True) + return df + + def to_fieldset(self): + db, fs = self._get_fields({}) + fs._db = db + return fs + + def get_longname_and_units(self, shortName, paramId): + return "", "" + + +class TrackConf: + def __init__(self, name, conf, data_dir, dataset): + self.name = name + self.dataset = dataset + self.label = self.name + self.path = conf.get("dir", "").replace( + IndexDb.ROOTDIR_PLACEHOLDER_TOKEN, data_dir + ) + self.file_name_pattern = conf.get("fname", "") + # self.conf_dir = os.path.join("_conf", self.name) + self.data_files = [] + self.conf = conf + + def load_data_file_list(self): + if len(self.data_files) == 0: + self.data_files = utils.get_file_list( + self.path, file_name_pattern=self.file_name_pattern + ) + + def select(self, name): + tr = self._make(name) + if tr is None: + raise Exception(f"No track is available with name={name}!") + return tr + + def describe(self): + self.load_data_file_list() + init_pandas_options() + t = {"Name": [], "Suffix": []} + for f in self.data_files: + n, s = os.path.splitext(os.path.basename(f)) + t["Name"].append(n) + t["Suffix"].append(s) + df = pd.DataFrame.from_dict(t) + df.set_index("Name", inplace=True) + return df + + def _make(self, name): + self.load_data_file_list() + for f in self.data_files: + if name == os.path.basename(f).split(".")[0]: + c = { + x: self.conf.get(x, None) + for x in [ + "skiprows", + "date_index", + "time_index", + "lon_index", + "lat_index", + ] + } + return Track(f, **c) + return None + + +class Dataset: + """ + Represents a Dataset + """ + + URL = "https://get.ecmwf.int/repository/test-data/metview/dataset" + LOCAL_ROOT = os.getenv( + "MPY_DATASET_ROOT", os.path.join(os.getenv("HOME", ""), "mpy_dataset") + ) + COMPRESSION = "bz2" + + def __init__(self, name_or_path, load_style=True): + self.field_conf = {} + self.track_conf = {} + + self.path = name_or_path + if any(x in self.path for x in ["/", "\\", "..", "./"]): + self.name = os.path.basename(self.path) + else: + self.name = self.path + self.path = "" + + assert self.name + # print(f"name={self.name}") + + # set local path + if self.path == "": + self.path = os.path.join(self.LOCAL_ROOT, self.name) + + # If the path does not exists, it must be a built-in dataset. Data will be + # downloaded into path. + if not os.path.isdir(self.path): + if self.check_remote(): + self.fetch(forced=True) + else: + raise Exception( + f"Could not find dataset={self.name} either under path={self.path} or on data server" + ) + # WARN: we do not store the dataset in the CACHE any more. Check code versions before + # 09082021 to see how the CACHE was used. + + if load_style: + self.load_style() + + self.load() + + for _, c in self.field_conf.items(): + LOG.debug(f"{c}") + + @staticmethod + def load_dataset(*args, **kwargs): + return Dataset(*args, **kwargs) + + def check_remote(self): + try: + return ( + requests.head( + f"{self.URL}/{self.name}/conf.tar", allow_redirects=True + ).status_code + == 200 + ) + except: + return False + + def load(self): + data_dir = os.path.join(self.path, "data") + index_dir = os.path.join(self.path, "index") + data_conf_file = os.path.join(self.path, "data.yaml") + with open(data_conf_file, "rt") as f: + d = yaml.safe_load(f) + regrid_conf = d.get("regrid", {}) + for item in d["experiments"]: + ((name, conf),) = item.items() + if conf.get("type") == "track": + self.track_conf[name] = TrackConf(name, conf, data_dir, self) + else: + c = ExperimentDb.make_from_conf( + name, conf, data_dir, index_dir, regrid_conf, self + ) + self.field_conf[c.name] = c + + def scan(self, name=None): + # indexer = ExperimentIndexer() + if name: + if name in self.field_conf: + self.field_conf[name].scan() + # indexer.scan(self.field_conf[name], to_disk=True) + else: + for name, c in self.field_conf.items(): + LOG.info("-" * 40) + c.scan() + # indexer.scan(c, to_disk=True) + + def find(self, name, comp="field"): + if comp == "all": + f = self.field_conf.get(name, None) + if f is not None: + return f + else: + return self.track_conf.get(name, None) + elif comp == "field": + return self.field_conf.get(name, None) + elif comp == "track": + return self.track_conf.get(name, None) + else: + return None + + def describe(self): + init_pandas_options() + print("Dataset components:") + t = {"Component": [], "Description": []} + for _, f in self.field_conf.items(): + t["Component"].append(f.name) + t["Description"].append(f.desc) + for _, f in self.track_conf.items(): + t["Component"].append(f.name) + t["Description"].append("Storm track data") + df = pd.DataFrame.from_dict(t) + df.set_index("Component", inplace=True) + return df + + def fetch(self, forced=False): + if not os.path.isdir(self.path): + Path(self.path).mkdir(0o755, parents=True, exist_ok=True) + + files = { + "conf.tar": ["data.yaml", "conf"], + f"index.tar.{self.COMPRESSION}": ["index"], + f"data.tar.{self.COMPRESSION}": ["data"], + } + + checked = False + for src, targets in files.items(): + # if forced or not utils.CACHE.all_exists(targets, self.path): + if forced: + if not checked and not self.check_remote(): + raise Exception( + f"Could not find dataset={self.name} on data server" + ) + else: + checked = True + + remote_file = os.path.join(self.URL, self.name, src) + target_file = os.path.join(self.path, src) + # LOG.debug(f"target_file={target_file}") + try: + # print("Download data ...") + utils.download(remote_file, target_file) + print(f"Unpack data ... {src}") + utils.unpack(target_file, remove=True) + # TODO: we skip the reference creation to make things faster. Enable it when + # it is needed! + # utils.CACHE.make_reference(targets, self.path) + except: + # if os.exists(target_file): + # os.remove(target_file) + raise Exception(f"Failed to download file={remote_file}") + + def __getitem__(self, key): + if isinstance(key, str): + r = self.find(key, comp="all") + if r is None: + raise Exception(f"No component={key} found in {self}") + return r + return None + + def load_style(self): + conf_dir = os.path.join(self.path, "conf") + mv.style.load_custom_config(conf_dir, force=True) + + def __str__(self): + return f"{self.__class__.__name__}[name={self.name}]" + + +def create_dataset_template(name_or_path): + path = name_or_path + if any(x in path for x in ["/", "\\", "..", "./"]): + name = os.path.basename(path) + else: + name = path + path = "" + + if path == "": + path = os.path.join(Dataset.LOCAL_ROOT, name) + + if not os.path.exists(path): + os.mkdir(path) + else: + if not os.path.isdir(path): + raise Exception(f"path must be a directory!") + if os.path.exists(os.path.join("path", "data.yaml")): + raise Exception( + f"The specified dataset directory={path} already exists and is not empty!" + ) + + # create dirs + for dir_name in ["conf", "data", "index"]: + os.mkdir(os.path.join(path, dir_name)) + + # copy files + files = { + "params.yaml": ("conf", ""), + "param_styles.yaml": ("conf", ""), + "areas.yaml": ("conf", ""), + "map_styles_template.yaml": ("conf", "map_styles.yaml"), + "dataset_template.yaml": ("", "data.yaml"), + } + for src_name, target in files.items(): + target_dir = os.path.join(path, target[0]) if target[0] else path + target_name = target[1] if target[1] else src_name + shutil.copyfile( + os.path.join(ETC_PATH, src_name), os.path.join(target_dir, target_name) + ) diff --git a/metview/etc/areas.yaml b/metview/etc/areas.yaml new file mode 100644 index 00000000..e63fd1d1 --- /dev/null +++ b/metview/etc/areas.yaml @@ -0,0 +1,3 @@ +- base: + map_area_definition: "corners" + area: [-90,-180,90,180] \ No newline at end of file diff --git a/metview/etc/dataset_template.yaml b/metview/etc/dataset_template.yaml new file mode 100644 index 00000000..530ec300 --- /dev/null +++ b/metview/etc/dataset_template.yaml @@ -0,0 +1,14 @@ +experiments: + - myexp: + label: "myexp" + dir: __ROOTDIR__/myexp + fname : "*.grib" + - track: + type: track + dir: __ROOTDIR__/track + fname: "*.csv" + skiprows: 10 + date_index: 0 + time_index: 1 + lat_index: 3 + lon_index: 4 \ No newline at end of file diff --git a/metview/etc/map_styles.yaml b/metview/etc/map_styles.yaml new file mode 100644 index 00000000..8734c129 --- /dev/null +++ b/metview/etc/map_styles.yaml @@ -0,0 +1,149 @@ +base: + map_coastline_land_shade: "off" + map_coastline_sea_shade: 'off' +base_diff: + map_coastline_colour: CHARCOAL + map_coastline_land_shade: "on" + map_coastline_land_shade_colour: "RGB(0.8549,0.8549,0.8549)" + map_coastline_sea_shade: 'off' + map_coastline_land_shade_colour: WHITE +grey_light_blue: + map_coastline_land_shade: "on" + map_coastline_land_shade_colour: "grey" + map_coastline_sea_shade: "on" + map_coastline_sea_shade_colour: "RGB(0.86,0.94,1)" + map_boundaries: "off" + map_boundaries_colour: "RGB(0.21,0.21,0.21)" + map_disputed_boundaries: "off" + map_administrative_boundaries: "off" + map_grid_colour: "RGB(0.294,0.294,0.2941)" + map_label_colour: "RGB(0.294,0.294,0.2941)" +black_coast_only: + map_coastline: 'on' + map_coastline_land_shade: 'off' + map_coastline_sea_shade: 'off' + # map_coastline_colour: BLACK +black_grey: + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: black + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.7767,0.7932,0.8076) + map_coastline: 'on' + map_coastline_colour: BACKGROUND + map_coastline_style: solid + map_coastline_thickness: 1 +blue_sea_only: + map_coastline: 'on' + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.7373,0.9294,0.9608) +brown: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.7216,0.6863,0.6471) + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.9059,0.9176,0.9373) + map_coastline_colour: TAN +cream_blue: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: CREAM + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: SKY +cream_land_only: + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: CREAM + map_coastline_sea_shade: 'off' + map_coastline_sea_shade_colour: WHITE + map_coastline_colour: BLACK + map_coastline_style: solid + map_coastline_thickness: 1 +dark_1: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: CHARCOAL + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.3677,0.3677,0.4009) + map_coastline_colour: WHITE +dark_2: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: CHARCOAL + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.2471,0.2471,0.2471) + map_coastline_colour: WHITE +dark_3: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.4431,0.4431,0.4431) + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.3176,0.3176,0.3176) +green_blue: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.2233,0.7101,0.3369) + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.4923,0.6884,0.8489) + map_coastline_colour: BLACK +green_land_only: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.1020,0.4000,0.0000) + map_coastline_sea_shade: 'off' +green_navy: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: EVERGREEN + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.0000,0.0000,0.5020) +grey_1: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.3725,0.4157,0.5451) + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.7767,0.7932,0.8076) + map_coastline_colour: BLACK + map_coastline_style: solid + map_coastline_thickness: 1 +grey_2: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: GREY + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.9059,0.9176,0.9373) + map_coastline_colour: CHARCOAL +grey_3: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.5373,0.5373,0.5373) + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.8392,0.8510,0.8706) + map_coastline_colour: CHARCOAL +grey_4: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.5569,0.5569,0.5569) + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.8000,0.8000,0.8000) + map_coastline_colour: CHARCOAL +grey_blue: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: GREY + map_coastline_sea_shade: 'on' + map_coastline_sea_shade_colour: RGB(0.7922,0.8431,0.9412) + map_coastline_colour: CHARCOAL +grey_darker_land_only: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.5020,0.5020,0.5020) + map_coastline_sea_shade: 'off' +grey_land_only: + map_coastline: 'on' + map_coastline_land_shade: 'on' + map_coastline_land_shade_colour: RGB(0.8549,0.8549,0.8549) + map_coastline_sea_shade: 'off' +white_coast_only: + map_coastline: 'on' + map_coastline_land_shade: 'off' + map_coastline_sea_shade: 'off' + map_coastline_colour: WHITE \ No newline at end of file diff --git a/metview/etc/map_styles_template.yaml b/metview/etc/map_styles_template.yaml new file mode 100644 index 00000000..ac36c89a --- /dev/null +++ b/metview/etc/map_styles_template.yaml @@ -0,0 +1,29 @@ +base: + map_coastline_resolution: "low" + map_coastline_land_shade: "on" + map_coastline_land_shade_colour: "grey" + map_coastline_sea_shade: "on" + map_coastline_sea_shade_colour: "RGB(0.86,0.94,1)" + map_boundaries: "on" + map_boundaries_colour: "RGB(0.21,0.21,0.21)" + map_disputed_boundaries: "off" + map_administrative_boundaries: "off" + map_grid_latitude_increment: 10 + map_grid_longitude_increment: 10 + map_grid_colour: "RGB(0.294,0.294,0.2941)" + map_label_colour: "RGB(0.294,0.294,0.2941)" +base_diff: + map_coastline_resolution: "low" + map_coastline_colour: CHARCOAL + map_coastline_land_shade: "on" + map_coastline_land_shade_colour: "RGB(0.8549,0.8549,0.8549)" + map_coastline_sea_shade: "off" + map_coastline_sea_shade_colour: "white" + map_boundaries: "on" + map_boundaries_colour: "RGB(0.21,0.21,0.21)" + map_disputed_boundaries: "off" + map_administrative_boundaries: "off" + map_grid_latitude_increment: 10 + map_grid_longitude_increment: 10 + map_grid_colour: "RGB(0.294,0.294,0.2941)" + map_label_colour: "RGB(0.294,0.294,0.2941)" \ No newline at end of file diff --git a/metview/etc/param_styles.yaml b/metview/etc/param_styles.yaml new file mode 100644 index 00000000..74e76bf0 --- /dev/null +++ b/metview/etc/param_styles.yaml @@ -0,0 +1,579 @@ +# ------------------------ +# DEFAULT STYLES +# ------------------------ +# this is the default style assigned to a scalar param! +default_mcont: + mcont: + contour_automatic_setting: ecmwf + legend: "on" +# this is the default style assigned to a vector param! +default_mwind: + mwind: + wind_thinning_method: "density" + wind_density: 4 + legend: "on" +# this is the default style for difference fields +default_diff: + mcont: + legend: "off" +# ------------------------ +# SCALAR STYLES +# ------------------------ +msl: + mcont: + legend: "off" + contour_highlight_colour: black + contour_highlight_thickness: 4 + contour_interval: 5 + contour_label: "on" + contour_label_frequency: 2 + contour_label_height: 0.4 + contour_level_selection_type: interval + contour_line_colour: black + contour_line_thickness: 2 + contour_reference_level: 1010 +# style for potential temperature +pt: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "interval" + contour_max_level: 380 + contour_min_level: 260 + contour_interval: 5 + contour_label: "off" + contour_shade: "on" + contour_shade_method: "area_fill" + contour_shade_colour_method: "palette" + contour_shade_palette_name: eccharts_rainbow_blue_purple_24 + grib_scaling_of_retrieved_fields: "off" +# style for pressure velocity +w: + mcont: + contour_automatic_setting: style_name + contour_style_name: sh_viobrn_fM5t5lst + legend: "on" +# style for relative and absolute vorticity +vo: + mcont: + contour_automatic_setting: style_name + contour_style_name: sh_blured_fM50t50lst_less + legend: "on" + grib_scaling_of_derived_fields: "on" +# style for potential vorticity on 320K +pv: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [1,2,4,6,8,10,20,30] + contour_label: "off" + contour_shade: "on" + contour_shade_method: "area_fill" + contour_shade_colour_method: "palette" + contour_shade_palette_name: matplotlib_Plasma_7_r +q_flux: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "interval" + contour_level_selection_type: "level_list" + contour_level_list: [-0.04, -0.01, -0.008, -0.004, -0.002, -0.001, -0.0005, -0.0002, 0.0002, 0.0005, 0.001, 0.002, 0.004, 0.008, 0.01, 0.04] + contour_label: "off" + contour_shade: "on" + contour_shade_method: "area_fill" + contour_shade_colour_method: "palette" + contour_shade_palette_name: "colorbrewer_PRGn_15_r" +t_flux: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "interval" + contour_level_selection_type: "level_list" + contour_level_list: [-120,-50,-20,-10,-5,-2,-1,-0.5, 0.5,1,2, 5, 10, 20, 50, 120] + contour_label: "off" + contour_shade: "on" + contour_shade_method: "area_fill" + contour_shade_colour_method: "palette" + contour_shade_palette_name: "colorbrewer_PuOr_15" +z: + mcont: + legend: "off" + contour_highlight_colour: black + contour_highlight_thickness: 4 + contour_interval: 5 + contour_label: "on" + contour_label_frequency: 2 + contour_label_height: 0.4 + contour_level_selection_type: interval + contour_line_colour: black + contour_line_thickness: 2 + contour_reference_level: 1010 +u_10m: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [-80,-30,-20,-10, -5, -2,0,2,5,10,20,30,80] + contour_label: "off" + contour_shade: "on" + contour_shade_colour_method: "palette" + contour_shade_method: "area_fill" + contour_shade_palette_name: "colorbrewer_PuOr" + contour_shade_colour_list_policy: "dynamic" +u_700: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [-90,-30,-25,-20,-15,-10,0,10,15,20,25,30,90] + contour_label: "off" + contour_shade: "on" + contour_shade_colour_method: "palette" + contour_shade_method: "area_fill" + contour_shade_palette_name: "colorbrewer_PuOr" + contour_shade_colour_list_policy: "dynamic" +u_500: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [-90,-50,-40,-30,-20,-10,0,10,20,30,40,50,90] + contour_label: "off" + contour_shade: "on" + contour_shade_colour_method: "palette" + contour_shade_method: "area_fill" + contour_shade_palette_name: "colorbrewer_PuOr" + contour_shade_colour_list_policy: "dynamic" +u_200: + mcont: + legend: "on" + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [-120,-80,-60,-40,-20,-10,0,10,20,40,60,80,120] + contour_label: "off" + contour_shade: "on" + contour_shade_colour_method: "palette" + contour_shade_method: "area_fill" + contour_shade_palette_name: "colorbrewer_PuOr" + contour_shade_colour_list_policy: "dynamic" +# ------------------------ +# VECTOR STYLES +# ------------------------ +# black wind arrows +# ------------------- +arrow_black: + mwind: &arrow_black + legend: "off" + wind_advanced_method: "off" + wind_thinning_method: "density" + wind_densisty: 2 + wind_arrow_colour: "black" + wind_arrow_min_speed: 0 +# 10m +arrow_black_10: + mwind: + <<: *arrow_black + wind_arrow_max_speed: 30 + wind_arrow_unit_velocity: 20.0 +# 850 hPa +arrow_black_850: + mwind: + <<: *arrow_black + wind_arrow_max_speed: 50 + wind_arrow_unit_velocity: 30.0 +# 700 hPa +arrow_black_700: + mwind: + <<: *arrow_black + wind_arrow_max_speed: 60 + wind_arrow_unit_velocity: 30.0 +# 500 hPa +arrow_black_500: + mwind: + <<: *arrow_black + wind_arrow_max_speed: 70 + wind_arrow_unit_velocity: 50.0 +# 200 hPa +arrow_black_200: + mwind: + <<: *arrow_black + wind_arrow_max_speed: 100 + wind_arrow_unit_velocity: 70.0 +# 100 hPa +arrow_black_100: + mwind: + <<: *arrow_black + wind_arrow_max_speed: 60 + wind_arrow_unit_velocity: 30.0 +# ---------------------- +# Blue-red wind arrows +# ---------------------- +# generic +arrow_blue_red: + mwind: &arrow_blue_red + legend: "on" + wind_advanced_method: "on" + wind_advanced_colour_selection_type: "interval" + wind_advanced_colour_max_level_colour: "red" + wind_advanced_colour_min_level_colour: "blue" + wind_advanced_colour_direction: "anticlockwise" + wind_thinning_method: "density" + wind_density: 2 + wind_advanced_colour_min_value: 0 +# 10m +arrow_blue_red_10: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 30 + wind_advanced_colour_level_interval: 5 + wind_arrow_unit_velocity: 20.0 +# style for wind arrows +arrow_blue_red: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 70 + wind_advanced_colour_min_value: 0 + wind_advanced_colour_level_interval: 10 + wind_arrow_unit_velocity: 50.0 +# 850 hPa +arrow_blue_red_850: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 50 + wind_arrow_unit_velocity: 30.0 +# 700 hPa +arrow_blue_red_700: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 60 + wind_arrow_unit_velocity: 30.0 +# 500 hPa +arrow_blue_red_500: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 70 + wind_arrow_unit_velocity: 50.0 +# 200 hPa +arrow_blue_red_200: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 100 + wind_arrow_unit_velocity: 70.0 +# 200 hPa +arrow_blue_red_100: + mwind: + <<: *arrow_blue_red + wind_advanced_colour_max_value: 60 + wind_arrow_unit_velocity: 30.0 +# ---------------------- +# Black wind flags +# ---------------------- +# 10m +flag_black_10: + mwind: + wind_field_type: "flags" + wind_flag_calm_indicator: "off" + wind_flag_length: 0.7 + wind_flag_origin_marker: "off" + wind_flag_colour: "black" + wind_thinning_method: "density" + wind_density: 1 +# upper levels +flag_black_upper: + mwind: + wind_field_type: "flags" + wind_flag_calm_indicator: "off" + wind_flag_length: 0.5 + wind_flag_origin_marker: "off" + wind_flag_colour: "black" + wind_thinning_method: "density" + wind_density: 1 + wind_flag_colour: "black" +# ------------------------ +# CROSS SECTION STYLES +# ------------------------ +xs_q: + mcont: + contour_automatic_setting: style_name + contour_style_name: sh_spechum_option1 + legend: "on" +xs_r: + mcont: + contour_automatic_setting: style_name + contour_style_name: sh_grnblu_f65t100i15_light + legend: "on" +xs_t: + mcont: + contour_automatic_setting: style_name + contour_style_name: sh_all_fM80t56i4_v2 + legend: "on" +# style for black wind arrows used in cross section +xs_arrow_black: + mwind: + legend: "off" + wind_advanced_method: "off" + wind_arrow_colour: "black" + #wind_advanced_colour_min_value: 0 + #wind_advanced_colour_level_interval: 10 + wind_arrow_unit_velocity: 30.0 + wind_thinning_factor: 1 +# ------------------------ +# SHAPES AND TRACKS +# ------------------------ +box: + mcont: + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [0.5,1.1] + contour_label: "off" + contour_shade: "on" + contour_shade_technique: grid_shading + contour_shade_colour_method: "list" + contour_shade_colour_list: "RGBA(1,0,0,128)" +frame: + mcont: + contour: "off" + contour_level_selection_type: "level_list" + contour_level_list: [0.5,1.1] + contour_label: "off" + contour_shade: "on" + contour_shade_technique: grid_shading + contour_shade_colour_method: "list" + contour_shade_colour_list: "red" +track: + - msymb: + symbol_type: "text" + symbol_table_mode: "advanced" + symbol_advanced_table_selection_type: "list" + symbol_advanced_table_text_font_colour: "black" + symbol_advanced_table_text_font_size: 0.5 + symbol_advanced_table_text_font_style: "normal" + symbol_advanced_table_text_display_type: "right" + - mgraph: + graph_line_colour: red + graph_line_thickness: 4 + graph_symbol: "on" + graph_symbol_colour: white + graph_symbol_height: 0.5 + graph_symbol_marker_index: 15 + graph_symbol_outline: "on" + graph_symbol_outline_colour: red +# ------------------------ +# DIFFERENCE STYLES +# ------------------------ +# core definition for the positive range +plus_core: + - mcont: &plus + legend: "on" + contour: "on" + contour_highlight: "off" + contour_line_colour: "black" + contour_level_selection_type: "level_list" + #contour_max_level: 40 + #contour_min_level: 0 + #contour_interval: 5 + contour_label: "off" + contour_shade: "on" + contour_shade_method: "area_fill" + #contour_shade_max_level_colour: red + contour_shade_max_level_colour: "RGB(110./255.,0,0)" + contour_shade_min_level_colour: "RGB(0.937,0.804,0.8041)" + contour_shade_colour_direction: anticlockwise + grib_scaling_of_derived_fields: "on" +# core definition for the negative range +minus_core: + - mcont: &minus + legend: "on" + contour: "on" + contour_highlight: "off" + contour_line_colour: black + contour_level_selection_type: level_list + #contour_max_level: 40 + #contour_min_level: 0 + #contour_interval: 5 + contour_label: "off" + contour_shade: "on" + contour_shade_method: area_fill + contour_shade_max_level_colour: "RGB(0.798,0.798,0.9192)" + contour_shade_min_level_colour: "RGB(0,25./255,51./255)" + #contour_shade_min_level_colour: blue + contour_shade_colour_direction: anticlockwise + grib_scaling_of_derived_fields: "on" +plus_core_shade_only: + - mcont: &plus_shade + <<: *plus + contour: "off" +minus_core_shade_only: + - mcont: &minus_shade + <<: *minus + contour: "off" +# diff_msl: +# mcont: +# <<: *diff_core +# # contour_level_list: [-50,-20,-10,-5,-0.5,0.5,5,10,20,50] +# contour_level_list: [-50,-30,-20,-10,-5,-0.5,0.5,5,10,20,30, 50] +diff_msl: + - mcont: + # contour_level_list: [-50,-20,-10,-5,-0.5,0.5,5,10,20,50] + <<: *minus_shade + contour_level_list: [-50,-15,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,15,50] +# temperature (K) +diff_t: + - mcont: + <<: *minus_shade + contour_level_list: [-30,-10,-8,-6,-4,-2] + grib_scaling_of_derived_fields: "off" + - mcont: + <<: *plus_shade + contour_level_list: [2,4,6,8,10,30] + grib_scaling_of_derived_fields: "off" +# equivalent potential temperature (K) +diff_pt: + - mcont: + <<: *minus_shade + contour_level_list: [-50,-25,-20,-15,-10,-5,-2] + grib_scaling_of_derived_fields: "off" + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,15,20,25,50] + grib_scaling_of_derived_fields: "off" +# potential vorticity +diff_pv: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-20,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,20,100] +# geopotential (dkm) +diff_z: + - mcont: + <<: *minus_shade + contour_level_list: [-50,-20,-15,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,15,20,50] +# precipitation (mm) +diff_tp: + - mcont: + <<: *minus_shade + contour_level_list: [-200,-20,-10,-5,-0.5] + - mcont: + <<: *plus_shade + contour_level_list: [0.5,5,10,20,200] +# wind u/v component (m/s) +diff_u_10: + - mcont: + <<: *minus_shade + contour_level_list: [-70,-10,-8,-5,-2,-1] + - mcont: + <<: *plus_shade + contour_level_list: [1,2,5,8,10,70] +# wind u/v component 700 hPa (m/s) +diff_u_700: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-20,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,20,100] +# wind u/v component 500 hPa (m/s) +diff_u_500: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-20,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,20,100] +# wind u/v component 200 hPa(m/s) +diff_u_200: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-30,-20,-10,-5] + - mcont: + <<: *plus_shade + contour_level_list: [5,10,20,30,100] +# wind speed (m/s) +diff_speed_10: + - mcont: + <<: *minus_shade + contour_level_list: [-70,-10,-8,-5,-2,-1] + - mcont: + <<: *plus_shade + contour_level_list: [1,2,5,8,10,70] +# wind speed 700 hPa (m/s) +diff_speed_700: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-20,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,20,100] +# wind speed 200 hPa(m/s) +diff_speed_200: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-30,-20,-10,-5] + - mcont: + <<: *plus_shade + contour_level_list: [5,10,20,30,100] +# wind gust (m/s) +diff_10fg3: + - mcont: + <<: *minus + contour_level_list: [-35,-20,-10,-5,-3] + - mcont: + <<: *plus + contour_level_list: [3,5,10,20,35] +# relative humidity (0-100) +diff_r: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-60,-30,-10,-5] + - mcont: + <<: *plus_shade + contour_level_list: [5,10,30,60,100] +# specific humidity (0-28) +diff_q: + - mcont: + <<: *minus_shade + contour_level_list: [-30,-10,-5,-2,-1,-0.5] + - mcont: + <<: *plus_shade + contour_level_list: [0.5,1,2,5,10,30] +# vorticity (1E-5 1/s) +diff_vo: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-20,-10,-5,-2] + - mcont: + <<: *plus_shade + contour_level_list: [2,5,10,20,100] +# vertical velocity (Pa/s) +diff_w: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-2,-1,-0.5,-0.2,-0.1] + - mcont: + <<: *plus_shade + contour_level_list: [0.1,0.2,0.5,1,2,100] +# t flux +diff_q_flux: + - mcont: + <<: *minus_shade + contour_level_list: [-0.1,-0.01,-0.005,-0.001,-0.0005] + - mcont: + <<: *plus_shade + contour_level_list: [0.0005, 0.001, 0.005, 0.01, 0.1] +# t flux +diff_t_flux: + - mcont: + <<: *minus_shade + contour_level_list: [-100,-20,-10,-5,-2,-1] + - mcont: + <<: *plus_shade + contour_level_list: [1,2,5,10,20,100] diff --git a/metview/etc/params.yaml b/metview/etc/params.yaml new file mode 100644 index 00000000..e0858b1d --- /dev/null +++ b/metview/etc/params.yaml @@ -0,0 +1,668 @@ +# -------------------- +# SCALAR PARAMETERS +# -------------------- +# 2m dewpoint temperature +- + styles: + diff: + - diff_t + match: + - + info_name: 2d + - + shortName: 2d + paramId: 168 +# 2m temperature +- + styles: + diff: + - diff_t + match: + - + info_name: 2t + - + shortName: 2t + paramId: 167 +# 10m wind gust 3h +- + styles: + diff: + - diff_10fg3 + match: + - + info_name: 10fg3 + - + shortName: 10fg3 + paramId: 228028 +# 10m wind speed +- + styles: + diff: + - diff_speed_10 + match: + - + info_name: 10si + - + shortName: 10si + paramId: 207 +# 10m u wind component +- + styles: + basic: + - u_10m + diff: + - diff_u_10 + match: + - + info_name: 10u + - + shortName: 10u + paramId: 165 +# 10m v wind component +- + styles: + basic: + - u_10m + diff: + - diff_u_10 + match: + - + info_name: 10v + - + shortName: 10v + paramId: 166 +# absolute vorticity +- + styles: + basic: + - vo + diff: + - diff_vo + scaling: + - xs + match: + - + info_name: absv + - + shortName: absv + paramId: 3041 +# equivalent potential temperature +- + styles: + basic: + - pt + diff: + - diff_pt + match: + - + info_name: eqpt + - + shortName: eqpt + paramId: 4 +# geopotential +- + styles: + basic: + - z + diff: + - diff_z + match: + - + info_name: z + - + shortName: z + paramId: 129 +# mean_sea_level_pressure: +- + styles: + basic: + - msl + diff: + - diff_msl + match: + - + info_name: msl + - + shortName: msl + paramId: 151 +# potential temperature +- + styles: + basic: + - pt + diff: + - diff_pt + match: + - + info_name: pt + - + shortName: pt + paramId: 3 +# potential vorticity +- + styles: + basic: + - pv + diff: + - diff_pv + scaling: + - xs + match: + - + info_name: pv +# relative humidity +- + styles: + diff: + - diff_r + xs: + - xs_r + match: + - + info_name: r + - + shortName: r + paramId: 157 +# relative vorticity +- + styles: + basic: + - vo + diff: + - diff_vo + match: + - + info_name: vo + - + shortName: vo + paramId: 138 + scaling: + - xs +# specific humidity +- + styles: + diff: + - diff_q + xs: + - xs_q + scaling: + - xs + match: + - + info_name: q + - + shortName: q + paramId: 133 +# sea surface temperature +- + styles: + diff: + - diff_t + match: + - + info_name: sst + - + shortName: sst + paramId: 34 +# temperature +- + styles: + diff: + - diff_t + xs: + - xs_t + scaling: + - xs + match: + - + info_name: t + - + shortName: t + paramId: 130 +# total precipitation +- + styles: + diff: + - diff_tp + match: + - + info_name: tp + - + shortName: tp + paramId: 228 +# vertical (pressure) velocity +- + styles: + basic: + - w + diff: + - diff_w + match: + - + info_name: w + - + shortName: w + paramId: 135 +# u wind component 1000hPa +- + styles: + basic: + - u_10m + diff: + - diff_u_10 + match: + - + info_name: u + levtype: pl + levelist: 1000 + - + shortName: u + paramId: 131 + levtype: pl + levelist: 1000 +# u wind component around 700 hPa +- + styles: + basic: + - u_700 + diff: + - diff_u_700 + match: + - + info_name: u + levtype: pl + levelist: [925, 850, 700] + - + shortName: u + paramId: 131 + levtype: pl + levelist: [925, 850, 700] +# u wind component around 500 hPa + 100 hPa +- + styles: + basic: + - u_500 + diff: + - diff_u_500 + match: + - + info_name: u + levtype: pl + levelist: [600, 500, 100] + - + shortName: u + paramId: 131 + levtype: pl + levelist: [600, 500, 100] +# u wind component around 200 hPa +- + styles: + basic: + - u_200 + diff: + - diff_u_200 + match: + - + info_name: u + levtype: pl + levelist: [400, 300, 250, 200, 150] + - + shortName: u + paramId: 131 + levtype: pl + levelist: [400, 300, 250, 200, 150] +# v wind component 1000hPa +- + styles: + basic: + - u_10m + diff: + - diff_u_10 + match: + - + info_name: v + levtype: pl + levelist: 1000 + - + shortName: v + paramId: 132 + levtype: pl + levelist: 1000 +# v wind component around 700 hPa +- + styles: + basic: + - u_700 + diff: + - diff_u_700 + match: + - + info_name: v + levtype: pl + levelist: [925, 850, 700] + - + shortName: v + paramId: 132 + levtype: pl + levelist: [925, 850, 700] +# v wind component around 500 hPa + 100 hPa +- + styles: + basic: + - u_500 + diff: + - diff_u_500 + match: + - + info_name: v + levtype: pl + levelist: [600, 500, 100] + - + shortName: v + paramId: 132 + levtype: pl + levelist: [600, 500, 100] +# v wind component around 200 hPa +- + styles: + basic: + - u_200 + diff: + - diff_u_200 + match: + - + info_name: v + levtype: pl + levelist: [400, 300, 250, 200, 150] + - + shortName: v + paramId: 132 + levtype: pl + levelist: [400, 300, 250, 200, 150] +# wind speed 1000hPa +- + styles: + diff: + - diff_speed_10 + match: + - + info_name: ws + levtype: pl + levelist: 1000 + - + shortName: ws + paramId: 10 + levtype: pl + levelist: 1000 +# wind speed around 700 hPa + 100 hPa +- + styles: + diff: + - diff_speed_700 + match: + - + info_name: ws + levtype: pl + levelist: [925, 850, 700, 600, 100] + - + shortName: ws + paramId: 10 + levtype: pl + levelist: [925, 850, 700, 600, 100] +# wind speed around 200 hPa +- + styles: + diff: + - diff_speed_200 + match: + - + info_name: ws + levtype: pl + levelist: [500, 400, 300, 250, 200, 150] + - + shortName: ws + paramId: 10 + levtype: pl + levelist: [500, 400, 300, 250, 200, 150] +# -------------------- +# VECTOR PARAMETERS +# -------------------- +# wind (generic) +- + type: vector + styles: + basic: + - arrow_blue_red + - arrow_black + match: + - + info_name: wind +# wind 10 m +- + type: vector + styles: + basic: + - arrow_blue_red_10 + - arrow_black_10 + - flag_black_10 + match: + - + info_name: wind10m +# wind 1000 hPa +- + type: vector + styles: + basic: + - arrow_blue_red_10 + - arrow_black_10 + - flag_black_10 + match: + - + info_name: wind + levtype: pl + levelist: 1000 +# wind 925-850 hPa +- + type: vector + styles: + basic: + - arrow_blue_red_850 + - arrow_black_850 + - flag_black_upper + match: + - + info_name: wind + levtype: pl + levelist: [925,850] +# wind 700-600 hPa +- + type: vector + styles: + basic: + - arrow_blue_red_700 + - arrow_black_700 + - flag_black_upper + match: + - + info_name: wind + levtype: pl + levelist: [700, 600] +# wind 500-400 hPa +- + type: vector + styles: + basic: + - arrow_blue_red_500 + - arrow_black_500 + - flag_black_upper + match: + - + info_name: wind + levtype: pl + levelist: [500, 400] +# wind 300-150 hPa +- + type: vector + styles: + basic: + - arrow_blue_red_200 + - arrow_black_200 + - flag_black_upper + match: + - + info_name: wind + levtype: pl + levelist: [300, 250, 200, 150] +# wind 100- hPa +- + type: vector + styles: + basic: + - arrow_blue_red_100 + - arrow_black_100 + - flag_black_upper + match: + - + info_name: wind + levtype: pl + levelist: 100 +# wind3d (generic) +- + type: vector + styles: + basic: + - xs_arrow_black + match: + - + info_name: wind3d +# ------------------------------------------- +# PARAMETERS REPRESENTING GEOMETRIC SHAPES +# ------------------------------------------- +- + styles: + basic: + - box + match: + - + info_name: box +- + styles: + basic: + - frame + match: + - + info_name: frame + +# ------------------------------------------- +# TRACKS +# ------------------------------------------- +- + styles: + basic: + - track + match: + - + info_name: track + +# --------------------------------- +# PARAMETERS WITHOUT A SHORTNAME +# --------------------------------- +- + styles: + basic: + - q_flux + diff: + - diff_q_flux + match: + - + info_name: qcl + - + paramId: 110 +- + styles: + basic: + - q_flux + diff: + - diff_q_flux + match: + - + info_name: qcon + - + paramId: 106 +- + styles: + basic: + - q_flux + diff: + - diff_q_flux + match: + - + info_name: qdyn + - + paramId: 94 +- + styles: + basic: + - q_flux + diff: + - diff_q_flux + match: + - + info_name: qvdiff + - + paramId: 99 +- + styles: + basic: + - t_flux + diff: + - diff_t_flux + match: + - + info_name: tcl + - + paramId: 109 +- + styles: + basic: + - t_flux + diff: + - diff_t_flux + match: + - + info_name: tcon + - + paramId: 105 +- + styles: + basic: + - t_flux + diff: + - diff_t_flux + match: + - + info_name: tdyn + - + paramId: 93 +- + styles: + basic: + - t_flux + diff: + - diff_t_flux + match: + - + info_name: trad + - + paramId: 95 +- + styles: + basic: + - t_flux + diff: + - diff_t_flux + match: + - + info_name: tvdiff + - + paramId: 98 diff --git a/metview/etc/scaling_ecmwf.yaml b/metview/etc/scaling_ecmwf.yaml new file mode 100644 index 00000000..c5fc9447 --- /dev/null +++ b/metview/etc/scaling_ecmwf.yaml @@ -0,0 +1,252 @@ +10**5/sec: + match: [] + paramId: + - '138' + - '155' + - '3041' + shortName: + - vo + - absv + - d +C: + match: + - centre: '7' + paramId: '11' + - centre: '46' + paramId: '128' + - centre: '46' + paramId: '129' + - centre: '46' + paramId: '158' + - centre: '46' + paramId: '175' + - centre: '46' + paramId: '187' + - centre: '46' + paramId: '188' + - centre: '46' + paramId: '189' + - centre: '46' + paramId: '190' + - centre: '46' + paramId: '191' + - centre: '46' + paramId: '232' + - centre: '46' + paramId: '234' + - centre: '46' + paramId: '243' + - centre: '46' + paramId: '251' + - centre: '46' + paramId: '253' + paramId: + - '34' + - '35' + - '36' + - '37' + - '38' + - '51' + - '52' + - '121' + - '122' + - '130' + - '139' + - '167' + - '168' + - '170' + - '183' + - '201' + - '202' + - '208' + - '235' + - '236' + - '238' + - '228004' + - '228008' + - '228010' + - '228011' + - '228013' + - '228026' + - '228027' + - '260510' + shortName: + - mn2t6 + - ltlt + - lmlt + - mn2t24 + - skt + - 2d + - mx2t24 + - clbt + - lict + - mn2t + - mn2t3 + - istl4 + - lblt + - mx2t3 + - istl3 + - stl1 + - istl2 + - tsn + - mx2t6 + - mean2t + - stl2 + - mx2t + - sst + - stl4 + - stl3 + - 2t + - istl1 + - t +Dob: + match: [] + paramId: + - '206' + shortName: [] +cm: + match: [] + paramId: + - '44' + - '45' + - '141' + - '149' + - '160' + - '173' + shortName: + - smlt + - es +dam: + match: [] + paramId: + - '129' + - '156' + - '206' + - '171129' + shortName: + - gh + - za + - z +g/kg: + match: [] + paramId: + - '133' + - '203' + - '233' + - '246' + - '247' + shortName: + - q +μg/kg: + match: [] + paramId: + - '217004' + - '210123' + - '210203' + - '210124' + - '217006' + - '217028' + - '210121' + - '217027' + - '217030' + shortName: + - ch4_c + - co + - go3 + - hcho + - hno3 + - ho2 + - no2 + - no + - oh + - so2 +hPa: + match: [] + paramId: + - '54' + - '134' + - '151' + shortName: + - msl + - sp +kg m**-2 h**-1: + match: [] + paramId: + - '3059' + - '3064' + - '174142' + - '174143' + - '228218' + - '228219' + - '228220' + - '228221' + - '228222' + - '228223' + - '228224' + - '228225' + - '228226' + - '228227' + - '260048' + shortName: + - mxtpr + - crr + - mxtpr3 + - mntpr6 + - srweq + - mntpr3 + - lsrr + - mntpr + - tprate + - crfrate + - lsrrate + - lssfr + - prate + - mxtpr6 + - csfr +mm: + match: [] + paramId: + - '8' + - '9' + - '140' + - '142' + - '143' + - '144' + - '171' + - '182' + - '184' + - '198' + - '205' + - '228' + - '237' + - '239' + - '240' + - '244' + - '174008' + - '174009' + - '228216' + shortName: + - ssro + - fzra + - lsp + - csf + - sf + - tp + - cp + - sro + - ro + - lsf +ppbv: + match: [] + paramId: + - '210062' + - '217004' + shortName: + - ch4_c + - ch4 +pv units: + match: [] + paramId: + - '60' + shortName: + - pv diff --git a/metview/etc/units-rules.json b/metview/etc/units-rules.json new file mode 100644 index 00000000..98a59373 --- /dev/null +++ b/metview/etc/units-rules.json @@ -0,0 +1,137 @@ +{ + "C": [ + { + "from" : "K", + "to" : "C", + "offset" : -273.15, + "scaling" : 1 + } + ], + "pv units" : [ + { + "from" : "K m**2 kg**-1 s**-1", + "offset" : -0.0, + "to" : "pv units", + "scaling" : 1000000.0 + } + ], + "10**5/sec" : [ + { + "from" : "s**-1", + "to" : "10**5/sec", + "offset" : -0.0, + "scaling" : 100000.0 + } + ], + "g/kg" : [ + { + "from" : "kg kg**-1", + "to" : "g/kg", + "offset" : -0.0, + "scaling" : 1000.0 + }, + { + "from" : "1", + "to" : "gr kg-1", + "offset" : 0.0, + "scaling" : 1000.0 + } + ], + "μg/kg" : [ + { + "from" : "kg kg**-1", + "to" : "μg/kg", + "offset" : -0.0, + "scaling" : 1e9 + }, + { + "from" : "1", + "to" : "μgr kg-1", + "offset" : 0.0, + "scaling" : 1e9 + } + ], + "mm" : [ + { + "from" : "m", + "to" : "mm", + "offset" : -0.0, + "scaling" : 1000.0 + }, + { + "from" : "m of water equivalent", + "to" : "mm", + "offset" : -0.0, + "scaling" : 1000.0 + }, + { + "from" : "m of water", + "to" : "mm", + "offset" : -0.0, + "scaling" : 1000.0 + } + ], + "cm" : [ + { + "from" : "m", + "to" : "cm", + "offset" : -0.0, + "scaling" : 100.0 + }, + { + "from" : "m of water equivalent", + "to" : "cm", + "offset" : -0.0, + "scaling" : 100.0 + } + ], + "dam" : [ + { + "from" : "m**2 s**-2", + "offset" : -0.0, + "to" : "dam", + "scaling" : 0.0101971621297793 + }, + { + "from" : "m", + "offset" : -0.0, + "to" : "dam", + "scaling" : 10.0 + }, + { + "from" : "gpm", + "offset" : -0.0, + "to" : "dam", + "scaling" : 10.0 + } + ], + "hPa" : [ + { + "from" : "Pa", + "offset" : -0.0, + "to" : "hPa", + "scaling" : 0.01 + } + ], + "kg m**-2 h**-1" : [ + { + "from" : "kg m**-2 s**-1", + "offset" : -0.0, + "to" : "kg m**-2 h**-1", + "scaling" : 3600.0 + } + ], + "Dob" : [ + { + "from" : "kg m**-2", + "offset" : -0.0, + "to" : "Dob", + "scaling" : 46696.2409526033 + } + ] +} + + + + + diff --git a/metview/gallery.py b/metview/gallery.py index 27f91c7d..bc0d4c11 100644 --- a/metview/gallery.py +++ b/metview/gallery.py @@ -7,20 +7,42 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. +import os import zipfile import metview as mv -def load_dataset(filename): - base_url = "http://download.ecmwf.org/test-data/metview/gallery/" +def load_dataset(filename, check_local=False): + def _simple_download(url, target): + import requests + + r = requests.get(url, allow_redirects=True) + r.raise_for_status() + open(target, "wb").write(r.content) + + if check_local and os.path.exists(filename): + try: + return mv.read(filename) + except: + return None + + base_url = "https://get.ecmwf.int/test-data/metview/gallery/" try: - d = mv.download(url=base_url + filename, target=filename) - if filename.endswith(".zip"): - with zipfile.ZipFile(filename, "r") as f: - f.extractall() - return d - except: + # d = mv.download(url=base_url + filename, target=filename) + _simple_download(os.path.join(base_url, filename), filename) + except Exception as e: raise Exception( - "Could not download file " + filename + " from the download server" + f"Could not download file={filename} from the download server. {e}" ) + + d = None + if filename.endswith(".zip"): + with zipfile.ZipFile(filename, "r") as f: + f.extractall() + else: + try: + d = mv.read(filename) + except: + pass + return d diff --git a/metview/layout.py b/metview/layout.py new file mode 100644 index 00000000..a22ff46f --- /dev/null +++ b/metview/layout.py @@ -0,0 +1,306 @@ +# +# (C) Copyright 2020- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import logging +import math + +import metview as mv + +LOG = logging.getLogger(__name__) + + +class Layout: + + GRID_DEF = { + 1: [1, 1], + 2: [2, 1], + 3: [3, 1], + 4: [2, 2], + 8: [3, 3], + 9: [3, 3], + 10: [4, 3], + } + + def _grid_row_col(self, page_num=0, rows=None, columns=None, layout=None): + if rows is None and columns is None: + if layout is not None and layout != "": + v = layout.split("x") + if len(v) == 2: + try: + r = int(v[0]) + c = int(v[1]) + # LOG.debug(f"r{r} c={c}") + except: + raise Exception(f"Invalid layout specification ({layout}") + if page_num > 0 and r * c < page_num: + raise Exception( + f"Layout specification={layout} does not match number of scenes={page_num}" + ) + else: + if page_num in self.GRID_DEF: + r, c = self.GRID_DEF[page_num] + elif page_num >= 1: + r, c = self._compute_row_col(page_num) + else: + raise Exception(f"Cannot create a layout for {page_num} pages!") + return r, c + + def build_grid(self, page_num=0, rows=None, columns=None, layout=None, view=None): + r, c = self._grid_row_col( + page_num=page_num, rows=rows, columns=columns, layout=layout + ) + + if isinstance(view, mv.style.GeoView): + v = view.to_request() + else: + v = view + + return self._build_grid(rows=r, columns=c, view=v) + + def _build_grid(self, rows=1, columns=1, view=None): + assert rows >= 1 and columns >= 1 + if rows == 1 and columns == 1: + return mv.plot_superpage(pages=[mv.plot_page(view=view)]) + else: + return mv.plot_superpage( + pages=mv.mvl_regular_layout( + view, columns, rows, 1, 1, [5, 100, 15, 100] + ) + ) + + def _compute_row_col(self, page_num): + # has been checked for 100 + r = int(math.sqrt(page_num)) + c = r + if c * r < page_num: + if (page_num - c * r) % c == 0: + c += (page_num - c * r) // c + else: + c += 1 + (page_num - c * r) // c + return (r, c) + + def build_diff(self, view): + page_1 = mv.plot_page(top=5, bottom=50, left=25, right=75, view=view) + + page_2 = mv.plot_page(top=52, bottom=97, right=50, view=view) + + page_3 = mv.plot_page(top=52, bottom=97, left=50, right=100, view=view) + + return mv.plot_superpage(pages=[page_1, page_2, page_3]) + + def build_xs(self, line, map_view): + xs_view = mv.mxsectview( + line=line, + bottom_level=1000, + top_level=150 + # wind_perpendicular : "off", + # wind_parallel :"on", + # wind_intensity :"off" + ) + + page_1 = mv.plot_page(top=35 if map_view is not None else 5, view=xs_view) + + if map_view is not None: + page = mv.plot_page(top=5, bottom=35, view=map_view) + return mv.plot_superpage(pages=[page_1, page]) + else: + return mv.plot_superpage(pages=[page_1]) + + def build_xs_avg( + self, + area, + direction, + bottom_level, + top_level, + vertical_scaling, + axis_label_height=0.4, + ): + axis = mv.maxis(axis_tick_label_height=axis_label_height) + + return mv.maverageview( + area=area, + direction=direction, + bottom_level=bottom_level, + top_level=top_level, + vertical_scaling=vertical_scaling, + horizontal_axis=axis, + vertical_axis=axis, + ) + + def build_stamp(self, page_num=0, layout="", view=None): + + if True: + coast_empty = mv.mcoast( + map_coastline="off", map_grid="off", map_label="off" + ) + + empty_view = mv.geoview( + page_frame="off", + subpage_frame="off", + coastlines=coast_empty, + subpage_x_position=40, + subpage_x_length=10, + subpage_y_length=10, + ) + + title_page = mv.plot_page( + top=0, bottom=5, left=30, right=70, view=empty_view + ) + + r, c = self._grid_row_col(page_num=page_num, layout=layout) + + pages = mv.mvl_regular_layout(view, c, r, 1, 1, [5, 100, 15, 100]) + pages.append(title_page) + return mv.plot_superpage(pages=pages) + + # g = self.build_grid(page_num=page_num, layout=layout, view=view) + # return g + + def build_rmse(self, xmin, xmax, ymin, ymax, xtick, ytick, xtitle, ytitle): + + horizontal_axis = mv.maxis( + axis_type="date", + axis_orientation="horizontal", + axis_position="bottom", + # axis_title : 'on', + # axis_title_height : 0.4, + axis_grid="on", + axis_grid_colour="grey", + # axis_title_text : xtitle, + # axis_title_quality : 'high', + axis_tick_interval=xtick, + axis_tick_label_height=0.6, + axis_date_type="days", + axis_years_label="off", + axis_months_label="off", + axis_days_label_height="0.3", + ) + + vertical_axis = mv.maxis( + axis_orientation="vertical", + axis_position="left", + axis_grid="on", + axis_grid_colour="grey", + axis_title="on", + axis_title_height=0.4, + axis_title_text=ytitle, + # axis_title_quality : 'high', + axis_tick_interval=ytick, + axis_tick_label_height=0.3, + ) + + cview = mv.cartesianview( + page_frame="off", + x_axis_type="date", + y_axis_type="regular", + y_automatic="off", + x_automatic="off", + y_min=ymin, + y_max=ymax, + x_date_min=xmin, + x_date_max=xmax, + horizontal_axis=horizontal_axis, + vertical_axis=vertical_axis, + ) + + return cview + + def build_xy(self, xmin, xmax, ymin, ymax, xtick, ytick, xtitle, ytitle): + + horizontal_axis = mv.maxis( + axis_orientation="horizontal", + axis_position="bottom", + axis_title="on", + axis_title_height=0.5, + axis_title_text=xtitle, + # axis_title : 'on', + # axis_title_height : 0.4, + axis_grid="on", + axis_grid_colour="grey", + # axis_title_text : xtitle, + # axis_title_quality : 'high', + axis_tick_interval=xtick, + axis_tick_label_height=0.6, + axis_date_type="days", + axis_years_label="off", + axis_months_label="off", + axis_days_label_height="0.3", + ) + + vertical_axis = mv.maxis( + axis_orientation="vertical", + axis_position="left", + axis_grid="on", + axis_grid_colour="grey", + axis_title="on", + axis_title_height=0.6, + axis_title_text=ytitle, + # axis_title_quality : 'high', + axis_tick_interval=ytick, + axis_tick_label_height=0.3, + ) + + cview = mv.cartesianview( + page_frame="off", + x_axis_type="regular", + y_axis_type="regular", + y_automatic="off", + x_automatic="off", + y_min=ymin, + y_max=ymax, + x_min=xmin, + x_max=xmax, + horizontal_axis=horizontal_axis, + vertical_axis=vertical_axis, + ) + + return cview + + @staticmethod + def compute_axis_range(v_min, v_max): + count = 15 + d = math.fabs(v_max - v_min) + if d > 0: + b = d / count + n = math.floor(math.log10(b)) + v = b / math.pow(10, n) + # print("d={} b={} n={} v={}".format(d,b,n,v)) + + if v <= 1: + v = 1 + elif v <= 2: + v = 2 + elif v <= 5: + v = 5 + else: + v = 10 + + bin_size = v * math.pow(10, n) + bin_start = math.ceil(v_min / bin_size) * bin_size + if bin_start >= v_min and math.fabs(bin_start - v_min) < bin_size / 10000: + bin_start = bin_start - bin_size / 100000 + av = v_min + else: + bin_start = bin_start - bin_size + av = bin_start + + max_iter = 100 + act_iter = 0 + while av < v_max: + av += bin_size + act_iter += 1 + if act_iter > max_iter: + return (0, v_min, v_max) + + bin_end = av # + bin_size / 100000 + return bin_size, bin_start, bin_end + else: + return (0, v_min, v_max) diff --git a/metview/metviewpy/__init__.py b/metview/metviewpy/__init__.py new file mode 100644 index 00000000..ecc7c909 --- /dev/null +++ b/metview/metviewpy/__init__.py @@ -0,0 +1,29 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import os + + +# if we're running pytest, it will need the fieldset functionality, so detect if it +# is running and if so, import the user-facing functions + + +def running_from_pytest(): + from inspect import stack + + call_stack = [s.function for s in stack()] + return "pytest_collection" in call_stack + + +if "METVIEW_PYTHON_ONLY" in os.environ or running_from_pytest(): + + from . import fieldset + + fieldset.bind_functions(globals(), module_name=__name__) diff --git a/metview/metviewpy/fieldset.py b/metview/metviewpy/fieldset.py new file mode 100644 index 00000000..c280186f --- /dev/null +++ b/metview/metviewpy/fieldset.py @@ -0,0 +1,890 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + + +from inspect import signature +import sys + +import numpy as np +import eccodes + +from . import maths +from .temporary import temp_file + +from . import indexdb as indexdb +from . import utils + + +BITS_PER_VALUE_FOR_WRITING = 24 + + +class CodesHandle: + """Wraps an ecCodes handle""" + + MISSING_VALUE = 1e34 + + STRING = ("s", eccodes.codes_get_string, eccodes.codes_set_string) + LONG = ("l", eccodes.codes_get_long, eccodes.codes_set_long) + DOUBLE = ("d", eccodes.codes_get_double, eccodes.codes_set_double) + LONG_ARRAY = ("la", eccodes.codes_get_long_array, eccodes.codes_set_long_array) + DOUBLE_ARRAY = ( + "da", + eccodes.codes_get_double_array, + eccodes.codes_set_double_array, + ) + NATIVE = ("n", eccodes.codes_get, eccodes.codes_set) + NATIVE_ARRAY = ("na", eccodes.codes_get_array, eccodes.codes_set_array) + + _STR_TYPES = { + v[0]: tuple(v[1:3]) + for v in [STRING, LONG, DOUBLE, LONG_ARRAY, DOUBLE_ARRAY, NATIVE, NATIVE_ARRAY] + } + + def __init__(self, handle, path, offset): + self.handle = handle + self.path = path + self.offset = offset + eccodes.codes_set(handle, "missingValue", CodesHandle.MISSING_VALUE) + + def __del__(self): + # print("CodesHandle:release ", self) + eccodes.codes_release(self.handle) + + def clone(self): + new_handle = eccodes.codes_clone(self.handle) + return CodesHandle(new_handle, None, None) + + # def __str__(self): + # s = "CodesHandle(" + # return s + str(self.handle) + "," + self.path + "," + str(self.offset) + ")" + + def call_get_func(self, func, key, default=False): + try: + result = func(self.handle, key) + except Exception as exp: + if default != False: + return default + else: + raise exp + return result + + def get_any(self, keys, key_type=None): + # single key with key_type specified + if key_type is not None: + assert isinstance(keys, str) + func = key_type[1] + return func(self.handle, keys) + # list of keys + else: + result = [] + for key in keys: + key_type_str = "s" + parts = key.split(":") + if len(parts) == 2: + key, key_type_str = parts + if ( + key_type_str == "n" + and eccodes.codes_get_size(self.handle, key) > 1 + ): + key_type_str = "na" + func = CodesHandle._STR_TYPES.get(key_type_str, None)[0] + result.append(self.call_get_func(func, key, default=None)) + return result + + def get_string(self, key): + return eccodes.codes_get_string(self.handle, key) + + def get_long(self, key): + return eccodes.codes_get_long(self.handle, key) + + def get_double(self, key): + return eccodes.codes_get_double(self.handle, key) + + def get_long_array(self, key): + return eccodes.codes_get_long_array(self.handle, key) + + def get_double_array(self, key): + return eccodes.codes_get_double_array(self.handle, key) + + def get_values(self): + vals = eccodes.codes_get_values(self.handle) + if self.get_long("bitmapPresent"): + vals[vals == CodesHandle.MISSING_VALUE] = np.nan + return vals + + def set_any(self, keys_and_vals, key_type=None): + for key, v in zip(keys_and_vals[0::2], keys_and_vals[1::2]): + if key_type is not None: + func = key_type[2] + else: + key_type_str = "s" # default is str + parts = key.split(":") + if len(parts) == 2: + key, key_type_str = parts + func = CodesHandle._STR_TYPES.get(key_type_str, None)[1] + func(self.handle, key, v) + + def set_string(self, key, value): + eccodes.codes_set_string(self.handle, key, value) + + def set_long(self, key, value): + eccodes.codes_set_long(self.handle, key, value) + + def set_double(self, key, value): + eccodes.codes_set_double(self.handle, key, value) + + def set_double_array(self, key, value): + eccodes.codes_set_double_array(self.handle, key, value) + + def set_values(self, value): + # replace nans with missing values + values_nans_replaced = np.nan_to_num( + value, copy=True, nan=CodesHandle.MISSING_VALUE + ) + self.set_long("bitsPerValue", BITS_PER_VALUE_FOR_WRITING) + self.set_double_array("values", values_nans_replaced) + + def write(self, fout, path): + self.offset = fout.tell() + eccodes.codes_write(self.handle, fout) + if path: + self.path = path + + +class GribFile: + """Encapsulates a GRIB file, giving access to an iteration of CodesHandles""" + + def __init__(self, path): + self.path = path + self.file = open(path, "rb") + self.num_messages = eccodes.codes_count_in_file(self.file) + + def __del__(self): + try: + # print("GribFile:close") + self.file.close() + except Exception: + pass + + def __len__(self): + return self.num_messages + + def __iter__(self): + return self + + def __next__(self): + handle = self._next_handle() + if handle is None: + raise StopIteration() + return handle + + def _next_handle(self): + offset = self.file.tell() + handle = eccodes.codes_new_from_file(self.file, eccodes.CODES_PRODUCT_GRIB) + if not handle: + return None + return CodesHandle(handle, self.path, offset) + + +class Field: + """Encapsulates single GRIB message""" + + def __init__(self, handle, gribfile, keep_values_in_memory=False, temp=None): + self.handle = handle + self.gribfile = gribfile + self.temp = temp + self.vals = None + self.keep_values_in_memory = keep_values_in_memory + + def grib_get(self, *args, **kwargs): + return self.handle.get_any(*args, **kwargs) + + def values(self): + if self.vals is None: + vals = self.handle.get_values() + if self.keep_values_in_memory: + self.vals = vals + else: + vals = self.vals + return vals + + def latitudes(self): + return self.handle.get_double_array("latitudes") + + def longitudes(self): + return self.handle.get_double_array("longitudes") + + def grib_set(self, *args, **kwargs): + result = self.clone() + result.handle.set_any(*args, **kwargs) + return result + + def encode_values(self, value): + self.handle.set_long("bitmapPresent", 1) + self.handle.set_values(value) + if not self.keep_values_in_memory: + self.vals = None + + def write(self, fout, path, temp=None): + self.temp = temp # store a reference to the temp file object for persistence + self.handle.write(fout, path) + + def grib_index(self): + return (self.handle.path, self.handle.offset) + # return (self.handle.path, self.grib_get("offset", key_type=CodesHandle.LONG)) + + def clone(self): + c = Field( + self.handle.clone(), + self.gribfile, + self.keep_values_in_memory, + ) + c.vals = None + return c + + def field_func(self, func): + """Applies a function to all values, returning a new Field""" + result = self.clone() + result.vals = func(self.values()) + result.encode_values(result.vals) + return result + + def field_other_func(self, func, other, reverse_args=False): + """Applies a function with something to all values, returning a new Field""" + result = self.clone() + if isinstance(other, Field): + other = other.values() + if reverse_args: + result.vals = func(other, self.values()) + else: + result.vals = func(self.values(), other) + result.encode_values(result.vals) + return result + + +# decorator to implement math functions in Fieldset +def wrap_maths(cls): + def wrap_single_method(fn): + def wrapper(self): + return self.field_func(fn) + + return wrapper + + def wrap_double_method(fn, **opt): + def wrapper(self, *args): + return self.fieldset_other_func(fn, *args, **opt) + + return wrapper + + for name, it in cls.WRAP_MATHS_ATTRS.items(): + if not isinstance(it, tuple): + it = (it, {}) + fn, opt = it + n = len(signature(fn).parameters) + if n == 1: + setattr(cls, name, wrap_single_method(fn)) + elif n == 2: + setattr(cls, name, wrap_double_method(fn, **opt)) + return cls + + +@wrap_maths +class Fieldset: + """A set of Fields, each of which can come from different GRIB files""" + + WRAP_MATHS_ATTRS = { + "abs": maths.abs, + "acos": maths.acos, + "asin": maths.asin, + "atan": maths.atan, + "atan2": maths.atan2, + "cos": maths.cos, + "div": maths.floor_div, + "exp": maths.exp, + "log": maths.log, + "log10": maths.log10, + "mod": maths.mod, + "sgn": maths.sgn, + "sin": maths.sin, + "square": maths.square, + "sqrt": maths.sqrt, + "tan": maths.tan, + "__neg__": maths.neg, + "__pos__": maths.pos, + "__invert__": maths.not_func, + "__add__": maths.add, + "__radd__": (maths.add, {"reverse_args": True}), + "__sub__": maths.sub, + "__rsub__": (maths.sub, {"reverse_args": True}), + "__mul__": maths.mul, + "__rmul__": (maths.mul, {"reverse_args": True}), + "__truediv__": maths.div, + "__rtruediv__": (maths.div, {"reverse_args": True}), + "__pow__": maths.pow, + "__rpow__": (maths.pow, {"reverse_args": True}), + "__ge__": maths.ge, + "__gt__": maths.gt, + "__le__": maths.le, + "__lt__": maths.lt, + "__eq__": maths.eq, + "__ne__": maths.ne, + "__and__": maths.and_func, + "__or__": maths.or_func, + "bitmap": (maths.bitmap, {"use_first_from_other": True}), + "nobitmap": maths.nobitmap, + } + + # QUALIFIER_MAP = {"float": "d"} + + # INT_KEYS = ["Nx", "Ny", "number"] + + def __init__( + self, path=None, fields=None, keep_values_in_memory=False, temporary=False + ): + self.fields = [] + self.count = 0 + self.temporary = None + self._db = None + + if (path is not None) and (fields is not None): + raise ValueError("Fieldset cannot take both path and fields") + + if path: + if isinstance(path, list): + v = [] + for p in path: + v.extend(utils.get_file_list(p)) + path = v + else: + path = utils.get_file_list(path) + + for p in path: + g = GribFile(p) + self.count = len(g) + for handle in g: + self.fields.append(Field(handle, p, keep_values_in_memory)) + if temporary: + self.temporary = temp_file() + if fields: + self.fields = fields + + def __len__(self): + return len(self.fields) + + def __str__(self): + n = len(self) + s = "s" if n > 1 else "" + return f"Fieldset ({n} field{s})" + + def _grib_get(self, *args, as_list=False, **kwargs): + ret = [x.grib_get(*args, **kwargs) for x in self.fields] + return ret if as_list else Fieldset._list_or_single(ret) + + def grib_get_string(self, key): + return self._grib_get(key, key_type=CodesHandle.STRING) + + def grib_get_long(self, key): + return self._grib_get(key, key_type=CodesHandle.LONG) + + def grib_get_double(self, key): + return self._grib_get(key, key_type=CodesHandle.DOUBLE) + + def grib_get_long_array(self, key): + return self._grib_get(key, key_type=CodesHandle.LONG_ARRAY) + + def grib_get_double_array(self, key): + return self._grib_get(key, key_type=CodesHandle.DOUBLE_ARRAY) + + def grib_get(self, keys, grouping="field"): + if grouping not in ["field", "key"]: + raise ValueError(f"grib_get: grouping must be field or key, not {grouping}") + ret = self._grib_get(keys, as_list=True) + if grouping == "key": + ret = list(map(list, zip(*ret))) # transpose lists of lists + return ret + + def _grib_set(self, *args, **kwargs): + result = Fieldset(temporary=True) + path = result.temporary.path + with open(path, "wb") as fout: + for f in self.fields: + result._append_field(f.grib_set(*args, **kwargs)) + result.fields[-1].write(fout, path, temp=result.temporary) + return result + + def grib_set_string(self, keys_and_vals): + return self._grib_set(keys_and_vals, key_type=CodesHandle.STRING) + + def grib_set_long(self, keys_and_vals): + return self._grib_set(keys_and_vals, key_type=CodesHandle.LONG) + + def grib_set_double(self, keys_and_vals): + return self._grib_set(keys_and_vals, key_type=CodesHandle.DOUBLE) + + def grib_set(self, keys_and_vals): + return self._grib_set(keys_and_vals) + + # def grib_set_double_array(self, key, value): + # return self._grib_set_any(key, value, "grib_set_double_array") + + def values(self): + v = [x.values() for x in self.fields] + return self._make_2d_array(v) + + def set_values(self, values): + if isinstance(values, list): + list_of_arrays = values + else: + if len(values.shape) > 1: + list_of_arrays = [a for a in values] + else: + list_of_arrays = [values] + if len(list_of_arrays) != len(self.fields): + msg = str(len(list_of_arrays)) + " instead of " + str(len(self.fields)) + raise ValueError("set_values has the wrong number of arrays: " + msg) + + return self.fieldset_other_func( + maths.set_from_other, list_of_arrays, index_other=True + ) + + def write(self, path): + with open(path, "wb") as fout: + for f in self.fields: + f.write(fout, path) + + def grib_index(self): + return [i.grib_index() for i in self.fields] + + def _append_field(self, field): + self.fields.append(field) + self._db = None + + def __getitem__(self, index): + try: + if isinstance(index, np.ndarray): + return Fieldset(fields=[self.fields[i] for i in index]) + # # GRIB key + # elif isinstance(index, str): + # keyname = index + # qualifier = "n" + # parts = keyname.split(":") + # if len(parts) > 1: + # keyname = parts[0] + # qualifier = Fieldset.QUALIFIER_MAP[parts[1]] + + # native_key = keyname + ":" + qualifier + # # print("native_key", native_key) + # value = self.grib_get([native_key])[0][0] + # # print(value) + # if value is None: + # raise IndexError + # if index in Fieldset.INT_KEYS: + # # print("int value:", int(value)) + # return int(value) + # # if isinstance(value, float): + # # return int(value) + # return value + elif isinstance(index, str): + return self._get_db().select_with_name(index) + else: + return Fieldset(fields=self._always_list(self.fields[index])) + except IndexError as ide: + # print("This Fieldset contains", len(self), "fields; index is", index) + raise ide + + def append(self, other): + self.fields = self.fields + other.fields + self._db = None + + def merge(self, other): + result = Fieldset(fields=self.fields, temporary=True) + result.append(other) + return result + + # def items(self): + # its = [] + # for i in range(len(self)): + # its.append((i, Fieldset(fields=[self.fields[i]]))) + # return its + + # def to_dataset(self, via_file=True, **kwarg): + # # soft dependency on cfgrib + # try: + # import xarray as xr + # except ImportError: # pragma: no cover + # print("Package xarray not found. Try running 'pip install xarray'.") + # raise + # # dataset = xr.open_dataset(self.url(), engine="cfgrib", backend_kwargs=kwarg) + # if via_file: + # dataset = xr.open_dataset(self.url(), engine="cfgrib", backend_kwargs=kwarg) + # else: + # print("Using experimental cfgrib interface to go directly from Fieldset") + # dataset = xr.open_dataset(self, engine="cfgrib", backend_kwargs=kwarg) + # return dataset + + def field_func(self, func): + """Applies a function to all values in all fields""" + result = Fieldset(temporary=True) + with open(result.temporary.path, "wb") as fout: + for f in self.fields: + result._append_field(f.field_func(func)) + result.fields[-1].write( + fout, result.temporary.path, temp=result.temporary + ) + return result + + def fieldset_other_func( + self, + func, + other, + reverse_args=False, + index_other=False, + use_first_from_other=False, + ): + """Applies a function to a fieldset and a scalar/fieldset, e.g. F+5""" + # print( + # f"Fieldset.fieldset_other_func() func={func}, other={other}, reverse_args={reverse_args}, index_other={index_other}, use_first_from_other={use_first_from_other}" + # ) + + def _process_one(f, g, result): + new_field = f.field_other_func(func, g, reverse_args=reverse_args) + result._append_field(new_field) + result.fields[-1].write(fout, result.temporary.path, temp=result.temporary) + + result = Fieldset(temporary=True) + with open(result.temporary.path, "wb") as fout: + if isinstance(other, Fieldset): + if len(other) == len(self.fields): + for f, g in zip(self.fields, other.fields): + _process_one(f, g, result) + elif use_first_from_other: + for f in self.fields: + _process_one(f, other.fields[0], result) + else: + raise Exception( + f"Fieldsets must have the same number of fields for this operation! {len(self.fields)} != {len(other)}" + ) + elif index_other: + for f, g in zip(self.fields, other): + _process_one(f, g, result) + else: + for f in self.fields: + _process_one(f, other, result) + return result + + def base_date(self): + if len(self.fields) > 0: + result = [] + for f in self.fields: + md = f.grib_get(["dataDate", "dataTime"]) + result.append( + utils.date_from_ecc_keys(md[0], md[1]) if len(md) == 2 else None + ) + return Fieldset._list_or_single(result) + + def valid_date(self): + if len(self.fields) > 0: + result = [] + for f in self.fields: + md = f.grib_get(["validityDate", "validityTime"]) + result.append( + utils.date_from_ecc_keys(md[0], md[1]) if len(md) == 2 else None + ) + return Fieldset._list_or_single(result) + + def accumulate(self): + if len(self.fields) > 0: + result = [np.nan] * len(self.fields) + for i, f in enumerate(self.fields): + result[i] = np.sum(f.values()) + return Fieldset._list_or_single(result) + + def average(self): + if len(self.fields) > 0: + result = [np.nan] * len(self.fields) + for i, f in enumerate(self.fields): + result[i] = f.values().mean() + return Fieldset._list_or_single(result) + + def maxvalue(self): + if len(self.fields) > 0: + result = np.array([np.nan] * len(self.fields)) + for i, f in enumerate(self.fields): + result[i] = f.values().max() + return result.max() + + def minvalue(self): + if len(self.fields) > 0: + result = np.array([np.nan] * len(self.fields)) + for i, f in enumerate(self.fields): + result[i] = f.values().min() + return result.min() + + def _make_single_result(self, v): + assert len(self) > 0 + result = Fieldset(temporary=True) + with open(result.temporary.path, "wb") as fout: + f = self.fields[0].clone() + f.encode_values(v) + result._append_field(f) + result.fields[-1].write(fout, result.temporary.path, temp=result.temporary) + return result + + def mean(self): + if len(self.fields) > 0: + v = self.fields[0].values() + for i in range(1, len(self.fields)): + v += self.fields[i].values() + v = v / len(self.fields) + return self._make_single_result(v) + else: + return None + + def rms(self): + if len(self.fields) > 0: + v = np.square(self.fields[0].values()) + for i in range(1, len(self.fields)): + v += np.square(self.fields[i].values()) + v = np.sqrt(v / len(self.fields)) + return self._make_single_result(v) + else: + return None + + def stdev(self): + v = self._compute_var() + return self._make_single_result(np.sqrt(v)) if v is not None else None + + def sum(self): + if len(self.fields) > 0: + v = self.fields[0].values() + for i in range(1, len(self.fields)): + v += self.fields[i].values() + return self._make_single_result(v) + else: + return None + + def var(self): + v = self._compute_var() + return self._make_single_result(v) if v is not None else None + + def _compute_var(self): + if len(self.fields) > 0: + v2 = self.fields[0].values() + v1 = np.square(v2) + for i in range(1, len(self.fields)): + v = self.fields[i].values() + v1 += np.square(v) + v2 += v + return v1 / len(self.fields) - np.square(v2 / len(self.fields)) + + def latitudes(self): + v = [x.latitudes() for x in self.fields] + return Fieldset._make_2d_array(v) + + def longitudes(self): + v = [x.longitudes() for x in self.fields] + return Fieldset._make_2d_array(v) + + def _lat_func(self, func, bitmap_poles=False): + result = Fieldset(temporary=True) + pole_limit = 90.0 - 1e-06 + with open(result.temporary.path, "wb") as fout: + for f in self.fields: + lat = f.latitudes() + if bitmap_poles: + lat[np.fabs(lat) > pole_limit] = np.nan + v = func(np.deg2rad(lat)) + c = f.clone() + c.encode_values(v) + result._append_field(c) + result.fields[-1].write( + fout, result.temporary.path, temp=result.temporary + ) + return result + + def coslat(self): + return self._lat_func(np.cos) + + def sinlat(self): + return self._lat_func(np.sin) + + def tanlat(self): + return self._lat_func(np.tan, bitmap_poles=True) + + @staticmethod + def _list_or_single(lst): + return lst if len(lst) != 1 else lst[0] + + @staticmethod + def _always_list(items): + return items if isinstance(items, list) else [items] + + @staticmethod + def _make_2d_array(v): + """Forms a 2D ndarray from a list of 1D ndarrays""" + v = Fieldset._list_or_single(v) + return np.stack(v, axis=0) if isinstance(v, list) else v + + # TODO: add all the field_func functions + + # TODO: add all field_field_func functions + + # TODO: add all field_scalar_func functions + + # TODO: add tests for all fieldset-fieldset methods: + # **, ==, !=, &, | + + # TODO: allow these methods to be called as functions (?) + + # TODO: function to write to single file if fields from different files + + # TODO: optimise write() to copy file if exists + + # TODO: to_dataset() + + # TODO: pickling + + # TODO: gribsetbits, default=24 + + def _get_db(self): + if self._db is None: + self._db = indexdb.FieldsetDb(fs=self) + assert self._db is not None + return self._db + + def _unique_metadata(self, key): + return self._get_db().unique(key) + + def select(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + return self._get_db().select(**args[0]) + else: + return self._get_db().select(**kwargs) + + def describe(self, *args, **kwargs): + return self._get_db().describe(*args, **kwargs) + + def ls(self, **kwargs): + return self._get_db().ls(**kwargs) + + def sort(self, *args, **kwargs): + return self._get_db().sort(*args, **kwargs) + + def deacc(self, **kwargs): + return utils.deacc(self, **kwargs) + + def speed(self, *args): + if len(args) == 0: + u = self[0::2] + v = self[1::2] + if len(u) != len(v): + raise Exception( + f"Fieldsets must contain an even number of fields for this operation! len={len(self.fields)} is not even!" + ) + return u.speed(v) + elif len(args) == 1: + other = args[0] + result = Fieldset(temporary=True) + param_ids = { + 131: 10, # atmospheric wind + 165: 207, # 10m wind + 228246: 228249, # 100m wind + 228239: 228241, # 200m wind + } + if len(self.fields) != len(other): + raise Exception( + f"Fieldsets must have the same number of fields for this operation! {len(self.fields)} != {len(other)}" + ) + with open(result.temporary.path, "wb") as fout: + for f, g in zip(self.fields, other.fields): + sp = np.sqrt(np.square(f.values()) + np.square(g.values())) + c = f.clone() + c.encode_values(sp) + param_id_u = f.grib_get("paramId", CodesHandle.LONG) + param_id_sp = param_ids.get(param_id_u, None) + if param_id_sp is not None: + c = c.grib_set(["paramId", param_id_sp], CodesHandle.LONG) + result._append_field(c) + result.fields[-1].write( + fout, result.temporary.path, temp=result.temporary + ) + return result + + +class FieldsetCF: + def __init__(self, fs): + self.fs = fs + + def items(self): + its = [] + for i in range(len(self.fs)): + its.append((i, FieldCF(self.fs[i]))) + # print("FieldsetCF items: ", d) + return its + + def __getitem__(self, index): + return FieldCF(self.fs[index]) + + +class FieldCF: + + QUALIFIER_MAP = {"float": "d"} + + INT_KEYS = ["Nx", "Ny", "number"] + + def __init__(self, fs): + self.fs = fs + + def __getitem__(self, key): + keyname = key + qualifier = "n" + parts = key.split(":") + if len(parts) > 1: + keyname = parts[0] + qualifier = FieldCF.QUALIFIER_MAP[parts[1]] + + native_key = keyname + ":" + qualifier + # print("native_key", native_key) + value = self.fs.grib_get([native_key])[0][0] + # print(value) + if value is None: + raise IndexError + if key in FieldCF.INT_KEYS: + # print("int value:", int(value)) + return int(value) + # if isinstance(value, float): + # return int(value) + return value + + +def read(p): + return Fieldset(path=p) + + +# expose all Fieldset functions as a module level function +def _make_module_func(name): + def wrapped(fs, *args): + return getattr(fs, name)(*args) + + return wrapped + + +module_obj = sys.modules[__name__] +for fn in dir(Fieldset): + if callable(getattr(Fieldset, fn)) and not fn.startswith("_"): + setattr(module_obj, fn, _make_module_func(fn)) + + +def bind_functions(namespace, module_name=None): + """Add to the module globals all metview functions except operators like: +, &, etc.""" + namespace["read"] = read + for fn in dir(Fieldset): + if callable(getattr(Fieldset, fn)) and not fn.startswith("_"): + namespace[fn] = _make_module_func(fn) + namespace["Fieldset"] = Fieldset diff --git a/metview/metviewpy/indexdb.py b/metview/metviewpy/indexdb.py new file mode 100644 index 00000000..860102dc --- /dev/null +++ b/metview/metviewpy/indexdb.py @@ -0,0 +1,581 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import copy +import logging +import os + +import pandas as pd + +from .indexer import GribIndexer, FieldsetIndexer +from .param import ( + ParamInfo, + ParamNameDesc, + ParamIdDesc, + init_pandas_options, +) +from .ipython import is_ipython_active + + +# logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") +# logging.basicConfig(level=logging.DEBUG, format="%(levelname)s - %(message)s") +LOG = logging.getLogger(__name__) + + +class IndexDb: + ROOTDIR_PLACEHOLDER_TOKEN = "__ROOTDIR__" + + def __init__( + self, + name, + label="", + desc="", + path="", + rootdir_placeholder_value="", + file_name_pattern="", + db_dir="", + blocks=None, + data_files=None, + merge_conf=None, + mapped_params=None, + regrid_from=None, + dataset=None, + ): + self.name = name + self.dataset = dataset + self.label = label + if self.label is None or self.label == "": + self.label = self.name + self.desc = desc + self.path = path + + self.rootdir_placeholder_value = rootdir_placeholder_value + self.file_name_pattern = file_name_pattern + if self.file_name_pattern == "": + self.path = os.path.dirname(self.path) + self.file_name_pattern = os.path.basename(self.path) + + self.db_dir = db_dir + self.mapped_params = {} if mapped_params is None else mapped_params + self.regrid_from = [] if regrid_from is None else regrid_from + self.blocks = {} if blocks is None else blocks + self.vector_loaded = False + self._param_types = {} + self.data_files = [] if data_files is None else data_files + self.merge_conf = [] if merge_conf is None else merge_conf + self._params = {} + + def select_with_name(self, name): + """ + Perform a select operation where selection options are derived + from the specified name. + """ + # print(f"select_with_name blocks: {self.blocks.keys()}") + # print(f"vector_loaded: {self.vector_loaded}") + + if "wind" in name and not self.vector_loaded: + self.load(vector=True) + pnf = ParamInfo.build_from_name(name, param_level_types=self.param_types) + if pnf is not None: + fs = self._select_fs( + **pnf.make_filter(), _named_vector_param=(not pnf.scalar) + ) + if fs is not None: + pnf.update_meta(fs._db._first_index_row()) + fs._ds_param_info = pnf + return fs + return None + + def select(self, **kwargs): + return self._select_fs(**kwargs) + + def _select_fs(self, **kwargs): + """ + Create a fieldset with the specified filter conditions. The resulting fieldset + will contain an index db. + """ + LOG.debug(f"kwargs={kwargs}") + + vector = kwargs.pop("_named_vector_param", False) + max_count = kwargs.pop("_max_count", -1) + + # print(f"kwargs={kwargs}") + # LOG.debug(f"blocks={self.blocks}") + dims = self._make_dims(kwargs) + + # We can only have a vector param when the fs["wind"]-like operator is + # invoked and we deduce the shortName from the specified name + if vector: + short_name = dims.get("shortName", []) + if isinstance(short_name, list): + assert len(short_name) == 1 + + # print(f"dims={dims}") + self.load(keys=list(dims.keys()), vector=vector) + + db, fs = self._get_fields(dims, max_count=max_count, vector=vector) + fs._db = db + # LOG.debug(f"fs={fs}") + # print(f"blocks={fs._db.blocks}") + return fs + + def _get_fields(self, dims, max_count=-1, vector=False): + + res = self.fieldset_class() + dfs = {} + LOG.debug(f"dims={dims}") + + name_filter = "shortName" in dims or "paramId" in dims + if not vector and name_filter: + # in this case filtering can only be done on the scalar block + if "scalar" in self.blocks.keys(): + self._get_fields_for_block("scalar", dims, dfs, res, max_count) + else: + for key in self.blocks.keys(): + self._get_fields_for_block(key, dims, dfs, res, max_count) + if max_count != -1 and len(res) >= max_count: + break + + # LOG.debug(f"len_res={len(res)}") + # LOG.debug(f"dfs={dfs}") + # LOG.debug(f"res={res}") + c = FieldsetDb( + res, + name=self.name, + blocks=dfs, + label=self.label, + mapped_params=self.mapped_params, + regrid_from=self.regrid_from, + ) + return c, res + + def _get_meta(self, dims): + LOG.debug(f"dims={dims}") + key = "scalar" + if key in self.blocks: + if self.blocks[key] is None: + self._load_block(key) + df = self._filter_df(df=self.blocks[key], dims=dims) + # LOG.debug(f"df={df}") + return df + return None + + def _build_query(self, dims, df): + q = "" + for column, v in dims.items(): + # print(f"v={v}") + if v: + col_type = None + if q: + q += " and " + + # datetime columns + if column in GribIndexer.DATETIME_KEYS: + name_date = GribIndexer.DATETIME_KEYS[column][0] + name_time = GribIndexer.DATETIME_KEYS[column][1] + # here we should simply say: name_date*10000 + name_time. However, + # pandas cannot handle it in the query because the column types are + # Int64. np.int64 would work but that cannot handle missing values. So + # we need to break down the condition into individual logical components! + s = [] + for x in v: + # print(f"x={x}") + # print(" date=" + x.strftime("%Y%m%d")) + # print(" time=" + x.strftime("%H%M")) + s.append( + f"(`{name_date}` == " + + str(int(x.strftime("%Y%m%d"))) + + " and " + + f"`{name_time}` == " + + str(int(x.strftime("%H%M"))) + + ")" + ) + q += "(" + " or ".join(s) + " ) " + else: + col_type = df.dtypes[column] + column = f"`{column}`" + if not isinstance(v, list): + q += f"{column} == {GribIndexer._convert_query_value(v, col_type)}" + else: + v = [GribIndexer._convert_query_value(x, col_type) for x in v] + q += f"{column} in {v}" + return q + + def _filter_df(self, df=None, dims={}): + if len(dims) == 0: + return df + else: + df_res = None + if df is not None: + # print("types={}".format(df.dtypes)) + q = self._build_query(dims, df) + # print("query={}".format(q)) + if q != "": + df_res = df.query(q, engine="python") + df_res.reset_index(drop=True, inplace=True) + # print(f"df_res={df_res}") + # LOG.debug(f"df_res={df_res}") + else: + return df + return df_res + + def _get_fields_for_block(self, key, dims, dfs, res, max_count): + # print(f"key={key} dims={dims}") + # LOG.debug(f"block={self.blocks[key]}") + if self.blocks[key] is None: + self._load_block(key) + df = self._filter_df(df=self.blocks[key], dims=dims) + # print(f"df={df}") + # LOG.debug(f"df={df}") + if df is not None and not df.empty: + df_fs = self._extract_fields(df, res, max_count) + # LOG.debug(f"len_res={len(res)}") + if df_fs is not None: + # LOG.debug(f"df_fs={df_fs}") + dfs[key] = df_fs + + def _make_param_info(self): + m = self._first_index_row() + if m: + name = m["shortName"] + pnf = ParamInfo( + name, meta=dict(m), scalar=not name in ParamInfo.VECTOR_NAMES + ) + return pnf + return None + + def _first_index_row(self): + if self.blocks: + df = self.blocks[list(self.blocks.keys())[0]] + if df is not None and not df.empty: + row = df.iloc[0] + return dict(row) + return {} + + def _make_dims(self, options): + dims = {} + GribIndexer._check_datetime_in_filter_input(options) + for k, v in options.items(): + name = str(k) + vv = copy.deepcopy(v) + r = GribIndexer._convert_filter_value(name, self._to_list(vv)) + for name, vv in r: + if len(r) > 1 or vv: + dims[name] = vv + return dims + + def _to_list(self, v): + if not isinstance(v, list): + v = [v] + return v + + @property + def param_types(self): + if len(self._param_types) == 0: + self.load() + for k, df in self.blocks.items(): + df_u = df[["shortName", "typeOfLevel"]].drop_duplicates() + for row in df_u.itertuples(name=None): + if not row[1] in self._param_types: + self._param_types[row[1]] = [row[2]] + else: + self._param_types[row[1]].append(row[2]) + # print(self._param_types) + return self._param_types + + def unique(self, key): + self.load([key]) + for _, v in self.blocks.items(): + if key in v.columns: + return list(v[key].unique()) + return [] + + @property + def param_meta(self): + if len(self._params) == 0: + self.load() + for par in sorted(self.unique("shortName")): + self._params[par] = ParamNameDesc(par) + self._params[par].load(self) + return self._params + + def param_id_meta(self, param_id): + self.load() + p = ParamIdDesc(param_id) + p.load(self) + return p + + def describe(self, *args, **kwargs): + param = args[0] if len(args) == 1 else None + if param is None: + param = kwargs.pop("param", None) + return ParamNameDesc.describe(self, param=param, **kwargs) + + def to_df(self): + return pd.concat([p for _, p in self.blocks.items()]) + + def __str__(self): + return "{}[name={}]".format(self.__class__.__name__, self.name) + + +class FieldsetDb(IndexDb): + def __init__(self, fs, name="", **kwargs): + super().__init__(name, **kwargs) + self.fs = fs + self.fieldset_class = fs.__class__ + self._indexer = None + + @property + def indexer(self): + if self._indexer is None: + self._indexer = FieldsetIndexer(self) + return self._indexer + + def scan(self, vector=False): + self.indexer.scan(vector=vector) + self.vector_loaded = vector + + def load(self, keys=[], vector=False): + # print(f"blocks={self.blocks}") + if self.indexer.update_keys(keys): + self.blocks = {} + self._param_types = {} + self.scan(vector=self.vector_loaded) + elif not self.blocks: + self._param_types = {} + self.scan(vector=vector) + self.vector_loaded = vector + elif vector and not self.vector_loaded: + self._param_types = {} + self.indexer._scan_vector() + self.vector_loaded = True + + def _extract_fields(self, df, fs, max_count): + if df.empty: + return None + + # print(f"cols={df.columns}") + if "_msgIndex3" in df.columns: + comp_num = 3 + elif "_msgIndex2" in df.columns: + comp_num = 2 + elif "_msgIndex1" in df.columns: + comp_num = 1 + else: + return None + + # print(f"comp_num={comp_num}") + idx = [[] for k in range(comp_num)] + comp_lst = list(range(comp_num)) + cnt = 0 + for row in df.itertuples(): + # print(f"{row}") + if max_count == -1 or len(fs) < max_count: + for comp in comp_lst: + fs.append(self.fs[row[-1 - (comp_num - comp - 1)]]) + idx[comp].append(len(fs) - 1) + cnt += 1 + else: + break + + # generate a new dataframe + if max_count == -1 or cnt == df.shape[0]: + df = df.copy() + else: + df = df.head(cnt).copy() + + for k, v in enumerate(idx): + df[f"_msgIndex{k+1}"] = v + return df + + def _extract_scalar_fields(self, df): + if df.empty: + return None, None + + assert "_msgIndex1" in df.columns + assert "_msgIndex2" not in df.columns + assert "_msgIndex3" not in df.columns + + fs = self.fieldset_class() + for row in df.itertuples(): + fs.append(self.fs[row[-1]]) + + assert len(fs) == len(df.index) + # generate a new dataframe + df = df.copy() + df["_msgIndex1"] = list(range(len(df.index))) + return df, fs + + def _clone(self): + db = FieldsetDb(self.name, label=self.label, regrid_from=self.regrid_from) + + if self._indexer is not None: + db.indexer.update_keys(self._indexer.keys_ecc) + db.blocks = {k: v.copy() for k, v in self.blocks.items()} + db.vector_loaded = self.vector_loaded + return db + + @staticmethod + def make_param_info(fs): + if fs._db is not None: + return fs._db._make_param_info() + else: + return ParamInfo.build_from_fieldset(fs) + + def ls(self, extra_keys=None, filter=None, no_print=False): + default_keys = [ + "centre", + "shortName", + "typeOfLevel", + "level", + "dataDate", + "dataTime", + "stepRange", + "dataType", + "number", + "gridType", + ] + ls_keys = default_keys + extra_keys = [] if extra_keys is None else extra_keys + if extra_keys is not None: + [ls_keys.append(x) for x in extra_keys if x not in ls_keys] + keys = list(ls_keys) + + # add keys appearing in the filter to the full list of keys + dims = {} if filter is None else filter + dims = self._make_dims(dims) + [keys.append(k) for k, v in dims.items() if k not in keys] + + # get metadata + self.load(keys=keys, vector=False) + + # performs filter + df = self._get_meta(dims) + + # extract results + keys = list(ls_keys) + keys.append("_msgIndex1") + df = df[keys] + df = df.sort_values("_msgIndex1") + df = df.rename(columns={"_msgIndex1": "Message"}) + df = df.set_index("Message") + + # only show the column for number in the default set of keys if + # there are any valid values in it + if "number" not in extra_keys: + r = df["number"].unique() + skip = False + if len(r) == 1: + skip = r[0] in ["0", None] + if skip: + df.drop("number", axis=1, inplace=True) + + init_pandas_options() + + # test whether we're in the Jupyter environment + if is_ipython_active(): + return df + elif not no_print: + print(df) + return df + + def sort(self, *args, **kwargs): + # handle arguments + keys = [] + asc = None + if len(args) >= 1: + keys = args[0] + if not isinstance(keys, list): + keys = [keys] + # optional positional argument - we only implement it to provide + # backward compability for the sort() Macro function + if len(args) == 2: + asc = args[1] + if isinstance(asc, list): + if len(keys) != len(asc): + raise ValueError( + f"sort(): when order is specified as a list it must have the same number of elements as keys! {len(keys)} != {len(asc)}" + ) + for i, v in enumerate(asc): + if v not in [">", "<"]: + raise ValueError( + f"sort(): invalid value={v} in order! Only " + > " and " + < " are allowed!" + ) + asc[i] = True if v == "<" else False + else: + if asc not in [">", "<"]: + raise ValueError( + f"sort(): invalid value={asc} in order! Only " + > " and " + < " are allowed!" + ) + asc = True if asc == "<" else False + + if "ascending" in kwargs: + if asc is not None: + raise ValueError( + "sort(): cannot take both a second positional argument and the ascending keyword argument!" + ) + asc = kwargs.pop("ascending") + + if asc is None: + asc = True + + if len(keys) == 0: + keys = self.indexer.DEFAULT_SORT_KEYS + + # print(f"keys={keys} asc={asc}") + + # get metadata + self.load(keys=keys, vector=False) + + scalar_df = self.blocks.get("scalar") + if scalar_df is not None: + dfs = self.indexer._sort_dataframe(scalar_df, columns=keys, ascending=asc) + # print(f"dfs={dfs.iloc[0:5]}") + # print(dfs) + + df, res = self._extract_scalar_fields(dfs) + # print(f"df={df.iloc[0:5]}") + + # LOG.debug(f"len_res={len(res)}") + # LOG.debug(f"dfs={dfs}") + # LOG.debug(f"res={res}") + c = FieldsetDb( + res, + name=self.name, + blocks={"scalar": df}, + label=self.label, + mapped_params=self.mapped_params, + regrid_from=self.regrid_from, + ) + + res._db = c + return res + + def get_longname_and_units(self, short_name, param_id): + # The name and units keys are not included in the default set of keys for the + # indexer. When we need them (primarily in ParamDesc) we simply get the first + # grib message and extract them from it. + a = {} + if short_name: + a["shortName"] = short_name + if param_id: + a["paramId"] = param_id + if a: + a["_max_count"] = 1 + r = self.select(**a) + if r is not None and len(r) > 0: + md = r[0].grib_get(["name", "units"]) + if md and len(md[0]) == 2: + return md[0][0], md[0][1] + return "", "" diff --git a/metview/metviewpy/indexer.py b/metview/metviewpy/indexer.py new file mode 100644 index 00000000..e94ed72a --- /dev/null +++ b/metview/metviewpy/indexer.py @@ -0,0 +1,655 @@ +# +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import copy +import datetime +import logging +import os +from pathlib import Path + +from . import utils +import numpy as np +import pandas as pd +import yaml + +# logging.basicConfig(level=logging.DEBUG) +# logging.basicConfig(level=logging.INFO) +LOG = logging.getLogger(__name__) + +NEWER = True + + +class GribIndexer: + VECTOR_PARAMS = { + "wind10m": ["10u", "10v"], + "wind100m": ["100u", "100v"], + "wind200m": ["200u", "200v"], + "wind": ["u", "v"], + "wind3d": ["u", "v", "w"], + } + + # tuple-> 0: ecCodes type, 1: pandas type, 2: Python type 3: use in duplicate check + DEFAULT_KEYS = { + "shortName": ("s", str, str, False), + "paramId": ("l", "Int32", int, False), + "date": ("l", "Int64", int, True), + "time": ("l", "Int64", int, True), + "step": ("l", "Int32", int, True), + "level": ("l", "Int32", int, True), + "typeOfLevel": ("s", str, str, False), + "number": ("s", str, str, True), + "experimentVersionNumber": ("s", str, str, False), + "marsClass": ("s", str, str, False), + "marsStream": ("s", str, str, False), + "marsType": ("s", str, str, False), + } + + DEFAULT_ECC_KEYS = [f"{k}:{v[0]}" for k, v in DEFAULT_KEYS.items()] + BLOCK_KEYS = ["shortName", "typeOfLevel"] + + DEFAULT_SORT_KEYS = [ + "date", + "time", + "step", + "number", + "level", + "paramId", + ] + DATE_KEYS = { + k: ("l", "Int64", int) + for k in ["date", "dataDate", "validityDate", "mars.date", "marsDate"] + } + TIME_KEYS = { + k: ("l", "Int64", int) + for k in ["time", "dataTime", "validityTime", "mars.time", "marsTime"] + } + + DATETIME_KEYS = { + "_dateTime": ("date", "time"), + "_dataDateTime": ("dataDate", "dataTime"), + "_validityDateTime": ("validityDate", "validityTime"), + } + + KEYS_TO_REPLACE = { + ("type", "mars.type"): "marsType", + ("stream", "mars.stream"): "marsStream", + ("class", "mars.class", "class_"): "marsClass", + ("perturbationNumber"): "number", + ("mars.date", "marsDate"): "date", + ("mars.time", "marsTime"): "time", + } + + PREDEF_KEYS = copy.deepcopy(DEFAULT_KEYS) + PREDEF_KEYS.update(DATE_KEYS) + PREDEF_KEYS.update(TIME_KEYS) + + PREDEF_PD_TYPES = {k: v[1] for k, v in PREDEF_KEYS.items()} + PREDEF_PT_TYPES = {k: v[2] for k, v in PREDEF_KEYS.items()} + + def __init__(self, db): + self.db = db + assert self.db is not None + self.ref_column_count = None + self.keys = [] + self.keys_ecc = [] + for k, v in GribIndexer.DEFAULT_KEYS.items(): + name = k + self.keys.append(name) + if v[0]: + name = f"{k}:{v[0]}" + self.keys_ecc.append(name) + + self.keys_duplicate_check = [ + k for k, v in GribIndexer.DEFAULT_KEYS.items() if v[3] == True + ] + + self.shortname_index = self.keys.index("shortName") + self.levtype_index = self.keys.index("typeOfLevel") + self.type_index = self.keys.index("marsType") + self.number_index = self.keys.index("number") + self.param_id_index = self.keys.index("paramId") + + self.wind_check_index = [] + for v in [ + "date", + "time", + "step", + "level", + "typeOfLevel", + "level", + "number", + "experimentVersionNumber", + "marsClass", + "marsStream", + "marsType", + ]: + self.wind_check_index.append(self.keys.index(v) + 1) + + # self.block_key_index = [self.keys.index(v) for v in GribIndexer.BLOCK_KEYS] + + self.pd_types = {k: v[1] for k, v in GribIndexer.DEFAULT_KEYS.items()} + self.pt_types = {k: v[2] for k, v in GribIndexer.DEFAULT_KEYS.items()} + + def update_keys(self, keys): + ret = False + for k in keys: + name = k + # we do not add datetime keys (they are pseudo keys, and their value + # is always generated on the fly) + if name not in self.keys and name not in GribIndexer.DATETIME_KEYS: + self.keys.append(name) + p = GribIndexer.PREDEF_KEYS.get(name, ("", str, str)) + ecc_name = name if p[0] == "" else name + ":" + p[0] + self.keys_ecc.append(ecc_name) + self.pd_types[name] = p[1] + self.pt_types[name] = p[2] + ret = True + return ret + + def _check_duplicates(self, name, df): + dup = df.duplicated(subset=self.keys_duplicate_check) + first_dup = True + cnt = 0 + for i, v in dup.items(): + if v: + if first_dup: + LOG.error( + f"{name}: has duplicates for key group: {self.keys_duplicate_check}!" + ) + first_dup = False + LOG.error(f" first duplicate: {df.iloc[i]}") + cnt += 1 + + if cnt > 1: + LOG.error(f" + {cnt-1} more duplicate(s)!") + + def _build_vector_index(self, df, v_name, v_comp): + # LOG.debug(f"v_name={v_name} v_comp={v_comp}") + comp_num = len(v_comp) + + # filter components belonging together + comp_df = [] + for i, comp_name in enumerate(v_comp): + query = f"shortName == '{comp_name}'" + r = df.query(query, engine="python") + # if we do not use copy, the assignment below as: + # comp_df[0].loc[... + # generates the SettingWithCopyWarning warning!!! + if i == 0: + r = df.query(query, engine="python").copy() + else: + r = df.query(query, engine="python") + if r.empty: + return [] + else: + comp_df.append(r) + + assert comp_num == len(comp_df) + + # pair up components within a 2D vector field. This + # version proved to be the fastest! + # LOG.debug(" pair up collected components:") + + # print(f"v_name={v_name} {len(comp_df[1].index)}") + # print("view=", comp_df[0]._is_view) + r = [] + used1 = np.full(len(comp_df[1].index), False, dtype="?") + comp_df[0].loc[:, "shortName"] = v_name + # 2D + if comp_num == 2: + for row0 in comp_df[0].itertuples(name=None): + i = 0 + for row1 in comp_df[1].itertuples(name=None): + if not used1[i]: + b = True + for x in self.wind_check_index: + if row0[x] != row1[x]: + b = False + break + if b: + d = list(row0[1:]) + d.extend(row1[-self.ref_column_count :]) + r.append(d) + used1[i] = True + break + i += 1 + # 3D + elif comp_num == 3: + used2 = np.full(len(comp_df[2].index), False, dtype="?") + for row0 in comp_df[0].itertuples(name=None): + i = 0 + for row1 in comp_df[1].itertuples(name=None): + if not used1[i]: + b = True + for x in self.wind_check_index: + if row0[x] != row1[x]: + b = False + break + if b: + j = 0 + for row2 in comp_df[2].itertuples(name=None): + if not used2[j]: + b = True + for x in self.wind_check_index: + if row0[x] != row2[x]: + b = False + break + if b: + d = list(row0[1:]) + d.extend(row1[-self.ref_column_count :]) + d.extend(row2[-self.ref_column_count :]) + r.append(d) + used1[i] = True + used2[j] = True + j = -1 + break + j += 1 + if j == -1: + break + i += 1 + return r + + def _make_dataframe(self, data, sort=False, columns=None): + if columns is not None: + df = pd.DataFrame(data, columns=columns) + else: + df = pd.DataFrame(data) + + for c in df.columns: + if self.pd_types.get(c, "") in ["Int32", "Int64"]: + df.fillna(value={c: np.nan}, inplace=True) + df = df.astype(self.pd_types) + if sort: + df = GribIndexer._sort_dataframe(df) + + return df + + @staticmethod + def _sort_dataframe(df, columns=None, ascending=True): + if columns is None: + columns = list(df.columns) + elif not isinstance(columns, list): + columns = [columns] + + # mergesoft is a stable sorting algorithm + df = df.sort_values(by=columns, ascending=ascending, kind="mergesort") + df = df.reset_index(drop=True) + return df + + def _write_dataframe(self, df, name, out_dir): + f_name = os.path.join(out_dir, f"{name}.csv.gz") + df.to_csv(path_or_buf=f_name, header=True, index=False, compression="gzip") + + @staticmethod + def read_dataframe(key, dir_name): + # assert len(key) == len(GribIndexer.BLOCK_KEYS) + name = key + f_name = os.path.join(dir_name, f"{name}.csv.gz") + # LOG.debug("f_name={}".format(f_name)) + return pd.read_csv(f_name, index_col=None, dtype=GribIndexer.PREDEF_PD_TYPES) + + @staticmethod + def get_storage_key_list(dir_name): + r = [] + # LOG.debug(f"dir_name={dir_name}") + suffix = ".csv.gz" + for f in utils.get_file_list(os.path.join(dir_name, f"*{suffix}")): + name = os.path.basename(f) + # LOG.debug(f"name={name}") + r.append(name[: -len(suffix)]) + return r + + @staticmethod + def is_key_wind(key): + return key in GribIndexer.VECTOR_PARAMS + + @staticmethod + def _convert_query_value(v, col_type): + # print(f"v={v} {type(v)} {col_type}") + return v if col_type != "object" else str(v) + + @staticmethod + def _check_datetime_in_filter_input(keys): + for k, v in GribIndexer.DATETIME_KEYS.items(): + name = k[1:] + name_date = v[0] + name_time = v[1] + if keys.get(name, []) and ( + keys.get(name_date, []) or keys.get(name_time, []) + ): + raise Exception( + f"Cannot specify {name} together with {name_date} and {name_time}!" + ) + + @staticmethod + def _convert_filter_value(name, val): + """ + Analyse the filter key-value pairs and perform the necessary conversions + """ + valid_name = name.split(":")[0] if ":" in name else name + + # datetime keys are pseudo keys, they start with _. Their value is converted to + # datetime. The key itself is not added to the scan! + if ("_" + valid_name) in GribIndexer.DATETIME_KEYS: + valid_name = "_" + valid_name + name_date = GribIndexer.DATETIME_KEYS[valid_name][0] + name_time = GribIndexer.DATETIME_KEYS[valid_name][1] + for i, t in enumerate(val): + val[i] = GribIndexer._to_datetime(name, t) + # print(f"t={t} -> {val[i]}") + # We add the date and time components with an empty value. So they will be + # added to the scan, but they will be ignored by the query. Conversely, + # the datetime key itself will be ignored in the scan, but will be used + # in the query. + return [("_" + name, val), (name_date, []), (name_time, [])] + # we convert dates to int + elif valid_name in GribIndexer.DATE_KEYS: + for i, t in enumerate(val): + d = GribIndexer._to_date(name, t) + # for daily climatologies dates where the year is missing the + # the a tuple is returned + if not isinstance(d, tuple): + val[i] = int(d.strftime("%Y%m%d")) + else: + val[i] = d[0] * 100 + d[1] + # we convert times to int + elif valid_name in GribIndexer.TIME_KEYS: + for i, t in enumerate(val): + val[i] = int(GribIndexer._to_time(name, t).strftime("%H%M")) + # print(f"t={t} -> {val[i]}") + else: + pt_type = GribIndexer.PREDEF_PT_TYPES.get(name, None) + # print(f"name={name} {pt_type}") + if pt_type is not None: + for i, t in enumerate(val): + val[i] = pt_type(t) + # print(f" t={t} -> {val[i]}") + + # remap some names to the ones already in the default set of indexer keys + for k, v in GribIndexer.KEYS_TO_REPLACE.items(): + if name in k: + name = v + + return [(name, val)] + + @staticmethod + def _to_datetime(param, val): + try: + if isinstance(val, datetime.datetime): + return val + elif isinstance(val, str): + return utils.date_from_str(val) + elif isinstance(val, (int, float)): + return utils.date_from_str(str(val)) + else: + raise + except: + raise Exception(f"Invalid datetime value={val} specified for key={param}") + + @staticmethod + def _to_date(param, val): + try: + if isinstance(val, datetime.datetime): + return val.date() + elif isinstance(val, datetime.date): + return val + elif isinstance(val, str): + d = utils.date_from_str(val) + return d.date() if not isinstance(d, tuple) else d + elif isinstance(val, (int, float)): + d = utils.date_from_str(str(val)) + return d.date() if not isinstance(d, tuple) else d + else: + raise + except: + raise Exception(f"Invalid date value={val} specified for key={param}") + + @staticmethod + def _to_time(param, val): + try: + if isinstance(val, (datetime.datetime)): + return val.time() + elif isinstance(val, datetime.time): + return val + elif isinstance(val, str): + return utils.time_from_str(val) + elif isinstance(val, int): + return utils.time_from_str(str(val)) + else: + raise + except: + raise Exception(f"Invalid time value={val} specified for key={param}") + + +class FieldsetIndexer(GribIndexer): + def __init__(self, *args): + super().__init__(*args) + self.ref_column_count = 1 + + def scan(self, vector=False): + data = self._scan(self.db.fs, mapped_params=self.db.mapped_params) + if data: + df = self._make_dataframe(data, sort=False) + self.db.blocks["scalar"] = df + if vector: + self._scan_vector() + + def _scan(self, fs, mapped_params={}): + LOG.info(f" scan fields ...") + data = {} + # print(f"fs_len={len(fs)}") + # print(f"keys_ecc={self.keys_ecc}") + if utils.is_fieldset_type(fs) and len(fs) > 0: + md_vals = fs.grib_get(self.keys_ecc, "key") + if mapped_params: + for i in range(len(fs)): + v = md_vals[self.param_id_index][i] + if v in mapped_params: + short_name = mapped_params[v] + md_vals[self.shortname_index][i] = short_name + + assert len(self.keys) == len(self.keys_ecc) + data = {k: md_vals[i] for i, k in enumerate(self.keys)} + data["_msgIndex1"] = list(range(len(fs))) + LOG.info(f" {len(fs)} GRIB messages processed") + return data + + def _scan_vector(self): + df = self.db.blocks["scalar"] + if df is not None and not df.empty: + for v_name, v_comp in GribIndexer.VECTOR_PARAMS.items(): + r = self._build_vector_index(df, v_name, v_comp) + comp_num = len(v_comp) + if r: + cols = [*self.keys] + for i in range(comp_num): + cols.extend([f"_msgIndex{i+1}"]) + w_df = self._make_dataframe(r, sort=False, columns=cols) + self.db.blocks[v_name] = w_df + # self._write_dataframe(w_df, v_name, out_dir) + else: + LOG.debug(" No paired fields found!") + continue + + +class ExperimentIndexer(GribIndexer): + def __init__(self, *args): + super().__init__(*args) + self.ref_column_count = 2 + + def scan(self): + out_dir = self.db.db_dir + Path(out_dir).mkdir(exist_ok=True, parents=True) + LOG.info(f"scan {self.db} out_dir={out_dir} ...") + + data = {k: [] for k in [*self.keys, "_msgIndex1", "_fileIndex1"]} + input_files = [] + + # print(f"out_dir={out_dir}") + # merge existing experiment objects + if self.db.merge_conf: + ds = [] + # simple merge + if isinstance(self.db.merge_conf, list): + for c_name in self.db.merge_conf: + ds.append( + {"data": db.dataset.find(c_name), "name": c_name, "ens": {}} + ) + # explicit ENS merge + else: + assert "pf" in self.db.merge_conf + # control forecast + c_name = self.db.merge_conf.get("cf", "") + if c_name != "": + ds.append( + { + "data": self.db.dataset.find(c_name, comp="field"), + "name": c_name, + "ens": {"type": "cf", "number": 0}, + } + ) + for i, c_name in enumerate(self.db.merge_conf.get("pf", [])): + ds.append( + { + "data": self.db.dataset.find(c_name, comp="field"), + "name": c_name, + "ens": {"type": "pf", "number": i + 1}, + } + ) + + for c in ds: + if c["data"] is None: + c_name = d["name"] + raise Exception( + f"Cannot merge experiments as {self.db}! Experiment {c_name} is not found!" + ) + else: + input_files = self._scan_one( + input_dir=c["data"].path, + file_name_pattern=c["data"].file_name_pattern, + input_files=input_files, + mapped_params=self.db.mapped_params, + ens=c["ens"], + data=data, + rootdir_placeholder_value=c["data"].rootdir_placeholder_value, + rootdir_placeholder_token=self.db.ROOTDIR_PLACEHOLDER_TOKEN, + ) + # index a single experiment + else: + input_files = self._scan_one( + input_dir=self.db.path, + file_name_pattern=self.db.file_name_pattern, + input_files=[], + mapped_params=self.db.mapped_params, + ens={}, + data=data, + rootdir_placeholder_value=self.db.rootdir_placeholder_value, + rootdir_placeholder_token=self.db.ROOTDIR_PLACEHOLDER_TOKEN, + ) + + # print(f"input_files={input_files}") + if len(input_files) > 0 and len(data["shortName"]) > 0: + # write config file for input file list + LOG.info(f"generate datafiles.yaml ...") + f_name = os.path.join(out_dir, "datafiles.yaml") + r = yaml.dump(input_files, default_flow_style=False) + with open(f_name, "w") as f: + f.write(r) + self.db.input_files = input_files + + # scalar + LOG.info(f"generate scalar fields index ...") + df = self._make_dataframe(data, sort=True) + self.db.blocks["scalar"] = df + self._write_dataframe(df, "scalar", out_dir) + + # vector (2D) + LOG.info(f"generate vector fields index ...") + for v_name, v_comp in GribIndexer.VECTOR_PARAMS.items(): + r = self._build_vector_index(df, v_name, v_comp) + comp_num = len(v_comp) + if r: + cols = [*self.keys] + for i in range(comp_num): + cols.extend([f"_msgIndex{i+1}", f"_fileIndex{i+1}"]) + w_df = self._make_dataframe(r, sort=True, columns=cols) + # print(f"wind_len={len(w_df.index)}") + self.db.blocks[v_name] = w_df + self._write_dataframe(w_df, v_name, out_dir) + else: + LOG.debug(" No paired fields found!") + continue + + def _scan_one( + self, + input_dir="", + file_name_pattern="", + input_files=[], + mapped_params={}, + ens={}, + data={}, + rootdir_placeholder_value="", + rootdir_placeholder_token=None, + ): + LOG.info("scan fields ...") + LOG.info(f" input_dir={input_dir} file_name_pattern={file_name_pattern}") + # print(f" input_dir={input_dir} file_name_pattern={file_name_pattern}") + + # for f_path in glob.glob(f_pattern): + cnt = 0 + input_files_tmp = [] + for f_path in utils.get_file_list( + input_dir, file_name_pattern=file_name_pattern + ): + # LOG.debug(f" f_path={f_path}") + fs = self.db.fieldset_class(path=f_path) + if utils.is_fieldset_type(fs) and len(fs) > 0: + cnt += 1 + input_files_tmp.append(f_path) + file_index = len(input_files) + len(input_files_tmp) - 1 + md_vals = fs.grib_get(self.keys_ecc, "key") + + if mapped_params: + for i in range(len(fs)): + v = md_vals[self.param_id_index][i] + if v in mapped_params: + short_name = mapped_params[v] + md_vals[self.shortname_index][i] = short_name + if ens: + for i in range(len(fs)): + md_vals[self.type_index][i] = ens["type"] + md_vals[self.number_index][i] = ens["number"] + + assert len(self.keys) == len(self.keys_ecc) + for i, c in enumerate(self.keys): + data[c].extend(md_vals[i]) + data["_msgIndex1"].extend(list(range(len(fs)))) + data["_fileIndex1"].extend([file_index] * len(fs)) + + # print({k: len(v) for k, v in data.items()}) + + if rootdir_placeholder_value: + input_files_tmp = [ + x.replace(rootdir_placeholder_value, rootdir_placeholder_token) + for x in input_files_tmp + ] + + input_files.extend(input_files_tmp) + + LOG.info(f" {cnt} GRIB files processed") + return input_files + + def allowed_keys(self): + r = list(self.keys) + r.extend(GribIndexer.DATE_KEYS) + r.extend(GribIndexer.TIME_KEYS) + r.extend(list(GribIndexer.DATETIME_KEYS.keys())) + return set(r) diff --git a/metview/metviewpy/ipython.py b/metview/metviewpy/ipython.py new file mode 100644 index 00000000..feb360ea --- /dev/null +++ b/metview/metviewpy/ipython.py @@ -0,0 +1,42 @@ +# +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +""" +ipython is not None when running a notebook +""" + +import logging +import sys + + +ipython_active = None + + +def is_ipython_active(): + global ipython_active + if ipython_active is None: + try: + from IPython import get_ipython + + ipython_active = get_ipython() is not None + except Exception: + ipython_active = False + return ipython_active + + +def import_widgets(): + try: + widgets = __import__("ipywidgets", globals(), locals()) + return widgets + except ImportError as imperr: + print("Could not import ipywidgets module - animation widget will not appear") + print("Call setoutput('jupyter', plot_widget=False) to suppress this message") + return None diff --git a/metview/metviewpy/maths.py b/metview/metviewpy/maths.py new file mode 100644 index 00000000..ba46c392 --- /dev/null +++ b/metview/metviewpy/maths.py @@ -0,0 +1,168 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +import numpy as np + + +def neg(x): + return -x + + +def pos(x): + return x + + +# def abs(x): +# return np.abs(x) + + +def not_func(x): + return (~(x != 0)).astype(int) + + +def add(x, y): + return x + y + + +def sub(x, y): + return x - y + + +def mul(x, y): + return x * y + + +def div(x, y): + return x / y + + +def pow(x, y): + return x**y + + +def ge(x, y): + return (x >= y).astype(int) + + +def gt(x, y): + return (x > y).astype(int) + + +def le(x, y): + return (x <= y).astype(int) + + +def lt(x, y): + return (x < y).astype(int) + + +def eq(x, y): + return (x == y).astype(int) + + +def ne(x, y): + return (x != y).astype(int) + + +def and_func(x, y): + return ne(x, 0) * ne(y, 0) + + +def or_func(x, y): + return np.clip(ne(x, 0) + ne(y, 0), 0, 1) + + +def set_from_other(x, y): + return y + + +# single argument functions + + +def abs(x): + return np.fabs(x) + + +def acos(x): + return np.arccos(x) + + +def asin(x): + return np.arcsin(x) + + +def atan(x): + return np.arctan(x) + + +def cos(x): + return np.cos(x) + + +def exp(x): + return np.exp(x) + + +def log(x): + return np.log(x) + + +def log10(x): + return np.log10(x) + + +def sgn(x): + return np.sign(x) + + +def square(x): + return np.square(x) + + +def sqrt(x): + return np.sqrt(x) + + +def sin(x): + return np.sin(x) + + +def tan(x): + return np.tan(x) + + +# double argument functions + + +def atan2(x, y): + return np.arctan2(x, y) + + +def floor_div(x, y): + return np.floor_divide(x, y) + + +def mod(x, y): + return np.mod(x, y) + + +# bitmapping + + +def bitmap(x, y): + if isinstance(y, (int, float)): + x[x == y] = np.nan + return x + elif isinstance(y, np.ndarray): + x[np.isnan(y)] = np.nan + return x + + +def nobitmap(x, y): + x[np.isnan(x)] = y + return x diff --git a/metview/metviewpy/param.py b/metview/metviewpy/param.py new file mode 100644 index 00000000..7ba5033b --- /dev/null +++ b/metview/metviewpy/param.py @@ -0,0 +1,582 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import logging +import re +import pandas as pd + +from .indexer import GribIndexer +from .ipython import is_ipython_active +from .utils import is_fieldset_type + +# logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") +# logging.basicConfig(level=logging.DEBUG, format="%(levelname)s - %(message)s") +LOG = logging.getLogger(__name__) + +PANDAS_ORI_OPTIONS = {} + + +def init_pandas_options(): + global PANDAS_ORI_OPTIONS + if len(PANDAS_ORI_OPTIONS) == 0: + opt = { + "display.max_colwidth": 300, + "display.colheader_justify": "center", + "display.max_columns": 100, + "display.max_rows": 500, + "display.width": None, + } + for k, _ in opt.items(): + PANDAS_ORI_OPTIONS[k] = pd.get_option(k) + for k, v in opt.items(): + pd.set_option(k, v) + + +def reset_pandas_options(): + global PANDAS_ORI_OPTIONS + if len(PANDAS_ORI_OPTIONS) > 0: + for k, v in PANDAS_ORI_OPTIONS.items(): + pd.set_option(k, v) + PANDAS_ORI_OPTIONS = {} + + +class ParamInfo: + SUFFIXES = { + "hPa": "isobaricInhPa", + "hpa": "isobaricInhPa", + "K": "theta", + "ml": "hybrid", + } + LEVEL_TYPES = {"sfc": "surface", "pl": "isobaricInhPa", "ml": "hybrid"} + LEVEL_RE = re.compile(r"(\d+)") + NUM_RE = re.compile(r"[0-9]+") + SURF_RE = re.compile(r"^\d+\w+") + # SURF_NAME_MAPPER = {"t2": "2t", "q2": "2q", "u10": "10u", "v10": "10v"} + KNOWN_SURF_NAMES = [ + "2t", + "2q", + "10u", + "10v", + "100u", + "100v", + "200u", + "200v", + "msl", + "wind10m", + "wind100m", + "wind200m", + ] + VECTOR_NAMES = [ + "wind100m", + "wind200m", + "wind10m", + "wind3d", + "wind", + ] # the longest ones first + + def __init__(self, name, meta=None, scalar=None): + self.name = name + self.scalar = scalar if scalar is not None else True + self.meta = {} if meta is None else meta + if len(self.meta) == 0: + self.meta["shortName"] = name + + def make_filter(self): + dims = {} + if self.name: + dims["shortName"] = [self.name] + for n in ["level", "typeOfLevel"]: + v = self.meta.get(n, None) + if v is not None: + dims[n] = [v] + return dims + + @staticmethod + def build_from_name(full_name, param_level_types=None): + full_name = full_name + name = full_name + level = None + level_type = "" + + # the name is a known param name + if param_level_types: + if name in param_level_types: + lev_t = param_level_types.get(name, []) + meta = {} + if len(lev_t) == 1: + meta = {"typeOfLevel": lev_t[0], "level": None} + scalar = not name in ParamInfo.VECTOR_NAMES + return ParamInfo(name, meta=meta, scalar=scalar) + + t = full_name + # surface fields + if t in ParamInfo.KNOWN_SURF_NAMES or ParamInfo.SURF_RE.match(t) is not None: + level_type = "surface" + + else: + # guess the level type from the suffix + for k, v in ParamInfo.SUFFIXES.items(): + if full_name.endswith(k): + level_type = v + t = full_name[: -(len(k))] + break + + # recognise vector params + for v in ParamInfo.VECTOR_NAMES: + if t.startswith(v): + name = v + t = t[len(v) :] + break + + # determine level value + m = ParamInfo.LEVEL_RE.search(t) + if m and m.groups() and len(m.groups()) == 1: + level = int(m.group(1)) + if level_type == "" and level > 10: + level_type = "isobaricInhPa" + if name == full_name: + name = ParamInfo.NUM_RE.sub("", t) + + # check param name in the conf + if param_level_types: + if not name in param_level_types: + raise Exception( + f"Param={name} (guessed from name={full_name}) is not found in dataset!" + ) + + lev_t = param_level_types.get(name, []) + if lev_t: + if not level_type and len(lev_t) == 1: + level_type = lev_t[0] + elif level_type and level_type not in lev_t: + raise Exception( + f"Level type cannot be guessed from param name={full_name}!" + ) + + if level_type == "": + level = None + scalar = not name in ParamInfo.VECTOR_NAMES + + LOG.debug(f"scalar={scalar}") + meta = {"level": level, "typeOfLevel": level_type} + return ParamInfo(name, meta=meta, scalar=scalar) + + @staticmethod + def build_from_fieldset(fs): + assert is_fieldset_type(fs) + f = fs[0:3] if len(fs) >= 3 else fs + m = ParamInfo._grib_get(f, GribIndexer.DEFAULT_ECC_KEYS) + name = level = lev_type = "" + scalar = True + + meta_same = True + for x in m.keys(): + if x != "shortName" and m[x].count(m[x][0]) != len(m[x]): + same = False + break + if meta_same: + if len(m["shortName"]) == 3 and m["shortName"] == ["u", "v", "w"]: + name = "wind3d" + scalar = False + elif len(m["shortName"]) >= 2: + if m["shortName"][0:2] == ["u", "v"]: + name = "wind" + scalar = False + elif m["shortName"][0:2] == ["10u", "10v"]: + name = "wind10m" + m["level"][0] = 0 + m["typeOfLevel"][0] = "sfc" + scalar = False + if not name: + name = m["shortName"][0] + + if name: + return ParamInfo(name, meta={k: v[0] for k, v in m.items()}, scalar=scalar) + else: + return None + + def _meta_match(self, meta, key): + local_key = key if key != "levelist" else "level" + if ( + key in meta + and meta[key] is not None + and meta[key] + and local_key in self.meta + ): + # print(f"local={self.meta[local_key]} other={meta[key]}") + if isinstance(meta[key], list): + return str(self.meta[local_key]) in meta[key] + else: + return meta[key] == str(self.meta[local_key]) + else: + return False + + def match(self, name, meta): + # print(f"{self}, name={name}, meta={meta}") + r = 0 + if self.name == name: + r += 3 + for n in ["shortName", "paramId"]: + if self._meta_match(meta, n): + r += 1 + # we only check the rest if the param is ok + if r > 0: + if self._meta_match(meta, "typeOfLevel"): + r += 1 + if self._meta_match(meta, "levelist"): + r += 1 + + return r + + def update_meta(self, meta): + self.meta = {**meta, **self.meta} + + @staticmethod + def _grib_get(f, keys, single_value_as_list=True): + md = f.grib_get(keys, "key") + m = {} + for k, v in zip(keys, md): + key_val = k.split(":")[0] + val = v + if k.endswith(":l"): + val = [] + for x in v: + try: + val.append(int(x)) + except: + val.append(None) + if not single_value_as_list and len(val) == 1: + val = val[0] + m[key_val] = val + return m + + def __str__(self): + return "{}[name={}, scalar={}, meta={}]".format( + self.__class__.__name__, self.name, self.scalar, self.meta + ) + + +class ParamDesc: + def __init__(self, name): + self.db = None + # self.name = name + self.md = {} + self.levels = {} + self._short_name = None + self._param_id = None + self._long_name = None + self._units = None + + def load(self, db): + raise NotImplementedError + + def _parse(self, md): + if "level" in md and len(md["level"]) > 0: + df = pd.DataFrame(md) + md.pop("typeOfLevel") + md.pop("level") + + for md_key in list(md.keys()): + d = list(df[md_key].unique()) + self.md[md_key] = d + + lev_types = list(df["typeOfLevel"].unique()) + for t in lev_types: + # print(f" t={t}") + self.levels[t] = [] + q = f"typeOfLevel == '{t}'" + # print(q) + dft = df.query(q) + if dft is not None: + self.levels[t] = list(dft["level"].unique()) + + for k, v in self.md.items(): + self.md[k] = sorted(v) + + for k, v in self.levels.items(): + self.levels[k] = sorted(v) + + @property + def short_name(self): + if self._short_name is None: + self._short_name = "" + if self.md["shortName"]: + self._short_name = self.md["shortName"][0] + return self._short_name + + @property + def param_id(self): + if self._param_id is None: + self._param_id = "" + if self.md["paramId"]: + self._param_id = self.md["paramId"][0] + return self._param_id + + @property + def long_name(self): + if self._long_name is None: + self._long_name = "" + if self.db is not None: + self._long_name, self._units = self.db.get_longname_and_units( + self.short_name, self.param_id + ) + return self._long_name + + @property + def units(self): + if self._units is None: + self._units = "" + if self.db: + self._long_name, self._units = self.db.get_longname_and_units( + self.short_name, self.param_id + ) + return self._units + + @staticmethod + def describe(db, param=None, no_print=False): + labels = {"marsClass": "class", "marsStream": "stream", "marsType": "type"} + in_jupyter = is_ipython_active() + + # describe all the params + if param is None: + t = {"parameter": [], "typeOfLevel": [], "level": []} + need_number = False + for k, v in db.param_meta.items(): + if not v.md.get("number", None) in [["0"], [None]]: + need_number = True + break + + for k, v in db.param_meta.items(): + t["parameter"].append(k) + if len(v.levels) > 1: + lev_type = "" + level = "" + cnt = 0 + for md_k, md in v.levels.items(): + if in_jupyter: + lev_type += md_k + "
" + level += str(ParamDesc.format_list(md)) + "
" + else: + prefix = " " if cnt > 0 else "" + lev_type += prefix + f"[{cnt+1}]:" + md_k + level += ( + prefix + f"[{cnt+1}]:" + str(ParamDesc.format_list(md)) + ) + cnt += 1 + t["typeOfLevel"].append(lev_type) + t["level"].append(level) + else: + for md_k, md in v.levels.items(): + t["typeOfLevel"].append(md_k) + t["level"].append(ParamDesc.format_list(md)) + + for md_k, md in v.md.items(): + if md_k != "number" or need_number: + md_k = labels.get(md_k, md_k) + if not md_k in t: + t[md_k] = [] + t[md_k].append(ParamDesc.format_list(md)) + + if in_jupyter: + txt = ParamDesc._make_html_table(t) + from IPython.display import HTML + + return HTML(txt) + else: + df = pd.DataFrame.from_dict(t) + df = df.set_index(["parameter"]) + init_pandas_options() + if not no_print: + print(df) + return df + + # specific param + else: + v = None + if isinstance(param, str): + v = db.param_meta.get(param, None) + elif isinstance(param, int): + v = db.param_id_meta(param) + + if v is None or len(v.md) == 0: + print(f"No shortName/paramId={param} found in data!") + return + + # if v is not None: + t = { + "key": ["shortName"], + "val": [v.short_name], + } + + if v.long_name != "" or v.units != "": + t["key"].append("name") + t["val"].append(v.long_name) + + t["key"].append("paramId") + t["val"].append(v.param_id) + # ParamDesc.format_list(v.md["shortName"], full=True), + + if v.long_name != "" or v.units != "": + t["key"].append("units") + t["val"].append(v.units) + + add_cnt = len(v.levels) > 1 + cnt = 0 + for md_k, md in v.levels.items(): + t["key"].append("typeOfLevel" + (f"[{cnt+1}]" if add_cnt else "")) + t["val"].append(md_k) + t["key"].append("level" + (f"[{cnt+1}]" if add_cnt else "")) + t["val"].append(ParamDesc.format_list(md, full=True)) + cnt += 1 + + for kk, md_v in v.md.items(): + if kk == "number" and md_v == ["0"]: + continue + if not kk in ["shortName", "paramId"]: + t["key"].append(labels.get(kk, kk)) + t["val"].append(ParamDesc.format_list(md_v, full=True)) + + if in_jupyter: + from IPython.display import HTML + + txt = ParamDesc._make_html_table(t, header=False) + return HTML(txt) + else: + df = pd.DataFrame.from_dict(t) + df = df.set_index("key") + init_pandas_options() + if not no_print: + print(df) + return df + + @staticmethod + def _make_html_table(d, header=True): + if len(d) > 1: + first_column_name = list(d.keys())[0] + txt = """ + + {} + {} +
""".format( + "" if not header else "".join([f"{k}" for k in d.keys()]), + "".join( + [ + "" + + d[first_column_name][i] + + "" + + "".join( + [ + f"{ParamDesc.format_list(d[k][i], full=True)}" + for k in list(d.keys())[1:] + ] + ) + + "" + for i in range(len(d[first_column_name])) + ] + ), + ) + return txt + else: + return "" + + @staticmethod + def format_list(v, full=False): + if isinstance(v, list): + if full is True: + return ",".join([str(x) for x in v]) + else: + if len(v) == 1: + return v[0] + if len(v) > 2: + return ",".join([str(x) for x in [v[0], v[1], "..."]]) + else: + return ",".join([str(x) for x in v]) + else: + return v + + +class ParamNameDesc(ParamDesc): + def __init__(self, name): + super().__init__(name) + self._short_name = name + + def load(self, db): + md = { + "typeOfLevel": [], + "level": [], + "date": [], + "time": [], + "step": [], + "number": [], + "paramId": [], + "marsClass": [], + "marsStream": [], + "marsType": [], + "experimentVersionNumber": [], + } + + self.db = db + self.md = {} + self.levels = {} + + # print(f"par={par}") + for b_name, b_df in db.blocks.items(): + if b_name == "scalar": + q = f"shortName == '{self.short_name}'" + dft = b_df.query(q) + elif b_name == self.short_name: + dft = b_df + else: + dft = None + + if dft is not None: + for k in md.keys(): + # print(f"{self.name}/{k}") + md[k].extend(dft[k].tolist()) + # print(f" df[{k}]={df[k]}") + # print(df) + + self._parse(md) + + +class ParamIdDesc(ParamDesc): + def __init__(self, param_id): + super().__init__("") + self._param_id = param_id + + def load(self, db): + md = { + "shortName": [], + "typeOfLevel": [], + "level": [], + "date": [], + "time": [], + "step": [], + "number": [], + "paramId": [], + "marsClass": [], + "marsStream": [], + "marsType": [], + "experimentVersionNumber": [], + } + + self.db = db + self.md = {} + self.levels = {} + + # print(f"par={par}" + b_df = db.blocks.get("scalar", None) + if b_df is not None: + q = f"paramId == {self._param_id}" + dft = b_df.query(q, engine="python") + if dft is not None: + for k in md.keys(): + md[k].extend(dft[k].tolist()) + self._parse(md) diff --git a/metview/metviewpy/temporary.py b/metview/metviewpy/temporary.py new file mode 100644 index 00000000..130e5f81 --- /dev/null +++ b/metview/metviewpy/temporary.py @@ -0,0 +1,71 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +# code copied from ecmwf/climetlab + +import os +import tempfile + + +class TmpFile: + """The TmpFile objets are designed to be used for temporary files. + It ensures that the file is unlinked when the object is + out-of-scope (with __del__). + Parameters + ---------- + path : str + Actual path of the file. + """ + + def __init__(self, path: str): + self.path = path + + def __del__(self): + self.cleanup() + + # def __enter__(self): + # return self.path + + # def __exit__(self, *args, **kwargs): + # self.cleanup() + + def cleanup(self): + if self.path is not None: + os.unlink(self.path) + self.path = None + + +def temp_file(extension=".tmp"): + """Create a temporary file with the given extension . + Parameters + ---------- + extension : str, optional + By default ".tmp" + Returns + ------- + TmpFile + """ + + fd, path = tempfile.mkstemp(suffix=extension) + os.close(fd) + return TmpFile(path) + + +def is_temp_file(path): + return tempfile.gettempdir() in path + + +# class TmpDirectory(tempfile.TemporaryDirectory): +# @property +# def path(self): +# return self.name + + +# def temp_directory(): +# return TmpDirectory() diff --git a/metview/metviewpy/utils.py b/metview/metviewpy/utils.py new file mode 100644 index 00000000..dbf0a448 --- /dev/null +++ b/metview/metviewpy/utils.py @@ -0,0 +1,350 @@ +# +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import datetime +import getpass +import glob +import logging +import math +from pathlib import Path +import shutil +import os +import re +import tempfile + +import numpy as np + +LOG = logging.getLogger(__name__) + + +CACHE_DIR = os.path.join(tempfile.gettempdir(), f"mpy-{getpass.getuser()}") + + +def deacc(fs, key=None, skip_first=False, mark_derived=False): + r = None + if key is None or key == "": + if len(fs) > 1: + v = fs[1:] - fs[:-1] + if not skip_first: + r = fs[0] * 0 + r = r.merge(v) + else: + r = v + else: + if not isinstance(key, str): + raise TypeError(f"deacc(): key must be a str (got {type(key)})!") + fs._get_db().load([key]) + key_vals = fs._unique_metadata(key) + if key_vals: + v = fs.select({key: key_vals[0]}) + gr_num = len(v) + r = None + if not skip_first: + r = v * 0 + for i in range(1, len(key_vals)): + v_next = fs.select({key: key_vals[i]}) + if len(v_next) != gr_num: + raise ValueError( + f"deacc(): unexpected number of fields (={len(v_next)}) found for {key}={key_vals[i]}! For each {key} value the number of fields must be the same as for {key}={key_vals[0]} (={gr_num})!" + ) + if r is None: + r = v_next - v + else: + # print(f"i={i}") + # v.ls() + # v_next.ls() + r.append(v_next - v) + v = v_next + if not mark_derived: + r = r.grib_set_long(["generatingProcessIdentifier", 148]) + return r + + +def date_from_str(d_str): + # yyyymmdd + if len(d_str) == 8: + return datetime.datetime.strptime(d_str, "%Y%m%d") + # yyyy-mm-dd .... + elif len(d_str) >= 10 and d_str[4] == "-" and d_str[7] == "-": + # yyyy-mm-dd + if len(d_str) == 10: + return datetime.datetime.strptime(d_str, "%Y-%m-%d") + # yyyy-mm-dd hh + elif len(d_str) == 13: + return datetime.datetime.strptime(d_str, "%Y-%m-%d %H") + # yyyy-mm-dd hh:mm + elif len(d_str) == 16: + return datetime.datetime.strptime(d_str, "%Y-%m-%d %H:%M") + # yyyy-mm-dd hh:mm:ss + elif len(d_str) == 19: + return datetime.datetime.strptime(d_str, "%Y-%m-%d %H:%M:%S") + # yyyymmdd.decimal_day + elif len(d_str) > 8 and d_str[8] == ".": + f = float("0" + d_str[8:]) + if f >= 1: + raise ValueError + else: + return datetime.datetime.strptime(d_str[:8], "%Y%m%d") + datetime.timedelta( + seconds=int(f * 86400) + ) + # mmdd or mdd (as in daily climatologies) + elif len(d_str) in [3, 4]: + # try to convert to datatime to see if it is valid date + d = datetime.datetime.strptime("0004" + d_str.rjust(4, "0"), "%Y%m%d") + # we just return a tuple since datetime cannot have an invalid date + return (d.month, d.day) + # b-dd e.g. apr-02 (as in daily climatologies) + elif len(d_str) == 6 and d_str[3] == "-": + months = [ + "jan", + "feb", + "mar", + "apr", + "may", + "jun", + "jul", + "aug", + "sep", + "nov", + "dec", + ] + m = d_str[0:3].lower() + try: + m_num = months.index(m) + 1 + except: + raise ValueError(f"Invalid month={m} specified in date={d_str}!") + # try to convert to datatime to see if it is valid date + d = datetime.datetime.strptime("0004" + f"{m_num:02}" + d_str[4:6], "%Y%m%d") + # we just return a tuple since datetime cannot have an invalid date + return (d.month, d.day) + + +def time_from_str(t_str): + h = m = 0 + if not ":" in t_str: + # formats: h[mm], hh[mm] + if len(t_str) in [1, 2]: + h = int(t_str) + elif len(t_str) in [3, 4]: + r = int(t_str) + h = int(r / 100) + m = int(r - h * 100) + else: + raise Exception(f"Invalid time={t_str}") + else: + # formats: h:mm, hh:mm + lst = t_str.split(":") + if len(lst) >= 2: + h = int(lst[0]) + m = int(lst[1]) + else: + raise Exception(f"Invalid time={t_str}") + + return datetime.time(hour=h, minute=m) + + +def date_from_ecc_keys(d, t): + try: + return datetime.datetime.combine( + date_from_str(str(d)).date(), time_from_str(str(t)) + ) + except: + return None + + +def is_fieldset_type(thing): + # will return True for binary or python fieldset objects + return "Fieldset" in thing.__class__.__name__ + + +def get_file_list(path, file_name_pattern=None): + m = None + # if isinstance(file_name_pattern, re.Pattern): + # m = file_name_pattern.match + # elif isinstance(file_name_pattern, str): + if isinstance(file_name_pattern, str): + if file_name_pattern.startswith('re"'): + m = re.compile(file_name_pattern[3:-1]).match + + # print(f"path={path} file_name_pattern={file_name_pattern}") + + if m is not None: + return [os.path.join(path, f) for f in filter(m, os.listdir(path=path))] + else: + if isinstance(file_name_pattern, str) and file_name_pattern != "": + path = os.path.join(path, file_name_pattern) + if not has_globbing(path): + return [path] + else: + return sorted(glob.glob(path)) + + +def has_globbing(text): + for x in ["*", "?"]: + if x in text: + return True + if "[" in text and "]" in text: + return True + else: + return False + + +def unpack(file_path, remove=False): + if any(file_path.endswith(x) for x in [".tar", ".tar.gz", ".tar.bz2"]): + target_dir = os.path.dirname(file_path) + LOG.debug(f"file_path={file_path} target_dir={target_dir}") + shutil.unpack_archive(file_path, target_dir) + if remove: + os.remove(file_path) + + +def download(url, target): + from tqdm import tqdm + import requests + + resp = requests.get(url, stream=True) + total = int(resp.headers.get("content-length", 0)) + with open(target, "wb") as file, tqdm( + desc=target, + total=total, + unit="iB", + unit_scale=True, + unit_divisor=1024, + ) as bar: + for data in resp.iter_content(chunk_size=1024): + size = file.write(data) + bar.update(size) + + +def simple_download(url, target): + import requests + + r = requests.get(url, allow_redirects=True) + r.raise_for_status() + open(target, "wb").write(r.content) + + +def _smooth_core(fs, repeat, m_func, m_arg, **kwargs): + """ + Performs spatial smoothing on each field in fs with the given callable + """ + # extract metadata for each field + meta = fs.grib_get(["gridType", "Ni", "generatingProcessIdentifier"]) + + # the resulting fieldset. We cannot use the Fieldset constructor here! + res = type(fs)() + + # build result in a loop + for fld, fld_meta in zip(fs, meta): + # smoothing only works for regular latlon grids + if fld_meta[0] == "regular_ll": + ncol = int(fld_meta[1]) + val = fld.values() + val = np.reshape(val, (-1, ncol)) + for _ in range(repeat): + val = m_func(val, m_arg, **kwargs) + r = fld.set_values(val.flatten()) + if fld_meta[2] is not None: + try: + r = r.grib_set_long( + ["generatingProcessIdentifier", int(fld_meta[2])] + ) + except: + pass + res.append(r) + else: + raise ValueError( + f"Unsupported gridType={fld_meta[0]} in field={i}. Only regular_ll is accepted!" + ) + + return res + + +def convolve(fs, weight, repeat=1, **kwargs): + """ + Performs spatial smoothing on each field in fs with a convolution + """ + from scipy.ndimage.filters import convolve + + m_arg = weight + return _smooth_core(fs, repeat, convolve, m_arg, **kwargs) + + +def smooth_n_point(fs, n=9, repeat=1, **kwargs): + """ + Performs spatial smoothing on each field in fs with an n-point averaging + """ + from scipy.ndimage.filters import convolve + + if n == 9: + weights = np.array( + [[0.0625, 0.125, 0.0625], [0.125, 0.25, 0.125], [0.0625, 0.125, 0.0625]], + dtype=float, + ) + m_arg = weights + elif n == 5: + weights = np.array( + [[0.0, 0.125, 0.0], [0.125, 0.5, 0.125], [0.0, 0.125, 0.0]], dtype=float + ) + m_arg = weights + else: + raise ValueError("smooth_n_point: n must be either 5 or 9!") + + return _smooth_core(fs, repeat, convolve, m_arg, **kwargs) + + +def smooth_gaussian(fs, sigma=1, repeat=1, **kwargs): + """ + Performs spatial smoothing on each field in fs with a Gaussian filter + """ + from scipy.ndimage.filters import gaussian_filter + + m_arg = sigma + return _smooth_core(fs, repeat, gaussian_filter, m_arg, **kwargs) + + +class Cache: + ROOT_DIR = os.path.join(tempfile.gettempdir(), f"mpy_ds_{getpass.getuser()}") + + def all_exists(self, items, path): + for name in items: + p = os.path.join(path, name) + # print(f"p={p}") + if not os.path.exists(p): + return False + elif os.path.isdir(p): + cont_file = os.path.join(path, f".content_{name}") + if os.path.exists(cont_file): + with open(cont_file, "r") as f: + try: + for item in f.read().split("\n"): + if item and not os.path.exists( + os.path.join(path, item) + ): + return False + except: + return False + else: + return False + return True + + def make_reference(self, items, path): + for name in items: + p = os.path.join(path, name) + if os.path.isdir(p): + cont_file = os.path.join(path, f".content_{name}") + with open(cont_file, "w") as f: + for item in Path(p).rglob("*"): + if item.is_file(): + f.write(item.relative_to(path).as_posix() + "\n") + + +CACHE = Cache() diff --git a/metview/plotting.py b/metview/plotting.py new file mode 100644 index 00000000..3dea0b60 --- /dev/null +++ b/metview/plotting.py @@ -0,0 +1,1044 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + + +import logging +import math + +import numpy as np +import pandas as pd + +import metview as mv +from metview.layout import Layout +from metview.style import Visdef, Style, GeoView +from metview.title import Title +from metview.track import Track +from metview.scaling import Scaling + +# logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") +# logging.basicConfig(level=logging.DEBUG, format="%(levelname)s - %(message)s") +# logging.basicConfig(level=logging.DEBUG) +LOG = logging.getLogger(__name__) + + +def _make_layers(*args, form_layout=False): + # in the positional arguments we have two options: + # 1. we only have non-list items. They belong to a single plot page. + # 2. we only have list items. Each list item defines a separate plot page. + plot_def = [] + if form_layout: + plot_def = list(args) + lst_cnt = sum([1 for x in plot_def if isinstance(x, list)]) + if not lst_cnt in [0, len(plot_def)]: + raise Exception( + f"Invalid plot arguments! Cannot mix list and non-list positional arguments." + ) + + if lst_cnt == 0: + plot_def = [plot_def] + else: + plot_def = [list(args)] + + # plot def at this point a list of lists. Each list item describes a + # map page item! + layers = [] + for d in plot_def: + data = None + layer = [] + for item in d: + if isinstance(item, (mv.Fieldset, Track)): + data = item + layer.append({"data": data, "vd": []}) + elif data is not None: + assert len(layer) > 0 + layer[-1]["vd"].append(item) + layers.append(layer) + + if form_layout: + return layers + else: + return layers[0] if layers else [] + + +def _make_visdef( + data, + vd, + use_eccharts=False, + style_db="param", + plot_type="map", + data_id=None, + pos_values=None, +): + if isinstance(data, mv.Fieldset): + if len(vd) == 0: + if use_eccharts: + return [mv.style.make_eccharts_mcont()] + else: + vd = mv.style.get_db(name=style_db).visdef( + data, plot_type=plot_type, data_id=data_id + ) + if plot_type == "diff" and pos_values is not None and pos_values: + s = mv.style.get_db(name=style_db).style( + data, plot_type=plot_type, data_id=data_id + ) + if s is not None and len(s.visdefs) == 2: + neg_values = [-x for x in reversed(pos_values)] + s.visdefs[0].set_values_list(neg_values) + s.visdefs[1].set_values_list(pos_values) + vd = s.to_request() + else: + vd = mv.style.get_db(name=style_db).visdef( + data, plot_type=plot_type, data_id=data_id + ) + else: + for i, v in enumerate(vd): + if isinstance(v, Style): + v = v.set_data_id(data_id) + vd[i] = v.to_request() + elif ( + isinstance(v, mv.Request) and data_id is not None and data_id != "" + ): + v = Visdef.from_request(v) + v.set_data_id(data_id) + vd[i] = v.to_request() + if vd is not None: + return [x for x in vd if x is not None] + else: + return [] + else: + return [] + + +def _get_data_area(data): + bb = [] + for d in data: + b = d.bounding_box() + if len(b) == 4: + if math.fabs(b[1] - b[0]) > 160: + b[0] = -90 + b[2] = 90 + if math.fabs(b[3] - b[1]) > 340: + b[1] = -180 + b[3] = 180 + if len(bb) == 0: + bb = b.tolist() + else: + bb = [ + min(b[0], bb[0]), + min(b[1], bb[1]), + max(b[2], bb[2]), + max(b[3], bb[3]), + ] + return bb + + +def _make_view(view, area, plot_type="map", data=None): + data = [] if data is None else data + + if area == "data" and data: + area = _get_data_area(data) + if len(area) != 4: + area = "base" + + if view is not None: + if area is not None: + view = mv.make_geoview(area=area, style=view["coastlines"]) + else: + if area is not None: + view = mv.make_geoview(area=area, plot_type=plot_type) + else: + view = mv.make_geoview(area="base", plot_type=plot_type) + return view + + +def _prepare_grid(d1, d2): + if d1._db is not None and d2._db is not None: + # interpolate from d2 to d1 + if d1._db.name in d2._db.regrid_from: + # print(f"regrid: {d1.label} -> {d2.label}") + # d11 = d1 + 0 + d = mv.regrid(data=d1, grid_definition_mode="template", template_data=d2[0]) + d = mv.grib_set_long(d, ["generatingProcessIdentifier", 148]) + d._db = d1._db._clone() + # print(f"{ len(d)}") + # print(f"{ len(d2)}") + # print(" {}".format(mv.grib_get(d[0], ["numberOfDataPoints"]))) + # print(" {}".format(mv.grib_get(d2[0], ["numberOfDataPoints"]))) + return (d, d2) + # interpolate from d1 to d2 + elif d2._db.name in d1._db.regrid_from: + # print(f"regrid: {d2.label} -> {d1.label}") + # d22 = d2 + 0 + d = mv.regrid(data=d2, grid_definition_mode="template", template_data=d1[0]) + d = mv.grib_set_long(d, ["generatingProcessIdentifier", 148]) + d._db = d2._db._clone() + return (d1, d) + + return (d1, d2) + + +def _y_max(data): + return max([max(d) for d in data]) + + +def _y_min(data): + return min([min(d) for d in data]) + + +def _scale_data(data, style_db="param", plot_type="map"): + if isinstance(data, mv.Fieldset): + scaler = mv.style.get_db(name=style_db).units_scaler(data, plot_type=plot_type) + if scaler is not None: + return scaler.scale_value(data) + return data + + +def _scale_xs(name, data): + if isinstance(data, mv.Fieldset): + scaler = None + for f in data: + if f.grib_get_string("shortName") != "lnsp": + scaler = mv.style.get_db(name="param").units_scaler(f, plot_type="xs") + if scaler is not None: + r = mv.Fieldset() + for f in data: + if f.grib_get_string("shortName") != "lnsp": + r.append(scaler.scale_value(f)) + else: + r.append(f) + return r + return data + + +def plot_maps( + *args, + layout=None, + view=None, + area=None, + use_eccharts=False, + title_font_size=0.4, + legend_font_size=0.35, + frame=-1, + animate="auto", +): + """ + Plot maps with generic contents + """ + + # in the positional arguments we have two options: + # 1. we only have non-list items. They belong to a single plot page. + # 2. we only have list items. Each list item defines a separate plot page. + plot_def = _make_layers(*args, form_layout=True) + + # collect the data items + data_items = [] + for i, sc_def in enumerate(plot_def): + for layer in sc_def: + data = layer["data"] + if isinstance(data, mv.Fieldset): + data_items.append(data[0]) + + # define the view + view = _make_view(view, area, data=data_items) + + # build the layout + num_plot = len(plot_def) + dw = Layout().build_grid(page_num=num_plot, layout=layout, view=view) + + # the plot description + desc = [] + + title = Title(font_size=title_font_size) + + # build each scene + data_id = ("d0", 0) + for i, sc_def in enumerate(plot_def): + desc.append(dw[i]) + # define layers + data_items = [] + use_data_id = ( + sum([1 for layer in sc_def if isinstance(layer["data"], mv.Fieldset)]) > 1 + ) + for layer in sc_def: + data = layer["data"] + vd = layer["vd"] + if isinstance(data, mv.Fieldset): + if use_data_id: + data_items.append((data, data_id[0])) + else: + data_items.append(data) + if frame != -1: + if data.ds_param_info.scalar: + data = data[frame] + else: + data = data[2 * frame : 2 * frame + 2] + + data = _scale_data(data, style_db="param", plot_type="map") + elif isinstance(data, Track): + data = data.build(style=vd) + + desc.append(data) + + if isinstance(data, mv.Fieldset): + vd = _make_visdef( + data, + vd, + use_eccharts=use_eccharts, + style_db="param", + plot_type="map", + data_id=data_id[0] if use_data_id else None, + ) + if vd: + desc.extend(vd) + + data_id = (f"d{data_id[1]+1}", data_id[1] + 1) + + if data_items: + legend = mv.mlegend(legend_text_font_size=legend_font_size) + desc.append(legend) + t = title.build(data_items) + # LOG.debug(f"t={t}") + desc.append(t) + + for i in range(len(plot_def), len(dw)): + desc.append(dw[i]) + + LOG.debug(f"desc={desc}") + + return mv.plot(desc, animate=animate) + + +def plot_diff_maps( + *args, + view=None, + area=None, + overlay=None, + diff_style=None, + pos_values=None, + title_font_size=0.4, + legend_font_size=0.35, + frame=-1, + animate="auto", +): + """ + Plot difference maps + """ + + # handle default arguments + pos_values = [] if pos_values is None else pos_values + diff_style = [] if diff_style is None else diff_style + if not isinstance(diff_style, list): + diff_style = [diff_style] + + # define the view + view = _make_view(view, area, plot_type="diff") + + # build the layout + dw = Layout().build_diff(view=view) + + data = {} + vd = {} + + # the positional arguments has the following order: + # data1, visdef1.1 visdef1.2 ... data2, visdef2.1, visdef2.2 ... + # the visdef objects are optional! + assert len(args) >= 2 + assert isinstance(args[0], mv.Fieldset) + layers = _make_layers(*args, form_layout=False) + assert len(layers) == 2 + # LOG.debug(f"layers={layers}") + data["0"] = layers[0]["data"] + data["1"] = layers[1]["data"] + vd["0"] = _make_visdef(data["0"], layers[0]["vd"]) + vd["1"] = _make_visdef(data["1"], layers[1]["vd"]) + + # overlay data + ov_data = {} + ov_vd = {} + if overlay is not None: + # single value, list or tuple: a data item that will be plotted into each map + if not isinstance(overlay, dict): + if isinstance(overlay, tuple): + ov_args = list(overlay) + else: + ov_args = [overlay] if not isinstance(overlay, list) else overlay + # print(ov_args) + ov_layers = _make_layers(*ov_args, form_layout=False) + # print(ov_layers) + assert len(ov_layers) == 1 + d = ov_layers[0]["data"] + if isinstance(d, Track): + d = d.build(style=ov_layers[0]["vd"]) + for k in ["d", "0", "1"]: + ov_data[k] = d + ov_vd[k] = _make_visdef(d, ov_layers[0]["vd"]) + else: + pass + + # LOG.debug("len_0={}".format(len(data["0"]))) + # LOG.debug("len_1={}".format(len(data["0"]))) + + # the plot description + desc = [] + + title = Title(font_size=title_font_size) + + # compute diff + data["0"], data["1"] = _prepare_grid(data["0"], data["1"]) + data["d"] = data["0"] - data["1"] + + data["d"]._ds_param_info = data["1"].ds_param_info + if data["0"].label and data["1"].label: + data["d"]._label = "{}-{}".format(data["0"].label, data["1"].label) + else: + data["d"]._label = "" + vd["d"] = _make_visdef( + data["d"], diff_style, plot_type="diff", pos_values=pos_values + ) + + # LOG.debug("len_d={}".format(len(data["d"]))) + + for i, k in enumerate(["d", "0", "1"]): + desc.append(dw[i]) + if frame == -1: + d = data[k] + else: + d = data[k][frame] + d._ds_param_info = data[k]._ds_param_info + d._label = data[k]._label + + desc.append(d) + if vd[k]: + desc.append(vd[k]) + + # add overlay + if k in ov_data: + if isinstance(ov_data[k], mv.Fieldset): + dd = ov_data[k] if frame == -1 else ov_data[k][frame] + else: + dd = ov_data[k] + desc.append(dd) + if k in ov_vd and ov_vd[k]: + desc.append(ov_vd[k]) + + t = title.build(data[k]) + legend = mv.mlegend(legend_text_font_size=legend_font_size) + desc.append(legend) + desc.append(t) + + # print(desc) + return mv.plot(desc, animate=animate) + + +def plot_xs( + *args, + map_line=True, + map_data=None, + line=[], + layout="", + view=None, + area=None, + title_font_size=0.3, + legend_font_size=0.2, + frame=-1, + animate="auto", +): + """ + Plot cross section with map + """ + + assert len(line) == 4 + assert len(args) >= 1 + assert isinstance(args[0], mv.Fieldset) + layers = _make_layers(*args, form_layout=False) + assert len(layers) > 0 + + # build the layout - if no map_data is specified no map view is + # added to the layout + if not map_line and map_data is None: + view = None + else: + view = _make_view(view, area) + dw = Layout().build_xs(line=line, map_view=view) + + # the plot description + desc = [] + + title = Title(font_size=title_font_size) + data_items = [] + + # build cross section plot + desc.append(dw[0]) + for layer in layers: + data = layer["data"] + vd = _make_visdef(data, layer["vd"], plot_type="xs") + param_info = data.ds_param_info + # print(f"param_info={param_info}") + data_items.append(data) + # print(f"data={len(data)}") + if param_info is not None and param_info.name == "wind3d": + xs_d = mv.mcross_sect( + data=data, + line=line, + wind_parallel="on", + w_wind_scaling_factor_mode="compute", + w_wind_scaling_factor="100", + # bottom_level=1015, + # top_level=250, + # vertical_scaling = "log", + # vertical_axis= vertical_axis + ) + desc.append(xs_d) + else: + if param_info is not None: + data = _scale_xs(param_info.name, data) + desc.append(data) + + if vd: + desc.extend(vd) + # print(f"vd={vd}") + + t = title.build_xs(data_items) + desc.append(t) + + # LOG.debug(f"desc={desc}") + + # build side map plot + if map_line or map_data is not None: + desc.append(dw[1]) + t = None + if map_data is not None and len(map_data) > 0: + layers = _make_layers(map_data, form_layout=False) + data_items = [] + for layer in layers: + data = layer["data"] + vd = _make_visdef(data, layer["vd"]) + + if isinstance(data, mv.Fieldset): + data_items.append(data) + if frame != -1: + data = data[frame] + + desc.append(data) + if vd: + desc.extend(vd) + + if data_items: + t = title.build(data_items) + + if map_line: + # define xsection line graph + graph = mv.mgraph( + graph_line_colour="red", graph_line_thickness=3, graph_symbol="off" + ) + + lv = mv.input_visualiser( + input_plot_type="geo_points", + input_longitude_values=[line[1], line[3]], + input_latitude_values=[line[0], line[2]], + ) + + desc.extend([lv, graph]) + + if t is not None: + desc.append(t) + + LOG.debug(f"desc={desc}") + return mv.plot(desc, animate=animate) + + +def plot_stamp( + *args, + an=[], + fc=[], + layout=None, + view=None, + area=None, + title_font_size=0.4, + frame=-1, + animate="auto", + diff_base=None, +): + """ + Plot ENS stamp maps + """ + + # define the view + view = _make_view(view, area) + + desc = [] + data = {} + vd = {} + + if diff_base is not None: + assert isinstance(diff_base, mv.Fieldset) + + if len(args) > 0: + assert isinstance(args[0], mv.Fieldset) + layers = _make_layers(*args, form_layout=False) + assert len(layers) == 1 + + # prepare ens + data["ens"] = layers[0]["data"] + assert data["ens"] is not None + if diff_base is not None: + vd["ens"] = _make_visdef(data["ens"], [], style_db="diff") + else: + vd["ens"] = _make_visdef(data["ens"], layers[0]["vd"]) + + # prepare an and fc + d = {"an": an, "fc": fc} + for k, v in d.items(): + if v: + layers = _make_layers(v, form_layout=False) + if layers: + data[k] = layers[0]["data"] + vd[k] = layers[0]["vd"] + if diff_base is not None: + vd[k] = vd["ens"] + else: + if len(vd[k]) == 0 and "ens" in vd: + vd[k] = vd["ens"] + else: + vd[k] = _make_visdef(data[k], vd[k]) + + # determine ens number + members = [] + if "ens" in data: + members = data["ens"]._unique_metadata("number") + LOG.debug(f"members={members}") + if len(members) == 0: + raise Exceception("No ENS data found in input!") + + # determine number of maps + num = len(members) + sum([1 for x in ["an", "fc"] if x in data]) + + # build the layout + dw = Layout().build_stamp(num, layout=layout, view=view) + + if len(dw) < num + 1: + raise Exception(f"Layout has less maps (={len(dw)}) than expected (={num})") + + title = Title(font_size=title_font_size) + + # ens members + for i, m in enumerate(members): + desc.append(dw[i]) + d = data["ens"].select(number=m) + if diff_base is not None: + d = d - diff_base + desc.append(d) + + if vd["ens"]: + desc.extend(vd["ens"]) + + t = title.build_stamp(d, member=str(i)) + desc.append(t) + + # add an and fc + n = len(members) + for t in ["an", "fc"]: + if t in data: + desc.append(dw[n]) + d = data[t] + if diff_base is not None: + d = d - diff_base + desc.append(d) + if vd[t]: + desc.append(vd[t]) + t = title.build_stamp(data[t], member="") + desc.append(t) + n += 1 + + for i in range(n, len(dw)): + desc.append(dw[i]) + + if len(members) > 0 and "ens" in data: + cont = mv.mcont(contour="off", contour_label="off") + dummy = d = data["ens"].select(number=members[0]) + t = title.build(dummy) + desc.extend([dw[-1], t, dummy, cont]) + + return mv.plot(desc, animate=animate) + + +def plot_rmse(*args, ref=None, area=None, title_font_size=0.4, y_max=None): + """ + Plot RMSE curve + """ + + desc = [] + + if not isinstance(ref, mv.Fieldset): + raise Exception(f"Missing or invalid ref argument!") + + layers = _make_layers(*args, form_layout=False) + + # compute the rmse for each input layer + data = [] # list of tuples + rmse_data = [] + title_data = [] + has_ef = False + + for layer in layers: + if isinstance(layer["data"], mv.Fieldset): + # determine ens number + members = layer["data"]._unique_metadata("number") + # print(f"members={members}") + # ens forecast + if len(members) > 1: + if has_ef: + raise Exception("Only one ENS fieldset can be used in plot_rmse()!") + has_ef = True + em_d = None # ens mean + for m in members: + pf_d = layer["data"].select(number=m) + ref_d, pf_d = _prepare_grid(ref, pf_d) + data.append(("cf" if m == "0" else "pf", layer["data"])) + rmse_data.append(mv.sqrt(mv.average((pf_d - ref_d) ** 2))) + em_d = pf_d if em_d is None else em_d + pf_d + + # compute rmse for ens mean + data.append(("em", layer["data"])) + rmse_data.append( + mv.sqrt(mv.average((em_d / len(members) - ref_d) ** 2)) + ) + + # deterministic forecast + else: + ref_d, dd = _prepare_grid(ref, layer["data"]) + data.append(("fc", layer["data"])) + rmse_data.append(mv.sqrt(mv.average((dd - ref_d) ** 2))) + + title_data.append(layer["data"]) + + # define x axis params + dates = ref.valid_date() + x_min = dates[0] + x_max = dates[-1] + x_tick = 1 + x_title = "" + + # define y axis params + y_min = 0 + if y_max is None: + y_tick, _, y_max = Layout.compute_axis_range(0, _y_max(rmse_data)) + else: + y_tick, _, _ = Layout.compute_axis_range(0, y_max) + y_title = "RMSE [" + mv.grib_get_string(ref[0], "units") + "]" + + # print(f"y_tick={y_tick} y_max={y_max}") + + # define the view + view = Layout().build_rmse( + x_min, x_max, y_min, y_max, x_tick, y_tick, x_title, y_title + ) + desc.append(view) + + # define curves + ef_label = {"cf": "ENS cf", "pf": "ENS pf", "em": "ENS mean"} + ef_colour = {"cf": "black", "pf": "red", "em": "kelly_green"} + fc_colour = ["red", "blue", "green", "black", "cyan", "evergreen", "gold", "pink"] + if has_ef: + fc_colour = [x for x in fc_colour if x not in list(ef_colour.values())] + + pf_label_added = False + colour_idx = -1 + legend_item_count = 0 + for i, d in enumerate(rmse_data): + vis = mv.input_visualiser( + input_x_type="date", input_date_x_values=dates, input_y_values=d + ) + + vd = {"graph_type": "curve"} + line_colour = "black" + line_width = 1 + + if data[i][0] == "fc": + line_width = 3 + colour_idx = (colour_idx + 1) % len(fc_colour) + line_colour = fc_colour[colour_idx] + # print(f"label={data[i][1][0].label}") + vd["legend_user_text"] = data[i][1].label + vd["legend"] = "on" + legend_item_count += 1 + elif data[i][0] == "pf": + line_width = 1 + line_colour = ef_colour["pf"] + if not pf_label_added: + pf_label_added = True + vd["legend_user_text"] = ef_label.get("pf", "") + vd["legend"] = "on" + legend_item_count += 1 + elif data[i][0] in ["cf", "em"]: + line_width = 3 + line_colour = ef_colour[data[i][0]] + vd["legend_user_text"] = ef_label.get(data[i][0], "") + vd["legend"] = "on" + legend_item_count += 1 + + vd["graph_line_colour"] = line_colour + vd["graph_line_thickness"] = line_width + + desc.append(vis) + desc.append(mv.mgraph(**vd)) + + # add title + title = Title(font_size=title_font_size) + t = title.build_rmse(ref, title_data) + if t is not None: + desc.append(t) + + # add legend + leg_left = 3.5 + # legY = 14 + leg_height = legend_item_count * (0.35 + 0.5) + (legend_item_count + 1) * 0.1 + leg_bottom = 17.5 - leg_height + + # Legend + legend = mv.mlegend( + legend_display_type="disjoint", + legend_entry_plot_direction="column", # "row", + legend_text_composition="user_text_only", + legend_border="on", + legend_border_colour="black", + legend_box_mode="positional", + legend_box_x_position=leg_left, + legend_box_y_position=leg_bottom, + legend_box_x_length=4, + legend_box_y_length=leg_height, + legend_text_font_size=0.35, + legend_box_blanking="on", + ) + desc.append(legend) + + mv.plot(desc, animate=False) + + +def plot_cdf(*args, location=None, title_font_size=0.4, x_range=None): + """ + Plot CDF curve + """ + + # check x range + x_range = [] if x_range is None else x_range + if x_range and len(x_range) not in [2, 3]: + raise Exception( + f"plot_cdf: invalid x_range specified. Format [x_min, x_max, [x_tick]]" + ) + if len(x_range) == 2 and x_range[1] <= x_range[0]: + raise Exception( + f"plot_cdf: invalid x_range specified. x_min={x_range[0]} >= x_max={x_range[1]}" + ) + + layers = _make_layers(*args, form_layout=False) + + desc = [] + cdf_data = [] + cdf_label = [] + title_data = [] + y_values = np.arange(0, 101) + plot_units = "" + units_scaler = None + + # compute the cdf for each input layer + for layer in layers: + if isinstance(layer["data"], mv.Fieldset): + # we assume each field has the same units and paramId + if plot_units == "": + meta = mv.grib_get(layer["data"][0], ["units", "paramId"]) + if meta and len(meta[0]) == 2: + meta = {"units": meta[0][0], "paramId": meta[0][1]} + units_scaler = Scaling.find_item(meta) + if units_scaler is not None: + plot_units = units_scaler.to_units + else: + plot_units = meta.get("units", "") + + # determine ens number and steps + members = layer["data"]._unique_metadata("number") + steps = layer["data"]._unique_metadata("step") + # print(f"members={members}") + # ens forecast + if len(members) > 1: + for step in steps: + v = layer["data"].select(step=step) + v = mv.nearest_gridpoint(v, location) + # print(f"step={step}") + x = np.percentile(v, y_values) + if units_scaler is not None: + x = units_scaler.scale_value(x) + # print(f" x={x}") + cdf_data.append(x) + cdf_label.append(layer["data"].label + f" +{step}h") + + # deterministic forecast + else: + raise Exception(f"plot_cds: only ENS data accepted as input!") + + title_data.append(layer["data"]) + + # define x axis params + if not x_range: + x_tick, x_min, x_max = Layout.compute_axis_range( + _y_min(cdf_data), _y_max(cdf_data) + ) + elif len(x_range) == 2: + x_min = x_range[0] + x_max = x_range[1] + x_tick, _, _ = Layout.compute_axis_range(x_min, x_max) + elif len(x_range) == 3: + x_min = x_range[0] + x_max = x_range[1] + x_tick = x_range[2] + else: + raise Exception(f"plot_cdf: invalid x_range={x_range} specified!") + + # print(f"x_tick={x_tick} x_min={x_min} x_max={x_max}") + x_title = f"[{plot_units}]" + + # define y axis params + y_min = 0 + y_max = 100 + y_tick = 10 + y_title = "Percentage [%]" + + # define the view + view = Layout().build_xy( + x_min, x_max, y_min, y_max, x_tick, y_tick, x_title, y_title + ) + desc.append(view) + + # define curves + line_colours = [ + "red", + "blue", + "green", + "black", + "cyan", + "evergreen", + "gold", + "pink", + ] + line_styles = ["solid", "dash", "dotted"] + + colour_idx = -1 + style_idx = 0 + + for i, d in enumerate(cdf_data): + vis = mv.input_visualiser(input_x_values=d, input_y_values=y_values) + colour_idx = (colour_idx + 1) % len(line_colours) + + vd = mv.mgraph( + graph_type="curve", + graph_line_colour=line_colours[colour_idx], + graph_line_thickness=3, + legend_user_text=cdf_label[i], + legend="on", + ) + + desc.append(vis) + desc.append(vd) + + # add title + title = Title(font_size=title_font_size) + t = title.build_cdf(title_data) + if t is not None: + desc.append(t) + + # add legend + legX = 3.5 + legY = 14 + + # Legend + legend = mv.mlegend( + legend_display_type="disjoint", + legend_entry_plot_direction="column", + legend_text_composition="user_text_only", + legend_border="on", + legend_border_colour="black", + legend_box_mode="positional", + legend_box_x_position=legX, + legend_box_y_position=legY, + legend_box_x_length=4, + legend_box_y_length=3, + legend_text_font_size=0.35, + legend_box_blanking="on", + ) + desc.append(legend) + + mv.plot(desc, animate=False) + + +def plot_xs_avg( + *args, + area=None, + direction="ew", + vertical_scaling="linear", + title_font_size=0.3, + legend_font_size=0.2, + axis_font_size=0.3, + animate="auto", +): + """ + Plot average cross section + """ + + if area is None: + area = [90, -180, -90, 180] + + assert isinstance(args[0], mv.Fieldset) + layers = _make_layers(*args, form_layout=False) + assert len(layers) > 0 + + view = Layout().build_xs_avg( + area=area, + direction=direction, + bottom_level=1000, + top_level=50, + vertical_scaling=vertical_scaling, + axis_label_height=axis_font_size, + ) + + # the plot description + desc = [] + + title = Title(font_size=title_font_size) + data_items = [] + + # build cross section plot + desc.append(view) + for layer in layers: + data = layer["data"] + + vd = _make_visdef(data, layer["vd"], plot_type="xs") + param_info = data.ds_param_info + # print(f"param_info={param_info}") + data_items.append(data) + # print(f"data={len(data)}") + + if param_info is not None: + data = _scale_xs(param_info.name, data) + desc.append(data) + + if vd: + desc.extend(vd) + # print(f"vd={vd}") + + t = title.build_xs(data_items) + desc.append(t) + + legend = mv.mlegend(legend_text_font_size=legend_font_size) + desc.append(legend) + desc.append(t) + + LOG.debug(f"desc={desc}") + + return mv.plot(desc, animate=animate) diff --git a/metview/scaling.py b/metview/scaling.py new file mode 100644 index 00000000..08259969 --- /dev/null +++ b/metview/scaling.py @@ -0,0 +1,144 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import json +import logging +import os +import sys + +import yaml + +LOG = logging.getLogger(__name__) + +ETC_PATH = os.path.join(os.path.dirname(__file__), "etc") + + +class UnitsScalingMethod: + """ + Performs units scaling of values + """ + + def __init__(self, scaling=1.0, offset=0.0, from_units="", to_units=""): + self.scaling = scaling + self.offset = offset + self.from_units = from_units + self.to_units = to_units + + def scale_value(self, value): + return self.scaling * value + self.offset + + def inverse_scale_value(self, value): + return (value - self.offset) / self.scaling + + def need_scaling(self, meta, scaling_retrieved, scaling_derived): + gen_id = meta.get("generatingProcessIdentifier", 0) + try: + gen_id = int(gen_id) + except: + gen_id = 0 + + return (scaling_retrieved and gen_id != 254) or ( + scaling_derived and gen_id == 254 + ) + + def __str__(self): + return "scaling: {} offset:{} from_units:{} to_units:{}".format( + self.scaling, self.offset, self.from_units, self.to_units + ) + + +class ScalingRule: + """ + Defines what countour scaling should be applied to a given field based + on its metadata + """ + + def __init__(self, to_units, conf): + self.to_units = to_units + self.units_methods = [ + it for it in Scaling.methods if it.to_units == self.to_units + ] + self.conf = conf + + def find_method(self, meta): + from_units = meta.get("units", "") + if from_units == "": + return None + + method = None + for item in self.units_methods: + if item.from_units == from_units: + method = item + break + + if method is None: + return None + + for m in self.conf.get("match", []): + match = False + for k, v in m.items(): + if meta.get(k, None) == v: + match = True + else: + break + + if match: + return method + + param_id = meta.get("paramId", "") + if param_id and param_id in self.conf.get("paramId", []): + return method + + short_name = meta.get("shortName", "") + if short_name and short_name in self.conf.get("shortName", []): + return method + + return None + + def __str__(self): + return "to_units:{}".format(self.to_units) + + +class Scaling: + methods = [] + rules = [] + loaded = False + + @staticmethod + def find_item(meta): + if not Scaling.loaded: + Scaling._load_def() + Scaling.loaded = True + + for item in Scaling.rules: + d = item.find_method(meta) + if d: + return d + return None + + @staticmethod + def _load_def(): + # load units conversion definition + file_name = os.path.join(ETC_PATH, "units-rules.json") + with open(file_name) as f: + data = json.load(f) + for k, v in data.items(): + for item in v: + item["to_units"] = item.pop("to") + item["from_units"] = item.pop("from") + Scaling.methods.append(UnitsScalingMethod(**item)) + + # load rules defining when to apply scaling on a parameter + file_name = os.path.join(ETC_PATH, "scaling_ecmwf.yaml") + # print(f"file_name={file_name}") + with open(file_name) as f: + data = yaml.load(f, Loader=yaml.SafeLoader) + for to_units, item in data.items(): + Scaling.rules.append(ScalingRule(to_units, item)) diff --git a/metview/style.py b/metview/style.py new file mode 100644 index 00000000..3625f4d1 --- /dev/null +++ b/metview/style.py @@ -0,0 +1,854 @@ +# +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import copy +import json +import logging +import os +from pathlib import Path + +import yaml +import metview as mv +from metview.metviewpy.param import ParamInfo +from metview.metviewpy.ipython import is_ipython_active +from metview.scaling import Scaling + +# logging.basicConfig(level=logging.DEBUG) +LOG = logging.getLogger(__name__) + +_DB = { + "param": [None, "params.yaml", "param_styles.yaml"], + "map": [None, "", "map_styles.yaml"], +} + +_MAP_CONF = None +ETC_PATH = os.path.join(os.path.dirname(__file__), "etc") +CUSTOM_CONF_PATH = [] +LOCAL_CONF_PATH = "" + + +PARAM_VISDEF_VERBS = ["mcont", "mwind", "mcoast", "msymb", "mgraph"] + + +# _DB = None + + +# def cont(): +# global _DB +# if not _DB: +# _DB = StyleDb() +# return _DB + + +class ContourStyleDbItem: + def __init__(self, name, db): + self.name = name + self.db = db + self.keywords = [] + self.colours = [] + self.layers = [] + self.description = "" + + def list_match(self, lst, pattern): + p = pattern.lower() + for t in lst: + if p in t.lower(): + return True + return False + + def keyword_match(self, pattern): + return self.list_match(self.keywords, pattern) + + def layer_match(self, pattern): + return self.list_match(self.layers, pattern) + + def colour_match(self, pattern): + return self.list_match(self.colours, pattern) + + def preview_file(self): + if self.db.preview_path: + return os.path.join(self.db.preview_path, self.name + ".png") + else: + return str() + + +class ContourStyleDb: + def __init__(self): + self.SHARE_DIR = os.path.join( + mv.version_info()["metview_dir"], "share", "metview", "eccharts" + ) + self.preview_path = os.path.join(self.SHARE_DIR, "style_previews") + self.items = [] + self.keywords = [] + self.colours = [] + self.layers = [] + self._load() + self._load_magics_def() + + def _load(self): + file_path = os.path.join(self.SHARE_DIR, "styles.json") + # print(f"file_path={file_path}") + try: + with open(file_path, "rt") as f: + conf = json.load(f) + col_set = set() + key_set = set() + layer_set = set() + for name, c in conf.items(): + item = ContourStyleDbItem(name, self) + item.colours = c.get("colours", []) + item.keywords = c.get("keywords", []) + item.layers = c.get("layers", []) + col_set.update(item.colours) + key_set.update(item.keywords) + layer_set.update(item.layers) + self.items.append(item) + self.colours = sorted(list(col_set)) + self.keywords = sorted(list(key_set)) + self.layers = sorted(list(layer_set)) + except: + pass + # LOG.exception("Failed to read eccharts styles", popup=True) + + self._load_magics_def() + + def find_by_name(self, name): + for i, item in enumerate(self.items): + if item.name == name: + return i, item + return -1, None + + def _load_magics_def(self): + pass + + def names(self): + return sorted([item.name for item in self.items]) + + +class Visdef: + # BUILDER = { + # "mcont": mv.mcont, + # "mwind": mv.mwind, + # "mcoast": mv.mcoast, + # "msymb": mv.msymb, + # "mgraph": mv.mgraph + # } + + def __init__(self, verb, params): + self.verb = verb.lower() + self.params = params + + self.BUILDER = { + "mcont": mv.mcont, + "mwind": mv.mwind, + "mcoast": mv.mcoast, + "msymb": mv.msymb, + "mgraph": mv.mgraph, + } + + def clone(self): + return Visdef(self.verb, copy.deepcopy(self.params)) + + def change(self, verb, param, value): + if verb == self.verb: + self.params[param] = value + + def change_symbol_text_list(self, text_value, idx_value): + assert self.verb == "msymb" + if self.verb == "msymb": + if self.params.get("symbol_type", "").lower() == "text": + self.params["symbol_advanced_table_text_list"] = text_value + self.params["symbol_advanced_table_level_list"] = idx_value + + def set_data_id(self, data_id): + if self.verb in ["mcont", "mwind"] and data_id is not None and data_id != "": + self.params["grib_id"] = data_id + + def set_values_list(self, values): + if self.verb == "mcont": + if self.params.get("contour_level_list", []): + self.params["contour_level_list"] = values + + @staticmethod + def from_request(req): + params = {k: v for k, v in req.items() if not k.startswith("_")} + vd = Visdef(req.verb.lower(), copy.deepcopy(params)) + return vd + + def to_request(self): + fn = self.BUILDER.get(self.verb, None) + if fn is not None: + return fn(**(self.params)) + else: + raise Exception(f"{self} unsupported verb!") + + def __str__(self): + return f"Visdef[verb={self.verb}, params={self.params}]" + + def __repr__(self): + return f"Visdef(verb={self.verb}, params={self.params})" + + +class Style: + def __init__(self, name, visdefs): + self.name = name + self.visdefs = visdefs + if not isinstance(visdefs, list): + self.visdefs = [self.visdefs] + + def clone(self): + return Style(self.name, [vd.clone() for vd in self.visdefs]) + + def to_request(self): + return [vd.to_request() for vd in self.visdefs] + + def update(self, *args, inplace=False, verb=None): + s = self if inplace == True else self.clone() + if isinstance(verb, str): + verb = [verb] + if verb: + for v in args: + if isinstance(v, dict): + v_vals = {v_key.lower(): v_val for v_key, v_val in v.items()} + for i, vd in enumerate(s.visdefs): + if vd.verb in verb: + vd.params.update(v_vals) + else: + for i, v in enumerate(args): + if isinstance(v, dict) and i < len(s.visdefs): + v = {v_key.lower(): v_val for v_key, v_val in v.items()} + s.visdefs[i].params.update(v) + return s + + def set_data_id(self, data_id): + allowed_verbs = ["mcont", "mwind"] + if isinstance(data_id, str) and any(x in allowed_verbs for x in self.verbs()): + return self.update({"grib_id": data_id}, verb=allowed_verbs) + else: + return self + + def verbs(self): + return [vd.verb for vd in self.visdefs] + + def __str__(self): + t = f"{self.__class__.__name__}[name={self.name}] " + for vd in self.visdefs: + t += f"{vd} " + return t + + +class ParamMatchCondition: + def __init__(self, cond): + self.name = cond.pop("info_name", "") + self.cond = cond + if "levelist" in self.cond: + if not isinstance(self.cond["levelist"], list): + self.cond["levelist"] = [self.cond["levelist"]] + # ParamInfo use "typeOfLevel" instead of "levtype" so we need to + # change the key a remap values + if "levtype" in self.cond: + v = self.cond.pop("levtype") + self.cond["typeOfLevel"] = v + self.cond["typeOfLevel"] = ParamInfo.LEVEL_TYPES.get(v, v) + + for k, v in self.cond.items(): + if isinstance(v, list): + self.cond[k] = [str(x) for x in v] + else: + self.cond[k] = str(v) + + def match(self, param_info): + return param_info.match(self.name, self.cond) + + +class ParamStyle: + def __init__(self, conf, db): + self.cond = [] + for d in conf["match"]: + if "info_name" in d: + self.info_name = d["info_name"] + self.cond.append(ParamMatchCondition(d)) + self.param_type = conf.get("param_type", "scalar") + + if self.param_type == "vector": + default_style = db.VECTOR_DEFAULT_STYLE_NAME + else: + default_style = db.SCALAR_DEFAULT_STYLE_NAME + + s = conf.get("styles", {}) + self.style = s.get("basic", [default_style]) + self.xs_style = s.get("xs", self.style) + self.diff_style = s.get("diff", [db.DIFF_DEFAULT_STYLE_NAME]) + + self.scaling = conf.get("scaling", []) + + def match(self, param): + return max([d.match(param) for d in self.cond]) + + def find_style(self, plot_type): + if plot_type == "" or plot_type == "map": + return self.style[0] + elif plot_type == "diff": + return self.diff_style[0] + elif plot_type == "xs": + return self.xs_style[0] + else: + return None + + def styles(self, plot_type): + if plot_type == "" or plot_type == "map": + return self.style + elif plot_type == "diff": + return self.diff_style + elif plot_type == "xs": + return self.xs_style + else: + return [] + + def __str__(self): + return "{}[name={},type={}]".format( + self.__class__.__name__, self.info_name, self.param_type + ) + + +class StyleDb: + SCALAR_DEFAULT_STYLE_NAME = "default_mcont" + VECTOR_DEFAULT_STYLE_NAME = "default_mwind" + DIFF_DEFAULT_STYLE_NAME = "default_diff" + + def __init__(self, param_file_name, style_file_name): + self.params = [] + self.styles = {} + self.param_file_name = param_file_name + self.style_file_name = style_file_name + + self._load_system_config() + self._load_custom_config() + self._load_local_config() + + def _load_system_config(self): + self._load( + os.path.join(ETC_PATH, self.param_file_name) + if self.param_file_name + else "", + os.path.join(ETC_PATH, self.style_file_name), + ) + + def _load_custom_config(self): + if CUSTOM_CONF_PATH and CUSTOM_CONF_PATH[-1]: + self._load( + os.path.join(CUSTOM_CONF_PATH[-1], self.param_file_name) + if self.param_file_name + else "", + os.path.join(CUSTOM_CONF_PATH[-1], self.style_file_name), + ) + + def _load_local_config(self): + if LOCAL_CONF_PATH: + self._load( + os.path.join(LOCAL_CONF_PATH, self.param_file_name) + if self.param_file_name + else "", + os.path.join(LOCAL_CONF_PATH, self.style_file_name), + ) + + def get_style(self, style): + if style in self.styles: + return self.styles[style] + else: + return self.styles.get("default", None) + + def _best_param_match(self, param_info): + r = 0 + p_best = None + for p in self.params: + m = p.match(param_info) + # print(f"p={p} m={m}") + if m > r: + r = m + p_best = p + return p_best + + def get_param_style_list(self, param_info, scalar=True, plot_type="map"): + p_best = self._best_param_match(param_info) + s = [] + # print(f"param_info={param_info}") + if p_best is not None: + s = p_best.styles(plot_type) + # print(f" -> style_name={style_name}") + if scalar: + s.append(self.SCALAR_DEFAULT_STYLE_NAME) + else: + s.append(self.VECTOR_DEFAULT_STYLE_NAME) + return s + + def get_param_style(self, param_info, scalar=True, plot_type="map", data_id=None): + p_best = self._best_param_match(param_info) + s = None + # print(f"param_info={param_info}") + if p_best is not None: + style_name = p_best.find_style(plot_type) + # print(f" -> style_name={style_name}") + s = self.styles.get(style_name, None) + else: + if scalar: + s = self.styles.get(self.SCALAR_DEFAULT_STYLE_NAME, None) + else: + s = self.styles.get(self.VECTOR_DEFAULT_STYLE_NAME, None) + + if s is not None: + # print(f"data_id={data_id}") + if data_id is not None: + s = s.set_data_id(data_id) + return s + else: + return s.clone() + return None + + def style(self, fs, plot_type="map", data_id=None): + param_info = fs.ds_param_info + if param_info is not None: + vd = self.get_param_style( + param_info, + scalar=param_info.scalar, + plot_type=plot_type, + data_id=data_id, + ) + # LOG.debug(f"vd={vd}") + return vd + return None + + def visdef(self, fs, plot_type="map", data_id=None): + vd = self.style(fs, plot_type=plot_type, data_id=data_id) + return vd.to_request() if vd is not None else None + + def style_list(self, fs, plot_type="map"): + param_info = fs.ds_param_info + if param_info is not None: + return self.get_param_style_list( + param_info, scalar=param_info.scalar, plot_type=plot_type + ) + return [] + + def units_scaler(self, fs, plot_type="map"): + param_info = fs.ds_param_info + if param_info is not None: + p = self._best_param_match(param_info) + if p is not None and ("all" in p.scaling or plot_type in p.scaling): + return Scaling.find_item( + ParamInfo._grib_get( + fs[0], ["units", "shortName"], single_value_as_list=False + ) + ) + + def _make_defaults(self): + d = { + self.SCALAR_DEFAULT_STYLE_NAME: "mcont", + self.VECTOR_DEFAULT_STYLE_NAME: "mwind", + self.DIFF_DEFAULT_STYLE_NAME: "mcont", + } + for name, verb in d.items(): + if name not in self.styles: + self.styles[name] = Style(name, Visdef(verb, {})) + assert self.SCALAR_DEFAULT_STYLE_NAME in self.styles + assert self.VECTOR_DEFAULT_STYLE_NAME in self.styles + assert self.DIFF_DEFAULT_STYLE_NAME in self.styles + + def _load(self, param_path, style_path): + if os.path.exists(style_path): + with open(style_path, "rt") as f: + c = yaml.safe_load(f) + self._load_styles(c) + if os.path.exists(param_path): + with open(param_path, "rt") as f: + c = yaml.safe_load(f) + self._load_params(c, param_path) + + def _load_styles(self, conf): + for name, d in conf.items(): + vd = [] + # print(f"name={name} d={d}") + if not isinstance(d, list): + d = [d] + + # print(f"name={name} d={d}") + # for mcoast the verb can be missing + if ( + len(d) == 1 + and isinstance(d[0], dict) + and (len(d[0]) > 1 or not list(d[0].keys())[0] in PARAM_VISDEF_VERBS) + ): + vd.append(Visdef("mcoast", d[0])) + else: + for v in d: + ((verb, params),) = v.items() + vd.append(Visdef(verb, params)) + self.styles[name] = Style(name, vd) + + def _load_params(self, conf, path): + # TODO: review defaults for maps + self._make_defaults() + + for d in conf: + assert isinstance(d, dict) + # print(f"d={d}") + p = ParamStyle(d, self) + for v in [p.style, p.xs_style, p.diff_style]: + # print(f"v={v}") + for s in v: + if not s in self.styles: + raise Exception( + f"{self} Invalid style={s} specified in {d}! File={path}" + ) + + self.params.append(p) + + def is_empty(self): + return len(self.styles) == 0 + + def __str__(self): + return self.__class__.__name__ + + def print(self): + pass + # print(f"{self} params=") + # for k, v in self.params.items(): + # print(v) + # print(f"{self} styles=") + # for k, v in self.styles.items(): + # print(v) + + +class GeoView: + def __init__(self, params, style): + self.params = copy.deepcopy(params) + for k in list(self.params.keys()): + if k.lower() == "coastlines": + self.params.pop("coastlines", None) + self.style = style + + def to_request(self): + v = copy.deepcopy(self.params) + if self.style is not None and self.style: + v["coastlines"] = self.style.to_request() + return mv.geoview(**v) + + def __str__(self): + t = f"{self.__class__.__name__}[params={self.params}, style={self.style}]" + return t + + +class MapConf: + items = [] + areas = [] + BUILTIN_AREAS = [ + "AFRICA", + "ANTARCTIC", + "ARCTIC", + "AUSTRALASIA", + "CENTRAL_AMERICA", + "CENTRAL_EUROPE", + "EAST_TROPIC", + "EASTERN_ASIA", + "EURASIA", + "EUROPE", + "GLOBAL", + "MIDDLE_EAST_AND_INDIA", + "NORTH_AMERICA", + "NORTH_ATLANTIC", + "NORTH_EAST_EUROPE", + "NORTH_POLE", + "NORTH_WEST_EUROPE", + "NORTHERN_AFRICA", + "PACIFIC", + "SOUTH_AMERICA", + "SOUTH_ATLANTIC_AND_INDIAN_OCEAN", + "SOUTH_EAST_ASIA_AND_INDONESIA", + "SOUTH_EAST_EUROPE", + "SOUTH_POLE", + "SOUTH_WEST_EUROPE", + "SOUTHERN_AFRICA", + "SOUTHERN_ASIA", + "WEST_TROPIC", + "WESTERN_ASIA", + ] + + def __init__(self): + self.areas = {} + self.style_db = get_db(name="map") + + # load areas + self._load_areas(os.path.join(ETC_PATH, "areas.yaml")) + self._load_custom_config() + self._load_local_config() + + def _load_custom_config(self): + if CUSTOM_CONF_PATH and CUSTOM_CONF_PATH[-1]: + self._load_areas(os.path.join(CUSTOM_CONF_PATH[-1], "areas.yaml")) + + def _load_local_config(self): + if LOCAL_CONF_PATH: + self._load_areas(os.path.join(LOCAL_CONF_PATH, "areas.yaml")) + + def _load_areas(self, file_path): + if os.path.exists(file_path): + with open(file_path, "rt") as f: + # the file can be empty! + d = yaml.safe_load(f) + if isinstance(d, list): + for item in d: + ((name, conf),) = item.items() + self.areas[name] = conf + + def area_names(self): + r = list(self.areas.keys()) + r.extend([a.lower() for a in self.BUILTIN_AREAS]) + return r + + def find(self, area=None, style=None): + area_v = "base" if area is None else area + style_v = "base" if style is None else style + s = None + if isinstance(area_v, list): + if len(area_v) == 4: + a = { + "area_mode": "user", + "map_projection": "cylindrical", + "map_area_definition": "corners", + "area": area_v, + } + else: + raise Exception( + "Invalid list specified for area. Required format: [S,W,N,E]" + ) + else: + a = self.areas.get(area_v, {}) + if len(a) == 0 and area_v.upper() in self.BUILTIN_AREAS: + a = {"area_mode": "name", "area_name": area} + if isinstance(style_v, mv.Request): + s = style_v + else: + s = self.style_db.get_style(style_v) + return a, s + + def make_geo_view(self, area=None, style=None, plot_type=None): + if style is None and plot_type == "diff": + style = "base_diff" + + a, s = self.find(area=area, style=style) + if a is not None and a: + a = copy.deepcopy(a) + else: + a = {} + + if s is not None: + if isinstance(s, mv.Request): + s = Visdef.from_request(s) + else: + s = s.clone() + if plot_type == "stamp": + s = s.update({"map_grid": "off", "map_label": "off"}) + if s: + a["coastlines"] = s.to_request() + + return mv.geoview(**a) + + +class StyleGallery: + def __init__(self): + pass + + def to_base64(self, image_path): + import base64 + + with open(image_path, "rb") as img_file: + return base64.b64encode(img_file.read()).decode("utf-8") + + def build_gallery(self, names, images, row_height="150px"): + figures = [] + # print(len(names)) + # print(len(images)) + for name, image in zip(names, images): + src = f"data:image/png;base64,{image}" + caption = f'
{name}
' + figures.append( + f""" +
+ + {caption} +
+ """ + ) + return f""" +
+ {''.join(figures)} +
+ """ + + def build(self): + if not is_ipython_active(): + return + + img_size = 120 + img, names = self._build(img_size) + + # reset jupyter output settings + mv.setoutput("jupyter", **mv.plot.jupyter_args) + + if len(img) > 0: + from IPython.display import HTML + + return HTML(self.build_gallery(names, img, row_height=f"{img_size}px")) + + +class MapStyleGallery(StyleGallery): + def _build(self, img_size): + img = [] + names = [] + # img_size = 120 + tmp_dir = os.path.join(os.getenv("METVIEW_TMPDIR", ""), "_mapstyle_") + Path(tmp_dir).mkdir(exist_ok=True) + for name, d in map_styles().items(): + f_name = os.path.join(tmp_dir, name + ".png") + if not os.path.exists(f_name): + view = mv.make_geoview(area=[30, -30, 75, 45], style=name) + view.update({"MAP_COASTLINE_RESOLUTION": "low"}, sub="coastlines") + mv.setoutput( + mv.png_output( + output_name=f_name[:-4], + output_width=300, + output_name_first_page_number="off", + ) + ) + mv.plot(view) + if os.path.exists(f_name): + names.append(name) + img.append(self.to_base64(f_name)) + + return (img, names) + + +class MapAreaGallery(StyleGallery): + def _build(self, img_size): + img = [] + names = [] + # img_size = 120 + tmp_dir = os.path.join(os.getenv("METVIEW_TMPDIR", ""), "_maparea_") + Path(tmp_dir).mkdir(exist_ok=True) + for name in map_area_names(): + f_name = os.path.join(tmp_dir, name + ".png") + if not os.path.exists(f_name): + view = mv.make_geoview(area=name, style="grey_1") + mv.setoutput( + mv.png_output( + output_name=f_name[:-4], + output_width=300, + output_name_first_page_number="off", + ) + ) + mv.plot(view) + if os.path.exists(f_name): + names.append(name) + img.append(self.to_base64(f_name)) + + return (img, names) + + +def MAP_CONF(): + global _MAP_CONF + if _MAP_CONF is None: + _MAP_CONF = MapConf() + assert not _MAP_CONF is None + return _MAP_CONF + + +def map_styles(): + return MAP_CONF().style_db.styles + + +def map_style_gallery(): + g = MapStyleGallery() + return g.build() + + +def map_area_names(): + return MAP_CONF().area_names() + + +def map_area_gallery(): + g = MapAreaGallery() + return g.build() + + +def make_geoview(**argv): + return MAP_CONF().make_geo_view(**argv) + + +def make_eccharts_mcont(): + return mv.mcont(contour_automatic_settings="ecmwf", legend="on") + + +def get_db(name="param"): + global _DB + assert name in _DB + if _DB[name][0] is None: + _DB[name][0] = StyleDb(_DB[name][1], _DB[name][2]) + return _DB[name][0] + + +def find(name): + db = get_db() + s = db.styles.get(name, None) + if s is not None: + return s.clone() + else: + return None + + +def load_custom_config(conf_dir, force=False): + global CUSTOM_CONF_PATH + global _MAP_CONF + if CUSTOM_CONF_PATH: + if CUSTOM_CONF_PATH[-1] == conf_dir: + if force: + CUSTOM_CONF_PATH.pop() + else: + return + + CUSTOM_CONF_PATH.append(conf_dir) + for k, v in _DB.items(): + if v[0] is not None: + v[0]._load_custom_config() + if _MAP_CONF is not None: + _MAP_CONF = None + _MAP_CONF = MapConf() + + +def reset_config(): + global CUSTOM_CONF_PATH + if CUSTOM_CONF_PATH: + CUSTOM_CONF_PATH = [] + global _DB + global _MAP_CONF + for name, v in _DB.items(): + if v[0] is not None: + v[0] = None + get_db(name=name) + if _MAP_CONF is not None: + _MAP_CONF = None + _MAP_CONF = MapConf() + + +if __name__ == "__main__": + vd = StyleConf() + vd.print() +else: + pass diff --git a/metview/title.py b/metview/title.py new file mode 100644 index 00000000..31dcbda3 --- /dev/null +++ b/metview/title.py @@ -0,0 +1,208 @@ +# +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import logging + +import metview as mv + +LOG = logging.getLogger(__name__) + +FC_TIME_PART = "Step: +h Valid: " +LEVEL_PART = "Lev: ()" + + +class Title: + def __init__(self, font_size=0.4): + self.font_size = font_size + + def _make_font_size(self, font_size): + if font_size is None: + return self.font_size + else: + return font_size + + def build(self, data, font_size=None): + font_size = self._make_font_size(font_size) + label = "" + if data: + if not isinstance(data, list): + data = [data] + + lines = [] + lines = {"text_line_count": len(data)} + for i, d_item in enumerate(data): + # print(f"d_item={d_item}") + if isinstance(d_item, tuple): + d = d_item[0] + data_id = d_item[1] + else: + d = d_item + data_id = None + + param = d.ds_param_info + if param is not None: + if param.meta.get("typeOfLevel", "") == "surface": + # lines.append(self.build_surface_fc(d.experiment.label, d.param.name, condition=cond)) + lines[f"text_line_{i+1}"] = self.build_surface_fc( + d.label, param.name, data_id=data_id + ) + else: + # lines.append(self.build_upper_fc(d.experiment.label, d.param.name, condition=cond)) + lines[f"text_line_{i+1}"] = self.build_upper_fc( + d.label, param.name, data_id=data_id + ) + + # print(f"line={lines}") + # return mv.mtext(text_lines=lines, text_font_size=font_size) + return mv.mtext(**lines, text_font_size=font_size) + + return mv.mtext( + { + "text_line_1": f"""{label} Par: Lev: () Step: +h Valid: """, + "text_font_size": font_size, + } + ) + + def _build_condition_str(self, condition): + if condition: + t = "where=" + for k, v in condition.items(): + t += f"'{k}={v}'" + return t + return str() + + def _add_grib_info(self, t, data_id): + if data_id is not None and data_id != "": + t = t.replace("h Valid: """, + "text_font_size": self.font_size, + } + ) + + return mv.mtext( + { + "text_font_size": self.font_size, + } + ) + + def build_rmse(self, ref, data): + if data: + if not isinstance(data, list): + data = [data] + + lines = [] + for d in data: + line = "RMSE" + if ref.label: + line += f"(ref={ref.label})" + if d.label: + line += f" {d.label}" + + # print(f"label={d.label}") + param = d.ds_param_info + if param is not None: + # print(f"meta={param.meta}") + line += f" Par: {param.name}" + meta = param.meta + if meta.get("mars.type", "") == "an": + pass + else: + r = meta.get("date", 0) + h = meta.get("time", 0) + line += f" Run: {r} {h} UTC" + lines.append(line) + + return mv.mtext(text_lines=lines, text_font_size=self.font_size) + + return mv.mtext(text_font_size=self.font_size) + + def build_cdf(self, data): + if data: + if not isinstance(data, list): + data = [data] + lines = [] + for d in data: + line = f"CDF {d.label}" + # print(f"label={d.label}") + param = d.ds_param_info + if param is not None: + # print(f"meta={param.meta}") + line += f" Par: {param.name}" + meta = param.meta + if meta.get("mars.type", "") == "an": + pass + else: + r = meta.get("date", 0) + h = meta.get("time", 0) + line += f" Run: {r} {h} UTC" + lines.append(line) + + return mv.mtext(text_lines=lines, text_font_size=self.font_size) + + return mv.mtext(text_font_size=self.font_size) diff --git a/metview/track.py b/metview/track.py new file mode 100644 index 00000000..33c92c2e --- /dev/null +++ b/metview/track.py @@ -0,0 +1,90 @@ +# +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +from datetime import datetime + +import metview as mv +import pandas as pd + +from metview.style import Style + + +class Track: + def __init__( + self, + path, + skiprows=None, + sep=None, + date_index=None, + time_index=None, + lat_index=None, + lon_index=None, + ): + self.path = path + self.skiprows = 0 if skiprows is None else skiprows + self.sep = sep + if self.sep == " ": + self.sep = r"\s+" + self.date_index = 0 if date_index is None else date_index + self.time_index = 1 if time_index is None else time_index + self.lat_index = 3 if lat_index is None else lat_index + self.lon_index = 2 if lon_index is None else lon_index + + def style(self): + return mv.style.get_db().get_style("track").clone() + + def build(self, style=None): + style = [] if style is None else style + df = pd.read_csv( + filepath_or_buffer=self.path, + sep=self.sep, + skiprows=self.skiprows, + header=None, + engine="python", + ) + + # print(df) + + v_date = df.iloc[:, self.date_index] + v_time = df.iloc[:, self.time_index] + val = [" {}/{:02d}".format(str(d)[-2:], t) for d, t in zip(v_date, v_time)] + lon = df.iloc[:, self.lon_index].values + lat = df.iloc[:, self.lat_index].values + idx_val = list(range(len(val))) + + # print(f"lon={lon}") + # print(f"lat={lat}") + # print(f"val={val}") + # for x in style: + # print(f"style={x}") + + if len(style) == 0: + s = mv.style.get_db().get_style("track").clone() + if len(style) == 1 and isinstance(style[0], Style): + s = style[0].clone() + else: + assert all(not isinstance(x, Style) for x in style) + s = mv.style.get_db().get_style("track").clone() + + for vd in s.visdefs: + if vd.verb == "msymb": + vd.change_symbol_text_list(val, idx_val) + + r = s.to_request() + + vis = mv.input_visualiser( + input_plot_type="geo_points", + input_longitude_values=lon, + input_latitude_values=lat, + input_values=idx_val, + ) + + return [vis, *r] diff --git a/metview/ui.py b/metview/ui.py new file mode 100644 index 00000000..d516c676 --- /dev/null +++ b/metview/ui.py @@ -0,0 +1,41 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +import metview as mv + +# this module is meant to expose some ui specific functions + + +def dialog(*args): + res = mv._dialog(*args) + return {k: v for k, v in res.items() if not k.startswith("_")} + + +def any(**kwargs): + return mv._any(**kwargs) + + +def colour(**kwargs): + return mv._colour(**kwargs) + + +def icon(**kwargs): + return mv._icon(**kwargs) + + +def option_menu(**kwargs): + return mv._option_menu(**kwargs) + + +def slider(**kwargs): + return mv._slider(**kwargs) + + +def toggle(**kwargs): + return mv._toggle(**kwargs) diff --git a/requirements/requirements-docs.txt b/requirements/requirements-docs.txt index c1a277c6..df6ff92f 100644 --- a/requirements/requirements-docs.txt +++ b/requirements/requirements-docs.txt @@ -6,25 +6,25 @@ # alabaster==0.7.10 # via sphinx -babel==2.5.3 +babel==2.9.1 # via sphinx -certifi==2018.1.18 +certifi==2024.7.4 # via requests -chardet==3.0.4 +charset-normalizer==3.3.2 # via requests docutils==0.14 # via sphinx -idna==2.6 +idna==3.7 # via requests imagesize==1.0.0 # via sphinx -jinja2==2.11.3 +jinja2==3.1.4 # via sphinx -markupsafe==1.0 +markupsafe==2.1.5 # via jinja2 packaging==17.1 # via sphinx -pygments>=2.7.4 +pygments==2.18.0 # via sphinx pyparsing==2.2.0 # via packaging @@ -32,7 +32,7 @@ pytest-runner==4.2 # via -r requirements-docs.in pytz==2018.3 # via babel -requests==2.21.0 +requests==2.32.2 # via sphinx six==1.11.0 # via @@ -44,7 +44,5 @@ sphinx==1.7.2 # via -r requirements-docs.in sphinxcontrib-websupport==1.0.1 # via sphinx -typing==3.7.4.3 - # via sphinx -urllib3==1.24.2 +urllib3==1.26.19 # via requests diff --git a/requirements/requirements-tests.txt b/requirements/requirements-tests.txt index bf44fec0..b5e5c96a 100644 --- a/requirements/requirements-tests.txt +++ b/requirements/requirements-tests.txt @@ -12,15 +12,15 @@ cfgrib==0.9.4.2 cftime==1.0.3.4 # via netcdf4 coverage==4.5.1 # via pytest-cov execnet==1.5.0 # via pytest-cache -future==0.17.1 # via cfgrib +future==0.18.3 # via cfgrib mccabe==0.6.1 # via pytest-mccabe more-itertools==4.2.0 # via pytest netcdf4==1.4.2 -numpy==1.14.3 +numpy==1.22 pandas==0.23.0 pep8==1.7.1 # via pytest-pep8 pluggy==0.6.0 # via pytest -py==1.10.0 # via pytest +py==1.11.0 # via pytest pycparser==2.18 # via cffi pyflakes==2.0.0 # via pytest-flakes pytest-cache==1.0 # via pytest-flakes, pytest-mccabe, pytest-pep8 diff --git a/setup.py b/setup.py index 06eb875d..ca6f62cb 100644 --- a/setup.py +++ b/setup.py @@ -53,8 +53,11 @@ def read(fname): "cffi", "numpy", "pandas", + "PyYAML", + "requests", ], tests_require=[ + "eccodes", "pytest", ], test_suite="tests", diff --git a/tests/MVPROFILEVIEW.png b/tests/MVPROFILEVIEW.png new file mode 100644 index 00000000..dd706d23 Binary files /dev/null and b/tests/MVPROFILEVIEW.png differ diff --git a/tests/args.py b/tests/args.py new file mode 100644 index 00000000..f8b30b93 --- /dev/null +++ b/tests/args.py @@ -0,0 +1,4 @@ +import metview as mv + +a = mv.arguments() +print(a) diff --git a/tests/daily_clims.grib b/tests/daily_clims.grib new file mode 100644 index 00000000..05e11023 Binary files /dev/null and b/tests/daily_clims.grib differ diff --git a/tests/ds/an/10u_sfc.grib b/tests/ds/an/10u_sfc.grib new file mode 100644 index 00000000..1f57df16 Binary files /dev/null and b/tests/ds/an/10u_sfc.grib differ diff --git a/tests/ds/an/10v_sfc.grib b/tests/ds/an/10v_sfc.grib new file mode 100644 index 00000000..f9c8b2ed Binary files /dev/null and b/tests/ds/an/10v_sfc.grib differ diff --git a/tests/ds/an/msl_sfc.grib b/tests/ds/an/msl_sfc.grib new file mode 100644 index 00000000..8c9e8683 Binary files /dev/null and b/tests/ds/an/msl_sfc.grib differ diff --git a/tests/ds/an/pv_pt.grib b/tests/ds/an/pv_pt.grib new file mode 100644 index 00000000..73361068 Binary files /dev/null and b/tests/ds/an/pv_pt.grib differ diff --git a/tests/ds/an/t_pl.grib b/tests/ds/an/t_pl.grib new file mode 100644 index 00000000..c2263aad Binary files /dev/null and b/tests/ds/an/t_pl.grib differ diff --git a/tests/ds/an/u_pl.grib b/tests/ds/an/u_pl.grib new file mode 100644 index 00000000..368fddac Binary files /dev/null and b/tests/ds/an/u_pl.grib differ diff --git a/tests/ds/an/v_pl.grib b/tests/ds/an/v_pl.grib new file mode 100644 index 00000000..120f0841 Binary files /dev/null and b/tests/ds/an/v_pl.grib differ diff --git a/tests/ds/an/w_pl.grib b/tests/ds/an/w_pl.grib new file mode 100644 index 00000000..1b1d6674 Binary files /dev/null and b/tests/ds/an/w_pl.grib differ diff --git a/tests/ds/oper/10u_sfc.grib b/tests/ds/oper/10u_sfc.grib new file mode 100644 index 00000000..c3b364f4 Binary files /dev/null and b/tests/ds/oper/10u_sfc.grib differ diff --git a/tests/ds/oper/10v_sfc.grib b/tests/ds/oper/10v_sfc.grib new file mode 100644 index 00000000..a0624499 Binary files /dev/null and b/tests/ds/oper/10v_sfc.grib differ diff --git a/tests/ds/oper/msl_sfc.grib b/tests/ds/oper/msl_sfc.grib new file mode 100644 index 00000000..4bf85828 Binary files /dev/null and b/tests/ds/oper/msl_sfc.grib differ diff --git a/tests/ds/oper/pv_pt.grib b/tests/ds/oper/pv_pt.grib new file mode 100644 index 00000000..47ac77b1 Binary files /dev/null and b/tests/ds/oper/pv_pt.grib differ diff --git a/tests/ds/oper/t_pl.grib b/tests/ds/oper/t_pl.grib new file mode 100644 index 00000000..9ab53a17 Binary files /dev/null and b/tests/ds/oper/t_pl.grib differ diff --git a/tests/ds/oper/tp_sfc.grib b/tests/ds/oper/tp_sfc.grib new file mode 100644 index 00000000..8c67cf71 Binary files /dev/null and b/tests/ds/oper/tp_sfc.grib differ diff --git a/tests/ds/oper/u_pl.grib b/tests/ds/oper/u_pl.grib new file mode 100644 index 00000000..da6de0ea Binary files /dev/null and b/tests/ds/oper/u_pl.grib differ diff --git a/tests/ds/oper/v_pl.grib b/tests/ds/oper/v_pl.grib new file mode 100644 index 00000000..54741ca3 Binary files /dev/null and b/tests/ds/oper/v_pl.grib differ diff --git a/tests/ds/oper/w_pl.grib b/tests/ds/oper/w_pl.grib new file mode 100644 index 00000000..bcdfd77e Binary files /dev/null and b/tests/ds/oper/w_pl.grib differ diff --git a/tests/monthly_avg.grib b/tests/monthly_avg.grib new file mode 100644 index 00000000..721ae4e2 Binary files /dev/null and b/tests/monthly_avg.grib differ diff --git a/tests/request.req b/tests/request.req index 45081447..2ca5ca7d 100644 --- a/tests/request.req +++ b/tests/request.req @@ -6,4 +6,4 @@ mcont, CONTOUR_HIGHLIGHT_FREQUENCY = 2, contour_level_selection_type = LEVEL_LIST, CONTOUR_LEVEL_LIST = -10/0/10, - contour_colour_list = 'RGB(0.5,0.2,0.8)'/'RGB(0.8,0.7,0.3)'/'RGB(0.4,0.8,0.3)' + contour_shade_colour_list = 'RGB(0.5,0.2,0.8)'/'RGB(0.8,0.7,0.3)'/'RGB(0.4,0.8,0.3)' diff --git a/tests/rgg_small_subarea_cellarea_ref.grib b/tests/rgg_small_subarea_cellarea_ref.grib new file mode 100644 index 00000000..f095ab50 Binary files /dev/null and b/tests/rgg_small_subarea_cellarea_ref.grib differ diff --git a/tests/sort/date.csv.gz b/tests/sort/date.csv.gz new file mode 100644 index 00000000..14e2c468 Binary files /dev/null and b/tests/sort/date.csv.gz differ diff --git a/tests/sort/date_level.csv.gz b/tests/sort/date_level.csv.gz new file mode 100644 index 00000000..3911d84d Binary files /dev/null and b/tests/sort/date_level.csv.gz differ diff --git a/tests/sort/default.csv.gz b/tests/sort/default.csv.gz new file mode 100644 index 00000000..675ca217 Binary files /dev/null and b/tests/sort/default.csv.gz differ diff --git a/tests/sort/default_desc.csv.gz b/tests/sort/default_desc.csv.gz new file mode 100644 index 00000000..8445c990 Binary files /dev/null and b/tests/sort/default_desc.csv.gz differ diff --git a/tests/sort/keys.csv.gz b/tests/sort/keys.csv.gz new file mode 100644 index 00000000..a6eb0b11 Binary files /dev/null and b/tests/sort/keys.csv.gz differ diff --git a/tests/sort/level.csv.gz b/tests/sort/level.csv.gz new file mode 100644 index 00000000..666ce0b5 Binary files /dev/null and b/tests/sort/level.csv.gz differ diff --git a/tests/sort/level_asc.csv.gz b/tests/sort/level_asc.csv.gz new file mode 100644 index 00000000..92e1d7e1 Binary files /dev/null and b/tests/sort/level_asc.csv.gz differ diff --git a/tests/sort/level_desc.csv.gz b/tests/sort/level_desc.csv.gz new file mode 100644 index 00000000..a8d3f587 Binary files /dev/null and b/tests/sort/level_desc.csv.gz differ diff --git a/tests/sort/level_units.csv.gz b/tests/sort/level_units.csv.gz new file mode 100644 index 00000000..10456663 Binary files /dev/null and b/tests/sort/level_units.csv.gz differ diff --git a/tests/sort/multi_asc.csv.gz b/tests/sort/multi_asc.csv.gz new file mode 100644 index 00000000..4f6f3ac2 Binary files /dev/null and b/tests/sort/multi_asc.csv.gz differ diff --git a/tests/sort/multi_desc.csv.gz b/tests/sort/multi_desc.csv.gz new file mode 100644 index 00000000..64de9065 Binary files /dev/null and b/tests/sort/multi_desc.csv.gz differ diff --git a/tests/sort/multi_mixed.csv.gz b/tests/sort/multi_mixed.csv.gz new file mode 100644 index 00000000..6965ce22 Binary files /dev/null and b/tests/sort/multi_mixed.csv.gz differ diff --git a/tests/sort/number.csv.gz b/tests/sort/number.csv.gz new file mode 100644 index 00000000..c3b5b560 Binary files /dev/null and b/tests/sort/number.csv.gz differ diff --git a/tests/sort/paramId.csv.gz b/tests/sort/paramId.csv.gz new file mode 100644 index 00000000..ce75074e Binary files /dev/null and b/tests/sort/paramId.csv.gz differ diff --git a/tests/sort/shortName.csv.gz b/tests/sort/shortName.csv.gz new file mode 100644 index 00000000..ead5206e Binary files /dev/null and b/tests/sort/shortName.csv.gz differ diff --git a/tests/sort/sort_data.grib b/tests/sort/sort_data.grib new file mode 100644 index 00000000..3dda39cd Binary files /dev/null and b/tests/sort/sort_data.grib differ diff --git a/tests/sort/step.csv.gz b/tests/sort/step.csv.gz new file mode 100644 index 00000000..b519a0eb Binary files /dev/null and b/tests/sort/step.csv.gz differ diff --git a/tests/sort/time.csv.gz b/tests/sort/time.csv.gz new file mode 100644 index 00000000..1e04281b Binary files /dev/null and b/tests/sort/time.csv.gz differ diff --git a/tests/sort/units.csv.gz b/tests/sort/units.csv.gz new file mode 100644 index 00000000..2bdf8ff7 Binary files /dev/null and b/tests/sort/units.csv.gz differ diff --git a/tests/t1000_LL_2x2.grb b/tests/t1000_LL_2x2.grb new file mode 100644 index 00000000..de391553 Binary files /dev/null and b/tests/t1000_LL_2x2.grb differ diff --git a/tests/t1000_LL_7x7.grb b/tests/t1000_LL_7x7.grb new file mode 100644 index 00000000..bbda7aba Binary files /dev/null and b/tests/t1000_LL_7x7.grb differ diff --git a/tests/t_with_missing.grib b/tests/t_with_missing.grib new file mode 100644 index 00000000..c56f9d9c Binary files /dev/null and b/tests/t_with_missing.grib differ diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 00000000..e9ea2221 --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,209 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import os +import shutil +import tempfile + +import metview as mv + + +PATH = os.path.dirname(__file__) +DS_DIR = "" + + +def build_dataset(): + # make conf file + global DS_DIR + DS_DIR = tempfile.mkdtemp() + # DS_DIR = "/var/folders/ng/g0zkhc2s42xbslpsywwp_26m0000gn/T/tmpa_i2t_y6" + # print(f"DS_DIR={DS_DIR}") + data_dir = os.path.join(DS_DIR, "data") + an_dir = os.path.join("__ROOTDIR__", "an") + oper_dir = os.path.join("__ROOTDIR__", "oper") + + ds_def = rf""" +experiments: + - an: + label: "an" + dir: {an_dir} + fname : re"[A-Za-z0-9]+_[A-z]+\.grib" + - oper: + label: "oper" + dir: {oper_dir} + fname : re"[A-Za-z0-9]+_[A-z]+\.grib" +""" + with open(os.path.join(DS_DIR, "data.yaml"), "w") as f: + f.write(ds_def) + + shutil.copytree(os.path.join(PATH, "ds"), os.path.join(DS_DIR, "data")) + conf_dir = os.path.join(DS_DIR, "conf") + if not os.path.exists(conf_dir): + os.mkdir(conf_dir) + + +def remove_dataset(): + global DS_DIR + if DS_DIR and os.path.exists(DS_DIR) and not DS_DIR in ["/", "."]: + shutil.rmtree(DS_DIR) + DS_DIR = "" + + +def test_dataset(): + build_dataset() + + ds = mv.load_dataset(DS_DIR) + assert list(ds.field_conf.keys()) == ["an", "oper"] + assert ds.name == os.path.basename(DS_DIR) + + # indexing + ds.scan() + index_dir = os.path.join(DS_DIR, "index") + assert os.path.exists(index_dir) + for comp in ["an", "oper"]: + for f in [ + "datafiles.yaml", + "scalar.csv.gz", + "wind10m.csv.gz", + "wind.csv.gz", + "wind3d.csv.gz", + ]: + assert os.path.exists(os.path.join(index_dir, comp, f)) + + # reload ds + ds = mv.load_dataset(DS_DIR) + + # Analysis + d = ds["an"].select( + dateTime=[mv.date("2016-09-25 00:00"), mv.date("2016-09-26 00:00")] + ) + + assert isinstance(d, mv.Fieldset) + assert set(ds["an"].blocks.keys()) == set(["scalar", "wind10m", "wind", "wind3d"]) + assert set(d._db.blocks.keys()) == set(["scalar", "wind10m", "wind", "wind3d"]) + + v = d["msl"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 2 + assert mv.grib_get(v, ["shortName", "date:l", "time:l"]) == [ + ["msl", 20160925, 0], + ["msl", 20160926, 0], + ] + + v = d["t500"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 2 + assert mv.grib_get(v, ["shortName", "date:l", "time:l", "level:l"]) == [ + ["t", 20160925, 0, 500], + ["t", 20160926, 0, 500], + ] + + v = d["pv320K"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 2 + assert mv.grib_get( + v, ["shortName", "date:l", "time:l", "level:l", "typeOfLevel"] + ) == [["pv", 20160925, 0, 320, "theta"], ["pv", 20160926, 0, 320, "theta"]] + + v = d["wind850"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 4 + assert mv.grib_get( + v, ["shortName", "date:l", "time:l", "level:l", "typeOfLevel"] + ) == [ + ["u", 20160925, 0, 850, "isobaricInhPa"], + ["v", 20160925, 0, 850, "isobaricInhPa"], + ["u", 20160926, 0, 850, "isobaricInhPa"], + ["v", 20160926, 0, 850, "isobaricInhPa"], + ] + + v = d["wind10m"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 4 + assert mv.grib_get(v, ["shortName", "date:l", "time:l", "typeOfLevel"]) == [ + ["10u", 20160925, 0, "surface"], + ["10v", 20160925, 0, "surface"], + ["10u", 20160926, 0, "surface"], + ["10v", 20160926, 0, "surface"], + ] + + v = d["wind"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 8 + assert mv.grib_get( + v, ["shortName", "date:l", "time:l", "level:l", "typeOfLevel"] + ) == [ + ["u", 20160925, 0, 500, "isobaricInhPa"], + ["v", 20160925, 0, 500, "isobaricInhPa"], + ["u", 20160925, 0, 850, "isobaricInhPa"], + ["v", 20160925, 0, 850, "isobaricInhPa"], + ["u", 20160926, 0, 500, "isobaricInhPa"], + ["v", 20160926, 0, 500, "isobaricInhPa"], + ["u", 20160926, 0, 850, "isobaricInhPa"], + ["v", 20160926, 0, 850, "isobaricInhPa"], + ] + + v = d["wind3d"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 12 + assert mv.grib_get( + v, ["shortName", "date:l", "time:l", "level:l", "typeOfLevel"] + ) == [ + ["u", 20160925, 0, 500, "isobaricInhPa"], + ["v", 20160925, 0, 500, "isobaricInhPa"], + ["w", 20160925, 0, 500, "isobaricInhPa"], + ["u", 20160925, 0, 850, "isobaricInhPa"], + ["v", 20160925, 0, 850, "isobaricInhPa"], + ["w", 20160925, 0, 850, "isobaricInhPa"], + ["u", 20160926, 0, 500, "isobaricInhPa"], + ["v", 20160926, 0, 500, "isobaricInhPa"], + ["w", 20160926, 0, 500, "isobaricInhPa"], + ["u", 20160926, 0, 850, "isobaricInhPa"], + ["v", 20160926, 0, 850, "isobaricInhPa"], + ["w", 20160926, 0, 850, "isobaricInhPa"], + ] + + # Oper + run = mv.date("2016-09-25 00:00") + d = ds["oper"].select(date=run.date(), time=run.time(), step=[120]) + + assert isinstance(d, mv.Fieldset) + assert set(ds["oper"].blocks.keys()) == set(["scalar", "wind10m", "wind", "wind3d"]) + assert set(d._db.blocks.keys()) == set(["scalar", "wind10m", "wind", "wind3d"]) + + v = d["msl"] + assert isinstance(v, mv.Fieldset) + assert len(v) == 1 + assert mv.grib_get(v, ["shortName", "date:l", "time:l", "step:l"]) == [ + ["msl", 20160925, 0, 120] + ] + + remove_dataset() + + +def test_dataset_create_template(): + + global DS_DIR + DS_DIR = tempfile.mkdtemp() + + mv.dataset.create_dataset_template(DS_DIR) + + for d in ["conf", "data", "index"]: + d_path = os.path.join(DS_DIR, d) + assert os.path.isdir(d_path) + + for f in ["params.yaml", "param_styles.yaml", "areas.yaml", "map_styles"]: + f_path = os.path.join(DS_DIR, "conf", f) + assert os.path.exists(d_path) + + f_path = os.path.join(DS_DIR, "data.yaml") + assert os.path.exists(f_path) + + remove_dataset() diff --git a/tests/test_indexer.py b/tests/test_indexer.py new file mode 100644 index 00000000..396d9bb6 --- /dev/null +++ b/tests/test_indexer.py @@ -0,0 +1,1482 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import copy +import datetime +import os +from black import assert_equivalent + +import numpy as np +import pandas as pd + + +if "METVIEW_PYTHON_ONLY" not in os.environ: + import metview as mv +else: + import metview.metviewpy as mv + +from metview.metviewpy.param import ParamInfo +from metview.metviewpy.indexer import GribIndexer + +PATH = os.path.dirname(__file__) + +DB_COLUMNS = copy.deepcopy(GribIndexer.DEFAULT_KEYS) +DB_COLUMNS["_msgIndex1"] = ("l", np.int64, False) +DB_COLUMNS_WIND2 = copy.deepcopy(DB_COLUMNS) +DB_COLUMNS_WIND2["_msgIndex2"] = ("l", np.int64, False) +DB_DEFAULT_COLUMN_NAMES = list(GribIndexer.DEFAULT_KEYS.keys()) + + +def file_in_testdir(filename): + return os.path.join(PATH, filename) + + +def build_index_db_dataframe(column_data, key_def=None): + c = {v: column_data[i] for i, v in enumerate(list(key_def.keys()))} + pd_types = {k: v[1] for k, v in key_def.items()} + return pd.DataFrame(c).astype(pd_types) + + +def file_in_sort_dir(filename): + return os.path.join(PATH, "sort", filename) + + +def build_metadata_dataframe(fs, keys): + val = fs.grib_get(keys, "key") + md = {k: v for k, v in zip(keys, val)} + return pd.DataFrame.from_dict(md) + + +def read_sort_meta_from_csv(name): + f_name = file_in_sort_dir(f"{name}.csv.gz") + return pd.read_csv(f_name, index_col=None, dtype=str) + + +def write_sort_meta_to_csv(name, md): + f_name = file_in_sort_dir(f"{name}.csv.gz") + md.to_csv(path_or_buf=f_name, header=True, index=False, compression="gzip") + + +def test_fieldset_select_single_file(): + f = mv.read(file_in_testdir("tuv_pl.grib")) + assert f._db is None + + # ------------------------ + # single resulting field + # ------------------------ + g = f.select(shortName="u", level=700) + assert len(g) == 1 + assert mv.grib_get(g, ["shortName", "level:l"]) == [["u", 700]] + g1 = f[7] + d = g - g1 + assert np.allclose(d.values(), np.zeros(len(d.values()))) + + # check index db contents + assert g._db is not None + assert "scalar" in g._db.blocks + assert len(g._db.blocks) == 1 + md = [ + ["u"], + [131], + [20180801], + [1200], + [0], + [700], + ["isobaricInhPa"], + ["0"], + ["0001"], + ["od"], + ["oper"], + ["an"], + [0], + ] + df_ref = build_index_db_dataframe(md, key_def=DB_COLUMNS) + # print(df_ref.dtypes) + # print(g._db.blocks) + df = g._db.blocks["scalar"] + # print(df.dtypes) + if not df.equals(df_ref): + print(df.compare(df_ref)) + assert False + + # ------------------------------------ + # single resulting field - paramId + # ------------------------------------ + g = f.select(paramId=131, level=700) + assert len(g) == 1 + assert mv.grib_get(g, ["paramId:l", "level:l"]) == [[131, 700]] + g1 = f[7] + d = g - g1 + assert np.allclose(d.values(), np.zeros(len(d.values()))) + + # check index db contents + assert g._db is not None + assert "scalar" in g._db.blocks + assert len(g._db.blocks) == 1 + md = [ + ["u"], + [131], + [20180801], + [1200], + [0], + [700], + ["isobaricInhPa"], + ["0"], + ["0001"], + ["od"], + ["oper"], + ["an"], + [0], + ] + df_ref = build_index_db_dataframe(md, key_def=DB_COLUMNS) + df = g._db.blocks["scalar"] + if not df.equals(df_ref): + print(df.compare(df_ref)) + assert False + + # ------------------------- + # multiple resulting fields + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + assert f._db is None + + g = f.select(shortName=["t", "u"], level=[700, 500]) + assert len(g) == 4 + assert mv.grib_get(g, ["shortName", "level:l"]) == [ + ["t", 700], + ["u", 700], + ["t", 500], + ["u", 500], + ] + + assert g._db is not None + assert len(g._db.blocks) == 1 + assert "scalar" in g._db.blocks + md = [ + ["t", "u", "t", "u"], + [130, 131, 130, 131], + [20180801] * 4, + [1200] * 4, + [0] * 4, + [700, 700, 500, 500], + ["isobaricInhPa"] * 4, + ["0"] * 4, + ["0001"] * 4, + ["od"] * 4, + ["oper"] * 4, + ["an"] * 4, + [0, 1, 2, 3], + ] + + df_ref = build_index_db_dataframe(md, key_def=DB_COLUMNS) + df = g._db.blocks["scalar"] + if not df.equals(df_ref): + print(df.compare(df_ref)) + assert False + + # ------------------------- + # empty result + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + g = f.select(shortName="w") + assert isinstance(g, mv.Fieldset) + assert len(g) == 0 + + # ------------------------- + # invalid key + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + g = f.select(INVALIDKEY="w") + assert isinstance(g, mv.Fieldset) + assert len(g) == 0 + + # ------------------------- + # str or int values + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + assert f._db is None + + g = f.select(shortName=["t"], level=["500", 700], marsType="an") + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "marsType"]) == [ + ["t", 700, "an"], + ["t", 500, "an"], + ] + + f = mv.read(file_in_testdir("t_time_series.grib")) + assert f._db is None + + g = f.select(shortName=["t"], step=[3, 6]) + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "step:l"]) == [ + ["t", 1000, 3], + ["t", 1000, 6], + ] + + g = f.select(shortName=["t"], step=["3", "06"]) + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "step:l"]) == [ + ["t", 1000, 3], + ["t", 1000, 6], + ] + + # ------------------------- + # repeated use + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + assert f._db is None + + g = f.select(shortName=["t"], level=[500, 700], marsType="an") + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "marsType"]) == [ + ["t", 700, "an"], + ["t", 500, "an"], + ] + + g = f.select(shortName=["t"], level=[500], marsType="an") + assert len(g) == 1 + assert mv.grib_get(g, ["shortName", "level:l", "marsType"]) == [ + ["t", 500, "an"], + ] + + # ------------------------- + # mars keys + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + assert f._db is None + + g = f.select(shortName=["t"], level=[500, 700], marsType="an") + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "marsType"]) == [ + ["t", 700, "an"], + ["t", 500, "an"], + ] + + g = f.select(shortName=["t"], level=[500, 700], type="an") + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "type"]) == [ + ["t", 700, "an"], + ["t", 500, "an"], + ] + # check the index db contents. "type" must be mapped to the "marsType" column of the + # db so no rescanning should happen. The db should only contain the default set of columns. + assert g._db is not None + assert "scalar" in g._db.blocks + assert len(g._db.blocks) == 1 + assert list(g._db.blocks["scalar"].keys())[:-1] == DB_DEFAULT_COLUMN_NAMES + + g = f.select(shortName=["t"], level=[500, 700], type="fc") + assert len(g) == 0 + + g = f.select({"shortName": "t", "level": [500, 700], "mars.type": "an"}) + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "mars.type"]) == [ + ["t", 700, "an"], + ["t", 500, "an"], + ] + + # ------------------------- + # custom keys + # ------------------------- + f = mv.read(file_in_testdir("tuv_pl.grib")) + assert f._db is None + + g = f.select(shortName=["t"], level=[500, 700], gridType="regular_ll") + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "gridType"]) == [ + ["t", 700, "regular_ll"], + ["t", 500, "regular_ll"], + ] + + g = f.select({"shortName": ["t"], "level": [500, 700], "mars.param:s": "130.128"}) + assert len(g) == 2 + assert mv.grib_get(g, ["shortName", "level:l", "mars.param"]) == [ + ["t", 700, "130.128"], + ["t", 500, "130.128"], + ] + + assert g._db is not None + assert "scalar" in g._db.blocks + assert len(g._db.blocks) == 1 + assert list(g._db.blocks["scalar"].keys())[:-1] == [ + *DB_DEFAULT_COLUMN_NAMES, + "gridType", + "mars.param:s", + ] + + +def test_fieldset_select_date(): + # date and time + f = mv.read(file_in_testdir("t_time_series.grib")) + assert f._db is None + + g = f.select(date="20201221", time="12", step="9") + assert len(g) == 2 + + ref_keys = ["shortName", "date", "time", "step"] + ref = [ + ["t", "20201221", "1200", "9"], + ["z", "20201221", "1200", "9"], + ] + + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date=20201221, time="1200", step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date=20201221, time="12:00", step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date=20201221, time=12, step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date="2020-12-21", time=1200, step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select( + date=datetime.datetime(2020, 12, 21), + time=datetime.time(hour=12, minute=0), + step=9, + ) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + # dataDate and dataTime + g = f.select(dataDate="20201221", dataTime="12", step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(dataDate="2020-12-21", dataTime="12:00", step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + # validityDate and validityTime + g = f.select(validityDate="20201221", validityTime="21") + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(validityDate="2020-12-21", validityTime="21:00") + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + # dateTime + g = f.select(dateTime="2020-12-21 12:00", step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + # dataDateTime + g = f.select(dataDateTime="2020-12-21 12:00", step=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + # validityDateTime + g = f.select(validityDateTime="2020-12-21 21:00") + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + # ------------------------------------ + # check multiple dates/times + # ------------------------------------ + + ref = [ + ["t", "20201221", "1200", "3"], + ["z", "20201221", "1200", "3"], + ["t", "20201221", "1200", "9"], + ["z", "20201221", "1200", "9"], + ] + + # date and time + g = f.select(date="2020-12-21", time=12, step=[3, 9]) + assert len(g) == 4 + assert mv.grib_get(g, ref_keys) == ref + + # dateTime + g = f.select(dateTime="2020-12-21 12:00", step=[3, 9]) + assert len(g) == 4 + assert mv.grib_get(g, ref_keys) == ref + + # validityDate and validityTime + g = f.select(validityDate="2020-12-21", validityTime=[15, 21]) + assert len(g) == 4 + assert mv.grib_get(g, ref_keys) == ref + + # validityDateTime + g = f.select(validityDateTime=["2020-12-21 15:00", "2020-12-21 21:00"]) + assert len(g) == 4 + assert mv.grib_get(g, ref_keys) == ref + + # ------------------------------------ + # check times with 1 digit hours + # ------------------------------------ + + # we create a new fieldset + f = mv.merge(f[0], mv.grib_set_long(f[2:4], ["time", 600])) + + ref = [ + ["t", "20201221", "0600", "3"], + ["z", "20201221", "0600", "3"], + ] + + g = f.select(date="20201221", time="6", step="3") + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date=20201221, time="06", step=3) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date=20201221, time="0600", step=3) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(date=20201221, time="06:00", step=3) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(validityDate="2020-12-21", validityTime=9) + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(validityDate="2020-12-21", validityTime="09") + assert len(g) == 2 + assert mv.grib_get(g, ref_keys) == ref + + g = f.select(validityDate="2020-12-21", validityTime=18) + assert len(g) == 0 + + # ------------------------------------ + # daily climatology dates (no year) + # ------------------------------------ + + f = mv.read(file_in_testdir("daily_clims.grib")) + + g = f.select(date="apr-01") + assert len(g) == 1 + assert int(mv.grib_get_long(g, "date")) == 401 + + g = f.select(date="Apr-02") + assert len(g) == 1 + assert int(mv.grib_get_long(g, "date")) == 402 + + g = f.select(date="402") + assert len(g) == 1 + assert int(mv.grib_get_long(g, "date")) == 402 + + g = f.select(date="0402") + assert len(g) == 1 + assert int(mv.grib_get_long(g, "date")) == 402 + + g = f.select(date=401) + assert len(g) == 1 + assert int(mv.grib_get_long(g, "date")) == 401 + + g = f.select(date=[401, 402]) + assert len(g) == 2 + assert [int(v) for v in mv.grib_get_long(g, "date")] == [402, 401] + + g = f.select(dataDate="apr-01") + assert len(g) == 1 + assert int(mv.grib_get_long(g, "dataDate")) == 401 + + +def test_fieldset_select_multi_file(): + f = mv.read(file_in_testdir("tuv_pl.grib")) + f.append(mv.read(file_in_testdir("ml_data.grib"))) + assert f._db is None + + # single resulting field + g = f.select(shortName="t", level=61) + # print(f._db.blocks) + assert len(g) == 1 + assert mv.grib_get(g, ["shortName", "level:l", "typeOfLevel"]) == [ + ["t", 61, "hybrid"] + ] + + g1 = f[34] + d = g - g1 + assert np.allclose(d.values(), np.zeros(len(d.values()))) + + assert g._db is not None + assert len(g._db.blocks) == 1 + assert "scalar" in g._db.blocks + md = [ + ["t"], + [130], + [20180111], + [1200], + [12], + [61], + ["hybrid"], + [None], + ["0001"], + ["od"], + ["oper"], + ["fc"], + [0], + ] + df_ref = build_index_db_dataframe(md, key_def=DB_COLUMNS) + df = g._db.blocks["scalar"] + + if not df.equals(df_ref): + print(df.compare(df_ref)) + assert False + + +def test_param_info(): + # no extra info + p = ParamInfo.build_from_name("2t") + assert p.name == "2t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "surface" + assert p.meta["level"] is None + + p = ParamInfo.build_from_name("msl") + assert p.name == "msl" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "surface" + assert p.meta["level"] is None + + p = ParamInfo.build_from_name("t500") + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 500 + + p = ParamInfo.build_from_name("t500hPa") + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 500 + + p = ParamInfo.build_from_name("t") + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "" + assert p.meta["level"] is None + + p = ParamInfo.build_from_name("t320K") + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "theta" + assert p.meta["level"] == 320 + + p = ParamInfo.build_from_name("t72ml") + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "hybrid" + assert p.meta["level"] == 72 + + p = ParamInfo.build_from_name("wind10m") + assert p.name == "wind10m" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "surface" + assert p.meta["level"] is None + + p = ParamInfo.build_from_name("wind100") + assert p.name == "wind" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 100 + + p = ParamInfo.build_from_name("wind700") + assert p.name == "wind" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 700 + + p = ParamInfo.build_from_name("wind") + assert p.name == "wind" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "" + assert p.meta["level"] == None + + p = ParamInfo.build_from_name("wind3d") + assert p.name == "wind3d" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "" + assert p.meta["level"] == None + + p = ParamInfo.build_from_name("wind3d500") + assert p.name == "wind3d" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 500 + + # exta info + param_level_types = { + "2t": ["surface"], + "msl": ["surface"], + "wind10m": ["surface"], + "t": ["isobaricInhPa", "theta"], + "wind": ["isobaricInhPa"], + } + + p = ParamInfo.build_from_name("2t", param_level_types=param_level_types) + assert p.name == "2t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "surface" + assert p.meta["level"] == None + + try: + p = ParamInfo.build_from_name("22t", param_level_types=param_level_types) + assert False + except: + pass + + # p = ParamInfo.build_from_name("t2", param_level_types=param_level_types) + # assert p.name == "2t" + # assert p.level_type == "surface" + # assert p.level is None + + p = ParamInfo.build_from_name("msl", param_level_types=param_level_types) + assert p.name == "msl" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "surface" + assert p.meta["level"] == None + + p = ParamInfo.build_from_name("t500", param_level_types=param_level_types) + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 500 + + p = ParamInfo.build_from_name("t500hPa", param_level_types=param_level_types) + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 500 + + p = ParamInfo.build_from_name("t", param_level_types=param_level_types) + assert p.name == "t" + assert p.scalar == True + assert "typeOfLevel" not in p.meta + assert "level" not in p.meta + + p = ParamInfo.build_from_name("t320K", param_level_types=param_level_types) + assert p.name == "t" + assert p.scalar == True + assert p.meta["typeOfLevel"] == "theta" + assert p.meta["level"] == 320 + + try: + p = ParamInfo.build_from_name("t72ml", param_level_types=param_level_types) + assert False + except: + pass + + p = ParamInfo.build_from_name("wind10m", param_level_types=param_level_types) + assert p.name == "wind10m" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "surface" + assert p.meta["level"] is None + + p = ParamInfo.build_from_name("wind100", param_level_types=param_level_types) + assert p.name == "wind" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 100 + + p = ParamInfo.build_from_name("wind700", param_level_types=param_level_types) + assert p.name == "wind" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 700 + + p = ParamInfo.build_from_name("wind", param_level_types=param_level_types) + assert p.name == "wind" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] is None + + try: + p = ParamInfo.build_from_name("wind3d", param_level_types=param_level_types) + assert False + except: + pass + + param_level_types["wind3d"] = ["isobaricInhPa"] + p = ParamInfo.build_from_name("wind3d", param_level_types=param_level_types) + assert p.name == "wind3d" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] is None + + p = ParamInfo.build_from_name("wind3d500", param_level_types=param_level_types) + assert p.name == "wind3d" + assert p.scalar == False + assert p.meta["typeOfLevel"] == "isobaricInhPa" + assert p.meta["level"] == 500 + + +def test_param_info_from_fs_single_file(): + if "METVIEW_PYTHON_ONLY" in os.environ: + return + + f = mv.read(file_in_testdir("tuv_pl.grib")) + g = f["u700"] + p = g.ds_param_info + assert len(g) == 1 + assert p.name == "u" + assert p.scalar == True + md = { + "shortName": "u", + "paramId": 131, + "date": 20180801, + "time": 1200, + "step": 0, + "level": 700, + "typeOfLevel": "isobaricInhPa", + "number": "0", + "experimentVersionNumber": "0001", + "marsClass": "od", + "marsStream": "oper", + "marsType": "an", + "_msgIndex1": 0, + } + assert md == p.meta + + g = f["wind500"] + p = g.ds_param_info + assert len(g) == 2 + assert p.name == "wind" + assert p.scalar == False + md = { + "shortName": "wind", + "paramId": 131, + "date": 20180801, + "time": 1200, + "step": 0, + "level": 500, + "typeOfLevel": "isobaricInhPa", + "number": "0", + "experimentVersionNumber": "0001", + "marsClass": "od", + "marsStream": "oper", + "marsType": "an", + "_msgIndex1": 0, + "_msgIndex2": 1, + } + assert md == p.meta + + # we lose the db + g = g + 0 + p = g.ds_param_info + assert len(g) == 2 + assert p.name == "wind" + assert p.scalar == False + md = { + "shortName": "u", + "paramId": 131, + "date": 20180801, + "time": 1200, + "step": 0, + "level": 500, + "typeOfLevel": "isobaricInhPa", + "number": "0", + "experimentVersionNumber": "0001", + "marsClass": "od", + "marsStream": "oper", + "marsType": "an", + } + assert md == p.meta + + g = f["t"] + p = g.ds_param_info + assert len(g) == 6 + assert p.name == "t" + assert p.scalar == True + md = { + "shortName": "t", + "paramId": 130, + "date": 20180801, + "time": 1200, + "step": 0, + "level": None, + "typeOfLevel": "isobaricInhPa", + "number": "0", + "experimentVersionNumber": "0001", + "marsClass": "od", + "marsStream": "oper", + "marsType": "an", + "_msgIndex1": 0, + } + assert md == p.meta + + # we lose the db + g = g + 0 + p = g.ds_param_info + assert len(g) == 6 + assert p.name == "t" + assert p.scalar == True + md = { + "shortName": "t", + "paramId": 130, + "date": 20180801, + "time": 1200, + "step": 0, + "level": 1000, + "typeOfLevel": "isobaricInhPa", + "number": "0", + "experimentVersionNumber": "0001", + "marsClass": "od", + "marsStream": "oper", + "marsType": "an", + } + assert md == p.meta + + +def test_fieldset_select_operator_single_file(): + f = mv.read(file_in_testdir("tuv_pl.grib")) + + g = f["u700"] + assert f._db is not None + assert g._db is not None + assert len(g) == 1 + assert mv.grib_get(g, ["shortName", "level:l"]) == [["u", 700]] + + g1 = f[7] + d = g - g1 + assert np.allclose(d.values(), np.zeros(len(d.values()))) + + g = f["t"] + assert len(g) == 6 + assert mv.grib_get(g, ["shortName", "level:l"]) == [ + ["t", 1000], + ["t", 850], + ["t", 700], + ["t", 500], + ["t", 400], + ["t", 300], + ] + + try: + g = f["w"] + assert False + except: + pass + + +def test_fieldset_select_operator_multi_file(): + f = mv.read(file_in_testdir("tuv_pl.grib")) + f.append(mv.read(file_in_testdir("ml_data.grib"))) + assert f._db is None + + # single resulting field + g = f["t61ml"] + assert f._db is not None + assert g._db is not None + assert len(g) == 1 + assert mv.grib_get(g, ["shortName", "level:l", "typeOfLevel"]) == [ + ["t", 61, "hybrid"] + ] + + g1 = f[34] + d = g - g1 + assert np.allclose(d.values(), np.zeros(len(d.values()))) + + +def test_indexer_dataframe_sort_value_with_key(): + + md = { + "paramId": [1, 2, 1, 2, 3], + "level": [925, 850, 925, 850, 850], + "step": [12, 110, 1, 3, 1], + "rest": ["1", "2", "aa", "b1", "1b"], + } + + md_ref = { + "paramId": [1, 1, 2, 2, 3], + "level": [925, 925, 850, 850, 850], + "step": [1, 12, 3, 110, 1], + "rest": ["aa", "1", "b1", "2", "1b"], + } + + df = pd.DataFrame(md) + df = GribIndexer._sort_dataframe(df) + df_ref = pd.DataFrame(md_ref) + + if not df.equals(df_ref): + print(df.compare(df_ref)) + assert False + + +def test_describe(): + + f = mv.read(file_in_testdir("tuv_pl.grib")) + + # full contents + df = f.describe(no_print=True) + + ref = { + "typeOfLevel": { + "t": "isobaricInhPa", + "u": "isobaricInhPa", + "v": "isobaricInhPa", + }, + "level": {"t": "300,400,...", "u": "300,400,...", "v": "300,400,..."}, + "date": {"t": 20180801, "u": 20180801, "v": 20180801}, + "time": {"t": 1200, "u": 1200, "v": 1200}, + "step": {"t": 0, "u": 0, "v": 0}, + "paramId": {"t": 130, "u": 131, "v": 132}, + "class": {"t": "od", "u": "od", "v": "od"}, + "stream": {"t": "oper", "u": "oper", "v": "oper"}, + "type": {"t": "an", "u": "an", "v": "an"}, + "experimentVersionNumber": {"t": "0001", "u": "0001", "v": "0001"}, + } + ref_full = ref + + assert ref == df.to_dict() + + df = f.describe() + assert ref == df.to_dict() + + # single param by shortName + df = f.describe("t", no_print=True) + + ref = { + "val": { + "shortName": "t", + "name": "Temperature", + "paramId": 130, + "units": "K", + "typeOfLevel": "isobaricInhPa", + "level": "300,400,500,700,850,1000", + "date": "20180801", + "time": "1200", + "step": "0", + "class": "od", + "stream": "oper", + "type": "an", + "experimentVersionNumber": "0001", + } + } + + assert ref["val"] == df["val"].to_dict() + + df = f.describe(param="t", no_print=True) + assert ref == df.to_dict() + df = f.describe("t") + assert ref == df.to_dict() + df = f.describe(param="t") + assert ref == df.to_dict() + + # single param by paramId + df = f.describe(130, no_print=True) + + ref = { + "val": { + "shortName": "t", + "name": "Temperature", + "paramId": 130, + "units": "K", + "typeOfLevel": "isobaricInhPa", + "level": "300,400,500,700,850,1000", + "date": "20180801", + "time": "1200", + "step": "0", + "class": "od", + "stream": "oper", + "type": "an", + "experimentVersionNumber": "0001", + } + } + + assert ref == df.to_dict() + + df = f.describe(param=130, no_print=True) + assert ref == df.to_dict() + df = f.describe(130) + assert ref == df.to_dict() + df = f.describe(param=130) + assert ref == df.to_dict() + + # append + g = f + 0 + df = g.describe(no_print=True) + assert ref_full == df.to_dict() + + g.append(f[0].grib_set_long(["level", 25])) + df = g.describe(no_print=True) + + ref = { + "typeOfLevel": { + "t": "isobaricInhPa", + "u": "isobaricInhPa", + "v": "isobaricInhPa", + }, + "level": {"t": "25,300,...", "u": "300,400,...", "v": "300,400,..."}, + "date": {"t": 20180801, "u": 20180801, "v": 20180801}, + "time": {"t": 1200, "u": 1200, "v": 1200}, + "step": {"t": 0, "u": 0, "v": 0}, + "paramId": {"t": 130, "u": 131, "v": 132}, + "class": {"t": "od", "u": "od", "v": "od"}, + "stream": {"t": "oper", "u": "oper", "v": "oper"}, + "type": {"t": "an", "u": "an", "v": "an"}, + "experimentVersionNumber": {"t": "0001", "u": "0001", "v": "0001"}, + } + + assert ref == df.to_dict() + + +def test_ls(): + + f = mv.read(file_in_testdir("tuv_pl.grib")) + + # default keys + df = f[:4].ls(no_print=True) + + ref = { + "centre": {0: "ecmf", 1: "ecmf", 2: "ecmf", 3: "ecmf"}, + "shortName": {0: "t", 1: "u", 2: "v", 3: "t"}, + "typeOfLevel": { + 0: "isobaricInhPa", + 1: "isobaricInhPa", + 2: "isobaricInhPa", + 3: "isobaricInhPa", + }, + "level": {0: 1000, 1: 1000, 2: 1000, 3: 850}, + "dataDate": {0: 20180801, 1: 20180801, 2: 20180801, 3: 20180801}, + "dataTime": {0: 1200, 1: 1200, 2: 1200, 3: 1200}, + "stepRange": {0: "0", 1: "0", 2: "0", 3: "0"}, + "dataType": {0: "an", 1: "an", 2: "an", 3: "an"}, + "gridType": { + 0: "regular_ll", + 1: "regular_ll", + 2: "regular_ll", + 3: "regular_ll", + }, + } + + assert ref == df.to_dict() + + # extra keys + df = f[:2].ls(extra_keys=["paramId"], no_print=True) + + ref = { + "centre": {0: "ecmf", 1: "ecmf"}, + "shortName": {0: "t", 1: "u"}, + "typeOfLevel": {0: "isobaricInhPa", 1: "isobaricInhPa"}, + "level": {0: 1000, 1: 1000}, + "dataDate": {0: 20180801, 1: 20180801}, + "dataTime": {0: 1200, 1: 1200}, + "stepRange": {0: "0", 1: "0"}, + "dataType": {0: "an", 1: "an"}, + "gridType": {0: "regular_ll", 1: "regular_ll"}, + "paramId": {0: 130, 1: 131}, + } + + assert ref == df.to_dict() + + # filter + df = f.ls(filter={"shortName": ["t", "v"], "level": 850}, no_print=True) + + ref = { + "centre": {3: "ecmf", 5: "ecmf"}, + "shortName": {3: "t", 5: "v"}, + "typeOfLevel": {3: "isobaricInhPa", 5: "isobaricInhPa"}, + "level": {3: 850, 5: 850}, + "dataDate": {3: 20180801, 5: 20180801}, + "dataTime": {3: 1200, 5: 1200}, + "stepRange": {3: "0", 5: "0"}, + "dataType": {3: "an", 5: "an"}, + "gridType": {3: "regular_ll", 5: "regular_ll"}, + } + + assert ref == df.to_dict() + + # append + g = f[:2] + df = g.ls(no_print=True) + + ref = { + "centre": {0: "ecmf", 1: "ecmf"}, + "shortName": {0: "t", 1: "u"}, + "typeOfLevel": { + 0: "isobaricInhPa", + 1: "isobaricInhPa", + }, + "level": {0: 1000, 1: 1000}, + "dataDate": {0: 20180801, 1: 20180801}, + "dataTime": {0: 1200, 1: 1200}, + "stepRange": {0: "0", 1: "0"}, + "dataType": {0: "an", 1: "an"}, + "gridType": { + 0: "regular_ll", + 1: "regular_ll", + }, + } + + assert ref == df.to_dict() + + g.append(f[2].grib_set_long(["level", 500])) + df = g.ls(no_print=True) + ref = { + "centre": {0: "ecmf", 1: "ecmf", 2: "ecmf"}, + "shortName": {0: "t", 1: "u", 2: "v"}, + "typeOfLevel": { + 0: "isobaricInhPa", + 1: "isobaricInhPa", + 2: "isobaricInhPa", + }, + "level": {0: 1000, 1: 1000, 2: 500}, + "dataDate": {0: 20180801, 1: 20180801, 2: 20180801}, + "dataTime": {0: 1200, 1: 1200, 2: 1200}, + "stepRange": {0: "0", 1: "0", 2: "0"}, + "dataType": {0: "an", 1: "an", 2: "an"}, + "gridType": { + 0: "regular_ll", + 1: "regular_ll", + 2: "regular_ll", + }, + } + + assert ref == df.to_dict() + + +def test_sort(): + + # In each message the message index (1 based) is encoded + # into latitudeOfLastGridPoint! + fs = mv.read(file_in_sort_dir("sort_data.grib")) + + default_sort_keys = ["date", "time", "step", "number", "level", "paramId"] + assert GribIndexer.DEFAULT_SORT_KEYS == default_sort_keys + + # Note: shortName, units and latitudeOfLastGridPoint are non-default sort keys! + keys = [ + "date", + "time", + "step", + "number", + "level", + "paramId", + "shortName", + "units", + "latitudeOfLastGridPoint", + ] + + # the reference csv files were generated like this: + # write_sort_meta_to_csv(f_name, md) + # + # the correctness of the reference was tested by generating md_ref using + # the GribIndexer and comparing it to md. E.g.: + # md_ori = build_metadata_dataframe(fs, keys) + # md_ref = GribIndexer._sort_dataframe(md_ori, columns=default_sort_keys) + # assert md.equals(md_ref) + + # default sorting + f_name = "default" + r = fs.sort() + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False + + # ----------------------------- + # default sorting direction + # ----------------------------- + + sort_keys = { + k: k + for k in [ + "date", + "time", + "step", + "number", + "level", + "paramId", + "shortName", + "units", + ] + } + sort_keys["date_level"] = ["date", "level"] + sort_keys["level_units"] = ["level", "units"] + sort_keys["keys"] = keys + + for f_name, key in sort_keys.items(): + r = fs.sort(key) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + # single key as list + key = f_name = "level" + r = fs.sort([key]) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + # ----------------------------- + # custom sorting direction + # ----------------------------- + + # default keys + + # ascending + f_name = "default" + r = fs.sort(ascending=True) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, "default ascending" + + # descending + f_name = "default_desc" + r = fs.sort(ascending=False) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, "default descending" + + # single key + key = "level" + + # ascending + f_name = f"{key}_asc" + r = fs.sort(key, "<") + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, ["<"]) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, ascending=True) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + # descending + f_name = f"{key}_desc" + r = fs.sort(key, ">") + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, [">"]) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, ascending=False) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + # multiple keys + key = ["level", "paramId", "date"] + + f_name = "multi_asc" + r = fs.sort(key, "<") + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, ascending=True) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + f_name = "multi_desc" + r = fs.sort(key, ">") + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, ascending=False) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + f_name = "multi_mixed" + r = fs.sort(key, ["<", ">", "<"]) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + md_ref = read_sort_meta_from_csv(f_name) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + r = fs.sort(key, ascending=[True, False, True]) + assert len(fs) == len(r) + md = build_metadata_dataframe(r, keys) + if not md.equals(md_ref): + print(md.compare(md_ref)) + assert False, f"key={key}" + + # invalid arguments + try: + r = fs.sort(key, [">", "<"]) + assert False + except ValueError: + pass + + try: + r = fs.sort(key, "1") + assert False + except ValueError: + pass + + try: + r = fs.sort(key, ascending=["True", "False"]) + assert False + except ValueError: + pass + + try: + r = fs.sort(key, ">", ascending="True") + assert False + except ValueError: + pass + + try: + r = fs.sort(key, ">", ascending=["True", "False"]) + assert False + except ValueError: + pass + + +def test_speed(): + # test with grib written with write() function + fs = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + + fs_u = fs.select(shortName="u") + fs_v = fs.select(shortName="v") + + # single field + u = fs_u[0] + v = fs_v[0] + r = mv.speed(u, v) + assert len(r) == 1 + np.testing.assert_allclose( + r.values(), np.sqrt(np.square(u.values()) + np.square(v.values())), rtol=1e-05 + ) + assert r.grib_get_long("paramId") == 10 + + # multi fields + u = fs_u[:2] + v = fs_v[:2] + r = mv.speed(u, v) + assert len(r) == 2 + for i in range(len(r)): + np.testing.assert_allclose( + r[i].values(), + np.sqrt(np.square(u[i].values()) + np.square(v[i].values())), + rtol=1e-05, + ) + assert r.grib_get_long("paramId") == [10, 10] + + # no arguments + u = fs_u + v = fs_v + r = fs["wind"].speed() + assert len(r) == 6 + for i in range(len(r)): + np.testing.assert_allclose( + r[i].values(), + np.sqrt(np.square(u[i].values()) + np.square(v[i].values())), + rtol=1e-05, + ) + assert r.grib_get_long("paramId") == [10] * 6 + + +def test_deacc(): + f = mv.Fieldset(path=os.path.join(PATH, "t_time_series.grib"))[:3] + + r = f.deacc() + assert len(r) == len(f) + assert r.grib_get_long("generatingProcessIdentifier") == [148] * len(r) + for i in range(len(f)): + v_ref = f[0].values() * 0 if i == 0 else f[i].values() - f[i - 1].values() + np.testing.assert_allclose(r[i].values(), v_ref, rtol=1e-03) + + r = f.deacc(key="") + assert len(r) == len(f) + assert r.grib_get_long("generatingProcessIdentifier") == [148] * len(r) + for i in range(len(f)): + v_ref = f[0].values() * 0 if i == 0 else f[i].values() - f[i - 1].values() + np.testing.assert_allclose(r[i].values(), v_ref, rtol=1e-03) + + r = f.deacc(skip_first=True) + assert len(r) == len(f) - 1 + assert r.grib_get_long("generatingProcessIdentifier") == [148] * len(r) + for i in range(len(r)): + v_ref = f[i + 1].values() - f[i].values() + np.testing.assert_allclose(r[i].values(), v_ref, rtol=1e-03) + + r = f.deacc(skip_first=True, mark_derived=True) + assert len(r) == len(f) - 1 + # assert r.grib_get_long("generatingProcessIdentifier") == [254] * len(r) + for i in range(len(r)): + v_ref = f[i + 1].values() - f[i].values() + np.testing.assert_allclose(r[i].values(), v_ref, rtol=1e-03) + + # use grouping by key + f = mv.Fieldset(path=os.path.join(PATH, "t_time_series.grib"))[:6] + + # only "step" is part of the default set of indexing keys! + keys = ["step", "startStep", "stepRange"] + for key in keys: + r = f.deacc(key=key) + assert len(r) == 6 + assert r.grib_get_long("generatingProcessIdentifier") == [148] * len(r) + + v_ref = f[0].values() * 0 + np.testing.assert_allclose(r[0].values(), v_ref, rtol=1e-03) + v_ref = f[1].values() * 0 + np.testing.assert_allclose(r[1].values(), v_ref, rtol=1e-03) + + steps = {2: (2, 0), 3: (3, 1), 4: (4, 2), 5: (5, 3)} + + for idx, steps in steps.items(): + v_ref = f[steps[0]].values() - f[steps[1]].values() + np.testing.assert_allclose(r[idx].values(), v_ref, rtol=1e-03) diff --git a/tests/test_metview.py b/tests/test_metview.py index d612b087..6609d975 100644 --- a/tests/test_metview.py +++ b/tests/test_metview.py @@ -14,10 +14,12 @@ import pickle import pytest import pandas as pd +from setuptools import find_namespace_packages import xarray as xr import metview as mv from metview import bindings +from metview.metviewpy import utils PATH = os.path.dirname(__file__) @@ -46,6 +48,28 @@ def file_in_testdir(filename): return os.path.join(PATH, filename) +def get_test_data(filename): + d_path = os.path.join(PATH, "data") + os.makedirs(d_path, mode=0o755, exist_ok=True) + f_path = os.path.join(d_path, filename) + if not os.path.exists(f_path): + URL = "https://get.ecmwf.int/repository/test-data/metview/tests" + utils.simple_download(url=f"{URL}/{filename}", target=f_path) + return f_path + + +def field_equal(f1, f2, **kwargs): + v1 = f1.values() + v2 = f2.values() + np.testing.assert_allclose(v1, v2, **kwargs) + + +def fieldset_equal(fs1, fs2, rtol=1e-7): + assert len(fs1) == len(fs2) + for i in range(len(fs1)): + field_equal(fs1[i], fs2[i], rtol=rtol, err_msg=f"i={i}") + + def test_version_info(): out = mv.version_info() assert "metview_version" in out @@ -53,12 +77,15 @@ def test_version_info(): def test_version_info_python(): out = mv.version_info() + print(out) assert "metview_python_version" in out assert isinstance(out["metview_python_version"], str) -def test_describe(): - mv.describe("type") +# describe() is now a method/function on Fieldsets, and calling it +# on a string does not work - we may want to re-add this in the future +# def test_describe(): +# mv.describe("type") def test_definitions(): @@ -101,6 +128,99 @@ def test_definitions_as_dict(): assert mcont_dict["CONTOUR_MIN_LEVEL"] == 4.9 +def test_definitions_as_args(): + mcont_def = mv.mcont( + legend=True, + contour_line_thickness=7, + contour_level_selection_type="level_list", + contour_level_list=[5, 7, 8.8, 9], + contour_hi_text="My high text", + CONTOUR_MIN_LEVEL=4.9, + ) + mcont_dict = mcont_def.copy() + assert mcont_dict["LEGEND"] == "ON" + assert mcont_def["CONTOUR_LINE_THICKNESS"] == 7 + assert mcont_def["CONTOUR_LEVEL_SELECTION_TYPE"] == "LEVEL_LIST" + assert mcont_def["CONTOUR_LEVEL_LIST"] == [5, 7, 8.8, 9] + assert mcont_def["CONTOUR_HI_TEXT"] == "My high text" + assert mcont_def["CONTOUR_MIN_LEVEL"] == 4.9 + + +def test_modify_request(): + c = mv.mcont( + { + "CONTOUR_LINE_COLOUR": "PURPLE", + "CONTOUR_LINE_THICKNESS": 3, + "CONTOUR_HIGHLIGHT": False, + } + ) + c["CONTOUR_linE_COLOUR"] = "GREEN" + c["CONTOUR_linE_style"] = "dash" + assert c["CONTOUR_LINE_THICKNESS"] == 3 + assert c["CONTOUR_LINE_COLOUR"] == "GREEN" + assert c["CONTOUR_LINE_STYLE"] == "DASH" + assert c["CONTour_LINE_STYLE"] == "DASH" + # for visual testing + # mv.setoutput(mv.png_output(output_name="ff")) + # a = mv.read(os.path.join(PATH, "test.grib")) + # mv.plot(a, c) + + +def test_modify_request_via_update(): + c = mv.mcont( + { + "CONTOUR_LINE_COLOUR": "PURPLE", + "CONTOUR_LINE_THICKNESS": 3, + "CONTOUR_HIGHLIGHT": False, + } + ) + new_params = {"CONTOUR_linE_COLOUR": "OliVE", "CONTOUR_linE_style": "DOT"} + c.update(new_params) + assert c["CONTOUR_LINE_THICKNESS"] == 3 + assert c["CONTOUR_LINE_COLOUR"] == "OliVE" + assert c["CONTOUR_LINE_STYLE"] == "DOT" + assert c["CONTour_LINE_STYLE"] == "DOT" + # for visual testing + # mv.setoutput(mv.png_output(output_name="gg")) + # a = mv.read(os.path.join(PATH, "test.grib")) + # mv.plot(a, c) + + +def test_modify_embedded_request_via_update(): + coastlines = mv.mcoast( + map_coastline_thickness=3, + map_coastline_land_shade="on", + map_coastline_land_shade_colour="cyan", + ) + + gv = mv.geoview( + map_area_definition="corners", + area=[17.74, -35.85, 81.57, 63.93], + coastlines=coastlines, + ) + + with pytest.raises(IndexError): + gv.update({"MAP_COASTLINE_land_SHADE_COLOUR": "yellow"}, sub="XXCOASTlines") + with pytest.raises(IndexError): + gv.update({"MAP_COASTLINE_land_SHADE_COLOUR": "yellow"}, sub=["a", 2]) + + gv.update({"MAP_COASTLINE_land_SHADE_COLOUR": "green"}, sub="COASTlines") + assert gv["map_area_definition"] == "CORNERS" + c = gv["COAStLINES"] + assert c["MAP_COASTLINE_LAND_SHADE_COLOUR"] == "green" + assert c["map_coastline_land_shade_colour"] == "green" + + # for visual testing + # mv.setoutput(mv.png_output(output_name="gg")) + # a = mv.read(os.path.join(PATH, "test.grib")) + # mv.plot(a, gv) + + +def test_read_png(): + p = mv.read(file_in_testdir("MVPROFILEVIEW.png")) + assert p.get_verb() == "PNG" + + def test_print(): mv.print("Start ", 7, 1, 3, " Finished!") mv.print(6, 2, " Middle ", 6) @@ -341,19 +461,19 @@ def test_division_fieldsets(): def test_power(): - raised_two = TEST_FIELDSET ** 2 + raised_two = TEST_FIELDSET**2 maximum = mv.maxvalue(raised_two) - assert np.isclose(maximum, MAX_VALUE ** 2) + assert np.isclose(maximum, MAX_VALUE**2) def test_power_reverse(): mask = TEST_FIELDSET > 290 FS_3_AND_4 = (mask * 3) + (1 - mask) * 4 - raised = 2 ** FS_3_AND_4 + raised = 2**FS_3_AND_4 minimum = mv.minvalue(raised) maximum = mv.maxvalue(raised) - assert np.isclose(minimum, 2 ** 3) - assert np.isclose(maximum, 2 ** 4) + assert np.isclose(minimum, 2**3) + assert np.isclose(maximum, 2**4) def test_mod(): @@ -397,24 +517,130 @@ def test_distance(): assert np.isclose(maximum, SEMI_EQUATOR) -def test_fieldset_and(): +def test_fieldset_and_fieldset(): t = (TEST_FIELDSET > 290) & (TEST_FIELDSET < 310) s = mv.accumulate(t) assert s == 43855 -def test_fieldset_or(): +def test_fieldset_and_bool(): + t = (TEST_FIELDSET > 290) & (True) + s = mv.accumulate(t) + assert s == 45423 + assert s == mv.accumulate(TEST_FIELDSET > 290) + t = (TEST_FIELDSET > 290) & (False) + s = mv.accumulate(t) + assert s == 0 + + +def test_bool_and_fieldset(): + t = (True) & (TEST_FIELDSET > 290) + s = mv.accumulate(t) + assert s == 45423 + assert s == mv.accumulate(TEST_FIELDSET > 290) + t = (False) & (TEST_FIELDSET > 290) + s = mv.accumulate(t) + assert s == 0 + + +def test_fieldset_or_fieldset(): t = (TEST_FIELDSET < 250) | (TEST_FIELDSET > 310) s = mv.accumulate(t) assert s == 12427 +def test_fieldset_or_bool(): + t = (TEST_FIELDSET < 250) | (True) + s = mv.accumulate(t) + assert s == 115680 + assert s == len(TEST_FIELDSET.values()) + t = (TEST_FIELDSET < 250) | (False) + s = mv.accumulate(t) + assert s == 10859 + assert s == mv.accumulate(TEST_FIELDSET < 250) + + +def test_bool_or_fieldset(): + t = (True) | (TEST_FIELDSET < 250) + s = mv.accumulate(t) + assert s == 115680 + assert s == len(TEST_FIELDSET.values()) + t = (False) | (TEST_FIELDSET < 250) + s = mv.accumulate(t) + assert s == 10859 + assert s == mv.accumulate(TEST_FIELDSET < 250) + + def test_fieldset_not(): t = ~(TEST_FIELDSET < 250) s = mv.accumulate(t) assert s == 104821 +def test_valid_date_nogrib(): + vd = mv.valid_date(base=mv.date("20201231"), step=33) + assert isinstance(vd, list) + assert vd[0] == datetime.datetime(2021, 1, 1, 9, 0, 0) + + vd = mv.valid_date(base=mv.date("20201231"), step=[11, 33]) + assert isinstance(vd, list) + assert vd == [ + datetime.datetime(2020, 12, 31, 11, 0, 0), + datetime.datetime(2021, 1, 1, 9, 0, 0), + ] + + vd = mv.valid_date(base=mv.date("2020-12-31 15:00"), step=[11, 34]) + assert isinstance(vd, list) + assert vd == [ + datetime.datetime(2021, 1, 1, 2, 0, 0), + datetime.datetime(2021, 1, 2, 1, 0, 0), + ] + + vd = mv.valid_date( + base=mv.date("2020-12-31 15:00"), + step=[11, 34], + step_units=datetime.timedelta(hours=1), + ) + assert isinstance(vd, list) + assert vd == [ + datetime.datetime(2021, 1, 1, 2, 0, 0), + datetime.datetime(2021, 1, 2, 1, 0, 0), + ] + + vd = mv.valid_date( + base=mv.date("2020-12-31 15:00"), + step=[1, 2], + step_units=datetime.timedelta(hours=10), + ) + assert isinstance(vd, list) + assert vd == [ + datetime.datetime(2021, 1, 1, 1, 0, 0), + datetime.datetime(2021, 1, 1, 11, 0, 0), + ] + + vd = mv.valid_date( + base=mv.date("2020-12-31 15:00"), + step=[11, 34], + step_units=datetime.timedelta(minutes=2), + ) + assert isinstance(vd, list) + assert vd == [ + datetime.datetime(2020, 12, 31, 15, 22, 0), + datetime.datetime(2020, 12, 31, 16, 8, 0), + ] + + vd = mv.valid_date( + base=mv.date("2020-12-31 15:00"), + step=[11, 34], + step_units=datetime.timedelta(minutes=2), + ) + assert isinstance(vd, list) + assert vd == [ + datetime.datetime(2020, 12, 31, 15, 22, 0), + datetime.datetime(2020, 12, 31, 16, 8, 0), + ] + + def test_valid_date_len_1(): vd = mv.valid_date(TEST_FIELDSET) assert isinstance(vd, datetime.datetime) @@ -729,6 +955,14 @@ def test_datainfo(): assert di3["proportion_missing"] == 0 +def test_grib_get_key_not_exist(): + a = mv.read(os.path.join(PATH, "tuv_pl.grib"))[0:2] + kv = a.grib_get(["silly"]) + assert kv == [[None], [None]] + with pytest.raises(Exception): + kv = a.grib_get_long(["silly"]) + + def test_empty_fieldset_contructor(): f = mv.Fieldset() assert type(f) == mv.Fieldset @@ -748,6 +982,83 @@ def test_fieldset_contructor_with_path(): assert ni0["index"] == 21597 +def test_fieldset_create_from_list_of_paths(): + paths = [os.path.join(PATH, "t_for_xs.grib"), os.path.join(PATH, "ml_data.grib")] + f = mv.Fieldset(path=paths) + assert f.count() == 42 + assert f[0:2].grib_get_long("level") == [1000, 850] + assert f[5:9].grib_get_long("level") == [300, 1, 1, 5] + assert f[40:42].grib_get_long("level") == [133, 137] + + +def test_fieldset_create_from_glob_path_single(): + f = mv.Fieldset(path=os.path.join(PATH, "test.g*ib")) + assert type(f) == mv.Fieldset + assert len(f) == 1 + ni = mv.nearest_gridpoint_info(f, 57.193, -2.360) + ni0 = ni[0] + assert np.isclose(ni0["latitude"], 57.0) + assert np.isclose(ni0["longitude"], 357.75) + assert np.isclose(ni0["distance"], 22.4505) + assert np.isclose(ni0["value"], 282.436) + assert ni0["index"] == 21597 + + +def test_fieldset_create_from_glob_path_multi(): + f = mv.Fieldset(path=os.path.join(PATH, "t_*.grib")) + assert type(f) == mv.Fieldset + assert len(f) == 17 + par_ref = [ + ["t", "1000"], + ["t", "850"], + ["t", "700"], + ["t", "500"], + ["t", "400"], + ["t", "300"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ] + assert par_ref == mv.grib_get(f, ["shortName", "level"]) + + +def test_fieldset_create_from_glob_paths(): + f = mv.Fieldset( + path=[os.path.join(PATH, "test.g*ib"), os.path.join(PATH, "t_*.grib")] + ) + assert type(f) == mv.Fieldset + assert len(f) == 18 + par_ref = [ + ["2t", "0"], + ["t", "1000"], + ["t", "850"], + ["t", "700"], + ["t", "500"], + ["t", "400"], + ["t", "300"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ] + assert par_ref == mv.grib_get(f, ["shortName", "level"]) + + def test_fieldset_append_from_empty(): f = mv.Fieldset() g = mv.read(os.path.join(PATH, "tuv_pl.grib")) @@ -851,6 +1162,13 @@ def test_fieldset_merge_4(): assert g3.grib_get_long("level") == [1000, 850, 700] +def test_fieldset_str(): + grib = mv.read(os.path.join(PATH, "t_for_xs.grib")) + assert str(grib) == "Fieldset (6 fields)" + assert str(grib[4]) == "Fieldset (1 field)" + assert str(mv.Fieldset()) == "Fieldset (0 fields)" + + def test_fieldset_pickling(): pickled_fname = file_in_testdir("pickled_fieldset.p") g = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) @@ -890,6 +1208,317 @@ def test_fieldset_pickling(): os.remove(pickled_fname) +def test_fieldset_basic_mean(): + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + m = mv.mean(alldata) # as function + assert len(m) == 1 + assert np.isclose(m.values()[0], 8619.0555) + assert np.isclose(m.values()[2], 8588.4003) + m = alldata.mean() # as method + assert len(m) == 1 + assert np.isclose(m.values()[0], 8619.0555) + assert np.isclose(m.values()[2], 8588.4003) + + +def test_fieldset_basic_mean_with_missing_vals(): + # replace first field with all missing values + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + f1vals = alldata[0].values() + f1vals[:] = np.nan + alldata[0] = alldata[0].set_values(f1vals) + m = mv.mean(alldata) # function + assert len(m) == 1 + assert np.isnan(m.values()[0]) + assert np.isnan(m.values()[2]) + m = mv.mean(alldata, missing=True) # function + assert len(m) == 1 + assert np.isclose(m.values()[0], 8644.7534) + assert np.isclose(m.values()[2], 8614.7797) + m = alldata.mean() # method + assert len(m) == 1 + assert np.isnan(m.values()[0]) + assert np.isnan(m.values()[2]) + m = alldata.mean(missing=True) # method + assert len(m) == 1 + assert np.isclose(m.values()[0], 8644.7534) + assert np.isclose(m.values()[2], 8614.7797) + + +def test_fieldset_mean_over_dim_number(): + # compute and check ensemble means + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + num_ens_members = len(mv.unique(alldata.grib_get_long("number"))) + assert num_ens_members == 6 + ens_mean = alldata.mean(dim="number") + # check general structure of the result + assert len(ens_mean) == len(alldata) / num_ens_members + assert mv.unique(ens_mean.grib_get_long("level")) == mv.unique( + alldata.grib_get_long("level") + ) + assert mv.unique(ens_mean.grib_get_long("step")) == mv.unique( + alldata.grib_get_long("step") + ) + assert mv.unique(ens_mean.grib_get_string("shortName")) == mv.unique( + alldata.grib_get_string("shortName") + ) + assert len(mv.unique(ens_mean.grib_get_long("number"))) == 1 + # check values for specific means #1 + mean1_computed = ens_mean.select(shortName="z", level=1000, step=3) + mean1_verified = alldata.select(shortName="z", level=1000, step=3).mean() + assert np.array_equal(mean1_computed.values(), mean1_verified.values()) + assert np.isclose(mean1_computed.values()[0], 1233.7) # via calculator + # check values for specific means #2 + mean2_computed = ens_mean.select(shortName="t", level=700, step=9) + mean2_verified = alldata.select(shortName="t", level=700, step=9).mean() + assert np.array_equal(mean2_computed.values(), mean2_verified.values()) + assert np.isclose(mean2_computed.values()[2], 276.208) # via calculator + # check values for specific means #3 + mean3_computed = ens_mean.select(shortName="u", level=500, step=6) + mean3_verified = alldata.select(shortName="u", level=500, step=6).mean() + assert np.array_equal(mean3_computed.values(), mean3_verified.values()) + + +def test_fieldset_mean_over_dim_step(): + # compute and check means over steps + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + num_steps = len(mv.unique(alldata.grib_get_long("step"))) + assert num_steps == 4 + step_mean = alldata.mean(dim="step") + # check general structure of the result + assert len(step_mean) == len(alldata) / num_steps + assert mv.unique(step_mean.grib_get_long("level")) == mv.unique( + alldata.grib_get_long("level") + ) + assert mv.unique(step_mean.grib_get_long("number")) == mv.unique( + alldata.grib_get_long("number") + ) + assert mv.unique(step_mean.grib_get_string("shortName")) == mv.unique( + alldata.grib_get_string("shortName") + ) + assert len(mv.unique(step_mean.grib_get_long("step"))) == 1 + # check values for specific means #1 + mean1_computed = step_mean.select(shortName="z", level=1000, number=3) + mean1_verified = alldata.select(shortName="z", level=1000, number=3).mean() + assert np.array_equal(mean1_computed.values(), mean1_verified.values()) + # check values for specific means #2 + mean2_computed = step_mean.select(shortName="t", level=850, number=5) + mean2_verified = alldata.select(shortName="t", level=850, number=5).mean() + assert np.array_equal(mean2_computed.values(), mean2_verified.values()) + assert np.isclose(mean2_computed.values()[1], 285.035) # via calculator + # check values for specific means #3 + mean3_computed = step_mean.select(shortName="u", level=500, number=1) + mean3_verified = alldata.select(shortName="u", level=500, number=1).mean() + assert np.array_equal(mean3_computed.values(), mean3_verified.values()) + + +def test_fieldset_mean_over_dim_step_with_expver(): + # compute and check step means when we have a dimension that's not + # in our default list + alldata = mv.read(file_in_testdir("zt_multi_expvers.grib")) + num_steps = len(mv.unique(alldata.grib_get_long("step"))) + assert num_steps == 4 + step_mean = alldata.mean( + dim="step", preserve_dims=["shortName", "experimentVersionNumber"] + ) + # check general structure of the result + assert len(step_mean) == len(alldata) / num_steps + assert mv.unique(step_mean.grib_get_long("level")) == mv.unique( + alldata.grib_get_long("level") + ) + assert mv.unique(step_mean.grib_get_long("number")) == mv.unique( + alldata.grib_get_long("number") + ) + assert mv.unique(step_mean.grib_get_string("shortName")) == mv.unique( + alldata.grib_get_string("shortName") + ) + assert len(mv.unique(step_mean.grib_get_long("step"))) == 1 + # check values for specific means #1 + mean1_computed = step_mean.select(shortName="z", experimentVersionNumber="0001") + mean1_verified = alldata.select( + shortName="z", experimentVersionNumber="0001" + ).mean() + assert np.array_equal(mean1_computed.values(), mean1_verified.values()) + # check values for specific means #2 + mean2_computed = step_mean.select(shortName="t", experimentVersionNumber="xxyz") + mean2_verified = alldata.select( + shortName="t", experimentVersionNumber="xxyz" + ).mean() + assert np.array_equal(mean2_computed.values(), mean2_verified.values()) + assert np.isclose(mean2_computed.values()[18], 312.2153) + + +def test_fieldset_mean_over_dim_with_empty_fieldset(): + alldata = mv.read(file_in_testdir("zt_multi_expvers.grib")) + with pytest.raises(ValueError): + alldata.select(shortName="t", level=850, number=5).mean() + + +def test_fieldset_mean_over_dim_number_with_holes(): + # compute and check ensemble means when there are holes in the hypercube + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + a = alldata.select(shortName=["z", "u"]) + t = alldata.select(shortName="t") + t = t.select(level=[850, 500]) # remove some levels from 't' + alldata = mv.Fieldset(fields=[a, t]) + + num_ens_members = len(mv.unique(alldata.grib_get_long("number"))) + assert num_ens_members == 6 + ens_mean = alldata.mean(dim="number") + # check general structure of the result + assert len(ens_mean) == len(alldata) / num_ens_members + assert mv.unique(ens_mean.grib_get_long("level")) == mv.unique( + alldata.grib_get_long("level") + ) + assert mv.unique(ens_mean.grib_get_long("step")) == mv.unique( + alldata.grib_get_long("step") + ) + assert mv.unique(ens_mean.grib_get_string("shortName")) == mv.unique( + alldata.grib_get_string("shortName") + ) + assert len(mv.unique(ens_mean.grib_get_long("number"))) == 1 + # check values for specific means #1 + mean1_computed = ens_mean.select(shortName="z", level=1000, step=3) + mean1_verified = alldata.select(shortName="z", level=1000, step=3).mean() + assert np.array_equal(mean1_computed.values(), mean1_verified.values()) + assert np.isclose(mean1_computed.values()[0], 1233.7) # via calculator + # check values for specific means #2 + mean2_computed = ens_mean.select(shortName="t", level=850, step=9) + mean2_verified = alldata.select(shortName="t", level=850, step=9).mean() + assert np.array_equal(mean2_computed.values(), mean2_verified.values()) + assert np.isclose(mean2_computed.values()[2], 285.548) # via calculator + # check values for specific means #3 + mean3_computed = ens_mean.select(shortName="u", level=500, step=6) + mean3_verified = alldata.select(shortName="u", level=500, step=6).mean() + assert np.array_equal(mean3_computed.values(), mean3_verified.values()) + # check that missing means are not present + mean_missing = ens_mean.select(shortName="t", level=1000) + assert len(mean_missing) == 0 + + +def test_fieldset_basic_sum(): + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + m = mv.sum(alldata) # as function + assert len(m) == 1 + assert np.isclose(m.values()[0], 2482287.9922) + assert np.isclose(m.values()[2], 2473459.2813) + m = alldata.sum() # as method + assert len(m) == 1 + assert np.isclose(m.values()[0], 2482287.9922) + assert np.isclose(m.values()[2], 2473459.2813) + + +def test_fieldset_basic_sum_with_missing_vals(): + # replace first field with all missing values + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + f1vals = alldata[0].values() + f1vals[:] = np.nan + alldata[0] = alldata[0].set_values(f1vals) + m = mv.sum(alldata) # function + print(m.values()[0], m.values()[2]) + assert len(m) == 1 + assert np.isnan(m.values()[0]) + assert np.isnan(m.values()[2]) + m = mv.sum(alldata, missing=True) # function + print(m.values()[0], m.values()[2]) + assert len(m) == 1 + assert np.isclose(m.values()[0], 2481044.2422) + assert np.isclose(m.values()[2], 2472441.8047) + m = alldata.sum() # method + print(m.values()[0], m.values()[2]) + assert len(m) == 1 + assert np.isnan(m.values()[0]) + assert np.isnan(m.values()[2]) + m = alldata.sum(missing=True) # method + print(m.values()[0], m.values()[2]) + assert len(m) == 1 + assert np.isclose(m.values()[0], 2481044.2422) + assert np.isclose(m.values()[2], 2472441.8047) + + +def test_fieldset_sum_over_dim_number(): + # compute and check ensemble sums + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + num_ens_members = len(mv.unique(alldata.grib_get_long("number"))) + assert num_ens_members == 6 + ens_sum = alldata.sum(dim="number") + # check general structure of the result + assert len(ens_sum) == len(alldata) / num_ens_members + assert mv.unique(ens_sum.grib_get_long("level")) == mv.unique( + alldata.grib_get_long("level") + ) + assert mv.unique(ens_sum.grib_get_long("step")) == mv.unique( + alldata.grib_get_long("step") + ) + assert mv.unique(ens_sum.grib_get_string("shortName")) == mv.unique( + alldata.grib_get_string("shortName") + ) + assert len(mv.unique(ens_sum.grib_get_long("number"))) == 1 + # check values for specific sums #1 + sum1_computed = ens_sum.select(shortName="z", level=1000, step=3) + sum1_verified = alldata.select(shortName="z", level=1000, step=3).sum() + assert np.array_equal(sum1_computed.values(), sum1_verified.values()) + # assert np.isclose(sum1_computed.values()[0], 1233.7) # via calculator + # check values for specific sums #2 + sum2_computed = ens_sum.select(shortName="t", level=700, step=9) + sum2_verified = alldata.select(shortName="t", level=700, step=9).sum() + assert np.array_equal(sum2_computed.values(), sum2_verified.values()) + # assert np.isclose(sum2_computed.values()[2], 276.208) # via calculator + # check values for specific sums #3 + sum3_computed = ens_sum.select(shortName="u", level=500, step=6) + sum3_verified = alldata.select(shortName="u", level=500, step=6).sum() + assert np.array_equal(sum3_computed.values(), sum3_verified.values()) + + +def test_fieldset_basic_stdev(): + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + z = alldata.select(shortName="z") + m = mv.stdev(z) # as function + # for testing + # val5 = [] + # for f in z: + # val5.append(f.values()[5]) + # print(np.std(h)) + assert len(m) == 1 + assert np.isclose(m.values()[5], 20948.418) + m = z.stdev() # as method + assert len(m) == 1 + assert np.isclose(m.values()[5], 20948.418) + + +def test_fieldset_stdev_over_dim_number(): + # compute and check ensemble stdevs + alldata = mv.read(file_in_testdir("ztu_multi_dim.grib")) + num_ens_members = len(mv.unique(alldata.grib_get_long("number"))) + assert num_ens_members == 6 + ens_stdev = alldata.stdev(dim="number") + # check general structure of the result + assert len(ens_stdev) == len(alldata) / num_ens_members + assert mv.unique(ens_stdev.grib_get_long("level")) == mv.unique( + alldata.grib_get_long("level") + ) + assert mv.unique(ens_stdev.grib_get_long("step")) == mv.unique( + alldata.grib_get_long("step") + ) + assert mv.unique(ens_stdev.grib_get_string("shortName")) == mv.unique( + alldata.grib_get_string("shortName") + ) + assert len(mv.unique(ens_stdev.grib_get_long("number"))) == 1 + # check values for specific stdevs #1 + stdev1_computed = ens_stdev.select(shortName="z", level=1000, step=3) + stdev1_verified = alldata.select(shortName="z", level=1000, step=3).stdev() + assert np.array_equal(stdev1_computed.values(), stdev1_verified.values()) + assert np.isclose(stdev1_computed.values()[0], 6.855, 0.0001) # via spreadsheet + # check values for specific stdevs #2 + stdev2_computed = ens_stdev.select(shortName="t", level=700, step=9) + stdev2_verified = alldata.select(shortName="t", level=700, step=9).stdev() + assert np.array_equal(stdev2_computed.values(), stdev2_verified.values()) + assert np.isclose(stdev2_computed.values()[0], 0.02276, 0.001) # via spreadsheet + # check values for specific stdevs #3 + stdev3_computed = ens_stdev.select(shortName="u", level=500, step=6) + stdev3_verified = alldata.select(shortName="u", level=500, step=6).stdev() + assert np.array_equal(stdev3_computed.values(), stdev3_verified.values()) + + def test_read_bufr(): bufr = mv.read(file_in_testdir("obs_3day.bufr")) assert mv.type(bufr) == "observations" @@ -1033,6 +1662,49 @@ def test_geopoints_nonequality_operator(): assert vdiff[15] == 1 +def _optional_check_create_geo_inline(): + # creates geopoints using numpy arrays and lists + # values are valid, np.nan and None + # - the ability to handle np.nan in lists is only available + # in Metview 5.15.0 and above + g = mv.create_geo( + type="ncols", + latitudes=np.array([4, 5, np.nan]), + longitudes=[2.3, np.nan, 6.5], + levels=850, # all rows will have 850 as their level + dates=[20180808, 20170707, 20160606], # dates as numbers + times=None, + stnids=["aberdeen", "aviemore", "edinburgh"], + elevations=np.array([np.nan, 14.1, np.nan]), + temp=[273.15, np.nan, 281.45], + precip=[None, np.nan, 1], + speed=np.array([2, 3, 5]), + ) + + def check_columns(gpt): + aeq = np.testing.assert_array_equal + aeq(mv.latitudes(gpt), np.array([4, 5, np.nan])) + aeq(mv.longitudes(gpt), np.array([2.3, np.nan, 6.5])) + aeq(mv.levels(gpt), np.array([850, 850, 850])) + aeq(mv.elevations(gpt), np.array([np.nan, 14.1, np.nan])) + aeq(gpt["temp"], np.array([273.15, np.nan, 281.45])) + aeq(gpt["precip"], np.array([np.nan, np.nan, 1])) + + check_columns(g) + temp_file = file_in_testdir("created_geo.gpts") + # check that it's written to disk ok + g.write(temp_file) + h = mv.read(temp_file) + check_columns(h) + os.remove(temp_file) + + +def test_create_geo_inline(): + version = mv.version_info() + if version["metview_version"] >= 51500.0: + _optional_check_create_geo_inline() + + def test_geopoints_set_dates(): g = mv.create_geo(3) @@ -1223,6 +1895,12 @@ def test_date_second(): assert mv.second(d1) == 0 +def test_date_string(): + npd1 = np.datetime64("2017-04-27T06:18:02") + assert mv.string(npd1, "yyyy-mm-dd") == "2017-04-27" + assert mv.string(npd1, "dd-mm-yyyy") == "27-04-2017" + + def test_numpy_numeric_args(): # we don't test them all here, but hopefully a representative sample assert mv.call("+", np.int32(5), 1) == 6 @@ -1588,21 +2266,6 @@ def test_temporary_file_deletion(file_name): assert not (os.path.isfile(temp_filepath)) -def test_mvl_ml2hPa(): - ml_data = mv.read(file_in_testdir("ml_data.grib")) - assert mv.type(ml_data) == "fieldset" - ml_t = mv.read(data=ml_data, param="t") - ml_lnsp = mv.read(data=ml_data, param="lnsp") - desired_pls = [1000, 900, 850, 500, 300, 100, 10, 1, 0.8, 0.5, 0.3, 0.1] - pl_data = mv.mvl_ml2hPa(ml_lnsp, ml_t, desired_pls) - assert mv.type(pl_data) == "fieldset" - pls = mv.grib_get_long(pl_data, "level") - lev_types = mv.grib_get_string(pl_data, "typeOfLevel") - lev_divisors = [1 if x == "isobaricInhPa" else 100 for x in lev_types] - pl_in_hpa = [a / b for a, b in zip(pls, lev_divisors)] - assert pl_in_hpa == desired_pls - - def test_push_nil(): n = mv.nil() assert n is None @@ -1662,6 +2325,14 @@ def test_get_vector_from_multi_field_grib(): assert v.shape == (6, 2664) +def test_2d_numpy_to_vector(): + np2d = np.array([[-1, -2, -3, -4], [-5, -6, -7, -8]]) + absnp = mv.abs(np2d) + # we expect the result to be flattened, but all the values used + assert absnp.shape == (8,) + assert np.array_equal(absnp, np.array([1, 2, 3, 4, 5, 6, 7, 8])) + + def test_vector_sort(): v1 = np.array([5, 3, 4, 9, 1, 4.2]) vsortasc = mv.sort(v1) @@ -2001,8 +2672,8 @@ def test_request(): r = mv.Request(d) assert isinstance(r, mv.Request) assert isinstance(r, dict) - r.set_verb("MSYMB") - assert r.get_verb() == "MSYMB" + r.set_verb("DVERB1") + assert r.get_verb() == "DVERB1" assert r["param1"] == "value1" assert r["param2"] == 180 r["param3"] = 108 @@ -2010,10 +2681,10 @@ def test_request(): # set verb in constructor d = {"param1": "value1", "param2": 180} - r = mv.Request(d, "MCOAST") + r = mv.Request(d, "DVERB2") assert isinstance(r, mv.Request) assert isinstance(r, dict) - assert r.get_verb() == "MCOAST" + assert r.get_verb() == "DVERB2" assert r["param1"] == "value1" assert r["param2"] == 180 r["param3"] = 108 @@ -2023,17 +2694,20 @@ def test_request(): r2 = mv.Request(r) assert isinstance(r2, mv.Request) assert isinstance(r2, dict) - assert r2.get_verb() == "MCOAST" + assert r2.get_verb() == "DVERB2" assert r2["param1"] == "value1" assert r2["param2"] == 180 r2["param3"] = 108 assert r2["param3"] == 108 + r2["param3"] = 179 + assert r2["param3"] == 179 + assert r["param3"] == 108 # original request is unchanged # destroy r and check that r2 is still ok r = 0 assert isinstance(r2, mv.Request) assert isinstance(r2, dict) - assert r2.get_verb() == "MCOAST" + assert r2.get_verb() == "DVERB2" assert r2["param1"] == "value1" assert r2["param2"] == 180 r2["param3"] = 108 @@ -2041,8 +2715,20 @@ def test_request(): # check the string representation d = {"param1": "value1", "param2": 180} - r = mv.Request(d, "MCOAST") - assert str(r) == "VERB: MCOAST{'param1': 'value1', 'param2': 180}" + r = mv.Request(d, "DVERB3") + rstr = str(r) + assert rstr == "VERB: DVERB3{'param1': 'value1', 'param2': 180}" + + # test the copying of an mcont + m1 = mv.mcont(contour_line_style="dot") + assert m1["contour_line_style"] == "DOT" + m2 = mv.Request(m1) + assert m1.val_pointer != m2.val_pointer + assert isinstance(m2, mv.Request) + assert m2["contour_line_style"] == "DOT" + m2["contour_line_style"] = "DASH" + assert m2["contour_line_style"] == "DASH" + assert m1["contour_line_style"] == "DOT" def test_read_request(): @@ -2069,9 +2755,9 @@ def test_read_request(): assert isinstance(req["CONTOUR_LEVEL_LIST"], list) assert isinstance(req["CONTOUR_LEVEL_LIST"][0], float) assert req["CONTOUR_LEVEL_LIST"] == [-10, 0, 10] - assert isinstance(req["contour_COLOUR_list"], list) - assert isinstance(req["contour_colour_list"][0], str) - assert req["CONTOUR_coLour_liSt"] == [ + assert isinstance(req["contour_sHade_COLOUR_list"], list) + assert isinstance(req["contour_shade_colour_list"][0], str) + assert req["CONTOUR_sHADe_coLour_liSt"] == [ "RGB(0.5,0.2,0.8)", "RGB(0.8,0.7,0.3)", "RGB(0.4,0.8,0.3)", @@ -2095,13 +2781,15 @@ def test_file(): def test_download(): - url = "http://download.ecmwf.org/test-data/metview/gallery/city_loc.gpt" + url = "http://get.ecmwf.int/test-data/metview/gallery/city_loc.gpt" g = mv.download(url=url) assert mv.type(g) == "geopoints" def test_download_gallery_data(): fname = "z_for_spectra.grib" + if os.path.exists(fname): + os.remove(fname) assert not os.path.isfile(fname) g = mv.gallery.load_dataset(fname) assert mv.type(g) == "fieldset" @@ -2112,7 +2800,11 @@ def test_download_gallery_data(): def test_download_gallery_zipped_data(): fname = "major_basins_of_the_world_0_0_0.zip" subname = "Major_Basins_of_the_World.prj" + if os.path.exists(fname): + os.remove(fname) assert not os.path.isfile(fname) + if os.path.exists(subname): + os.remove(subname) assert not os.path.isfile(subname) g = mv.gallery.load_dataset(fname) assert os.path.isfile(fname) @@ -2127,3 +2819,249 @@ def test_download_gallery_data_bad_fname(): fname = "zzz_for_spectra.grib" with pytest.raises(Exception): g = mv.gallery.load_dataset(fname) + + +def test_speed(): + # test with grib written with write() function + fs = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + + fs_u = fs.select(shortName="u") + fs_v = fs.select(shortName="v") + + # single field + u = fs_u[0] + v = fs_v[0] + r = mv.speed(u, v) + assert len(r) == 1 + np.testing.assert_allclose( + r.values(), np.sqrt(np.square(u.values()) + np.square(v.values())), rtol=1e-05 + ) + assert r.grib_get_long("paramId") == 10 + + # multi fields + u = fs_u[:2] + v = fs_v[:2] + r = mv.speed(u, v) + assert len(r) == 2 + for i in range(len(r)): + np.testing.assert_allclose( + r[i].values(), + np.sqrt(np.square(u[i].values()) + np.square(v[i].values())), + rtol=1e-05, + ) + assert r.grib_get_long("paramId") == [10, 10] + + +def test_mvl_ml2hpa(): + # 1. test against reference fields + + fs = mv.Fieldset(path=get_test_data("tq_ml137.grib")) + t = fs.select(shortName="t") + lnsp = fs.select(shortName="lnsp") + + # this is not the same order as in t_ref! + pres = [100, 500, 925, 850, 1000] + t_ref = mv.Fieldset(path=get_test_data("ml2pl_ref.grib")) + + r = mv.mvl_ml2hPa(lnsp, t, pres) + + levels = r.grib_get_long("level") + np.testing.assert_allclose(levels, pres) + + sh = r.grib_get_string("shortName") + assert sh == ["t"] * len(pres) + + for p in pres: + np.testing.assert_allclose( + r.select(level=p).values(), + t_ref.select(level=p).values(), + rtol=1e-05, + err_msg=f"p={p}", + ) + + # 2. test pressure values + + ml_data = mv.read(file_in_testdir("ml_data.grib")) + assert mv.type(ml_data) == "fieldset" + ml_t = mv.read(data=ml_data, param="t") + ml_lnsp = mv.read(data=ml_data, param="lnsp") + desired_pls = [1000, 900, 850, 500, 300, 100, 10, 1, 0.8, 0.5, 0.3, 0.1] + + pl_data = mv.mvl_ml2hPa(ml_lnsp, ml_t, desired_pls) + + assert mv.type(pl_data) == "fieldset" + pls = pl_data.grib_get_long("level") + lev_types = pl_data.grib_get_string("typeOfLevel") + lev_divisors = [1 if x == "isobaricInhPa" else 100 for x in lev_types] + pl_in_hpa = [a / b for a, b in zip(pls, lev_divisors)] + assert pl_in_hpa == desired_pls + + +def test_smooth_n_point(): + f = mv.Fieldset(path=get_test_data("regular_ll_iPos_jNeg_iCons_multi.grib")) + f_ref = mv.Fieldset(path=get_test_data("ref_smooth_n_point.grib")) + + meta_ref = f.grib_get_long("generatingProcessIdentifier") + + r = f.smooth_n_point(n=5) + fieldset_equal(r, f_ref[:2]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + r = f.smooth_n_point(n=5, repeat=2) + fieldset_equal(r, f_ref[2:4]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + r = f.smooth_n_point(n=9, repeat=2, mode="nearest") + fieldset_equal(r, f_ref[4:6]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + +def test_smooth_gaussian(): + f = mv.Fieldset(path=get_test_data("regular_ll_iPos_jNeg_iCons_multi.grib")) + f_ref = mv.Fieldset(path=get_test_data("ref_smooth_gaussian.grib")) + + meta_ref = f.grib_get_long("generatingProcessIdentifier") + + r = f.smooth_gaussian(sigma=1) + fieldset_equal(r, f_ref[:2]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + r = f.smooth_gaussian(sigma=2, repeat=2, mode="nearest") + fieldset_equal(r, f_ref[2:4]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + +def test_convolve(): + f = mv.Fieldset(path=get_test_data("regular_ll_iPos_jNeg_iCons_multi.grib")) + f_ref = mv.Fieldset(path=get_test_data("ref_convolve.grib")) + + meta_ref = f.grib_get_long("generatingProcessIdentifier") + + weights = np.array( + [[0.0625, 0.125, 0.0625], [0.125, 0.25, 0.125], [0.0625, 0.125, 0.0625]] + ) + + r = f.convolve(weights) + fieldset_equal(r, f_ref[:2]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + r = f.convolve(weights, repeat=2, mode="nearest") + fieldset_equal(r, f_ref[2:4]) + meta = r.grib_get_long("generatingProcessIdentifier") + np.testing.assert_allclose(meta, meta_ref) + + +# subprocess.run does not work well with Metview, so we use another method +# to run a Python process that uses Metview and returns the output +def _run_external_python_command(cmd, args): + import subprocess + import sys + import tempfile + + fd, path = tempfile.mkstemp() + + run_args = [sys.executable] # the Python executable + run_args.append(cmd) + run_args.extend(args) + + # create a clean environment to avoid clashes with the current session + env = {"PATH": os.environ["PATH"]} + mv_start_cmd = os.environ.get("METVIEW_PYTHON_START_CMD", None) + if mv_start_cmd is not None: + env.update({"METVIEW_PYTHON_START_CMD": mv_start_cmd}) + ppath = os.environ.get("PYTHONPATH", None) + if ppath is not None: + env.update({"PYTHONPATH": ppath}) + + with os.fdopen(fd, "w") as f: + subprocess.call(run_args, stdout=f, env=env) + f = open(path, "r") + a = f.read() + f.close() + return a + + +def test_arguments(): + output = _run_external_python_command( + file_in_testdir("args.py"), ["arg1", "arg 2", "59.35", "99.9000"] + ) + assert output == "['arg1', 'arg 2', 59.35, 99.9]\n" + + +@pytest.mark.parametrize( + "arg,ref", + [ + (["a", "b"], "ab"), + (["a", "b", "c"], "abc"), + ([None, "a"], "a"), + ([None, "a", "b"], "ab"), + ([np.array([1, 2, 3]), np.array([4, 5, 6])], np.array([1, 2, 3, 4, 5, 6])), + ([None, np.array([1, 2, 3])], np.array([1, 2, 3])), + ( + [None, np.array([1, 2, 3]), np.array([4, 5, 6])], + np.array([1, 2, 3, 4, 5, 6]), + ), + ([[1, 2], [3, 4]], [1, 2, 3, 4]), + ([[1, 2], 3], [1, 2, 3]), + ([[1, 2], "a"], [1, 2, "a"]), + ([None, [1, 2]], [1, 2]), + ([None, [1, 2], [3, 4]], [1, 2, 3, 4]), + ([None, 1], 1), + ], +) +def test_macro_python_compat_concat(arg, ref): + res = mv.compat.concat(*arg) + if isinstance(ref, np.ndarray): + np.allclose(res, ref) + else: + assert res == ref + + +def test_macro_python_compat_concat_fieldset(): + f1 = mv.Fieldset(path=file_in_testdir("t1000_LL_2x2.grb")) + f2 = mv.Fieldset(path=file_in_testdir("t1000_LL_7x7.grb")) + + res = mv.compat.concat(f1, f2) + assert isinstance(res, mv.Fieldset) + assert len(res) == len(f1) + len(f2) + assert res[0].grib_get_long("iDirectionIncrementInDegrees") == 2 + assert res[1].grib_get_long("iDirectionIncrementInDegrees") == 7 + + res = mv.compat.concat(None, f1) + assert isinstance(res, mv.Fieldset) + assert len(res) == len(f1) + assert res[0].grib_get_long("iDirectionIncrementInDegrees") == 2 + + +def test_macro_python_compat_concat_gpt(): + f = mv.read(file_in_testdir("xyv.gpt")) + + res = mv.compat.concat(None, f) + assert isinstance(res, mv.bindings.Geopoints) + assert len(res) == len(f) + + +def test_macro_python_compat_concat_bufr(): + f = mv.read(file_in_testdir("obs_3day.bufr")) + + res = mv.compat.concat(None, f) + assert isinstance(res, mv.bindings.Bufr) + + +@pytest.mark.parametrize( + "start,stop,step,num,ref", + [ + (0, 11, 4, 2, [0, 1, 4, 5, 8, 9]), + (0, 6, 1, 1, [0, 1, 2, 3, 4, 5]), + ], +) +def test_macro_python_compat_index4(start, stop, step, num, ref): + vals = np.array(list(range(12))) + res = mv.compat.index4(vals, start, stop, step, num) + np.allclose(vals[res], np.array(ref)) diff --git a/tests/test_pp_metview.py b/tests/test_pp_metview.py new file mode 100644 index 00000000..eeab5153 --- /dev/null +++ b/tests/test_pp_metview.py @@ -0,0 +1,1390 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +from inspect import ArgInfo +import numpy as np +import os +import pytest + +import metview.metviewpy as mv +from metview.metviewpy import utils +from metview.metviewpy.temporary import is_temp_file + +PATH = os.path.dirname(__file__) + + +def test_empty_fieldset_contructor(): + f = mv.Fieldset() + assert type(f) is mv.Fieldset + assert len(f) == 0 + + +def test_fieldset_contructor_bad_file_path(): + with pytest.raises(FileNotFoundError): + f = mv.Fieldset(path="does/not/exist") + + +def test_non_empty_fieldset_contructor_len(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + assert type(f) is mv.Fieldset + assert len(f) == 1 + + +def test_non_empty_fieldset_contructor_len_18(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + assert type(f) is mv.Fieldset + assert len(f) == 18 + + +def test_fieldset_create_from_list_of_paths(): + paths = [os.path.join(PATH, "t_for_xs.grib"), os.path.join(PATH, "ml_data.grib")] + f = mv.Fieldset(path=paths) + assert len(f) == 42 + assert f[0:2].grib_get_long("level") == [1000, 850] + assert f[5:9].grib_get_long("level") == [300, 1, 1, 5] + assert f[40:42].grib_get_long("level") == [133, 137] + + +def test_fieldset_create_from_glob_path_single(): + f = mv.Fieldset(path=os.path.join(PATH, "test.g*ib")) + assert type(f) == mv.Fieldset + assert len(f) == 1 + + +def test_fieldset_create_from_glob_path_multi(): + f = mv.Fieldset(path=os.path.join(PATH, "t_*.grib")) + assert type(f) == mv.Fieldset + assert len(f) == 17 + par_ref = [ + ["t", "1000"], + ["t", "850"], + ["t", "700"], + ["t", "500"], + ["t", "400"], + ["t", "300"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ] + assert par_ref == f.grib_get(["shortName", "level"]) + + +def test_fieldset_create_from_glob_paths(): + f = mv.Fieldset( + path=[os.path.join(PATH, "test.g*ib"), os.path.join(PATH, "t_*.grib")] + ) + assert type(f) == mv.Fieldset + assert len(f) == 18 + par_ref = [ + ["2t", "0"], + ["t", "1000"], + ["t", "850"], + ["t", "700"], + ["t", "500"], + ["t", "400"], + ["t", "300"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ["z", "1000"], + ["t", "1000"], + ] + assert par_ref == f.grib_get(["shortName", "level"]) + + +def test_read_1(): + f = mv.read(os.path.join(PATH, "test.grib")) + assert type(f) is mv.Fieldset + assert len(f) == 1 + + +def test_grib_get_string_1(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + sn = f.grib_get_string("shortName") + assert sn == "2t" + + +def test_grib_get_string_18(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + sn = f.grib_get_string("shortName") + assert sn == ["t", "u", "v"] * 6 + + +def test_grib_get_long_1(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + lev = f.grib_get_long("level") + assert lev == 0 + + +def test_grib_get_long_18(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + lev = f.grib_get_long("level") + assert lev == ([1000] * 3) + ([850] * 3) + ([700] * 3) + ([500] * 3) + ( + [400] * 3 + ) + ([300] * 3) + + +def test_grib_get_double_1(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + m = f.grib_get_double("max") + assert np.isclose(m, 316.061) + + +def test_grib_get_double_18(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + m = f.grib_get_double("max") + ref_m = [ + 320.564, + 21.7131, + 19.8335, + 304.539, + 43.1016, + 28.661, + 295.265, + 44.1455, + 31.6385, + 275.843, + 52.74, + 47.0099, + 264.003, + 62.2138, + 55.9496, + 250.653, + 66.4555, + 68.9203, + ] + np.testing.assert_allclose(m, ref_m, 0.001) + + +def test_grib_get_long_array_1(): + f = mv.Fieldset(path=os.path.join(PATH, "rgg_small_subarea_cellarea_ref.grib")) + pl = f.grib_get_long_array("pl") + assert isinstance(pl, np.ndarray) + assert len(pl) == 73 + assert pl[0] == 24 + assert pl[1] == 28 + assert pl[20] == 104 + assert pl[72] == 312 + + +def test_grib_get_double_array_1(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + v = f.grib_get_double_array("values") + assert isinstance(v, np.ndarray) + assert len(v) == 115680 + assert np.isclose(v[0], 260.4356) + assert np.isclose(v[24226], 276.1856) + assert np.isclose(v[36169], 287.9356) + assert np.isclose(v[115679], 227.1856) + + +def test_grib_get_double_array_18(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + v = f.grib_get_double_array("values") + assert isinstance(v, list) + assert len(v) == 18 + assert isinstance(v[0], np.ndarray) + assert isinstance(v[17], np.ndarray) + assert len(v[0]) == 2664 + assert len(v[17]) == 2664 + eps = 0.001 + assert np.isclose(v[0][0], 272.5642, eps) + assert np.isclose(v[0][1088], 304.5642, eps) + assert np.isclose(v[17][0], -3.0797, eps) + assert np.isclose(v[17][2663], -11.0797, eps) + + +def test_grib_get_generic(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:4] + sn = f.grib_get(["shortName"]) + assert sn == [["t"], ["u"], ["v"], ["t"]] + cs = f.grib_get(["centre:s"]) + assert cs == [["ecmf"], ["ecmf"], ["ecmf"], ["ecmf"]] + cl = f.grib_get(["centre:l"]) + assert cl == [[98], [98], [98], [98]] + lg = f.grib_get(["level:d", "cfVarName"]) + assert lg == [[1000, "t"], [1000, "u"], [1000, "v"], [850, "t"]] + lgk = f.grib_get(["level:d", "cfVarName"], "key") + assert lgk == [[1000, 1000, 1000, 850], ["t", "u", "v", "t"]] + with pytest.raises(ValueError): + lgk = f.grib_get(["level:d", "cfVarName"], "silly") + ln = f.grib_get(["level:n"]) + assert ln == [[1000], [1000], [1000], [850]] + cn = f.grib_get(["centre:n"]) + assert cn == [["ecmf"], ["ecmf"], ["ecmf"], ["ecmf"]] + vn = f[0].grib_get(["longitudes:n"]) + assert vn[0][0][0] == 0 + assert vn[0][0][1] == 5 + assert vn[0][0][5] == 25 + + +def test_grib_get_generic_key_not_exist(): + a = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:2] + kv = a.grib_get(["silly"]) + assert kv == [[None], [None]] + with pytest.raises(Exception): + kv = a.grib_get_long(["silly"]) + + +def test_values_1(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + v = f.values() + assert isinstance(v, np.ndarray) + assert len(v) == 115680 + assert np.isclose(v[0], 260.4356) + assert np.isclose(v[24226], 276.1856) + assert np.isclose(v[36169], 287.9356) + assert np.isclose(v[115679], 227.1856) + + +def test_values_18(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + v = f.values() + assert isinstance(v, np.ndarray) + assert v.shape == (18, 2664) + assert isinstance(v[0], np.ndarray) + assert isinstance(v[17], np.ndarray) + assert len(v[0]) == 2664 + assert len(v[17]) == 2664 + eps = 0.001 + assert np.isclose(v[0][0], 272.5642, eps) + assert np.isclose(v[0][1088], 304.5642, eps) + assert np.isclose(v[17][0], -3.0797, eps) + assert np.isclose(v[17][2663], -11.0797, eps) + + +def test_values_with_missing(): + f = mv.Fieldset(path=os.path.join(PATH, "t_with_missing.grib")) + v = f.values() + assert isinstance(v, np.ndarray) + assert v.shape == (2664,) + eps = 0.001 + assert np.isclose(v[0], 272.5642, eps) + assert np.isnan(v[798]) + assert np.isnan(v[806]) + assert np.isnan(v[1447]) + assert np.isclose(v[2663], 240.5642, eps) + + +def test_grib_set_string(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:2] + g = f.grib_set_string(["pressureUnits", "silly"]) + assert g.grib_get_string("pressureUnits") == ["silly"] * 2 + assert f.grib_get_string("pressureUnits") == ["hPa"] * 2 + g = f.grib_set_string(["pressureUnits", "silly", "shortName", "q"]) + assert g.grib_get_string("pressureUnits") == ["silly"] * 2 + assert g.grib_get_string("shortName") == ["q", "q"] + assert f.grib_get_string("pressureUnits") == ["hPa"] * 2 + assert f.grib_get_string("shortName") == ["t", "u"] + + +def test_grib_set_long(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:2] + g = f.grib_set_long(["level", 95]) + assert g.grib_get_long("level") == [95] * 2 + assert f.grib_get_long("level") == [1000] * 2 + g = f.grib_set_long(["level", 95, "time", 1800]) + assert g.grib_get_long("level") == [95] * 2 + assert g.grib_get_long("time") == [1800] * 2 + assert f.grib_get_long("level") == [1000] * 2 + assert f.grib_get_long("time") == [1200] * 2 + + +def test_grib_set_double(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:2] + g = f.grib_set_double(["level", 95]) + assert g.grib_get_double("level") == [95] * 2 + assert f.grib_get_double("level") == [1000] * 2 + orig_point = f.grib_get_double("longitudeOfFirstGridPointInDegrees") + g = f.grib_set_double(["longitudeOfFirstGridPointInDegrees", 95.6]) + assert g.grib_get_double("longitudeOfFirstGridPointInDegrees") == [95.6] * 2 + assert f.grib_get_double("longitudeOfFirstGridPointInDegrees") == orig_point + + +def test_grib_set_generic(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:2] + g = f.grib_set(["shortName", "r"]) + assert g.grib_get_string("shortName") == ["r"] * 2 + assert f.grib_get_string("shortName") == ["t", "u"] + g = f.grib_set(["shortName:s", "q"]) + assert g.grib_get_string("shortName") == ["q"] * 2 + assert f.grib_get_string("shortName") == ["t", "u"] + + g = f.grib_set(["level:l", 500, "shortName", "z"]) + assert g.grib_get_long("level") == [500] * 2 + assert g.grib_get_string("shortName") == ["z"] * 2 + assert f.grib_get_long("level") == [1000] * 2 + assert f.grib_get_string("shortName") == ["t", "u"] + + g = f.grib_set(["level:d", 500]) + np.testing.assert_allclose( + np.array(g.grib_get_double("level")), np.array([500] * 2) + ) + np.testing.assert_allclose( + np.array(f.grib_get_double("level")), np.array([1000] * 2) + ) + + g = f.grib_set_double(["longitudeOfFirstGridPointInDegrees", 95.6]) + np.testing.assert_allclose( + np.array(g.grib_get_double("longitudeOfFirstGridPointInDegrees")), + np.array([95.6] * 2), + ) + np.testing.assert_allclose( + np.array(f.grib_get_double("longitudeOfFirstGridPointInDegrees")), [0, 0] + ) + + +def test_write_fieldset(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + temp_path = "written_tuv_pl.grib" + f.write(temp_path) + assert os.path.isfile(temp_path) + g = mv.Fieldset(path=temp_path) + assert type(g) == mv.Fieldset + assert len(g) == 18 + sn = g.grib_get_string("shortName") + assert sn == ["t", "u", "v"] * 6 + f = 0 + os.remove(temp_path) + + +def test_write_modified_fieldset_binop(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + fp20 = f + 20 + temp_path = "written_tuv_pl_binop.grib" + fp20.write(temp_path) + assert os.path.isfile(temp_path) + g = mv.Fieldset(path=temp_path) + assert type(g) == mv.Fieldset + assert len(g) == 18 + sn = g.grib_get_string("shortName") + assert sn == ["t", "u", "v"] * 6 + np.testing.assert_allclose(g.values(), f.values() + 20) + f = 0 + os.remove(temp_path) + + +def test_write_modified_fieldset_unop(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + negf = -f + temp_path = "written_tuv_pl_unop.grib" + negf.write(temp_path) + assert os.path.isfile(temp_path) + g = mv.Fieldset(path=temp_path) + assert type(g) == mv.Fieldset + assert len(g) == 18 + sn = g.grib_get_string("shortName") + assert sn == ["t", "u", "v"] * 6 + np.testing.assert_allclose(g.values(), -f.values(), 0.0001) + f = 0 + os.remove(temp_path) + + +def test_field_func(): + def sqr_func(x): + return x * x + + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + g = f.field_func(sqr_func) + assert type(g) == mv.Fieldset + assert len(g) == 18 + vf = f.values() + vg = g.values() + np.testing.assert_allclose(vg, vf * vf, 0.0001) + + +def test_field_func_neg(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + g = -f + assert type(g) == mv.Fieldset + assert len(g) == 18 + vf = f.values() + vg = g.values() + np.testing.assert_allclose(vg, -vf) + + +def test_field_func_pos(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + g = +f # should return values unaltered + assert type(g) == mv.Fieldset + assert len(g) == 18 + vf = f.values() + vg = g.values() + np.testing.assert_allclose(vg, vf) + + +def test_field_func_abs(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + g = f.abs() + assert type(g) == mv.Fieldset + assert len(g) == 18 + vf = f.values() + vg = g.values() + np.testing.assert_allclose(vg, np.abs(vf)) + + +def test_temporary_file(): + # create a temp file, then delete the fieldset - temp should be removed + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + g = -f + temp_path = g.temporary.path + assert os.path.isfile(temp_path) + g = None + assert not os.path.isfile(temp_path) + + +def test_permanent_file_not_accidentally_deleted(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + assert os.path.isfile(path) + f = None + assert os.path.isfile(path) + + +def test_single_index_0(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + f0 = f[0] + assert type(f0) is mv.Fieldset + assert len(f0) == 1 + assert f0.grib_get_string("shortName") == "t" + v = f0.values() + eps = 0.001 + assert len(v) == 2664 + assert np.isclose(v[1088], 304.5642, eps) + + +def test_single_index_17(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + f17 = f[17] + assert type(f17) is mv.Fieldset + assert len(f17) == 1 + assert f17.grib_get_string("shortName") == "v" + v = f17.values() + eps = 0.001 + assert len(v) == 2664 + assert np.isclose(v[2663], -11.0797, eps) + + +def test_single_index_minus_1(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + fm1 = f[-1] + assert type(fm1) is mv.Fieldset + assert len(fm1) == 1 + assert fm1.grib_get_string("shortName") == "v" + v = fm1.values() + eps = 0.001 + assert len(v) == 2664 + assert np.isclose(v[2663], -11.0797, eps) + + +def test_single_index_bad(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + with pytest.raises(IndexError): + fbad = f[27] + + +def test_slice_0_5(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + f05 = f[0:5] + assert type(f05) is mv.Fieldset + assert len(f05) == 5 + assert f05.grib_get_string("shortName") == ["t", "u", "v", "t", "u"] + v = f05.values() + assert v.shape == (5, 2664) + # check the original fieldset + assert len(f) == 18 + sn = f.grib_get_string("shortName") + assert sn == ["t", "u", "v"] * 6 + + +def test_array_indexing(): + path = os.path.join(PATH, "tuv_pl.grib") + f = mv.Fieldset(path=path) + indexes = np.array([1, 16, 5, 9]) + fv = f[indexes] + assert type(fv) is mv.Fieldset + assert len(fv) == 4 + assert fv.grib_get_string("shortName") == ["u", "u", "v", "t"] + # check with bad indexes + indexes = np.array([1, 36, 5, 9]) + with pytest.raises(IndexError): + fvbad = f[indexes] + + +def test_fieldset_iterator(): + grib = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + sn = grib.grib_get_string("shortName") + assert len(sn) == 18 + iter_sn = [] + for f in grib: + iter_sn.append(f.grib_get_string("shortName")) + assert len(iter_sn) == len(sn) + assert iter_sn == sn + iter_sn = [f.grib_get_string("shortName") for f in grib] + assert iter_sn == sn + + +def test_fieldset_iterator_multiple(): + grib = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + sn = grib.grib_get_string("shortName") + assert len(sn) == 18 + for i in [1, 2, 3]: + iter_sn = [] + for f in grib: + iter_sn.append(f.grib_get_string("shortName")) + assert len(iter_sn) == len(sn) + for i in range(0, 18): + assert sn[i] == iter_sn[i] + + +def test_fieldset_iterator_with_zip(): + # this tests something different with the iterator - this does not try to + # 'go off the edge' of the fieldset, because the length is determined by + # the list of levels + grib = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + ref_levs = grib.grib_get_long("level") + assert len(ref_levs) == 18 + levs1 = [] + levs2 = [] + for k, f in zip(grib.grib_get_long("level"), grib): + levs1.append(k) + levs2.append(f.grib_get_long("level")) + assert levs1 == ref_levs + assert levs2 == ref_levs + + +def test_fieldset_iterator_with_zip_multiple(): + # same as test_fieldset_iterator_with_zip() but multiple times + grib = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + ref_levs = grib.grib_get_long("level") + assert len(ref_levs) == 18 + for i in [1, 2, 3]: + levs1 = [] + levs2 = [] + for k, f in zip(grib.grib_get_long("level"), grib): + levs1.append(k) + levs2.append(f.grib_get_long("level")) + print(grib.grib_get_long("level")) + assert levs1 == ref_levs + assert levs2 == ref_levs + + +def test_fieldset_reverse_iterator(): + grib = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + sn = grib.grib_get_string("shortName") + sn_reversed = list(reversed(sn)) + assert sn_reversed[0] == "v" + assert sn_reversed[17] == "t" + gribr = reversed(grib) + iter_sn = [f.grib_get_string("shortName") for f in gribr] + assert len(iter_sn) == len(sn_reversed) + assert iter_sn == sn_reversed + assert iter_sn == ["v", "u", "t"] * 6 + + +def test_fieldset_append(): + g = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + h = mv.Fieldset(path=os.path.join(PATH, "all_missing_vals.grib")) + i = g[0:3] + i.append(h) + assert i.grib_get_string("shortName") == ["t", "u", "v", "z"] + + +def test_fieldset_merge(): + g = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + h = mv.Fieldset(path=os.path.join(PATH, "all_missing_vals.grib")) + i = g[0:3] + j = i.merge(h) # does not alter the original fieldset + assert i.grib_get_string("shortName") == ["t", "u", "v"] + assert j.grib_get_string("shortName") == ["t", "u", "v", "z"] + + +def test_field_scalar_func(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:3] + # fieldset op scalar + g = f + 10 + assert type(g) == mv.Fieldset + assert len(g) == 3 + np.testing.assert_allclose(g.values(), f.values() + 10) + q = f - 5 + np.testing.assert_allclose(q.values(), f.values() - 5) + m = f * 1.5 + np.testing.assert_allclose(m.values(), f.values() * 1.5) + d = f / 3.0 + np.testing.assert_allclose(d.values(), f.values() / 3.0, 0.0001) + p = f**2 + np.testing.assert_allclose(p.values(), f.values() ** 2) + first_val = f.values()[0][0] # 272 + ge = f >= first_val + v = ge.values() + assert v[0][0] == 1 # 272 + assert v[0][2645] == 0 # 240 + assert v[0][148] == 1 # 280 + assert v[1][0] == 0 # -6 + gt = f > first_val + v = gt.values() + assert v[0][0] == 0 # 272 + assert v[0][2645] == 0 # 240 + assert v[0][148] == 1 # 280 + assert v[1][0] == 0 # - 6 + lt = f < first_val + v = lt.values() + assert v[0][0] == 0 # 272 + assert v[0][2645] == 1 # 240 + assert v[0][148] == 0 # 280 + assert v[1][0] == 1 # - 6 + lt = f <= first_val + v = lt.values() + assert v[0][0] == 1 # 272 + assert v[0][2645] == 1 # 240 + assert v[0][148] == 0 # 280 + assert v[1][0] == 1 # - 6 + e = f == first_val + v = e.values() + assert v[0][0] == 1 # 272 + assert v[0][2645] == 0 # 240 + assert v[0][148] == 0 # 280 + assert v[1][0] == 0 # - 6 + ne = f != first_val + v = ne.values() + assert v[0][0] == 0 # 272 + assert v[0][2645] == 1 # 240 + assert v[0][148] == 1 # 280 + assert v[1][0] == 1 # - 6 + andd = (f > 270) & (f < 290) # and + v = andd.values() + assert v[0][0] == 1 # 272 + assert v[0][2645] == 0 # 240 + assert v[0][148] == 1 # 280 + assert v[1][0] == 0 # - 6 + orr = (f < 270) | (f > 279) # or + v = orr.values() + assert v[0][0] == 0 # 272 + assert v[0][2645] == 1 # 240 + assert v[0][148] == 1 # 280 + assert v[1][0] == 1 # - 6 + nott = ~((f > 270) & (f < 290)) # not + v = nott.values() + assert v[0][0] == 0 # 272 + assert v[0][2645] == 1 # 240 + assert v[0][148] == 0 # 280 + assert v[1][0] == 1 # - 6 + # scalar op fieldset + h = 20 + f + assert type(h) == mv.Fieldset + assert len(h) == 3 + np.testing.assert_allclose(h.values(), f.values() + 20) + r = 25 - f + np.testing.assert_allclose(r.values(), 25 - f.values()) + mr = 3 * f + np.testing.assert_allclose(mr.values(), f.values() * 3) + dr = 200 / f + np.testing.assert_allclose(dr.values(), 200 / f.values(), 0.0001) + pr = 2**f + np.testing.assert_allclose(pr.values(), 2 ** f.values(), 1) + + +def test_fieldset_fieldset_func(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[0:3] + g = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib"))[5:8] + h = f + g + np.testing.assert_allclose(h.values(), f.values() + g.values()) + i = g + f + np.testing.assert_allclose(i.values(), g.values() + f.values()) + q = f - g + np.testing.assert_allclose(q.values(), f.values() - g.values()) + r = g - f + np.testing.assert_allclose(r.values(), g.values() - f.values()) + t = g * f + np.testing.assert_allclose(t.values(), g.values() * f.values(), 0.0001) + d = g / f + np.testing.assert_allclose(d.values(), g.values() / f.values(), 0.0001) + gt = f > g + assert gt[0].values()[0] == 1 + assert gt[1].values()[0] == 0 + assert gt[2].values()[0] == 1 + assert gt[2].values()[22] == 0 + gt = f >= g + assert gt[0].values()[0] == 1 + assert gt[1].values()[0] == 0 + assert gt[2].values()[0] == 1 + assert gt[2].values()[22] == 0 + lt = f < g + assert lt[0].values()[0] == 0 + assert lt[1].values()[0] == 1 + assert lt[2].values()[0] == 0 + assert lt[2].values()[22] == 1 + lt = f <= g + assert lt[0].values()[0] == 0 + assert lt[1].values()[0] == 1 + assert lt[2].values()[0] == 0 + assert lt[2].values()[22] == 1 + + +def test_fieldset_multiple_funcs(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + g = 1 - ((f[0] + f[3]) - 5) + np.testing.assert_allclose(g.values(), 1 - ((f[0].values() + f[3].values()) - 5)) + + +def test_fieldset_funs_with_read(): + f = mv.read(os.path.join(PATH, "tuv_pl.grib")) + assert isinstance(f, mv.Fieldset) + g = f + 18 + assert isinstance(g, mv.Fieldset) + diff = g - f + assert isinstance(diff, mv.Fieldset) + assert np.isclose(diff.minvalue(), 18) + assert np.isclose(diff.maxvalue(), 18) + + +def test_field_maths_funcs(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + f = f[0] + v = f.values() + + # no arg + r = f.abs() + np.testing.assert_allclose(r.values(), np.fabs(v), rtol=1e-05) + + r = f.cos() + np.testing.assert_allclose(r.values(), np.cos(v), rtol=1e-05) + + f1 = f / 100 + r = f1.exp() + np.testing.assert_allclose(r.values(), np.exp(f1.values()), rtol=1e-05) + + r = f.log() + np.testing.assert_allclose(r.values(), np.log(v), rtol=1e-05) + + r = f.log10() + np.testing.assert_allclose(r.values(), np.log10(v), rtol=1e-05) + + r = f.sgn() + np.testing.assert_allclose(r.values(), np.sign(v), rtol=1e-05) + + r = f.sin() + np.testing.assert_allclose(r.values(), np.sin(v), rtol=1e-05) + + r = f.sqrt() + np.testing.assert_allclose(r.values(), np.sqrt(v), rtol=1e-05) + + r = f.tan() + np.testing.assert_allclose(r.values(), np.tan(v), rtol=1e-04) + + # inverse functions + # scale input between [-1, 1] + f1 = (f - 282) / 80 + v1 = f1.values() + r = f1.acos() + np.testing.assert_allclose(r.values(), np.arccos(v1), rtol=1e-05) + + r = f1.asin() + np.testing.assert_allclose(r.values(), np.arcsin(v1), rtol=1e-05) + + r = f1.atan() + np.testing.assert_allclose(r.values(), np.arctan(v1), rtol=1e-05) + + # 1 arg + f1 = f - 274 + v1 = f1.values() + + r = f.atan2(f1) + np.testing.assert_allclose(r.values(), np.arctan2(v, v1), rtol=1e-05) + + r = f.div(f1) + np.testing.assert_allclose(r.values(), np.floor_divide(v, v1), rtol=1e-05) + + r = f.mod(f1) + np.testing.assert_allclose(r.values(), np.mod(v, v1), rtol=1e-04) + + +def test_str(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + assert str(f) == "Fieldset (18 fields)" + + +def test_set_values_single_field(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + f0 = f[0] + f0_vals = f0.values() + vals_plus_10 = f0_vals + 10 + f0_modified = f0.set_values(vals_plus_10) + f0_mod_vals = f0_modified.values() + np.testing.assert_allclose(f0_mod_vals, vals_plus_10) + # write to disk, read and check again + testpath = "f0_modified.grib" + f0_modified.write(testpath) + f0_read = mv.Fieldset(path=testpath) + np.testing.assert_allclose(f0_read.values(), vals_plus_10) + os.remove(testpath) + + +def test_set_values_multiple_fields(): + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + f03 = f[0:3] + f47 = f[4:7] + f03_modified = f03.set_values(f47.values()) + np.testing.assert_allclose(f03_modified.values(), f47.values()) + # same, but with a list of arrays instead of a 2D array + list_of_arrays = [f.values() for f in f47] + f03_modified_2 = f03.set_values(list_of_arrays) + np.testing.assert_allclose(f03_modified_2.values(), f47.values()) + # wrong number of arrays + f48 = f[4:8] + with pytest.raises(ValueError): + f03_modified_3 = f03.set_values(f48.values()) + + +def test_set_values_with_missing_values(): + f = mv.Fieldset(path=os.path.join(PATH, "t_with_missing.grib")) + new_vals = f.values() + 40 + g = f.set_values(new_vals) + v = g.values() + assert v.shape == (2664,) + eps = 0.001 + assert np.isclose(v[0], 272.5642 + 40, eps) + assert np.isnan(v[798]) + assert np.isnan(v[806]) + assert np.isnan(v[1447]) + assert np.isclose(v[2663], 240.5642 + 40, eps) + + +def test_set_values_with_missing_values_2(): + f = mv.Fieldset(path=os.path.join(PATH, "t_with_missing.grib")) + g = f[0] + v = g.values() + v[1] = np.nan + h = g.set_values(v) + hv = h.values()[:10] + assert np.isclose(hv[0], 272.56417847) + assert np.isnan(hv[1]) + assert np.isclose(hv[2], 272.56417847) + + +def test_set_values_resize(): + # NOTE: the current change in behavour - in 'standard Metview' the user + # has to supply "resize" as an optional argument in order to allow an array + # of different size to be used; if not supplied, and the given array is not the + # same size as the original field, an error is thrown; here, we allow resizing + # without the need for an extra argument - do we want to do this check? + f = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + f0 = f[0] + f0_20vals = f0.values()[0:20] + f0_modified = f0.set_values(f0_20vals) + f0_mod_vals = f0_modified.values() + eps = 0.001 + np.testing.assert_allclose(f0_mod_vals, f0_20vals, eps) + + +def test_vals_destroyed(): + f = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + assert f.fields[0].vals is None + g = f.values() + assert isinstance(g, np.ndarray) + assert f.fields[0].vals is None + f = -f + assert f.fields[0].vals is None + g = f.values() + assert isinstance(g, np.ndarray) + assert f.fields[0].vals is None + f = f + 1 + assert f.fields[0].vals is None + g = f.values() + assert isinstance(g, np.ndarray) + assert f.fields[0].vals is None + + +def test_accumulate(): + f = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_7x7.grb")) + v = mv.accumulate(f) + assert isinstance(v, float) + assert np.isclose(v, 393334.244141) + + f = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + v = mv.accumulate(f) + assert isinstance(v, list) + v_ref = [ + 408058.256226, + 413695.059631, + 430591.282776, + 428943.981812, + 422329.622498, + 418016.024231, + 409755.097961, + 402741.786194, + ] + assert len(v) == len(f) + np.testing.assert_allclose(v, v_ref) + + +def test_average(): + fs = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + + # const fields + v = mv.average(fs * 0 + 1) + assert isinstance(v, float) + assert np.isclose(v, 1) + + # # single field + v = mv.average(fs) + assert isinstance(v, float) + assert np.isclose(v, 279.06647863) + + # multiple fields + f = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + v = mv.average(f) + assert isinstance(v, list) + v_ref = [ + 290.639783636, + 294.654600877, + 306.688947846, + 305.515656561, + 300.804574428, + 297.732210991, + 291.848360371, + 286.85312407, + ] + + assert len(v) == len(f) + np.testing.assert_allclose(v, v_ref) + + +def test_latitudes(): + fs = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_2x2.grb")) + + v = mv.latitudes(fs) + assert isinstance(v, np.ndarray) + assert len(v) == 16380 + assert np.isclose(v[0], 90) + assert np.isclose(v[1], 90) + assert np.isclose(v[8103], 0) + assert np.isclose(v[11335], -34) + assert np.isclose(v[16379], -90) + + f = fs.merge(fs) + lst = mv.latitudes(f) + assert len(lst) == 2 + for v in lst: + assert np.isclose(v[0], 90) + assert np.isclose(v[1], 90) + assert np.isclose(v[8103], 0) + assert np.isclose(v[11335], -34) + assert np.isclose(v[16379], -90) + + +def test_longitudes(): + fs = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_2x2.grb")) + + v = mv.longitudes(fs) + assert isinstance(v, np.ndarray) + assert len(v) == 16380 + assert np.isclose(v[0], 0) + assert np.isclose(v[1], 2) + assert np.isclose(v[8103], 6) + assert np.isclose(v[11335], 350) + assert np.isclose(v[16379], 358) + + f = fs.merge(fs) + lst = mv.longitudes(f) + assert len(lst) == 2 + for v in lst: + assert np.isclose(v[0], 0) + assert np.isclose(v[1], 2) + assert np.isclose(v[8103], 6) + assert np.isclose(v[11335], 350) + assert np.isclose(v[16379], 358) + + +def test_coslat(): + fs = mv.Fieldset(path=os.path.join(PATH, "t_time_series.grib")) + + # WARN: it is important that the data should be at least 16 bit + # to keep accuracy in resulting fields + + f = fs[0] + r = mv.coslat(f) + np.testing.assert_allclose( + r.values(), np.cos(np.deg2rad(f.latitudes())), rtol=1e-06 + ) + + f = fs[:2] + r = mv.coslat(f) + assert len(r) == 2 + for i in range(len(r)): + np.testing.assert_allclose( + r[i].values(), np.cos(np.deg2rad(f[i].latitudes())), rtol=1e-06 + ) + + +def test_mean(): + fs = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + + # single fields + f = fs + r = mv.mean(f) + v_ref = mv.values(fs) + assert len(r) == 1 + np.testing.assert_allclose(r.values(), v_ref, rtol=1e-05) + + # known mean + f = fs.merge(2 * fs) + f = f.merge(3 * fs) + r = f.mean() + v_ref = mv.values(fs) * 2 + assert len(r) == 1 + np.testing.assert_allclose(r.values(), v_ref, rtol=1e-05) + + +def test_maxvalue(): + fs = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + + f = fs + f = f.merge(3 * fs) + f = f.merge(2 * fs) + v = mv.maxvalue(f) + assert isinstance(v, float) + assert np.isclose(v, 948.1818237304688) + + +def test_minvalue(): + fs = mv.Fieldset(path=os.path.join(PATH, "test.grib")) + + f = 3 * fs + f = f.merge(fs) + f = f.merge(2 * fs) + v = mv.minvalue(f) + assert isinstance(v, float) + assert np.isclose(v, 206.93560791015625) + + +def test_sinlat(): + fs = mv.Fieldset(path=os.path.join(PATH, "t_time_series.grib")) + + # WARN: it is important that the data should be at least 16 bit + # to keep accuracy in resulting fields + + f = fs[0] + r = mv.sinlat(f) + np.testing.assert_allclose( + r.values(), np.sin(np.deg2rad(f.latitudes())), rtol=1e-06 + ) + + f = fs[:2] + r = mv.sinlat(f) + assert len(r) == 2 + for i in range(len(r)): + np.testing.assert_allclose( + r[i].values(), np.sin(np.deg2rad(f[i].latitudes())), rtol=1e-06 + ) + + +def test_tanlat(): + fs = mv.Fieldset(path=os.path.join(PATH, "t_time_series.grib")) + + # WARN: it is important that the data should be at least 16 bit + # to keep accuracy in resulting fields + + # TODO: use pole_limit value from fieldset + + pole_limit = 90.0 - 1e-06 + + f = fs[0] + r = mv.tanlat(f) + lat = f.latitudes() + lat[np.fabs(lat) > pole_limit] = np.nan + np.testing.assert_allclose( + r.values(), np.tan(np.deg2rad(lat)), rtol=1e-06, atol=1e-06 + ) + + f = fs[:2] + r = mv.tanlat(f) + assert len(r) == 2 + for i in range(len(r)): + lat = f[i].latitudes() + lat[np.fabs(lat) > pole_limit] = np.nan + np.testing.assert_allclose( + r[i].values(), np.tan(np.deg2rad(lat)), rtol=1e-06, atol=1e-06 + ) + + +def test_stdev(): + fs = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_7x7.grb")) + + # single field + r = mv.stdev(fs) + assert len(r) == 1 + np.testing.assert_allclose(r.values(), fs.values() * 0) + + # known variance + f = fs.merge(4 * fs) + f = f.merge(10 * fs) + r = mv.stdev(f) + assert len(r) == 1 + np.testing.assert_allclose(r.values(), np.sqrt(np.square(fs.values()) * 42 / 3)) + + # real life example + fs = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + r = mv.stdev(fs) + assert len(r) == 1 + + v_ref = np.ma.std(np.array([x.values() for x in fs]), axis=0) + np.testing.assert_allclose(r.values(), v_ref, rtol=1e-03) + + +def test_sum(): + fs = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_7x7.grb")) + + # single fields + f = fs + r = mv.sum(f) + assert len(r) == 1 + np.testing.assert_allclose(r.values(), fs.values()) + + # known sum + f = fs.merge(fs) + f = f.merge(fs) + r = f.sum() + assert len(r) == 1 + np.testing.assert_allclose(r.values(), fs.values() * 3) + + # real life example + f = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + r = f.sum() + assert len(r) == 1 + v_ref = r.values() * 0 + for g in f: + v_ref += g.values() + np.testing.assert_allclose(r.values(), v_ref) + + +def test_var(): + fs = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_7x7.grb")) + + # single field + r = mv.var(fs) + assert len(r) == 1 + np.testing.assert_allclose(r.values(), fs.values() * 0) + + # known variance + f = fs.merge(4 * fs) + f = f.merge(10 * fs) + r = mv.var(f) + assert len(r) == 1 + np.testing.assert_allclose(r.values(), np.square(fs.values()) * 42 / 3) + + # real life example + fs = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + r = mv.var(fs) + assert len(r) == 1 + + v_ref = np.ma.var(np.array([x.values() for x in fs]), axis=0) + np.testing.assert_allclose(r.values(), v_ref, rtol=1e-03) + + +def test_date(): + + fs = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + + # analysis, so valid=base + bdate_ref = [ + "2016-01-01 00:00:00", + "2016-02-01 00:00:00", + "2016-03-01 00:00:00", + "2016-04-01 00:00:00", + "2016-05-01 00:00:00", + "2016-06-01 00:00:00", + "2016-07-01 00:00:00", + "2016-08-01 00:00:00", + ] + vdate_ref = bdate_ref + + v = mv.base_date(fs) + assert len(v) == len(fs) + for i, d in enumerate(v): + assert d == utils.date_from_str(bdate_ref[i]) + + v = mv.valid_date(fs) + assert len(v) == len(fs) + for i, d in enumerate(v): + assert d == utils.date_from_str(vdate_ref[i]) + + +def test_bitmap(): + fs = mv.Fieldset(path=os.path.join(PATH, "t1000_LL_2x2.grb")) + + # -- const field + f = fs * 0 + 1 + + # non missing + r = mv.bitmap(f, 0) + np.testing.assert_allclose(r.values(), f.values()) + + # all missing + r = mv.bitmap(f, 1) + np.testing.assert_allclose(r.values(), f.values() * np.nan) + + # -- non const field + f = fs + + # bitmap with value + f_mask = f > 300 + r = mv.bitmap(f_mask, 1) + v_ref = f_mask.values() + v_ref[v_ref == 1] = np.nan + np.testing.assert_allclose(r.values(), v_ref) + + f_mask = f > 300 + r = mv.bitmap(f_mask * 2, 2) + v_ref = f_mask.values() * 2 + v_ref[v_ref == 2] = np.nan + np.testing.assert_allclose(r.values(), v_ref) + + # bitmap with field + f = mv.bitmap(fs > 300, 0) + r = mv.bitmap(fs, f) + v_ref = fs.values() * f.values() + np.testing.assert_allclose(r.values(), v_ref) + + # multiple fields + f = mv.Fieldset(path=os.path.join(PATH, "monthly_avg.grib")) + f = f[0:2] + + # with value + f_mask = f > 300 + r = mv.bitmap(f_mask, 1) + assert len(r) == len(f) + for i in range(len(r)): + v_ref = f_mask[i].values() + v_ref[v_ref == 1] = np.nan + np.testing.assert_allclose(r[i].values(), v_ref) + + # with field + f1 = mv.bitmap(f > 300, 0) + r = mv.bitmap(f, f1) + assert len(r) == len(f1) + for i in range(len(r)): + v_ref = f[i].values() * f1[i].values() + np.testing.assert_allclose(r[i].values(), v_ref) + + # with single field + f1 = mv.bitmap(f[0] > 300, 0) + r = mv.bitmap(f, f1) + assert len(r) == len(f) + for i in range(len(r)): + v_ref = f[i].values() * f1.values() + np.testing.assert_allclose(r[i].values(), v_ref) + + +def test_nobitmap(): + + fs = mv.Fieldset(path=os.path.join(PATH, "t_with_missing.grib")) + + # single field + f = fs + r = mv.nobitmap(f, 1) + assert len(r) == 1 + v_ref = f.values() + v_ref[np.isnan(v_ref)] = 1 + np.testing.assert_allclose(r.values(), v_ref) + + # multiple fields + f = fs.merge(2 * fs) + r = mv.nobitmap(f, 1) + assert len(r) == 2 + + for i in range(len(r)): + v_ref = f[i].values() + v_ref[np.isnan(v_ref)] = 1 + np.testing.assert_allclose(r[i].values(), v_ref) + + +def test_grib_index_0(): + # empty fieldset + fs = mv.Fieldset() + gi = fs.grib_index() + assert gi == [] + + +def test_grib_index_1(): + # single field + grib_path = os.path.join(PATH, "test.grib") + fs = mv.Fieldset(path=grib_path) + gi = fs.grib_index() + assert gi == [(grib_path, 0)] + + +def test_grib_index_2(): + # multiple fields + grib_path = os.path.join(PATH, "tuv_pl.grib") + fs = mv.Fieldset(path=grib_path) + gi = fs.grib_index() + assert isinstance(gi, list) + assert len(gi) == 18 + for f, g in zip(fs, gi): + assert g == (grib_path, f.grib_get_long("offset")) + assert gi[5] == (grib_path, 7200) + + +def test_grib_index_3(): + # merged fields from different files + gp1 = os.path.join(PATH, "tuv_pl.grib") + gp2 = os.path.join(PATH, "t_time_series.grib") + fs1 = mv.Fieldset(path=gp1) + fs2 = mv.Fieldset(path=gp2) + fs3 = fs1[4:7] + fs3.append(fs2[1]) + fs3.append(fs1[2]) + gi = fs3.grib_index() + assert isinstance(gi, list) + assert len(gi) == 5 + # assert gi == [(gp1, 5760), (gp1, 7200), (gp1, 8640), (gp2, 5520), (gp1, 2880)] + assert gi == [(gp1, 5760), (gp1, 7200), (gp1, 8640), (gp2, 5436), (gp1, 2880)] + + +def test_grib_index_4(): + # test with a derived fieldset + fs = mv.Fieldset(os.path.join(PATH, "t_time_series.grib"))[0:4] + fs1 = fs + 1 + gi = fs1.grib_index() + for i in gi: + assert is_temp_file(i[0]) + offsets = [i[1] for i in gi] + assert offsets == [0, 8440, 16880, 25320] + + +def test_grib_index_5(): + # test with grib written with write() function + f_orig = mv.Fieldset(path=os.path.join(PATH, "tuv_pl.grib")) + f = (f_orig[0:4]).merge(f_orig[7]) + p = "written_tuv_pl_5.grib" + f.write(p) + gi = f.grib_index() + assert gi == [(p, 0), (p, 1440), (p, 2880), (p, 4320), (p, 5760)] + f = 0 + os.remove(p) diff --git a/tests/test_style.py b/tests/test_style.py new file mode 100644 index 00000000..d36a5841 --- /dev/null +++ b/tests/test_style.py @@ -0,0 +1,164 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + + +import metview as mv +from metview import bindings +from metview.style import Style, Visdef + + +def test_visdef_object(): + vd = Visdef( + "msymb", + { + "symbol_type": "text", + "symbol_table_mode": "advanced", + "symbol_advanced_table_text_font_colour": "black", + }, + ) + + # clone + v = vd.clone() + assert v.verb == vd.verb + assert v.params == vd.params + assert id(v) != id(vd) + assert id(v.params) != id(vd.params) + + # change + vd.change("msymb", "symbol_advanced_table_text_font_colour", "red") + assert vd.params["symbol_advanced_table_text_font_colour"] == "red" + vd.change("mcont", "symbol_advanced_table_text_font_colour", "blue") + assert vd.params["symbol_advanced_table_text_font_colour"] == "red" + + # from request + r = mv.msymb( + { + "symbol_type": "text", + "symbol_table_mode": "advanced", + "symbol_advanced_table_text_font_colour": "black", + } + ) + + v = Visdef.from_request(r) + assert v.verb == vd.verb + for k, v in v.params.items(): + assert r[k.upper()].lower() == v.lower() + + +def test_style_object(): + vd = Visdef( + "msymb", + { + "symbol_type": "text", + "symbol_table_mode": "advanced", + "symbol_advanced_table_text_font_colour": "black", + }, + ) + + # create + s = Style("super", vd) + assert s.name == "super" + assert len(s.visdefs) == 1 + assert s.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + + # update + ns = s.update({"symbol_advanced_table_text_font_colour": "red"}) + assert id(ns) != id(s) + assert s.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + assert ns.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "red" + + ns = s.update({"symbol_advanced_table_text_font_colour": "red"}, verb="mcont") + assert id(ns) != id(s) + assert s.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + assert ns.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + + ns = s.update({"symbol_advanced_table_text_font_colour": "red"}, verb="msymb") + assert id(ns) != id(s) + assert s.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + assert ns.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "red" + + # multiple visdefs + + vd1 = Visdef( + "msymb", + { + "symbol_type": "text", + "symbol_table_mode": "advanced", + "symbol_advanced_table_text_font_colour": "black", + }, + ) + vd2 = Visdef("mgraph", {"graph_line_colour": "blue", "graph_line_thickness": 4}) + + s = Style("super", [vd1, vd2]) + assert s.name == "super" + assert len(s.visdefs) == 2 + assert s.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + assert s.visdefs[1].params["graph_line_colour"] == "blue" + + # update + ns = s.update({"symbol_advanced_table_text_font_colour": "red"}, verb="msymb") + assert id(ns) != id(s) + assert s.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "black" + assert ns.visdefs[0].params["symbol_advanced_table_text_font_colour"] == "red" + + ns = s.update( + { + "graph_line_colour": "yellow", + }, + verb="mgraph", + ) + assert id(ns) != id(s) + assert s.visdefs[1].params["graph_line_colour"] == "blue" + assert ns.visdefs[1].params["graph_line_colour"] == "yellow" + + +def test_style_set_grib_id(): + # mcont + vd = Visdef( + "mcont", + { + "legend": "on", + }, + ) + + s = Style("super", vd) + ns = s.set_data_id("11") + assert id(s) != id(ns) + assert ns.visdefs[0].verb == "mcont" + assert ns.visdefs[0].params["grib_id"] == "11" + assert s.visdefs[0].params.get("grib_id", None) is None + + # mwind + vd = Visdef( + "mwind", + { + "legend": "on", + }, + ) + + s = Style("super", vd) + ns = s.set_data_id("12") + assert id(s) != id(ns) + assert ns.visdefs[0].verb == "mwind" + assert ns.visdefs[0].params["grib_id"] == "12" + assert s.visdefs[0].params.get("grib_id", None) is None + + # other + vd = Visdef( + "msymb", + { + "legend": "on", + }, + ) + + s = Style("super", vd) + ns = s.set_data_id("10") + assert id(s) == id(ns) + assert s.visdefs[0].params.get("grib_id", None) is None diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..89ecce1d --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,104 @@ +# (C) Copyright 2017- ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# + +import datetime + +from metview.metviewpy import utils + + +def test_utils_date(): + + d = utils.date_from_str("20210204") + assert d == datetime.datetime(2021, 2, 4) + + d = utils.date_from_str("2021-02-04") + assert d == datetime.datetime(2021, 2, 4) + + d = utils.date_from_str("2021-02-04 06") + assert d == datetime.datetime(2021, 2, 4, 6, 0, 0) + + d = utils.date_from_str("2021-02-04 06:14") + assert d == datetime.datetime(2021, 2, 4, 6, 14, 0) + + d = utils.date_from_str("2021-02-04 06:14:32") + assert d == datetime.datetime(2021, 2, 4, 6, 14, 32) + + d = utils.date_from_str("20210204.25") + assert d == datetime.datetime(2021, 2, 4, 6, 0, 0) + + d = utils.date_from_str("402") + assert d == (4, 2) + + d = utils.date_from_str("0402") + assert d == (4, 2) + + try: + d = utils.date_from_str("0432") + assert False + except: + pass + + d = utils.date_from_str("apr-02") + assert d == (4, 2) + + d = utils.date_from_str("Apr-02") + assert d == (4, 2) + + try: + d = utils.date_from_str("afr-02") + assert False + except: + pass + + +def test_utils_time(): + + d = utils.time_from_str("6") + assert d == datetime.time(6, 0, 0) + + d = utils.time_from_str("06") + assert d == datetime.time(6, 0, 0) + + d = utils.time_from_str("14") + assert d == datetime.time(14, 0, 0) + + d = utils.time_from_str("600") + assert d == datetime.time(6, 0, 0) + + d = utils.time_from_str("612") + assert d == datetime.time(6, 12, 0) + + d = utils.time_from_str("1100") + assert d == datetime.time(11, 0, 0) + + d = utils.time_from_str("1457") + assert d == datetime.time(14, 57, 0) + + d = utils.time_from_str("6:12") + assert d == datetime.time(6, 12, 0) + + d = utils.time_from_str("06:12") + assert d == datetime.time(6, 12, 0) + + d = utils.time_from_str("14:57") + assert d == datetime.time(14, 57, 0) + + +def test_has_globbing(): + + assert utils.has_globbing("a*") == True + assert utils.has_globbing("a?") == True + assert utils.has_globbing("my_path/a*.grib") == True + assert utils.has_globbing("my_path/[Aa].grib") == True + assert utils.has_globbing("my_path/[a-m].grib") == True + assert utils.has_globbing("my_path/test.grib") == False + assert utils.has_globbing("my_path/te[st.grib") == False + assert utils.has_globbing("my_path/tes]t.grib") == False + # assert utils.has_globbing("my_path/t]e[st.grib") == False diff --git a/tests/zt_multi_expvers.grib b/tests/zt_multi_expvers.grib new file mode 100644 index 00000000..b47ed397 Binary files /dev/null and b/tests/zt_multi_expvers.grib differ diff --git a/tests/ztu_multi_dim.grib b/tests/ztu_multi_dim.grib new file mode 100644 index 00000000..7e66904c Binary files /dev/null and b/tests/ztu_multi_dim.grib differ