obsarray 1.0.2__py3-none-any.whl

obsarray/templater/dstemplater.py

@@ -0,0 +1,133 @@
+"""
+Handling for multiple ds templates
+"""
+
+from typing import Optional, Dict, List
+import xarray
+from obsarray import create_ds
+
+
+__author__ = "Sam Hunt <sam.hunt@npl.co.uk>"
+
+
+class DSTemplater:
+    """
+    Class to generate ``xarray.Dataset``'s from a set of templates.
+
+    :param templates: template dictionary for each product format
+    :param metadata_defs: metadata for each product format
+
+    Each dictionary has as keys ``"template_name"`` strings that define the names of the different available dataset
+    templates, e.g. "Level-0".
+
+    For the ``templates`` dictionary the corresponding entries should be a `variable definition dictionary <https://dsbuilder.readthedocs.io/en/latest/content/quickstart.html#defining-a-template-dataset>`_.
+
+    For ``metadata_defs`` the corresponding entries should be a dictionary of per product metadata.
+    """
+
+    def __init__(
+        self, templates: Optional[Dict] = None, metadata_defs: Optional[Dict] = None
+    ):
+        self.templates = templates if templates is not None else {}
+        self.metadata_defs = metadata_defs if metadata_defs is not None else {}
+
+    def create(self, template_name: str, size: Dict[str, int]) -> xarray.Dataset:
+        """
+        Returns template dataset
+
+        :param template_name: name of template to create (value returned by ``self.return_template_names()``)
+        :param size: entry per dataset dimension with value of size
+
+        :returns: Empty dataset
+        """
+
+        # Find template
+        if template_name in self.return_template_names():
+            template = self.templates[template_name]
+        else:
+            raise NameError(
+                "No template: "
+                + template_name
+                + " - must be one of "
+                + str(self.return_template_names())
+            )
+
+        # Find metadata def
+        metadata = (
+            self.metadata_defs[template_name]
+            if template_name in self.metadata_defs.keys()
+            else None
+        )
+
+        return create_ds(template, size, metadata=metadata)
+
+    def return_template_names(self) -> list:
+        """
+        Returns available ds template names
+
+        :returns: template names
+        """
+
+        return list(self.templates.keys())
+
+    def return_template_var_names(self, template_name: str) -> List[str]:
+        """
+        Returns variables for specified template
+
+        :param template_name: template name (value returned by ``self.return_template_names()``)
+
+        :returns: variable names for specified template
+        """
+
+        return list(self.templates[template_name].keys())
+
+    def return_ds_format_variable_dict(
+        self, template_name: str, variable_name: str
+    ) -> Dict:
+        """
+        Returns variable definition info for specified template variable
+
+        :param template_name: template name (value returned by ``self.return_template_names()``)
+        :param variable_name: variable name
+
+        :returns: variable definition info
+        """
+
+        return self.templates[template_name][variable_name]
+
+    def return_template_dim_names(self, template_name: str) -> List[str]:
+        """
+        Returns dims required for specified template
+
+        :param template_name: template name (value returned by ``self.return_template_names()``)
+
+        :returns: template dimensions
+        """
+
+        template = self.templates[template_name]
+
+        template_dims = set()
+
+        for var_name in template.keys():
+            template_dims.update(template[var_name]["dim"])
+
+        return list(template_dims)
+
+    def create_size_dict(self, ds_format: str) -> Dict[str, None]:
+        """
+        Returns empty size dictionary for specified template
+
+        :param template_name: template name (value returned by ``self.return_template_names()``)
+
+        :return: empty sizes dictionary
+        """
+
+        dim_sizes_dict = dict()
+        for dim in self.return_template_dim_names(ds_format):
+            dim_sizes_dict[dim] = None
+
+        return dim_sizes_dict
+
+
+if __name__ == "__main__":
+    pass

obsarray/templater/dswriter.py

