Source code for flowws_analysis.Colormap

import argparse
import functools

import flowws
from flowws import Argument as Arg
import matplotlib,
import numpy as np

[docs]@flowws.add_stage_arguments class Colormap(flowws.Stage): """Access and use matplotlib colormaps on scalar quantities. This module emits a `color` value, calculated using a given scalar argument and matplotlib colormap name. Valid scalars quantities can be provided to this module by saving them in the scope and adding their name to the `color_scalars` list. """ ARGS = [ Arg('colormap_name', '-c', str, 'viridis', help='Name of the matplotlib colormap to use'), Arg('argument', '-a', str, help='Name of the value to map to colors'), Arg('range', '-r', (float, float), help='Minimum and maximum values of the scalar to be mapped'), ] def run(self, scope, storage): """Generate an array of colors using the given color scalars.""" color_scalars = scope.setdefault('color_scalars', []) if 'type' not in color_scalars and 'type' in scope: color_scalars.append('type') argument = self.arguments.get('argument', None) if (argument is not None and argument not in color_scalars and argument in scope): color_scalars.append(argument) self.arg_specifications['argument'].valid_values = color_scalars if argument is None: self.arguments['argument'] = color_scalars[0] self.arg_specifications['colormap_name'].valid_values = \ sorted( N = len(scope['position']) try: values = scope[self.arguments['argument']].copy() except KeyError: values = np.full(N, 0.5) normalize = None if self.arguments.get('range', None): (vmin, vmax) = self.arguments['range'] normalize = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) cmap =['colormap_name']) smap =, cmap) scope['color'] = smap.to_rgba(values)