class SingleArrayExpressionBase(ExpressionMixin, Task):
"""Base class for tasks with a single array and expression as inputs."""
inputs = [
params.NdArrayParameter('array_in', required=True),
params.StringParameter('expression', required=True)
]
outputs = [params.NdArrayParameter('array_out')]
allow_extra_args = True
def get_context(self, arr, expr, context):
"""
Returns a context dictionary for use in evaluating the expression.
:param arr: The input array.
:param expr: The input expression.
:param context: Evaluation context.
"""
expression_names = [
x for x in self.get_expression_names(expr)
if x not in set(context.keys()).union(['i'])
]
if len(expression_names) != 1:
raise ValueError('The expression must have exactly one variable.')
return {expression_names[0]: arr}
class ReduceByExpression(ExpressionMixin, Task):
"""Iteratively reduces a list of arrays using an expression."""
name = 'raster:reduce_by_expression'
inputs = [
params.ListParameter(params.NdArrayParameter(''),
'arrays_in',
required=True),
params.StringParameter('expression', required=True),
params.NdArrayParameter('initial_array', required=False)
]
outputs = [params.NdArrayParameter('array_out')]
allow_extra_args = True
def execute(self, arrays_in, expression, initial_array=None, **kwargs):
expression_names = [
x for x in self.get_expression_names(expression)
if x not in set(kwargs.keys())
]
if len(expression_names) != 2:
raise ValueError("The expression must have exactly two variables.")
def reduce_fn(x, y):
context = {expression_names[0]: x, expression_names[1]: y}
context.update(kwargs)
return self.evaluate_expression(expression, context)
args = [reduce_fn, arrays_in]
if initial_array is not None:
args.append(initial_array)
return reduce(*args)
class LoadRasterDataset(Task):
"""Loads a raster dataset from a NetCDF file."""
name = 'raster:load_dataset'
inputs = [params.StringParameter('path', required=True)]
outputs = [params.RasterDatasetParameter('dataset_out')]
def execute(self, path):
return Dataset(path, 'r')
class ArrayFromDataset(Task):
"""Reads a variable from a raster dataset into an array."""
name = 'raster:array_from_dataset'
inputs = [
params.RasterDatasetParameter('dataset', required=True),
params.StringParameter('variable', required=True)
]
outputs = [params.NdArrayParameter('array_out')]
def execute(self, dataset, variable):
return dataset[variable][:]
class MapByExpression(SingleArrayExpressionBase):
"""Applies a given expression to a list of arrays, returning a list with new arrays."""
name = 'raster:map_by_expression'
inputs = [
params.ListParameter(params.NdArrayParameter(''),
'arrays_in',
required=True),
params.StringParameter('expression', required=True),
params.BooleanParameter('generator', required=False)
]
outputs = [params.ListParameter(params.NdArrayParameter(''), 'arrays_out')]
def execute(self, arrays_in, expression, generator=False, **kwargs):
result = (self.evaluate_expression(
expression,
dict(self.get_context(a, expression, kwargs), i=i, **kwargs))
for i, a in enumerate(arrays_in))
return result if generator else list(result)