Commit 68497c0a authored by Klaus Zimmermann's avatar Klaus Zimmermann
Browse files

Make spell fusing across periods optional (closes #74)

parent 11279451
from cf_units import Unit
import dask
import dask.array as da
import numpy as np
......@@ -10,10 +9,11 @@ from .support import (normalize_axis,
class SpellLength(ThresholdMixin, ReducerMixin, IndexFunction):
def __init__(self, threshold, condition, reducer):
def __init__(self, threshold, condition, reducer, fuse_periods=False):
super().__init__(threshold, condition, reducer, units=Unit('days'))
self.spanning_spells = True
self.kernels = make_spell_length_kernels(self.scalar_reducer)
self.fuse_periods = fuse_periods
def pre_aggregate_shape(self, *args, **kwargs):
return (4,)
......@@ -78,18 +78,23 @@ class SpellLength(ThresholdMixin, ReducerMixin, IndexFunction):
spell_length = self.lazy_reducer(stack, axis=-1)
return spell_length
stack = []
this = data[0]
slice_shape = this.shape[:-1]
previous_tail = da.ma.masked_array(
da.zeros(slice_shape, dtype=np.float32),
da.ma.getmaskarray(data[0, ..., 3]))
if self.fuse_periods and len(data) > 1:
stack = []
this = data[0]
slice_shape = this.shape[:-1]
previous_tail = da.ma.masked_array(
da.zeros(slice_shape, dtype=np.float32),
da.ma.getmaskarray(data[0, ..., 3]))
for next_chunk in data[1:]:
spell_length, previous_tail = fuse(this, next_chunk, previous_tail)
stack.append(spell_length)
this = next_chunk
for next_chunk in data[1:]:
spell_length, previous_tail = fuse(this,
next_chunk,
previous_tail)
stack.append(spell_length)
this = next_chunk
stack.append(fuse_last(next_chunk, previous_tail))
res_data = da.stack(stack, axis=0)
stack.append(fuse_last(next_chunk, previous_tail))
res_data = da.stack(stack, axis=0)
else:
res_data = self.lazy_reducer(data[..., 1:], axis=-1)
return cube, res_data
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment