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)