Source code for flowws_analysis.Center

import flowws
from flowws import Argument as Arg
import plato
import numpy as np

def center(box, positions, target=(0, 0, 0)):
    """Centers a set of positions in a periodic box."""
    fractions = plato.math.make_fractions(box, positions)
    thetas = 2*np.pi*fractions
    sums = np.sum(np.exp(1j*thetas), axis=0)
    fractions -= np.angle(sums)/2/np.pi + .5
    fractions -= target
    fractions %= 1.
    result = plato.math.fractions_to_coordinates(box, fractions)
    return np.ascontiguousarray(result)

def wrap(box, positions):
    fractions = plato.math.make_fractions(box, positions)
    fractions %= 1.
    return plato.math.fractions_to_coordinates(box, fractions)

[docs]@flowws.add_stage_arguments class Center(flowws.Stage): """Center the system through periodic boundary conditions. This module modifies the positions of the system to have either the center of mass of the system or a single indicated particle at (0, 0, 0). """ ARGS = [ Arg('particle', '-p', int, -1, help='Particle index to center with (default: use center of mass of the system)'), ] 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