Commit 1cda4d74 authored by Lars Bärring's avatar Lars Bärring Committed by Klaus Zimmermann
Browse files

Index function first occurrence

parent f43f715d
......@@ -20,6 +20,20 @@ index_functions:
condition:
kind: operator
first_occurrence:
description: |
Calculates the first time some condition is met.
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.
Finally, locate the first occurrence when condition is met.
parameters:
threshold:
kind: quantity
condition:
kind: operator
spell_length:
description: |
Calculates statistics on lengths of spells.
......
# -*- coding: utf-8 -*-
from datetime import datetime
from cf_units import Unit
import dask.array as da
import numpy as np
......@@ -82,6 +84,51 @@ class CountOccurrences:
return res.astype('float32')
class FirstOccurrence:
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.NO_OCCURRENCE = np.inf
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)
cond = self.condition(data, self.threshold.data)
res = np.ma.where(cond.any(axis=axis),
cond.argmax(axis=axis),
self.NO_OCCURRENCE)
return res.astype('float32')
def lazy_func(self, data, axis, **kwargs):
axis = normalize_axis(axis, data.ndim)
mask = da.ma.getmaskarray(data).any(axis=axis)
cond = self.lazy_condition(data, self.threshold.data)
res = da.where(cond.any(axis=axis),
cond.argmax(axis=axis),
self.NO_OCCURRENCE)
res = da.ma.masked_array(da.ma.getdata(res), mask)
return res.astype('float32')
def post_process(self, collapsed_cube, data_result, coords,
period, **kwargs):
time = collapsed_cube.coord('time')
calendar = time.units.calendar
offsets = np.empty_like(time.points, dtype=data_result.dtype)
for i, representative_date in enumerate(time.cells()):
year = representative_date.point.year
start_date = datetime(year, period.first_month_number, 1)
units = Unit('days since {}-01-01'.format(year), calendar=calendar)
offsets[i] = units.date2num(start_date)
collapsed_cube.data = (collapsed_cube.core_data()
+ offsets[:, None, None])
return collapsed_cube
class SpellLength:
def __init__(self, threshold, condition, reducer):
self.threshold = threshold
......
......@@ -50,6 +50,7 @@ setuptools.setup(
],
'climix.index_functions': [
'count_occurrences=climix.index_functions:CountOccurrences',
'first_occurrence=climix.index_functions:FirstOccurrence',
'spell_length=climix.index_functions:SpellLength',
'thresholded_statistics=climix.index_functions:ThresholdedStatistics',
],
......
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