Commit cb47acaf authored by Klaus Zimmermann's avatar Klaus Zimmermann
Browse files

Store parameters as scalar coords (closes #73)

parent e81369cc
...@@ -41,4 +41,8 @@ class PointLocalAggregator(Aggregator): ...@@ -41,4 +41,8 @@ class PointLocalAggregator(Aggregator):
cube.standard_name = standard_name cube.standard_name = standard_name
if proposed_standard_name is not None: if proposed_standard_name is not None:
cube.attributes['proposed_standard_name'] = proposed_standard_name cube.attributes['proposed_standard_name'] = proposed_standard_name
for extra_coord in self.index_function.extra_coords:
change_units(extra_coord,
self.output_metadata.units, unit_standard_name)
cube.add_aux_coord(extra_coord)
return cube return cube
...@@ -67,6 +67,7 @@ class CountOccurrences: ...@@ -67,6 +67,7 @@ class CountOccurrences:
self.lazy_condition = DASK_OPERATORS[condition] self.lazy_condition = DASK_OPERATORS[condition]
self.standard_name = None self.standard_name = None
self.units = Unit('days') self.units = Unit('days')
self.extra_coords = [threshold]
def prepare(self, input_cube): def prepare(self, input_cube):
change_units(self.threshold, change_units(self.threshold,
...@@ -75,13 +76,13 @@ class CountOccurrences: ...@@ -75,13 +76,13 @@ class CountOccurrences:
def call_func(self, data, axis, **kwargs): def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
cond = self.condition(data, self.threshold.data) cond = self.condition(data, self.threshold.points)
res = np.count_nonzero(cond, axis=axis) res = np.count_nonzero(cond, axis=axis)
return res.astype('float32') return res.astype('float32')
def lazy_func(self, data, axis, **kwargs): def lazy_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
cond = self.lazy_condition(data, self.threshold.data) cond = self.lazy_condition(data, self.threshold.points)
res = da.count_nonzero(cond, axis=axis) res = da.count_nonzero(cond, axis=axis)
return res.astype('float32') return res.astype('float32')
...@@ -94,6 +95,7 @@ class FirstOccurrence: ...@@ -94,6 +95,7 @@ class FirstOccurrence:
self.standard_name = None self.standard_name = None
self.units = Unit('days') self.units = Unit('days')
self.NO_OCCURRENCE = np.inf self.NO_OCCURRENCE = np.inf
self.extra_coords = [threshold]
def prepare(self, input_cube): def prepare(self, input_cube):
change_units(self.threshold, change_units(self.threshold,
...@@ -102,7 +104,7 @@ class FirstOccurrence: ...@@ -102,7 +104,7 @@ class FirstOccurrence:
def call_func(self, data, axis, **kwargs): def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
cond = self.condition(data, self.threshold.data) cond = self.condition(data, self.threshold.points)
res = np.ma.where(cond.any(axis=axis), res = np.ma.where(cond.any(axis=axis),
cond.argmax(axis=axis), cond.argmax(axis=axis),
self.NO_OCCURRENCE) self.NO_OCCURRENCE)
...@@ -111,7 +113,7 @@ class FirstOccurrence: ...@@ -111,7 +113,7 @@ class FirstOccurrence:
def lazy_func(self, data, axis, **kwargs): def lazy_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
mask = da.ma.getmaskarray(data).any(axis=axis) mask = da.ma.getmaskarray(data).any(axis=axis)
cond = self.lazy_condition(data, self.threshold.data) cond = self.lazy_condition(data, self.threshold.points)
res = da.where(cond.any(axis=axis), res = da.where(cond.any(axis=axis),
cond.argmax(axis=axis), cond.argmax(axis=axis),
self.NO_OCCURRENCE) self.NO_OCCURRENCE)
...@@ -141,6 +143,7 @@ class LastOccurrence: ...@@ -141,6 +143,7 @@ class LastOccurrence:
self.standard_name = None self.standard_name = None
self.units = Unit('days') self.units = Unit('days')
self.NO_OCCURRENCE = -np.inf self.NO_OCCURRENCE = -np.inf
self.extra_coords = [threshold]
def prepare(self, input_cube): def prepare(self, input_cube):
change_units(self.threshold, change_units(self.threshold,
...@@ -149,7 +152,7 @@ class LastOccurrence: ...@@ -149,7 +152,7 @@ class LastOccurrence:
def call_func(self, data, axis, **kwargs): def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
cond = self.condition(np.flip(data, axis=axis), self.threshold.data) cond = self.condition(np.flip(data, axis=axis), self.threshold.points)
ndays = data.shape[axis] ndays = data.shape[axis]
res = np.ma.where(cond.any(axis=axis), res = np.ma.where(cond.any(axis=axis),
ndays - cond.argmax(axis=axis), ndays - cond.argmax(axis=axis),
...@@ -159,7 +162,7 @@ class LastOccurrence: ...@@ -159,7 +162,7 @@ class LastOccurrence:
def lazy_func(self, data, axis, **kwargs): def lazy_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
mask = da.ma.getmaskarray(data).any(axis=axis) mask = da.ma.getmaskarray(data).any(axis=axis)
cond = self.lazy_condition(da.flip(data, axis), self.threshold.data) cond = self.lazy_condition(da.flip(data, axis), self.threshold.points)
ndays = data.shape[axis] ndays = data.shape[axis]
res = da.where(cond.any(axis=axis), res = da.where(cond.any(axis=axis),
ndays - cond.argmax(axis=axis), ndays - cond.argmax(axis=axis),
...@@ -189,6 +192,7 @@ class SpellLength: ...@@ -189,6 +192,7 @@ class SpellLength:
self.reducer = NUMPY_REDUCERS[reducer] self.reducer = NUMPY_REDUCERS[reducer]
self.standard_name = None self.standard_name = None
self.units = Unit('days') self.units = Unit('days')
self.extra_coords = [threshold]
def prepare(self, input_cube): def prepare(self, input_cube):
change_units(self.threshold, change_units(self.threshold,
...@@ -210,7 +214,7 @@ class SpellLength: ...@@ -210,7 +214,7 @@ class SpellLength:
return res.astype('float32') return res.astype('float32')
def __call__(self, raw_data): def __call__(self, raw_data):
data = self.condition(raw_data, self.threshold.data) data = self.condition(raw_data, self.threshold.points)
where = np.flatnonzero where = np.flatnonzero
x = data[1:] != data[:-1] x = data[1:] != data[:-1]
starts = where(x) + 1 starts = where(x) + 1
...@@ -244,6 +248,7 @@ class ThresholdedStatistics: ...@@ -244,6 +248,7 @@ class ThresholdedStatistics:
self.reducer = NUMPY_REDUCERS[reducer] self.reducer = NUMPY_REDUCERS[reducer]
self.lazy_condition = DASK_OPERATORS[condition] self.lazy_condition = DASK_OPERATORS[condition]
self.lazy_reducer = DASK_REDUCERS[reducer] self.lazy_reducer = DASK_REDUCERS[reducer]
self.extra_coords = [threshold]
def prepare(self, input_cube): def prepare(self, input_cube):
change_units(self.threshold, change_units(self.threshold,
...@@ -254,12 +259,12 @@ class ThresholdedStatistics: ...@@ -254,12 +259,12 @@ class ThresholdedStatistics:
def call_func(self, data, axis, **kwargs): def call_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
comb = self.condition(data, self.threshold.data) comb = self.condition(data, self.threshold.points)
res = self.reducer(np.ma.masked_where(~comb, data), axis=axis) res = self.reducer(np.ma.masked_where(~comb, data), axis=axis)
return res.astype('float32') return res.astype('float32')
def lazy_func(self, data, axis, **kwargs): def lazy_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim) axis = normalize_axis(axis, data.ndim)
comb = self.lazy_condition(data, self.threshold.data) comb = self.lazy_condition(data, self.threshold.points)
res = self.lazy_reducer(da.ma.masked_where(~comb, data), axis=axis) res = self.lazy_reducer(da.ma.masked_where(~comb, data), axis=axis)
return res.astype('float32') return res.astype('float32')
...@@ -80,27 +80,25 @@ class ParameterKind(Enum): ...@@ -80,27 +80,25 @@ class ParameterKind(Enum):
REDUCER = 'reducer' REDUCER = 'reducer'
@dataclass
class Parameter:
kind: ParameterKind
PARAMETER_KINDS = {} PARAMETER_KINDS = {}
@dataclass @dataclass
class ParameterQuantity(Parameter): class ParameterQuantity:
var_name: str
standard_name: str standard_name: str
data: Any data: Any
units: str units: str
long_name: str = None long_name: str = None
kind: ParameterKind = ParameterKind.QUANTITY
@property @property
def parameter(self): def parameter(self):
return iris.cube.Cube(self.data, return iris.coords.AuxCoord(self.data,
self.standard_name, self.standard_name,
self.long_name, self.long_name,
units=self.units) self.var_name,
units=self.units)
def instantiate(self, parameters): def instantiate(self, parameters):
data = self.data data = self.data
...@@ -111,7 +109,7 @@ class ParameterQuantity(Parameter): ...@@ -111,7 +109,7 @@ class ParameterQuantity(Parameter):
if ln is not None: if ln is not None:
ln = ln.format(**parameters) ln = ln.format(**parameters)
param = ParameterQuantity( param = ParameterQuantity(
ParameterKind.QUANTITY, format_var_name(self.var_name, parameters),
self.standard_name.format(**parameters), self.standard_name.format(**parameters),
data, data,
self.units, self.units,
...@@ -123,8 +121,10 @@ PARAMETER_KINDS['quantity'] = ParameterQuantity ...@@ -123,8 +121,10 @@ PARAMETER_KINDS['quantity'] = ParameterQuantity
@dataclass @dataclass
class ParameterOperator(Parameter): class ParameterOperator:
name: str
operator: str operator: str
kind: ParameterKind = ParameterKind.OPERATOR
@property @property
def parameter(self): def parameter(self):
...@@ -140,8 +140,10 @@ PARAMETER_KINDS['operator'] = ParameterOperator ...@@ -140,8 +140,10 @@ PARAMETER_KINDS['operator'] = ParameterOperator
@dataclass @dataclass
class ParameterReducer(Parameter): class ParameterReducer:
name: str
reducer: str reducer: str
kind: ParameterKind = ParameterKind.REDUCER
@property @property
def parameter(self): def parameter(self):
...@@ -156,6 +158,9 @@ class ParameterReducer(Parameter): ...@@ -156,6 +158,9 @@ class ParameterReducer(Parameter):
PARAMETER_KINDS['reducer'] = ParameterReducer PARAMETER_KINDS['reducer'] = ParameterReducer
Parameter = Union[ParameterQuantity, ParameterOperator, ParameterReducer]
@dataclass @dataclass
class IndexFunction: class IndexFunction:
name: str name: str
...@@ -188,8 +193,8 @@ class IndexDefinition: ...@@ -188,8 +193,8 @@ class IndexDefinition:
return idx return idx
def build_parameter(metadata): def build_parameter(name, metadata):
return PARAMETER_KINDS[metadata['kind']](**metadata) return PARAMETER_KINDS[metadata['kind']](name, **metadata)
def build_index(metadata, source=None): def build_index(metadata, source=None):
...@@ -210,7 +215,7 @@ def build_index(metadata, source=None): ...@@ -210,7 +215,7 @@ def build_index(metadata, source=None):
if params is None: if params is None:
parameters = {} parameters = {}
else: else:
parameters = {name: build_parameter(params[name]) parameters = {name: build_parameter(name, params[name])
for name in params} for name in params}
index_function = IndexFunction(metadata['index_function']['name'], index_function = IndexFunction(metadata['index_function']['name'],
parameters) parameters)
......
...@@ -4,11 +4,11 @@ from .change_pr_units import is_precipitation, change_pr_units ...@@ -4,11 +4,11 @@ from .change_pr_units import is_precipitation, change_pr_units
from .cube_diffs import cube_diffs from .cube_diffs import cube_diffs
def change_units(cube, new_units, new_standard_name=None): def change_units(cube_or_coord, new_units, new_standard_name=None):
if is_precipitation(new_standard_name): if is_precipitation(new_standard_name):
change_pr_units(cube, new_units, new_standard_name) change_pr_units(cube_or_coord, new_units, new_standard_name)
else: else:
cube.convert_units(new_units) cube_or_coord.convert_units(new_units)
__all__ = [ __all__ = [
......
...@@ -3,6 +3,7 @@ ...@@ -3,6 +3,7 @@
from collections import namedtuple, OrderedDict from collections import namedtuple, OrderedDict
from cf_units import Unit from cf_units import Unit
import iris
import six import six
...@@ -73,18 +74,18 @@ def get_mass_conversion(old_qty, new_qty, density): ...@@ -73,18 +74,18 @@ def get_mass_conversion(old_qty, new_qty, density):
return 1. return 1.
def change_pr_units(cube, new_units=None, new_standard_name=None, def change_pr_units(cube_or_coord, new_units=None, new_standard_name=None,
integration_time=Unit('1 day'), integration_time=Unit('1 day'),
density=DENSITY_WATER): density=DENSITY_WATER):
"""Converts precipitation units considering standard names. """Converts precipitation units considering standard names.
Converts precipitation units and cube.data to requested units. Converts precipitation units and data to requested units.
Standard names are converted too, if required. Standard names are converted too, if required.
Parameters Parameters
---------- ----------
cube: iris.cube.Cube cube_or_coord: iris.cube.Cube or iris.coords.Coord
cube with precipitation data precipitation data
new_units: cf_units.Unit new_units: cf_units.Unit
units to convert to units to convert to
new_standard_name: str new_standard_name: str
...@@ -92,8 +93,8 @@ def change_pr_units(cube, new_units=None, new_standard_name=None, ...@@ -92,8 +93,8 @@ def change_pr_units(cube, new_units=None, new_standard_name=None,
Returns Returns
------- -------
iris.cube.Cube iris.cube.Cube or iris.coords.Coord
cube with new units and standard name cube or coord with new units and standard name
""" """
new_units, new_standard_name = ensure_units_and_standard_name( new_units, new_standard_name = ensure_units_and_standard_name(
...@@ -101,8 +102,8 @@ def change_pr_units(cube, new_units=None, new_standard_name=None, ...@@ -101,8 +102,8 @@ def change_pr_units(cube, new_units=None, new_standard_name=None,
new_standard_name) new_standard_name)
old_units, old_standard_name = ensure_units_and_standard_name( old_units, old_standard_name = ensure_units_and_standard_name(
cube.units, cube_or_coord.units,
cube.standard_name) cube_or_coord.standard_name)
old_qty = PRECIPITATION_INFO[old_standard_name] old_qty = PRECIPITATION_INFO[old_standard_name]
new_qty = PRECIPITATION_INFO[new_standard_name] new_qty = PRECIPITATION_INFO[new_standard_name]
...@@ -111,6 +112,11 @@ def change_pr_units(cube, new_units=None, new_standard_name=None, ...@@ -111,6 +112,11 @@ def change_pr_units(cube, new_units=None, new_standard_name=None,
* get_mass_conversion(old_qty, new_qty, density)) * get_mass_conversion(old_qty, new_qty, density))
conv_factor = (old_units*conv).convert(1., new_units) conv_factor = (old_units*conv).convert(1., new_units)
cube.data = cube.core_data() * float(conv_factor) if isinstance(cube_or_coord, iris.cube.Cube):
cube.units = new_units cube_or_coord.data = cube_or_coord.core_data() * float(conv_factor)
cube.standard_name = new_standard_name elif isinstance(cube_or_coord, iris.coords.Coord):
cube_or_coord.points *= float(conv_factor)
if cube_or_coord.bounds is not None:
cube_or_coord.bounds *= float(conv_factor)
cube_or_coord.units = new_units
cube_or_coord.standard_name = new_standard_name
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