@@ -0,0 +1,113 @@
+"""
+DSWriter class
+"""
+
+import os
+from obsarray.templater.template_util import DatasetUtil
+
+__author__ = "Sam Hunt <sam.hunt@npl.co.uk>"
+
+
+class DSWriter:
+    """
+    Class to write datasets to output files
+    """
+
+    @staticmethod
+    def write(ds, path, overwrite=False, fmt="netCDF4", compression_level=None):
+        """
+        Write xarray dataset to file
+
+        :type ds: xarray.Dataset
+        :param ds: dataset
+
+        :type path: str
+        :param path: file path
+
+        :type overwrite: bool
+        :param overwrite: set to true to overwrite existing file
+
+        :type fmt: str
+        :param fmt: format to write to, may be 'netCDF4' (default) or 'csv'
+
+        :type compression_level: int
+        :param compression_level: the file compression level if 'netCDF4' fmt, 0 - 9 (default is 5)
+        """
+
+        if os.path.isfile(path):
+            if overwrite is True:
+                os.remove(path)
+            else:
+                raise IOError("The file already exists: " + path)
+
+        if (fmt.lower() == "netcdf4") or (fmt.lower() == "netcdf"):
+            DSWriter._write_netcdf(ds, path, compression_level=compression_level)
+
+        elif fmt == "csv":
+            DSWriter._write_csv(ds, path)
+
+        else:
+            raise NameError("Invalid fmt: " + fmt)
+
+    @staticmethod
+    def _write_netcdf(ds, path, compression_level=None):
+        """
+        Write xarray dataset to file to netcdf
+
+        :type ds: xarray.Dataset
+        :param ds: dataset
+
+        :type path: str
+        :param path: file path
+
+        :type compression_level: int
+        :param compression_level: the file compression level if 'netCDF4' fmt, 0 - 9 (default is 5)
+        """
+
+        if compression_level is None:
+            compression_level = 5
+
+        comp = dict(zlib=True, complevel=compression_level)
+
+        encoding = dict()
+        for var_name in ds.data_vars:
+            var_encoding = dict(comp)
+            var_encoding.update(ds[var_name].encoding)
+            if "dtype" in var_encoding.keys():
+                var_encoding.update(
+                    {
+                        "_FillValue": DatasetUtil.get_default_fill_value(
+                            var_encoding["dtype"]
+                        )
+                    }
+                )
+                if "_FillValue" in ds[var_name].attrs:
+                    ds[var_name].attrs.pop("_FillValue")
+            encoding.update({var_name: var_encoding})
+
+        ds.to_netcdf(path, format="netCDF4", engine="netcdf4", encoding=encoding)
+
+    @staticmethod
+    def _write_csv(ds, path):
+        """
+        Write xarray dataset to file to csv
+
+        :type ds: xarray.Dataset
+        :param ds: dataset
+
+        :type path: str
+        :param path: file path
+        """
+
+        # write data variables (via pandas dataframe)
+        ds.to_dataframe().to_csv(path)
+
+        # write metadata
+        metadata_path = os.path.splitext(path)[0] + "_meta.txt"
+        with open(metadata_path, "w") as f:
+            for meta_name in ds.attrs.keys():
+                f.write(meta_name + ": " + ds.attrs[meta_name] + "\n")
+
+
+if __name__ == "__main__":
+    pass

obsarray/templater/template_util.py

