Source code for flowws_analysis.SaveGarnett

import collections
import flowws
from flowws import Argument as Arg
import garnett
import numpy as np

FakeFrame = collections.namedtuple(
    'FakeFrame', [
        'box', 'position', 'orientation', 'types', 'typeid', 'data',
        'view_rotation', 'N',

[docs]@flowws.add_stage_arguments class SaveGarnett(flowws.Stage): """Save trajectory quantities using Garnett. This stage currently only saves an individual frame, but saving an entire trajectory is intended to work in the future. """ ARGS = [ Arg('filename', '-f', str, 'trajectory.sqlite', help='Name of file to save trajectory to'), ] def run(self, scope, storage): """Save trajectory quantities.""" filename = self.arguments['filename'] box = scope['box'] box = garnett.trajectory.Box( Lx=box[0], Ly=box[1], Lz=box[2], xy=box[3], xz=box[4], yz=box[5]) positions = scope['position'] if 'orientation' in scope: orientations = scope['orientation'] else: orientations = np.tile([(1, 0, 0, 0)], (len(positions), 1)) types = scope['type'] if 'type_names' not in scope: type_names = [chr(ord('A') + t) for t in sorted(set(types))] else: type_names = scope['type_names'] the_frame = FakeFrame( box, positions, orientations, type_names, types, None, None, len(positions)) mode = ('w' if any(filename.endswith(suf) for suf in ['.pos', '.cif']) else 'wb') with self.arguments['filename'], mode=mode, on_filesystem=True) as f: garnett.write([the_frame],