"""
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
natoms: int = 0
nactive: int = 0
_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))
self.natoms = cms.particle_total
self.nactive = cms.particle_total
[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 make_whole(topo, tr):
"""
In MD simulation, molecules can be broken due to the periodic boundary
condition, which makes some atoms be at one side of the simulation box and
the other atoms at the opposite side. This function will edit the atom
coordinates so to make broken molecules whole again for each frame of the
MD trajectory `tr`.
:type topo: List[List[int]], usually cms.Cms.glued_topology
:param topo: Lists of particles connected to each of the particles in
trajectory
:type tr: `list`
:param tr: The simulation trajectory to be modified
:rtype: `list`
:return: Modified simulation trajectory
"""
return tr and _pfx_apply(topo, 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))
[docs]def make_whole_cms(cms_model):
"""
Similar to `make_whole` (see above), but on a cms model instead of a
simulation trajectory.
:rtype: `cms.Cms`
:return: Modified cms model
"""
fr = Frame(cms_model)
_pfx_apply(cms_model.glued_topology, [fr])
cms_model.update_with_frame(fr)
return cms_model
[docs]def glue_cms(gids, cms_model):
"""
Similar to `glue` (see above), but on a cms model instead of a
simulation trajectory.
:rtype: `cms.Cms`
:return: Modified cms model
"""
fr = Frame(cms_model)
_pfx_apply(cms_model.glued_topology, [fr],
lambda pfx: _glue_action(pfx, gids))
cms_model.update_with_frame(fr)
return cms_model
[docs]def center_cms(gids, cms_model, dims=None):
"""
Similar to `center` (see above), but on a cms model instead of a
simulation trajectory.
:param dim: dimensions to be centered on
:type dim: `None` or `list` of `int`, e.g., [2] for z axis,
[0, 1] for x-y plane. If set to `None`, center on all three
spatial dimensions
:rtype: `cms.Cms`
:return: Modified cms model
"""
fr = Frame(cms_model)
if dims:
pos = fr.pos()
dims_to_keep = list({0, 1, 2} - set(dims))
cached = pos[:, dims_to_keep]
_pfx_apply(cms_model.glued_topology, [fr],
lambda pfx: _glue_action(pfx, gids),
lambda pfx: _align_action(pfx, gids))
if dims:
pos[:, dims_to_keep] = cached
cms_model.update_with_frame(fr)
return cms_model
[docs]def superimpose_cms(gids, cms_model, ref_pos, weights=None):
"""
Similar to `superimpose` (see above), but on a cms model instead of a
simulation trajectory.
:rtype: `cms.Cms`
:return: Modified cms model
"""
fr = Frame(cms_model)
_pfx_apply(
cms_model.glued_topology, [fr], lambda pfx: _glue_action(pfx, gids),
lambda pfx: _align_action(pfx, gids, ref_pos=ref_pos, weights=weights))
cms_model.update_with_frame(fr)
return cms_model