"""
Periodic boundary condition related functions. These functions use particle
connection maintained by cms.Cms (cms.Cms.glued_topology) for coordinates
transformation.
Copyright Schrodinger, LLC. All rights reserved.
"""
import dataclasses
import numpy
import schrodinger.application.desmond.packages.msys.pfx as pfx_module
[docs]@dataclasses.dataclass
class Frame:
"""
This a minimum interface to allow pfx functions for cms.Cms object.
"""
box: numpy.array = None
_pos: numpy.array = None
_vel: numpy.array = None
[docs] def __init__(self, cms):
self._pos, self._vel = cms.get_pos_vel()
self.box = numpy.reshape(cms.box, newshape=(3, 3))
[docs] def pos(self):
return self._pos
[docs] def vel(self):
return self._vel
def _glue_action(pfx, gids):
pfx.glue(gids)
def _align_action(pfx, gids, ref_pos=None, weights=None):
pfx.align(gids, coords=ref_pos, weights=weights)
def _pfx_apply(topo, tr, *actions):
"""
Given topology (topo), apply actions to each frames of trajectory (tr).
:type topo: List[List[int]]
:param topo: Lists of particles connected to each of the particles in
trajectory
:type tr: List[traj.Frame]
:param tr: Trajectory to be transformed by actions
:type actions: Iterable[function], where function takes two arguments,
msys.pfx and List[int]
:param actions: Transformactions to be applied to trajectory
:rtype: List
:return: Modified trajectory
"""
pfx = pfx_module.Pfx(topology=topo, fixbonds=True)
for a in actions:
a(pfx)
for e in tr:
pfx.apply(e.pos(), e.box, e.vel())
return tr
[docs]def glue(topo, gids, tr):
"""
Given topology (topo), glue gids together for each frames of trajectory (tr).
:type topo: List[List[int]]
:param topo: Lists of particles connected to each of the particles in
trajectory
:type gids: List[int]
:param gids: List of particle ids to be glued
:type tr: List[traj.Frame]
:param tr: Simulation trajectory to be modified
:rtype: List
:return: Modified simulation trajectory
"""
return tr and _pfx_apply(topo, tr, lambda pfx: _glue_action(pfx, gids))
[docs]def center(topo, gids, tr, dims=None):
"""
Given topology (topo), center around gids for each frames of trajectory (tr).
:type topo: List[List[int]]
:param topo: Lists of particles connected to each of the particles in
trajectory
:type gids: List[int]
:param gids: List of particle ids
:type tr: List[traj.Frame]
:param tr: simulation trajectory to be modified
:type dims: Set[int]
:param dims: Dimensions to be centered on
:rtype: List
:return: Modified simulation trajectory
"""
pfx = pfx_module.Pfx(topology=topo, fixbonds=True)
_glue_action(pfx, gids)
_align_action(pfx, gids)
if dims:
dims_to_keep = list({0, 1, 2} - set(dims))
for fr in tr:
if dims:
pos = fr.pos()
saved = pos[:, dims_to_keep]
pfx.apply(fr.pos(), fr.box, fr.vel())
if dims:
pos[:, dims_to_keep] = saved
return tr
[docs]def superimpose(topo, gids, tr, ref_pos, weights=None):
"""
Given topology (topo), superimpose gids for each frames of trajectory (tr)
to ref_pos.
:type topo: List[List[int]]
:param topo: Lists of particles connected to each of the particles in
trajectory
:type gids: List[int]
:param gids: List of particle ids
:type tr: List[traj.Frame]
:param tr: Trajectory to be modified
:type ref_pos: numpy.array, 3xlen(gids) matrix
:param ref_pos: reference positions
:type weigths: List or None
:param weigths: weights for each particle in gids
:rtype: List
:return: Modified trajectory
"""
return tr and _pfx_apply(
topo, tr, lambda pfx: _glue_action(pfx, gids),
lambda pfx: _align_action(pfx, gids, ref_pos=ref_pos, weights=weights))