def run(self, scope, storage):
"""Center the system."""
box = scope.get('box', None)
positions = scope['position']
self.arg_specifications['particle'].valid_values = flowws.Range(
-1, len(positions), (True, False))
index = self.arguments['particle']
index = index if index >= 0 else None
if index is not None:
center_point = positions[index]
positions -= center_point[np.newaxis]
if box is not None:
positions = wrap(box, positions)
elif box is None:
center_point = np.mean(positions, axis=0)
positions -= center_point[np.newaxis]
else:
positions = center(box, positions)
scope['position'] = positions
def run(self, scope, storage):
"""Load records found in a getar file into the scope."""
scope['filename'] = self.arguments['filename']
scope['frame'] = self.arguments['frame']
scope['cache_key'] = scope['filename'], scope['frame']
garnett_traj = self._get_traj(self.arguments['filename'], storage)
self.arg_specifications['frame'].valid_values = flowws.Range(
0, len(garnett_traj), (True, False))
frame = garnett_traj[self.arguments['frame']]
# account for changes in the garnett API around v0.7
try:
types = frame.typeid
positions = frame.position
# account for some types of frames, like those from CIF
# files, not exposing orientations and raising an
# AttributeError instead
try:
orientations = frame.orientation
except AttributeError:
orientations = None
except AttributeError:
type_map = {k: i for (i, k) in enumerate(sorted(set(frame.types)))}
types = np.array([type_map[t] for t in frame.types],
dtype=np.uint32)
positions = frame.positions
try:
orientations = frame.orientations
except AttributeError:
orientations = None
try:
type_shapes = [
shape.type_shape for shape in frame.shapedef.values()
]
type_shapes = json.dumps(type_shapes)
scope['type_shapes.json'] = type_shapes
except AttributeError: # no shapedefs
pass
try:
scope['diameter'] = frame.diameter
except AttributeError: # no diameters
pass
scope['position'] = positions
if orientations is not None:
scope['orientation'] = orientations
scope['type'] = types
scope['box'] = frame.box.get_box_array()
scope['dimensions'] = frame.box.dimensions
def _set_record_frames(self):
frame = self.arguments['frame']
index_to_find = index_sort_key(
self._cached_frame_indices[self.arguments['frame']][1])
self.arg_specifications['frame'].valid_values = flowws.Range(
0, len(self._cached_frame_indices), (True, False))
for (rec, indices) in self._cached_record_frames.items():
try:
index = find_le(indices, index_to_find)[1]
rec.setIndex(index)
except ValueError:
pass
return self._cached_record_frames.keys()
class Plato(flowws.Stage):
"""Render shapes via :std:doc:`plato<plato:index>`.
This module uses the `position`, `orientation`, `type`, `color`,
and `type_shapes.json` quantities found in the scope, if provided,
to produce a scene of plato shapes.
The `type_shapes.json` value should provide a string of a
json-encoded list containing one *shape description* object
(described below) for each type. These will be converted into
plato primitives in conjunction with the `type` and other arrays.
Shape description objects are JSON objects with the following keys:
- type: one of "ConvexPolyhedron", "Disk", "Mesh", "Polygon", "Sphere"
- rounding_radius (only if type is "ConvexPolyhedron" or "Polygon"): rounding radius of a rounded shape, or 0 for perfectly faceted shapes
- vertices (only if type is "ConvexPolyhedron", "Mesh", or "Polygon"): coordinates in the shape's reference frame for its vertices; 2D for polygon and 3D for other shapes
- indices (only if type is "Mesh"): Array of triangle indices associated with the given set of vertices
"""
ARGS = [
Arg('outline', '-o', float, 0,
help='High-quality outline for spheres and polyhedra'),
Arg('cartoon_outline', None, float, 0,
help='Cartoon-like outline mode for all shapes'),
Arg('color_scale', None, float, 1, valid_values=flowws.Range(0, 10, True),
help='Factor to scale color RGB intensities by'),
Arg('draw_scale', '-s', float, 1,
help='Scale to multiply particle size by'),
Arg('display_box', '-b', bool, True,
help='Display the system box'),
Arg('transparency', None, bool, False,
help='Enable special translucent particle rendering'),
Arg('additive_rendering', None, bool, False,
help='Use additive rendering for shapes'),
Arg('fast_antialiasing', None, bool, False,
help='Use Fast Approximate Antialiasing (FXAA)'),
Arg('ambient_occlusion', None, bool, False,
help='Use Screen Space Ambient Occlusion (SSAO)'),
Arg('disable_rounding', None, bool, False,
help='Disable spheropolyhedra and spheropolygons'),
Arg('disable_selection', None, bool, False,
help='Don\'t allow selection of particles for this scene'),
]
def run(self, scope, storage):
"""Generate a scene of plato primitives."""
scene_kwargs = {}
positions = np.asarray(scope['position'])
N = len(positions)
if 'type' in scope:
types = np.asarray(scope['type'])
else:
types = np.repeat(0, N)
unique_types = list(sorted(set(types)))
if 'type_shapes.json' in scope:
type_shapes = json.loads(scope['type_shapes.json'])
else:
type_shapes = []
if 'dimensions' in scope:
dimensions = scope['dimensions']
try:
dimensions = dimensions[0]
except TypeError:
pass
elif type_shapes and any(shape['type'].lower() in ('disk', 'polygon')
for shape in type_shapes):
dimensions = 2
elif np.allclose(positions[:, 2], 0):
dimensions = 2
else:
dimensions = 3
orientations = np.atleast_2d(scope.get('orientation', []))
if len(orientations) < N:
orientations = np.tile([[1, 0, 0, 0.]], (N, 1))
if 'color' not in scope or len(scope['color']) < N:
colors = np.empty((N, 4), dtype=np.float32)
colors[:, :3] = plato.cmap.cubeellipse_intensity(
types.astype(np.float32), h=1.2, s=-0.25, lam=.45)
colors[:, 3] = 1
else:
colors = scope['color']
colors[:, :3] *= self.arguments['color_scale']
diameters = np.atleast_1d(scope.get('diameter', 1))
if len(diameters) < N:
diameters = np.repeat(1, N)
diameters = diameters*self.arguments['draw_scale']
if self.arguments.get('disable_rounding', False):
for description in type_shapes:
if 'rounding_radius' in description:
description['rounding_radius'] = 0
while len(type_shapes) < len(unique_types):
if dimensions == 2:
type_shapes.append(dict(type='Disk'))
else:
type_shapes.append(dict(type='Sphere'))
primitives = list(scope.get('plato_primitives', []))
primitive_indices = [[]]*len(primitives)
for (t, description) in zip(unique_types, type_shapes):
filt = np.where(types == t)[0]
prim_type = description['type'].lower()
is_2d = prim_type in ('disk', 'polygon')
prim_class = PRIM_NAME_MAP[prim_type]
if prim_type == 'convexpolyhedron' and description.get('rounding_radius', 0):
prim_class = draw.ConvexSpheropolyhedra
elif prim_type == 'polygon' and description.get('rounding_radius', 0):
prim_class = draw.Spheropolygons
kwargs = dict(description)
kwargs.pop('type')
for key in list(kwargs):
new_key = TYPE_SHAPE_KWARG_MAP.get(key, key)
kwargs[new_key] = kwargs.pop(key)
prim = prim_class(**kwargs)
if is_2d:
prim.positions = positions[filt, :2]
else:
prim.positions = positions[filt]
prim.orientations = orientations[filt]*np.sqrt(self.arguments['draw_scale'])
prim.colors = colors[filt]
prim.diameters = diameters[filt]
prim.outline = self.arguments['outline']
primitive_indices.append(filt)
primitives.append(prim)
if 'box' in scope and self.arguments['display_box']:
prim = draw.Box.from_box(scope['box'])
prim.width = min(scope['box'][:dimensions])*5e-3
prim.color = (0, 0, 0, 1)
primitive_indices.append([])
primitives.append(prim)
scene_kwargs['size'] = 1.05*np.array(scope['box'][:2])
# use default size of 800px wide
scene_kwargs['pixel_scale'] = 800/scene_kwargs['size'][0]
self.scene = draw.Scene(primitives, **scene_kwargs)
if dimensions == 2:
self.scene.enable('pan')
for (argument_name, feature_name) in SCENE_FEATURE_REMAP.items():
if self.arguments[argument_name]:
self.scene.enable(feature_name)
if self.arguments['cartoon_outline']:
self.scene.enable('outlines', self.arguments['cartoon_outline'])
scope['primitive_indices'] = primitive_indices
scope.setdefault('visuals', []).append(self)
scope.setdefault('visual_link_rotation', []).append(self)
if not self.arguments['disable_selection']:
scope['selection_visual_target'] = self
def draw_plato(self):
return self.scene
def _make_config_widget(self, arg, stage):
callback = functools.partial(self._rerun, arg, stage)
result = None
if arg.type == int:
result = QtWidgets.QSpinBox()
if arg.name in stage.arguments:
val = stage.arguments[arg.name]
else:
val = 0
if arg.valid_values is not None:
range_ = arg.valid_values
result.setMinimum(range_.min + (not range_.inclusive[0]))
result.setMaximum(range_.max - (not range_.inclusive[1]))
else:
result.setMaximum(max(1024, val * 4))
result.setValue(val)
result.valueChanged[int].connect(callback)
elif arg.type == float:
result = QtWidgets.QDoubleSpinBox()
if arg.name in stage.arguments:
val = stage.arguments[arg.name]
else:
val = 0
if arg.valid_values is not None:
range_ = arg.valid_values
else:
range_ = flowws.Range(0, val * 4 if val else 8, True)
delta = range_.max - range_.min
result.setDecimals(6)
result.setMinimum(range_.min + 1e-2 * delta *
(not range_.inclusive[0]))
result.setMaximum(range_.max - 1e-2 * delta *
(not range_.inclusive[1]))
result.setStepType(
QtWidgets.QAbstractSpinBox.AdaptiveDecimalStepType)
result.setSingleStep(5e-2 * delta)
result.setValue(val)
result.valueChanged[float].connect(callback)
elif arg.type == str:
if arg.valid_values is not None:
result = QtWidgets.QComboBox()
result.addItems(arg.valid_values)
if arg.name in stage.arguments:
result.setCurrentText(stage.arguments[arg.name])
result.currentIndexChanged[str].connect(callback)
else:
result = QtWidgets.QLineEdit()
if arg.name in stage.arguments:
result.setText(stage.arguments[arg.name])
result.textChanged[str].connect(callback)
elif arg.type == bool:
result = QtWidgets.QCheckBox()
if arg.name in stage.arguments:
result.setChecked(stage.arguments[arg.name])
result.stateChanged.connect(_bool_checkbox_helper(callback))
return result