field.py 27.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12
# -*- coding: utf-8 -*-
"""
Classes for the handling the fields.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import warnings
import logging
import copy
from collections import OrderedDict
Jeff Piollé's avatar
Jeff Piollé committed
13 14
from typing import (Any, Dict, Hashable, Iterable, Iterator, List,
                    Mapping, Optional, Sequence, Set, Tuple, Union, cast)
15 16 17 18 19

import numpy
import xarray as xr

from cerbere.datamodel.variable import Variable
Jeff Piollé's avatar
Jeff Piollé committed
20
import cerbere.cfconvention as cf
21

Jeff Piollé's avatar
Jeff Piollé committed
22
__all__ = ['Field']
23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41


FIELD_ATTRS = [
    'standard_name',
    'authority',
    'units',
    'valid_min',
    'valid_max',
    '_FillValue'
]
FIELD_EXCL_ATTRS = [
    'long_name'
]


class Field(object):
    """A Field describes a scientific data array. It contains data and
    metadata attributes.

Jeff Piollé's avatar
Jeff Piollé committed
42 43
    This is an extension of xarray's DataArray with stricter requirements on
    attributes.
44

Jeff Piollé's avatar
Jeff Piollé committed
45
    A :class:`Field` object can be constructed with:
46

Jeff Piollé's avatar
Jeff Piollé committed
47
    * a xarray :class:`~xarray.DataArray` object, provided in ``array``
48

