Commit 785805fb by Klaus Zimmermann

### Index function first spell (closes #245)

parent ac2bcc73
 ... @@ -68,6 +68,23 @@ index_functions: ... @@ -68,6 +68,23 @@ index_functions: condition: condition: kind: operator kind: operator first_spell: description: | Calculates the beginning of the first spell that meets a minimium duration. First, the threshold is transformed to the same standard_name and units as the input data. Then the thresholding is performed as condition(data, threshold), ie if condition is <, data < threshold. After that spells of at least minimum length are identified. Finally, locate the first occurrence of such spells. parameters: threshold: kind: quantity duration: kind: quantity condition: kind: operator last_occurrence: last_occurrence: description: | description: | Calculates the last time some condition is met. Calculates the last time some condition is met. ... ...
 ... @@ -22,5 +22,6 @@ from .percentile_functions import ( # noqa: F401 ... @@ -22,5 +22,6 @@ from .percentile_functions import ( # noqa: F401 ) ) from .spell_functions import ( # noqa: F401 from .spell_functions import ( # noqa: F401 FirstSpell, SpellLength, SpellLength, ) )
 ... @@ -2,10 +2,60 @@ from cf_units import Unit ... @@ -2,10 +2,60 @@ from cf_units import Unit import dask.array as da import dask.array as da import numpy as np import numpy as np from .spell_kernels import make_spell_length_kernels from .spell_kernels import make_first_spell_kernels, make_spell_length_kernels from .support import normalize_axis, IndexFunction, ThresholdMixin, ReducerMixin from .support import normalize_axis, IndexFunction, ThresholdMixin, ReducerMixin class FirstSpell(ThresholdMixin, IndexFunction): def __init__(self, threshold, condition, duration, dead_period): super().__init__(threshold, condition, units=Unit("days")) self.duration = duration self.dead_period = dead_period self.kernels = make_first_spell_kernels(duration.points[0]) def pre_aggregate_shape(self, *args, **kwargs): return (4,) def call_func(self, data, axis, **kwargs): raise NotImplementedError def lazy_func(self, data, axis, **kwargs): axis = normalize_axis(axis, data.ndim) mask = da.ma.getmaskarray(data).any(axis=axis) data = da.moveaxis(data, axis, -1) offset = self.dead_period.points[0] data = data[..., offset:] first_spell_data = da.reduction( data, self.chunk, self.aggregate, keepdims=True, output_size=4, axis=-1, dtype=int, concatenate=False, meta=np.array((), dtype=int), ) res = first_spell_data[..., 2].copy() res = da.where(res >= 0, res - offset, res) res = da.ma.masked_array(da.ma.getdata(res), mask) return res.astype("float32") def chunk(self, raw_data, axis, keepdims, computing_meta=False): if computing_meta: return np.array((), dtype=int) data = self.condition(raw_data, self.threshold.points) chunk_res = self.kernels.chunk(data) return chunk_res def aggregate(self, x_chunk, axis, keepdims): if not isinstance(x_chunk, list): return x_chunk res = self.kernels.aggregate(np.array(x_chunk)) return res class SpellLength(ThresholdMixin, ReducerMixin, IndexFunction): class SpellLength(ThresholdMixin, ReducerMixin, IndexFunction): def __init__(self, threshold, condition, reducer, fuse_periods=False): def __init__(self, threshold, condition, reducer, fuse_periods=False): super().__init__(threshold, condition, reducer, units=Unit("days")) super().__init__(threshold, condition, reducer, units=Unit("days")) ... ...