support.py 2.05 KB
 Klaus Zimmermann committed Nov 21, 2019 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 `````` Klaus Zimmermann committed Nov 21, 2019 82 83 84 `````` threshold.points = threshold.points.astype(input_cube.dtype) if threshold.has_bounds(): threshold.bounds = threshold.bounds.astype(input_cube.dtype) `````` Klaus Zimmermann committed Nov 21, 2019 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]``````