Skip to content

GitLab

Commit 81da3e71 authored by Klaus Zimmermann's avatar Klaus Zimmermann
Browse files

Index functions refactoring (closes #137)

parent 29e95e28
Hide whitespace changes
Inline Side-by-side
climix/index_functions/__init__.py 0 → 100644
View file @ 81da3e71
# -*- coding: utf-8 -*-
from .index_functions import ( # noqa: F401
CountOccurrences,
FirstOccurrence,
LastOccurrence,
SpellLength,
Statistics,
ThresholdedStatistics,
TemperatureSum,
)
climix/index_functions.py → climix/index_functions/index_functions.py
View file @ 81da3e71
......@@ -6,73 +6,14 @@ 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 CountOccurrences:
def __init__(self, threshold, condition):
self.threshold = threshold
self.condition = NUMPY_OPERATORS[condition]
self.lazy_condition = DASK_OPERATORS[condition]
self.standard_name = None
self.units = Unit('days')
self.extra_coords = [threshold.copy()]
from .support import (normalize_axis,
IndexFunction,
ThresholdMixin, ReducerMixin)
def prepare(self, input_cube):
change_units(self.threshold,
input_cube.units,
input_cube.standard_name)
class CountOccurrences(ThresholdMixin, IndexFunction):
def __init__(self, threshold, condition):
super().__init__(threshold, condition, units=Unit('days'))
def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim)
......@@ -87,20 +28,10 @@ class CountOccurrences:
return res.astype('float32')
class FirstOccurrence:
class FirstOccurrence(ThresholdMixin, IndexFunction):
def __init__(self, threshold, condition):
self.threshold = threshold
self.condition = NUMPY_OPERATORS[condition]
self.lazy_condition = DASK_OPERATORS[condition]
self.standard_name = None
self.units = Unit('days')
super().__init__(threshold, condition, units=Unit('days'))
self.NO_OCCURRENCE = np.inf
self.extra_coords = [threshold.copy()]
def prepare(self, input_cube):
change_units(self.threshold,
input_cube.units,
input_cube.standard_name)
def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim)
......@@ -135,20 +66,10 @@ class FirstOccurrence:
return collapsed_cube
class LastOccurrence:
class LastOccurrence(ThresholdMixin, IndexFunction):
def __init__(self, threshold, condition):
self.threshold = threshold
self.condition = NUMPY_OPERATORS[condition]
self.lazy_condition = DASK_OPERATORS[condition]
self.standard_name = None
self.units = Unit('days')
super().__init__(threshold, condition, units=Unit('days'))
self.NO_OCCURRENCE = -np.inf
self.extra_coords = [threshold.copy()]
def prepare(self, input_cube):
change_units(self.threshold,
input_cube.units,
input_cube.standard_name)
def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim)
......@@ -185,19 +106,9 @@ class LastOccurrence:
return collapsed_cube
class SpellLength:
class SpellLength(ThresholdMixin, ReducerMixin, IndexFunction):
def __init__(self, threshold, condition, reducer):
self.threshold = threshold
self.condition = NUMPY_OPERATORS[condition]
self.reducer = NUMPY_REDUCERS[reducer]
self.standard_name = None
self.units = Unit('days')
self.extra_coords = [threshold.copy()]
def prepare(self, input_cube):
change_units(self.threshold,
input_cube.units,
input_cube.standard_name)
super().__init__(threshold, condition, reducer, units=Unit('days'))
def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim)
......@@ -241,13 +152,12 @@ class SpellLength:
return float(res)
class Statistics:
class Statistics(ReducerMixin, IndexFunction):
def __init__(self, reducer):
self.reducer = NUMPY_REDUCERS[reducer]
self.lazy_reducer = DASK_REDUCERS[reducer]
self.extra_coords = []
super().__init__(reducer)
def prepare(self, input_cube):
super().prepare(input_cube)
self.standard_name = input_cube.standard_name
self.units = input_cube.units
......@@ -262,19 +172,12 @@ class Statistics:
return res.astype('float32')
class ThresholdedStatistics:
class ThresholdedStatistics(ThresholdMixin, ReducerMixin, IndexFunction):
def __init__(self, threshold, condition, reducer):
self.threshold = threshold
self.condition = NUMPY_OPERATORS[condition]
self.reducer = NUMPY_REDUCERS[reducer]
self.lazy_condition = DASK_OPERATORS[condition]
self.lazy_reducer = DASK_REDUCERS[reducer]
self.extra_coords = [threshold.copy()]
super().__init__(threshold, condition, reducer, units=Unit('days'))
def prepare(self, input_cube):
change_units(self.threshold,
input_cube.units,
input_cube.standard_name)
super().prepare(input_cube)
self.standard_name = input_cube.standard_name
self.units = input_cube.units
......@@ -291,21 +194,18 @@ class ThresholdedStatistics:
return res.astype('float32')
class TemperatureSum:
class TemperatureSum(ThresholdMixin, IndexFunction):
def __init__(self, threshold, condition):
self.threshold = threshold
super().__init__(threshold, condition, units=Unit('days'))
if condition in ['>', '>=']:
self.fun = lambda d, t: np.maximum(d - t, 0)
self.lazy_fun = lambda d, t: da.maximum(d - t, 0)
else:
self.fun = lambda d, t: np.maximum(t - d, 0)
self.lazy_fun = lambda d, t: da.maximum(t - d, 0)
self.extra_coords = [threshold.copy()]
def prepare(self, input_cube):
change_units(self.threshold,
input_cube.units,
input_cube.standard_name)
super().prepare(input_cube)
self.standard_name = input_cube.standard_name
if input_cube.units.is_convertible('degC'):
self.units = 'degC days'
......
climix/index_functions/support.py 0 → 100644
View file @ 81da3e71
# -*- 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
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]
Supports Markdown
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