@@ -0,0 +1,257 @@
+"""
+Utilities for creating template xarray datasets
+"""
+
+from typing import Optional, Dict, List
+from obsarray.templater.dataset_util import DatasetUtil
+import xarray
+
+
+__author__ = "Sam Hunt <sam.hunt@npl.co.uk>"
+
+
+def create_ds(
+    template: Dict[str, Dict],
+    size: Dict[str, int],
+    metadata: Optional[Dict] = None,
+    append_ds: Optional[xarray.Dataset] = None,
+    propagate_ds: Optional[xarray.Dataset] = None,
+) -> xarray.Dataset:
+    """
+    Returns template dataset
+
+    :param template: dictionary defining ds variable structure, as defined below.
+    :param size: dictionary of dataset dimensions, entry per dataset dimension with value of size as int
+    :param metadata: dictionary of dataset metadata
+    :param append_ds: base dataset to append with template variables
+    :param propagate_ds: template dataset is populated with data from propagate_ds for their variables with
+    common names and dimensions. Useful for transferring common data between datasets at different processing levels
+    (e.g. times, etc.).
+    :returns: template dataset
+
+    For the ``template`` dictionary each key/value pair defines one variable, where the key is the variable name and the value is a dictionary with the following entries:
+
+    * ``"dtype"`` (*np.typecodes*/*str*) - variable data type, either a numpy data type or special value ``"flag"`` for
+      flag variable
+    * ``"dim"`` (*list*) - list of variable dimension names
+    * ``"attributes"`` (*dict*) - (optional) variable attributes
+    * ``"encoding"`` (*dict*) - (optional) variable encoding.
+
+    For more information on the required form of these entries, see the :ref:`variables definition section <variables dict>`
+    of the user guide.
+    """
+
+    # Create dataset
+    ds = append_ds if append_ds is not None else xarray.Dataset()
+
+    # Add variables
+    ds = TemplateUtil.add_variables(ds, template, size)
+
+    # Add metadata
+    if metadata is not None:
+        ds = TemplateUtil.add_metadata(ds, metadata)
+
+    # Propagate variable data
+    if propagate_ds is not None:
+        TemplateUtil.propagate_values(ds, propagate_ds)
+
+    return ds
+
+
+def create_var(var_name: str, var_attrs: Dict, size: Dict[str, int]) -> xarray.Variable:
+    """
+    Returns template variable
+
+    :param var_name: variable name
+    :param var_attrs: variable definition dictionary (as an entry to a template dictionary)
+    :param size: dictionary of dataset dimensions, entry per dataset dimension with value of size as int
+    :return:
+    """
+
+    return TemplateUtil._create_var(var_name, var_attrs, size)
+
+
+class TemplateUtil:
+    """
+    Class to create template xarray datasets
+    """
+
+    @staticmethod
+    def add_variables(
+        ds: xarray.Dataset, template: Dict[str, Dict], size: Dict[str, int]
+    ) -> xarray.Dataset:
+        """
+        Adds defined variables dataset
+
+        :param ds: dataset
+        :param template: dictionary defining variables, see the :ref:`variables definition section <variables dict>` of the user guide for more information.
+        :param size: dictionary of dataset dimensions, entry per dataset dimension with value of size as int
+
+        :returns: dataset with defined variables
+        """
+
+        for var_name in template.keys():
+
+            var = TemplateUtil._create_var(var_name, template[var_name], size)
+
+            ds[var_name] = var
+
+        return ds
+
+    @staticmethod
+    def _create_var(
+        var_name: str, var_attrs: dict, size: Dict[str, int]
+    ) -> xarray.Variable:
+
+        du = DatasetUtil()
+
+        # Check variable definition
+        TemplateUtil._check_variable_definition(var_name, var_attrs)
+
+        # Unpack variable attributes
+        dtype = var_attrs["dtype"]
+        dim_names = var_attrs["dim"]
+        attributes = var_attrs["attributes"] if "attributes" in var_attrs else None
+
+        err_corr = None
+        if attributes is not None:
+            if "err_corr" in attributes:
+                err_corr = attributes.pop("err_corr")
+
+        # Determine variable shape from dims
+        try:
+            dim_sizes = TemplateUtil._return_variable_shape(dim_names, size)
+        except KeyError:
+            raise KeyError(
+                "Dim Name Error - Variable "
+                + var_name
+                + " defined with dim not in dim_sizes_dict"
+            )
+
+        # Create variable and add to dataset
+        if isinstance(dtype, str):
+            if dtype == "flag":
+                flag_meanings = attributes.pop("flag_meanings")
+                variable = du.create_flags_variable(
+                    dim_sizes,
+                    meanings=flag_meanings,
+                    dim_names=dim_names,
+                    attributes=attributes,
+                )
+
+            else:
+                raise ValueError("unknown dtype - " + dtype)
+
+        else:
+            if err_corr is None:
+                variable = du.create_variable(
+                    dim_sizes, dim_names=dim_names, dtype=dtype, attributes=attributes
+                )
+
+            else:
+                variable = du.create_unc_variable(
+                    dim_sizes,
+                    dim_names=dim_names,
+                    dtype=dtype,
+                    attributes=attributes,
+                    err_corr=err_corr,
+                )
+
+            if "encoding" in var_attrs:
+                du.add_encoding(variable, **var_attrs["encoding"])
+                variable.attrs.pop("_FillValue")
+
+        return variable
+
+    @staticmethod
+    def _check_variable_definition(variable_name: str, variable_attrs: Dict):
+        """
+        Checks validity of variable definition, raising errors as appropriate
+
+        :param variable_name: variable name
+        :param variable_attrs: variable defining dictionary
+        """
+
+        # Variable name must be type str
+        if type(variable_name) != str:
+            raise TypeError(
+                "Invalid variable name: " + str(variable_name) + " (must be string)"
+            )
+
+        # todo - add more tests to check validity of variable definition
+
+    @staticmethod
+    def _return_variable_shape(dim_names: List[str], size: Dict[str, int]) -> List[int]:
+        """
+        Returns dimension sizes of specified dimensions
+
+        :param dim_names: dimension names
+        :param size: dictionary of dataset dimensions, entry per dataset dimension with value of size as int
+
+        :returns: dimension sizes
+        """
+
+        return [size[dim_name] for dim_name in dim_names]
+
+    @staticmethod
+    def add_metadata(ds: xarray.Dataset, metadata: Dict) -> xarray.Dataset:
+        """
+        Adds metadata to dataset
+
+        :param ds: dataset
+        :param metadata: dictionary of dataset metadata
+
+        :returns: dataset with updated metadata
+        """
+
+        ds.attrs.update(metadata)
+
+        return ds
+
+    @staticmethod
+    def propagate_values(target_ds, source_ds, exclude=None):
+        """
+        Populates target_ds in-place with data from source_ds for their variables with common names and dimensions.
+        Useful for transferring common data between datasets at different processing levels (e.g. times, etc.).
+
+        N.B. propagates data only, not variables as a whole with attributes etc.
+
+        :type target_ds: xarray.Dataset
+        :param target_ds: ds to populate (perhaps data at new processing level)
+
+        :type source_ds: xarray.Dataset
+        :param source_ds: ds to take data from (perhaps data at previous processing level)
+        """
+
+        # Find variable names common to target_ds and source_ds, excluding specified exclude variables
+        common_variable_names = list(set(target_ds).intersection(source_ds))
+        # common_variable_names = list(set(target_ds.variables).intersection(source_ds.variables))
+        # print(common_variable_names)
+
+        if exclude is not None:
+            common_variable_names = [
+                name for name in common_variable_names if name not in exclude
+            ]
+
+        # Remove any common variables that have different dimensions in target_ds and source_ds
+        common_variable_names = [
+            name
+            for name in common_variable_names
+            if target_ds[name].dims == source_ds[name].dims
+        ]
+
+        # Propagate data
+        for common_variable_name in common_variable_names:
+            if (
+                target_ds[common_variable_name].shape
+                == source_ds[common_variable_name].shape
+            ):
+                target_ds[common_variable_name].values = source_ds[
+                    common_variable_name
+                ].values
+
+        # to do - add method to propagate common unpopulated metadata
+
+
+if __name__ == "__main__":
+    pass