Jeff Piollé's avatar
Jeff Piollé committed
49 50 51
    A :class:`Field` object can be attached to a  :class:`~cerbere.dataset.dataset.Dataset`
    object (of any class inherited from a :class:`~cerbere.dataset.dataset.Dataset`,
    provided with the ``dataset`` argument.
52

Jeff Piollé's avatar
Jeff Piollé committed
53 54 55 56 57 58 59 60 61 62 63 64 65
    Args:

        data: the scientific data. Can be of the following types:

              * :class:`numpy.ma.MaskedArray`
              * :class:`xarray.DataArray`

            Optional (the field can be created with default values)

        name: the label of the field (don't use any white space). This
            corresponds to the variable name in a netcdf file

        dims (list): the scientific array dimension names
66

Jeff Piollé's avatar
Jeff Piollé committed
67 68
        datatype: the type of the data. Infer the type from the provided data by
            default.
69

Jeff Piollé's avatar
Jeff Piollé committed
70 71
        precision (int, optional): number of significant digits (when writing
            the data on file)
72

Jeff Piollé's avatar
Jeff Piollé committed
73 74 75 76 77 78
        fields : the subfields composing the main field.
            This is intended to group for instance components of the same
            variable (such as a vector's northward and eastward components
            for wind, currents,...). This allows to relate these components
            to the same physical variable (e.g wind) and to a single
            qc_levels and qc_details information.
79

Jeff Piollé's avatar
Jeff Piollé committed
80 81 82
        fillvalue : the default value to associate with missing data in the
            field's values. The fillvalue must be of the same type as
            `datatype` and `values`
83

Jeff Piollé's avatar
Jeff Piollé committed
84 85 86 87
        attrs (dict) : a dictionary of the metadata associated with the field's
            values.
        units : the units in which the data are given (if
            applicable)
88

Jeff Piollé's avatar
Jeff Piollé committed
89 90
        description (str) : full name of the phenomenon. This corresponds
            to a long_name in attribute in a netCDF file.
91

Jeff Piollé's avatar
Jeff Piollé committed
92 93
        authority (str): naming authority referencing the provided
            standard name
94

Jeff Piollé's avatar
Jeff Piollé committed
95
        standard_name (optional, str): standard label for a phenomenon, with
Jeff Piollé's avatar
Jeff Piollé committed
96 97 98
            respect to the convention stated in `authority` argument.
            This corresponds to a standard_name attribute in a CF compliant
            NetCDF file.
99

Jeff Piollé's avatar
Jeff Piollé committed
100
        quality_vars: list of related quality fields
101

Jeff Piollé's avatar
Jeff Piollé committed
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
    """

    def __init__(self,
                 data,
                 name: Optional[str] = None,
                 dims: Optional[Tuple] = None,
                 datatype: Optional[numpy.dtype] = None,
                 fields: Optional[Tuple['Field']] = None,
                 dataset: Optional['Dataset'] = None,
                 fillvalue: Optional[Any] = None,
                 precision: Optional[int] = None,
                 description: Optional[str] = None,
                 standard_name: Optional[Union[str, Tuple[str, str]]] = None,
                 units: Optional[str] = None,
                 quality_vars: Optional[List[str]] = None,
                 attrs: Optional[Mapping[str, Any]] = None,
Jeff Piollé's avatar
Jeff Piollé committed
118
                 **kwargs) -> None:
Jeff Piollé's avatar
Jeff Piollé committed
119
        """
120
        """
Jeff Piollé's avatar
Jeff Piollé committed
121 122
        if name is not None and not isinstance(name, str):
            raise TypeError('name must be a string')
123

Jeff Piollé's avatar
Jeff Piollé committed
124 125
        if isinstance(data, xr.DataArray):
            self.array = data
126 127 128 129
        else:
            # create the DataArray from the provided information

            # dimensions
Jeff Piollé's avatar
Jeff Piollé committed
130 131 132 133
            if isinstance(dims, (list, tuple)):
                dims = tuple(dims)
            elif dims is None:
                dims = ()
134 135 136
            else:
                raise TypeError("Wrong type for dimensions")

Jeff Piollé's avatar
Jeff Piollé committed
137
            if data is None:
138 139 140 141 142 143
                # create default array
                if datatype is None:
                    raise ValueError(
                        "If you don't provide any data, you must at least "
                        "provide a datatype"
                    )
Jeff Piollé's avatar
Jeff Piollé committed
144
                if not isinstance(dims, OrderedDict):
145 146 147 148 149
                    raise TypeError(
                        "dimensions should be provided with their size in a "
                        "OrderedDict"
                    )
                data = numpy.ma.masked_all(
Jeff Piollé's avatar
Jeff Piollé committed
150
                    tuple(dims.values()), datatype)
151
            else:
Jeff Piollé's avatar
Jeff Piollé committed
152
                data = data
153 154

            # instantiate the xarray representation
Jeff Piollé's avatar
Jeff Piollé committed
155 156
            kwargs['dims'] = list(dims)
            kwargs['attrs'] = attrs
157 158
            self.array = xr.DataArray(
                data,
Jeff Piollé's avatar
Jeff Piollé committed
159 160 161 162
               # dims=list(dims),
               # attrs=attrs,
                name=name,
                **kwargs
163 164 165 166
            )

        # Overrides DataArray object when conflicts with the superceding
        # arguments
Jeff Piollé's avatar
Jeff Piollé committed
167 168 169 170 171 172
        self.name = name
        self.standard_name = standard_name
        self.description = description
        self.fill_value = fillvalue
        self.units = units
        self.array.attrs['quality_vars'] = quality_vars
173 174 175 176 177 178 179 180 181 182

        # components for complex fields
        if fields is not None:
            # Add components in case of a composite field :
            # each of these components must be itself a field
            for fld in fields:
                if not isinstance(fld, Field):
                    raise TypeError("Components must by Field class object")
            self.components = fields

Jeff Piollé's avatar
Jeff Piollé committed
183 184 185 186
        # attachment to a parent Dataset object
        self.dataset = dataset
        if (self.dataset is not None
                and self.array.name not in dataset._varnames):
187 188 189 190 191 192 193 194
            raise ValueError(
                "Field {} not found in this mapper".format(self.array.name)
            )

        # @TODO self.handler ???
        self.array.encoding['cerbere_status'] = "changed"

    @classmethod
Jeff Piollé's avatar
Jeff Piollé committed
195
    def to_field(cls, data: xr.DataArray) -> 'Field':
196 197 198
        """Cast a xarray DataArray to a
        :class:`cerbere.datamodel.field.Field` object
        """
Jeff Piollé's avatar
Jeff Piollé committed
199
        return Field(data=data)
200

Jeff Piollé's avatar
Jeff Piollé committed
201
    @property
Jeff Piollé's avatar
Jeff Piollé committed
202
    def to_dataarray(self):
203
        """Return the field values a xarray DataArray"""
Jeff Piollé's avatar
Jeff Piollé committed
204
        if self.dataset is None:
205 206
            return self.array
        else:
Jeff Piollé's avatar
Jeff Piollé committed
207
            return self.dataset.get_values(
208 209 210 211 212 213 214 215 216 217 218 219 220
                self.array.name,
                as_masked_array=False
            )

    def __str__(self):
        result = "Field : '%s'\n" % self.array.name
        if ('long_name' in self.array.attrs and
                self.array.attrs['long_name'] is not None):
            result += '    {}\n'.format(self.array.attrs['long_name'])
        result += '\n'

        # dimensions
        result = result + '    dimensions :\n'
Jeff Piollé's avatar
Jeff Piollé committed
221
        if self.dataset is None:
222 223
            dims = OrderedDict(self.array.sizes.items())
        else:
Jeff Piollé's avatar
Jeff Piollé committed
224
            dims = self.dataset.get_field_dims(self.name)
225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255
        for dim, size in dims.items():
            result += '      # {} : {}\n'.format(dim, size)

        attrs = self.array.attrs.items()

        # standard attributes
        result = result + '    standard CF attributes :\n'
        for att, val in attrs:
            if att in FIELD_ATTRS:
                result += '      # {} : {}\n'.format(att, val)

        # free form attributes
        result = result + '    other attributes :\n'
        for att, val in attrs:
            if att not in FIELD_ATTRS and att not in FIELD_EXCL_ATTRS:
                result = result + '      # {} : {}\n'.format(att, val)

        return result

    @property
    def name(self):
        return self.array.name

    @name.setter
    def name(self, value):
        self.array.name = value

    @property
    def attrs(self):
        return self.array.attrs

256 257 258 259
    @attrs.setter
    def attrs(self, attrs):
        self.array.attrs = attrs

260 261
    @property
    def dims(self):
Jeff Piollé's avatar
Jeff Piollé committed
262
        if self.dataset is None:
263 264
            return tuple(self.array.dims)
        else:
PIOLLE's avatar
PIOLLE committed
265
            return self.dataset.get_field_dims(self.name)
266 267 268

    @dims.setter
    def dims(self, dims):
Jeff Piollé's avatar
Jeff Piollé committed
269
        if self.dataset is None:
270 271
            self.array.dims = dims
        else:
Jeff Piollé's avatar
Jeff Piollé committed
272
            self.dataset.set_dimensions(dims)
273

274 275 276 277
    @property
    def dimnames(self):
        return tuple(self.dims.keys())

278 279
    def get_dimsize(self, dimname) -> int:
        """Return the size of a field dimension"""
Jeff Piollé's avatar
Jeff Piollé committed
280
        if self.dataset is None:
281 282
            return self.array.sizes[dimname]
        else:
Jeff Piollé's avatar
Jeff Piollé committed
283
            return self.dataset.get_dimsize(dimname)
284 285 286 287 288 289 290 291 292 293 294 295 296 297

    @property
    def fill_value(self):
        """return the value for missing data"""
        try:
            return self.array.encoding['_FillValue']
        except KeyError:
            return None

    @fill_value.setter
    def fill_value(self, fill_value):
        """set the value for missing data"""
        self.array.encoding['_FillValue'] = fill_value

298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323
    @property
    def valid_min(self):
        """return the minimum valid value"""
        try:
            return self.array.attrs['valid_min']
        except KeyError:
            return None

    @valid_min.setter
    def valid_min(self, value):
        """set the minimum valid value"""
        self.array.attrs['valid_min'] = value

    @property
    def valid_max(self):
        """return the maximum valid value"""
        try:
            return self.array.attrs['valid_max']
        except KeyError:
            return None

    @valid_max.setter
    def valid_max(self, value):
        """set the maximum valid value"""
        self.array.attrs['valid_max'] = value

324
    @property
Jeff Piollé's avatar
Jeff Piollé committed
325 326
    def units(self) -> str:
        """return the field units (``units`` CF attribute)"""
327 328 329
        try:
            return self.array.attrs['units']
        except KeyError:
Jeff Piollé's avatar
Jeff Piollé committed
330
            return
331 332

    @units.setter
Jeff Piollé's avatar
Jeff Piollé committed
333 334
    def units(self, units: str):
        """set the variable units (``units`` CF attribute)"""
335 336
        self.array.attrs['units'] = units

PIOLLE's avatar
PIOLLE committed
337 338
    @property
    def description(self):
Jeff Piollé's avatar
Jeff Piollé committed
339
        """return the field description (``long_name`` CF attribute)"""
PIOLLE's avatar
PIOLLE committed
340 341 342 343 344 345
        try:
            return self.array.attrs['long_name']
        except KeyError:
            return None

    @units.setter
Jeff Piollé's avatar
Jeff Piollé committed
346 347
    def description(self, description: str) -> None:
        """set the field description (``long_name`` CF attribute)"""
PIOLLE's avatar
PIOLLE committed
348 349 350
        self.array.attrs['long_name'] = description

    @property
Jeff Piollé's avatar
Jeff Piollé committed
351 352
    def standard_name(self) -> str:
        """return the field standard name (``standard_name`` CF attribute)"""
PIOLLE's avatar
PIOLLE committed
353 354 355 356 357 358 359 360 361
        try:
            return (
                self.array.attrs['standard_name'],
                self.array.attrs['cf_authority']
            )
        except KeyError:
            return None

    @units.setter
Jeff Piollé's avatar
Jeff Piollé committed
362 363
    def standard_name(self, standard_name: str) -> None:
        """set the standard_name (``standard_name`` CF attribute)"""
PIOLLE's avatar
PIOLLE committed
364 365 366 367 368 369 370 371 372 373
        if isinstance(standard_name, tuple):
            self.array.attrs['standard_name'] = standard_name[0]
            self.array.attrs['authority'] = standard_name[1]
        elif standard_name is not None:
            self.array.attrs['standard_name'] = standard_name[0]
            self.array.attrs['authority'] = cf.CF_AUTHORITY
        else:
            self.array.attrs['standard_name'] = None
            self.array.attrs['authority'] = None

374 375
    @property
    def datatype(self):
Jeff Piollé's avatar
Jeff Piollé committed
376
        if self.dataset is None:
377 378
            return self.array.dtype
        else:
Jeff Piollé's avatar
Jeff Piollé committed
379
            return self.dataset.dataset[self.name].dtype
380 381 382 383 384 385 386 387 388 389 390 391 392

    @property
    def variable(self):
        """return the field variable definition"""
        var = Variable(self.array.name)
        if 'long_name' in self.array.attrs:
            var.description = self.array.attrs['long_name']
        if 'standard_name' in self.array.attrs:
            var.standardname = self.array.attrs['standard_name']
            try:
                var.authority = self.array.attrs['authority']
            except KeyError:
                logging.error(
393 394
                    "No authority attribute defined for standard name: {}"
                    .format(var.standardname)
395 396 397 398 399 400 401 402 403 404 405 406 407
                )
        return var

    @variable.setter
    def variable(self, variable):
        """set the field variable definition"""
        self.array.name = variable.shortname
        self.array.attrs['long_name'] = variable.description
        self.array.attrs['authority'] = variable.authority
        self.array.attrs['standard_name'] = variable.standardname

    def is_composite(self) -> bool:
        """
Jeff Piollé's avatar
Jeff Piollé committed
408
        True if the field is a composite field.
409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429

        A composite field is a composition of sub fields (vector components,
        real and imaginary part of a complex, ...)
        """
        return self.components is not None

    @property
    def components(self):
        """Return the list of components of the field.

        Components (or sub fields) are intended for non scalar fields (ex:
        vector like current or wind, real and imaginary part of a complex,...)
        """
        res = []
        if self.is_composite():
            for rec in self.components:
                res.extend(rec.get_components())
            return res
        else:
            return [self]

Jeff Piollé's avatar
Jeff Piollé committed
430 431 432 433 434
    def _attach_dataset(
            self,
            dataset: 'Dataset'
        ) -> None:
        """Attach a parent dataset to the field.
435

Jeff Piollé's avatar
Jeff Piollé committed
436 437 438 439
        cerbere uses lazy loading. Data are not actually loaded until they
        are explicitly requested (calling :func:`get_values`). This pointer to
        the parent dataset, set by this function, is used to locate where the
        data have to be read from (or saved later for a new or updated file).
440 441

        Args:
Jeff Piollé's avatar
Jeff Piollé committed
442
            handler(:class:`Dataset`): the parent dataset
443 444 445
                attached to the field.

        """
Jeff Piollé's avatar
Jeff Piollé committed
446
        self.dataset = dataset
447 448 449

    def get_values(
            self,
Jeff Piollé's avatar
Jeff Piollé committed
450 451 452
            index: Optional[Mapping[str, slice]] = None,
            padding: Optional[bool] = False,
            **kwargs) -> 'np.ma.MaskedArray':
453
        """
Jeff Piollé's avatar
Jeff Piollé committed
454
        Return the field values as a :class:`numpy.ma.MaskedArray` object.
455 456

        Args:
Jeff Piollé's avatar
Jeff Piollé committed
457 458
            index: any kind of xarray indexing compatible with xarray
                :func:`~xarray.DataArray.isel` selection method.
459

Jeff Piollé's avatar
Jeff Piollé committed
460 461
            padding: pad the result with fill values where slices are out of the
             field dimension limits. Default is False.
462

Jeff Piollé's avatar
Jeff Piollé committed
463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484
        Returns:
            The field values as a numpy MaskedArray

        """
        return self._get_values(
            index=index, padding=padding, as_masked_array=True, **kwargs
        )

    def get_xvalues(
            self,
            index: Optional[Mapping[str, slice]] = None,
            padding: Optional[bool] = False,
            **kwargs) -> 'xr.DataArray':
        """
        Return the field values as a :class:`xarray.DataArray` object.

        Args:
            index: any kind of xarray indexing compatible with
                :func:`xarray.DataArray.isel` selection method.

            padding: pad the result with fill values where slices are out of the
             field dimension limits. Default is False.
485 486

        Returns:
Jeff Piollé's avatar
Jeff Piollé committed
487
            The field values as a xarray DataArray
488 489

        """
Jeff Piollé's avatar
Jeff Piollé committed
490 491 492 493 494 495 496 497 498 499
        return self._get_values(
            index=index, padding=padding, as_masked_array=False, **kwargs
        )

    def _get_values(
            self,
            index=None,
            padding=False,
            as_masked_array=True,
            **kwargs):
500 501 502 503 504 505
        allkwargs = {
            'index': index,
            'padding': padding,
            'as_masked_array': as_masked_array,
            **kwargs
        }
Jeff Piollé's avatar
Jeff Piollé committed
506 507
        if self.dataset is None:
            return numpy.ma.array(self._read_dataarray(self.array, **allkwargs))
508
        else:
Jeff Piollé's avatar
Jeff Piollé committed
509
            return self.dataset.get_values(self.name, **allkwargs)
510 511

    @classmethod
Jeff Piollé's avatar
Jeff Piollé committed
512
    def _read_dataarray(
513 514
            cls,
            xrdata,
Jeff Piollé's avatar
Jeff Piollé committed
515 516 517
            index: Mapping[Hashable, Any]=None,
            padding: bool=False,
            as_masked_array: bool=True,
518 519 520 521 522 523 524 525 526 527 528 529
            **kwargs
    ):
        """
        Extract values as from a xarray DataArray object.

        Values are returned as a numpy ``MaskedArray`` or xarray ``DataArray``
        object.

        Args:
            xrdata (:class:`xarray.DataArray`): the xarray ``DataArray`` object
                from which to extract the values.

Jeff Piollé's avatar
Jeff Piollé committed
530
            index (dict, optional): any kind of xarray indexing compatible with
531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549
                ``isel`` selection method.

            as_masked_array (bool): return the result as a numpy masked array
                (by default), or as a xarray if set to False.

            padding (bool): if True, pad the result with fill values where
                slices are out of the field size.

        Returns:
            :class:`numpy.ma.MaskedArray` or :class:`xarray.DataArray`: the
                data as a numpy ``MaskedArray``, if ``as_masked_array`` is set,
                or as a xarray ``DataArray`` object
        """
        if as_masked_array:
            data = xrdata.isel(index).to_masked_array(copy=False)
        else:
            data = xrdata.isel(index).values()

        if padding:
Jeff Piollé's avatar
Jeff Piollé committed
550
            data = cls._pad_data(xrdata, data, index)
551 552 553 554

        return data

    @classmethod
Jeff Piollé's avatar
Jeff Piollé committed
555 556
    def _pad_data(
            cls,
Jeff Piollé's avatar
Jeff Piollé committed
557 558
            array: 'to_dataarray.core.dataset.Dataset',
            subset: 'to_dataarray.core.dataset.Dataset',
Jeff Piollé's avatar
Jeff Piollé committed
559 560
            index: Optional[Mapping[str, slice]]
    ) -> 'numpy.ndarray':
561 562 563 564 565
        """
        pad with fill values the ``subset`` array extracted from ``array``
        where ``index`` is beyond the limits of ``array``.
        """
        pad_edges = []
Jeff Piollé's avatar
Jeff Piollé committed
566 567 568
        for dim in list(array.dims):
            if dim in index:
                dslice = index[dim]
569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584
                pad_edges.append([
                    abs(min(0, dslice.start)),
                    abs(min(0, array.sizes[dim] - dslice.stop))
                ])
            else:
                pad_edges.append([0, 0])

        res = numpy.pad(
            subset, pad_edges, 'constant',
            constant_values=numpy.nan
        )
        if isinstance(subset, numpy.ma.MaskedArray):
            res = numpy.ma.fix_invalid(res, copy=False)

        return res

Jeff Piollé's avatar
Jeff Piollé committed
585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601
    def set_values(
            self,
            values: numpy.ndarray,
            index: Optional[Mapping[str, slice]] = None,
            **kwargs) -> None:
        """set the values of a field.

        It is  possible to set only a subset of the field data array, using
        ``index``:

        >>> import numpy as np
        >>> data = np.ma.zeros((100, 200))
        >>> field = Field(data, name='test', dims=('x', 'y'))
        >>> field.set_values(
        >>>        np.full((10, 5,), 2),
        >>>        {'x': slice(10, 20), 'y': slice(0, 5)}
        >>>        )
602

Jeff Piollé's avatar
Jeff Piollé committed
603 604 605 606
        Args:
            values: the values to replace the ones in the field
            index: a dict of slices or indices of the subset to replace in the
                current field data array
607

Jeff Piollé's avatar
Jeff Piollé committed
608 609 610
        """
        if self.dataset is None:
            self._set_xrvalues(self.array, values, index=index)
611
        else:
Jeff Piollé's avatar
Jeff Piollé committed
612
            self.dataset.set_values(self.name, values, index=index)
613 614 615 616 617 618 619 620 621

    @classmethod
    def _set_xrvalues(
            cls,
            xrdata,
            values,
            index=None
    ):
        if index is None:
Jeff Piollé's avatar
Jeff Piollé committed
622
            xrdata.values[:] = values
623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643

        else:
            tmp = []
            for dim in xrdata.dims:
                if dim not in index:
                    tmp.append(slice(None))
                else:
                    tmp.append(index[dim])
            subset = tmp

            xrdata[tuple(subset)] = values

    def is_saved(self):
        """
        Return True is the content of the field is saved on file and
        was not updated since
        """
        if self.handler is None:
            return False
        return self.handler.is_saved()

Jeff Piollé's avatar
Jeff Piollé committed
644 645 646 647 648
    def bitmask_or(
            self,
            meanings,
            index: Mapping[Hashable, Any]=None,
            **kwargs):
649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700
        """helper function to get a boolean mask from a bit field.

        Bit (or flag) fields are arrays of integers where each bit has a
        specific meaning, described in a ``flag_meaning`` field attribute.
        Providing a list of the meanings to be tested in ``meaning``, a boolean
        mask if built, using the `or` logical operator, with a True value
        everywhere at least one of the provided meanings is set.

        The field must defined as in CF convention, with ``flags_masks``
        and ``flag_meanings`` field attributes.

        Args:
            meanings(list or str): a list of the meanings that have to be set
                or a str if only one bit is tested. Available meanings are
                listed in the ``flag_meanings`` attribute of the field.

        Returns:
            A boolean array.
        """
        if ('flag_meanings' not in self.array.attrs or
                'flag_masks' not in self.array.attrs):
            raise ValueError(
                "This is not mask field. Either flag_meanings or flag_masks "
                "is missing."
            )

        # transform mask attributes to list if it is not the case
        allmeanings = self.array.attrs['flag_meanings']
        if isinstance(allmeanings, str):
            allmeanings = [_ for _ in allmeanings.split(' ') if _ != '']

        allmasks = self.array.attrs['flag_masks']
        if isinstance(allmasks, str):
            allmasks = [_ for _ in allmasks.split(' ') if _ != '']

        criteria = meanings
        if isinstance(meanings, str):
            criteria = [meanings]

        # calculate sum (and) of all mask bits requested to be set
        masksum = 0
        for criterium in criteria:
            try:
                bit = allmeanings.index(criterium)
            except:
                raise ValueError("Unknown flag meaning {}".format(criterium))
            masksum += allmasks[bit]

        # calculate mask
        return self.get_values(slices=index) & int(masksum) != 0

    @classmethod
Jeff Piollé's avatar
Jeff Piollé committed
701 702 703 704 705 706 707
    def compute(
            cls,
            func,
            field1: 'Field', field2: 'Field'=None,
            **kwargs) -> 'Field':
        """Apply a function to a field (possibly combining with a second one)
         and returns the result as a new field.
708

Jeff Piollé's avatar
Jeff Piollé committed
709
        The function may be for instance a numpy MaskedArray operator
710 711 712 713 714
        such as numpy.ma.anom, numpy.ma.corr,...

        To be used with caution.

        Args:
Jeff Piollé's avatar
Jeff Piollé committed
715 716 717 718 719 720
            func (function) : the function to be called (ex: numpy.ma.anom)
            field1 (Field) : the field argument to the function
            field2 (Field, optional) : an optional 2nd field argument to the
                function
            kwargs : any argument to Field creation further describing the
                returned Field (units, name, ...).
721 722 723 724

        Returns:
            Field: the result field
        """
Jeff Piollé's avatar
Jeff Piollé committed
725
        if 'name' not in kwargs:
726 727
            varname = 'result'
        if field2 is None:
Jeff Piollé's avatar
Jeff Piollé committed
728
            values = func(field1.get_values())
729
        else:
Jeff Piollé's avatar
Jeff Piollé committed
730 731 732 733
            values = func(field1.get_values(), field2.get_values())
        field = Field(data=values,
                      name=varname,
                      dims=copy.copy(field1.dims),
734
                      datatype=field1.datatype,
Jeff Piollé's avatar
Jeff Piollé committed
735 736
                      fillvalue=field1.fill_value,
                      **kwargs)
737 738
        return field

PIOLLE's avatar
PIOLLE committed
739 740
    def clone(
            self,
Jeff Piollé's avatar
Jeff Piollé committed
741 742 743 744 745
            index: Mapping[Hashable, Any] = None,
            padding: bool = False,
            prefix: str = None,
            **kwargs) -> 'Field':
        """Create a copy of a field, or a subset defined by index, and
PIOLLE's avatar
PIOLLE committed
746
        padding out as required.
Jeff Piollé's avatar
Jeff Piollé committed
747 748

        The returned field does not contain any attachment to the source file
PIOLLE's avatar
PIOLLE committed
749
        attached to the original field, if any.
750

Jeff Piollé's avatar
Jeff Piollé committed
751 752 753
        Args:
            index (dict, optional):any kind of xarray indexing compatible with
                xarray :func:`~xarray.DataArray.isel` selection method.
PIOLLE's avatar
PIOLLE committed
754 755 756
            padding (bool, optional): True to pad out feature with fill values
                to the extent of the dimensions.
            prefix (str, optional): add a prefix string to the field names of
Jeff Piollé's avatar
Jeff Piollé committed
757
                the extracted subset.
PIOLLE's avatar
PIOLLE committed
758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778
        """
        if index is None:
            new_field = Field(data=self.array.copy(deep=True))
        else:
            new_index = {
                dim: val
                for dim, val in index.items() if dim in self.array.dims
            }
            subarray = self.array[new_index]
            if padding:
                data = self.__pad_data(subarray.values, new_index)
                subarray.set_values(data)

            new_field = Field(data=subarray.copy(deep=True))

        # detach from any dataset
        new_field.dataset = None

        if prefix is not None:
            new_field.set_name(prefix + new_field.name)
        return new_field
779

Jeff Piollé's avatar
Jeff Piollé committed
780 781
    def rename(self, newname: str) -> None:
        """Rename the field inplace.
782 783

        Args:
Jeff Piollé's avatar
Jeff Piollé committed
784
            newname (str): new name of the field
785 786 787 788 789 790 791 792 793 794
        """

        if self._mapper is not None:
            if self.name not in self._mapper.get_fieldnames():
                raise ValueError("Field {} not existing".format(self.name))

            self.dataset = self.dataset.rename({self.name: newname})

        self.name = newname

Jeff Piollé's avatar
Jeff Piollé committed
795
    def __add__(self, other: 'Field') -> 'Field':
796 797 798 799 800
        """Return a new field with the sum of current and an other field."""
        res = Field.convert_from_xarray(self.xrdata + other.xrdata)
        res.xrdata.name = "{}_{}_sum".format(self.name, other.name)
        return res

Jeff Piollé's avatar
Jeff Piollé committed
801
    def __sub__(self, other: 'Field') -> 'Field':
802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834
        """Return a new field with the difference of current and an other
        field.
        """
        res = Field.convert_from_xarray(self.xrdata - other.xrdata)
        res.xrdata.name = "{}_{}_difference".format(self.name, other.name)
        return res


def module(u, v, variable=None):
    """Return the module field from its two components

    The module is sqrt(u² + v²)

    Args:
        u (Field) : the eastward component
        v (Field) : the northward component
        variable (Variable) : variable of the returned module field. If not
            provided, the returned field is created with a basic variable
            definition.

    Returns:
        Field: the module field
    """
    values = numpy.ma.sqrt(numpy.ma.power(u.get_values(), 2)
                           + numpy.ma.power(v.get_values(), 2)
                           )
    if variable is None:
        if 'eastward_' in u.variable.shortname:
            varname = u.variable.shortname.replace('eastward_', '')
        else:
            varname = 'module'
        variable = Variable(varname)
    field = Field(variable,
Jeff Piollé's avatar
Jeff Piollé committed
835
                  dims=copy.copy(u.dimensions),
836
                  datatype=u.datatype,
Jeff Piollé's avatar
Jeff Piollé committed
837
                  fillvalue=u.fill_value,
838 839 840
                  values=values,
                  units=u.units)
    return field