support.py 2.05 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
# -*- coding: utf-8 -*-

from cf_units import Unit
import dask.array as da
import numpy as np

from ..util import change_units


SUPPORTED_OPERATORS = [
    '<',
    '>',
    '<=',
    '>=',
]
SUPPORTED_REDUCERS = [
    'min',
    'max',
    'sum',
    'mean',
]


NUMPY_OPERATORS = {
    '<': np.less,
    '>': np.greater,
    '<=': np.less_equal,
    '>=': np.greater_equal,
}
NUMPY_REDUCERS = {
    'min': np.min,
    'max': np.max,
    'sum': np.sum,
    'mean': np.mean,
}


DASK_OPERATORS = {
    '<': da.less,
    '>': da.greater,
    '<=': da.less_equal,
    '>=': da.greater_equal,
}
DASK_REDUCERS = {
    'min': da.min,
    'max': da.max,
    'sum': da.sum,
    'mean': da.mean,
}


def normalize_axis(axis, ndim):
    if isinstance(axis, list) and len(axis) == 1:
        axis = axis[0]
    if axis < 0:
        # just cope with negative axis numbers
        axis += ndim
    return axis


class IndexFunction:
    def __init__(self, standard_name=None, units=Unit('no_unit')):
        super().__init__()
        self.standard_name = standard_name
        self.units = units
        self.extra_coords = []

    def prepare(self, input_cube):
        pass


class ThresholdMixin:
    def __init__(self, threshold, condition, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.threshold = threshold
        self.condition = NUMPY_OPERATORS[condition]
        self.lazy_condition = DASK_OPERATORS[condition]
        self.extra_coords.append(threshold.copy())

    def prepare(self, input_cube):
        threshold = self.threshold
82
83
84
        threshold.points = threshold.points.astype(input_cube.dtype)
        if threshold.has_bounds():
            threshold.bounds = threshold.bounds.astype(input_cube.dtype)
85
86
87
88
89
90
91
92
93
94
95
        change_units(threshold,
                     input_cube.units,
                     input_cube.standard_name)
        super().prepare(input_cube)


class ReducerMixin:
    def __init__(self, reducer, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.reducer = NUMPY_REDUCERS[reducer]
        self.lazy_reducer = DASK_REDUCERS[reducer]