"""
Simulation Trajectory Analysis Framework (STAF)
Copyright Schrodinger, LLC. All rights reserved.
"""
import copy
import numpy
from schrodinger.application.desmond.packages import msys
from schrodinger.application.desmond.packages import pfx
from schrodinger.application.desmond.packages import topo
[docs]class GeomAnalyzerBase(object):
"""
Base class of all geometry analyzer classes
At this level, we make a distinction of two types of analyzers: static and
dynamic analyzers, which we call staalyzers and dynalyzers, respectively.
A dynamic analyzer (or dynalyzer) redefines its calculations from frame
to frame, typically due to atom selection change, i.e., dynamic ASL.
Generally speaking, any analyses where the exact types and/or quantities
of calculations depend on the dynamics of the simulation system will fall
into the conceptual regime of dynamic analyzer.
A static analyzer (or staalyzer) is the opposite. The exact types and/or
quantities of calculations are predefined and doesn't depend on the
coordinates of the particles.
Also we regard dynamic analyzer as the more general concept.
This means that a static analyzer is a dynamic analyzer with trivial
(nil) dependency on the coordinates.
All subclasses should define two private methods:
- `_precalc`
- This method will be called by a `GeomCalc` object to register
geometry calculations.
- `_postcalc`
- This method will be called by a `GeomCalc` object to finish
the particular analysis calculation. And the results should are saved
in the `self._result`
In between the `_precalc` and `_postcalc` calls, the `GeomCalc` object
will be called (outside the analyzer class) to calculate all requested
geometry calculation frame by frame. Also see the docstring of
the `analysis.analyze` function.
For dynalyzers, the subclass is expected to define one more private method:
- `_dyncalc`
- This method will be called by a `GeomCalc` object to register
geometry calculations for each trajectory frame.
"""
def __new__(cls, *args, **kwargs):
# We force set `_is_dynamic` in case a subclass doesn't care about
# initializing this base class.
instance = super().__new__(cls)
instance._is_dynamic = False
instance._result = None
return instance
[docs] def __init__(self, is_dynamic=False):
self._is_dynamic = is_dynamic
def _precalc(self, calc):
"""
:type calc: `GeomCalc`
"""
pass
def _dyncalc(self, calc):
"""
:type calc: `GeomCalc`
"""
# Any staalyzer is a trivial dynalyzer. A staalyzer subclass never
# defines the `_dyncalc`, but this method should default to the
# `_precalc` in case a staalyzer is used as a dynalyzer, which is
# still conceptually correct and does happen in practice.
self._precalc(calc)
[docs] def disableDyncalc(self):
"""
Disable the execution of _dyncalc(). This is used to avoid redundant
_precalc() calculations delegated in _dyncalc().
"""
def noop(*args):
pass
self._dyncalc = noop
def _postcalc(self, calc, pbc, fr):
"""
:type calc: `GeomCalc`
:type pbc: `Pbc`
:type fr: `traj.Frame`
"""
pass
def __call__(self):
return self._result
@staticmethod
def _get_sites(cms_model, xids):
"""
This function takes a list of XIDs (AIDs or `CenterOf` objects) and
turns them into a list of sites (GIDs or `Positer`).
:type xids: `int`, `Com`, `Coc`, `Centroid`
:rtype: `list` of combinations of GIDs or `Positer`
"""
if cms_model.need_msys:
return [
topo.aids2gids(cms_model, [x], include_pseudoatoms=False)[0]
if isinstance(x, (numpy.integer, int)) else Positer([x], 1)
for x in xids
]
else:
return [
cms_model.convert_to_gids([x], with_pseudo=False)[0]
if isinstance(x, (numpy.integer, int)) else Positer([x], 1)
for x in xids
]
@staticmethod
def _sites2gids(calc, sites, gid_offset):
"""
This function returns GIDs for the provided sites. If `Positer` type is
detected, a new GIDs is registered.
:type calc: `GeomCalc`
:type particles: a `list` that contains either `int` and/or `Positer`
:param gid_offset: number of sites/GIDs in the msys_model, the GIDs
generated for all the particles are being offset by
this value, which is equal to the total number of
GIDs.
:type gid_offset: `int`
:rtype: `list` of GIDs
"""
return [
s if isinstance(s, (numpy.integer,
int)) else calc.addPosition(s, s.numPos()) +
gid_offset for s in sites
]
[docs] def isDynamic(self):
return self._is_dynamic
[docs]def center_fr(data, pbc, fr, custom):
"""
This function will copy the input trajectory frame (`fr`), and center
the selected particles in the copy.
:type data: `dict` where key = `(msys_model, cms_model, center_gids)`
and value = <centered-frame>.
where `msys_model` and `center_gids` are mandated, and `cms_model` can
be `None`. `center_gids` is the GIDs of the particles to be centered.
:type pbc: `analysis.Pbc`
:type fr: `traj.Frame`
:type custom: `_CustomCalc`
:return: Updated `data`, where values are updated for the given `fr`.
"""
# `pbc` and `custom` are not used in the current implementation.
for key in data:
msys_model, cms_model, center_gids = key
centered_fr = fr.copy()
pos = centered_fr.pos()
# Only center the original positions, not the positers.
centered_fr.hijackPos(pos[:centered_fr.natoms])
if cms_model.need_msys:
topo.center(msys_model, center_gids, [centered_fr])
else:
pfx.center(cms_model.glued_topology, center_gids, [centered_fr])
centered_fr.hijackPos(pos)
data[key] = centered_fr
return data
[docs]def center_ct(data, pbc, fr, custom):
"""
Center selected particles in the simulation box, and it will automatically
make all molecules whole. This function will create a centered frame and a
copy of the full system CT from the centered frame.
Caveat: `center_fr` should have been called on the same key and frame. This
is an implicit coupling (bad) between the two functions.
:type data: `dict` where key = `(msys_model, cms_model, center_gids)`
and value = <centered-fsys-ct> where `msys_model` and `center_gids` are
mandated, and `cms_model` can be `None`. `center_gids` is the GIDs of
the particles to be centered.
:type pbc: `analysis.Pbc`
:type fr: `traj.Frame`
:type custom: `_CustomCalc`
:return: Updated `data`, where values are updated for the given `fr`.
"""
# `pbc` and `fr` are not used in the current implementation.
for key in data:
_, cms_model, _ = key
centered_fr = custom[center_fr][key]
ct = data[key] or cms_model.fsys_ct.copy()
data[key] = topo.get_active_fsys_ct_from_frame(ct, cms_model,
centered_fr)
return data
[docs]def center_fr_along_z(data, pbc, fr, custom):
"""
This function will copy the input trajectory frame (`fr`), and center
the selected particles along the Z-axis in the copy.
:type data: `dict` where key = `(msys_model, cms_model, center_gids)` and
value = <centered-frame> where `msys_model` and `center_gids` are
mandated, and `cms_model` can be `None`. `center_gids` is the GIDs of
the particles to be centered.
:type pbc: `analysis.Pbc`
:type fr: `traj.Frame`
:type custom: `_CustomCalc`
:return: Updated `data`, where values are updated for the given `fr`.
"""
# `pbc` and `custom` are not used in the current implementation.
for key in data:
msys_model, cms_model, center_gids = key
centered_fr = fr.copy()
pos = centered_fr.pos()
# Only center the original positions, not the positers.
centered_fr.hijackPos(pos[:centered_fr.natoms])
if cms_model.need_msys:
topo.center(msys_model, center_gids, [centered_fr], dims=[2])
else:
pfx.center(cms_model.glued_topology,
center_gids, [centered_fr],
dims=[2])
centered_fr.hijackPos(pos)
data[key] = centered_fr
return data
[docs]class CenteredSoluteAnalysis(GeomAnalyzerBase):
"""
This class provides a temporary trajectory frame where the solute atoms are
centered. It helps resolve PBC wrapping issues for analyzers such as
`analysis.RMSD`, `analysis.PosAlign`.
"""
[docs] def __init__(self,
msys_model,
cms_model,
*arg,
asl_center='solute',
asl_exclude='ions or water or metals',
**kwarg):
"""
:param asl_center: ASL for the atoms to be centered
:param asl_exclude: ASL for the atoms to be excluded from centering.
For example, ASL 'protein' could select ion and
water molecules which may be too mobile for the
centering.
"""
super().__init__(*arg, **kwarg)
self._result = []
if cms_model.need_msys:
solute_gids = topo.asl2gids(cms_model,
asl_center,
include_pseudoatoms=False)
exclusion_gids = topo.asl2gids(cms_model,
asl_exclude,
include_pseudoatoms=False)
else:
solute_gids = cms_model.convert_to_gids(
cms_model.select_atom(asl_center), with_pseudo=False)
exclusion_gids = cms_model.convert_to_gids(
cms_model.select_atom(asl_exclude), with_pseudo=False)
solute_gids = tuple(sorted(set(solute_gids) - set(exclusion_gids)))
self._key = (msys_model, cms_model, solute_gids)
def _precalc(self, calc):
calc.addCustom(center_fr, self._key)
def _getCenteredFrame(self, calc):
"""
:return: Centered frame
:rtype: `traj.Frame`
"""
return calc.getCustom(center_fr)[self._key]
[docs]class MaestroAnalysis(CenteredSoluteAnalysis):
"""
Analyzer classes that perform calculations on the solute-centered
full-system CT could inherit this base class.
"""
def _precalc(self, calc):
"""
Reusable intermediate data saved in `calc`: We use the `center_fr` and
`center_ct` functions to create one copy of the trajectory frame and one
copy of the full-system CT. In both, the solute atoms are centered.
"""
# Ensures `center_fr` is called before `center_ct`.
super()._precalc(calc)
calc.addCustom(center_ct, self._key)
def _getCenteredCt(self, calc):
"""
:return: Centered full system CT
:rtype: `structure.Structure`
"""
return calc.getCustom(center_ct)[self._key]
[docs]class CustomMaestroAnalysis(MaestroAnalysis):
"""
Compute the result of a custom function on centered models. Under the hood,
this custom function serves as a CID for `_CustomCalc`. The same key of
`MaestroAnalysis` is used, and the value is a tuple of the custom function's
return and the centered fullsystem CT.
"""
[docs] def __init__(self, msys_model, cms_model, func):
"""
:type func: callable, with the same arguments as `center_fr`
"""
assert cms_model
super().__init__(msys_model, cms_model)
self._func = func
def _precalc(self, calc):
super()._precalc(calc)
calc.addCustom(self._func, self._key)
def _postcalc(self, calc, *_):
"""
The first entry is the result of the self._func for the current frame,
the second entry is a tuple of centered `traj.Frame`,
and `schrodinger.structure.Structure` objects.
"""
self._result = (calc.getCustom(self._func)[self._key],
self._getCenteredCt(calc))
[docs]class CompositeAnalyzer(GeomAnalyzerBase):
"""
A composite analyzer calls one or more other analyzers (subanalyzers) to
obtain (intermediate) results.
The subclass should define the subanalyzers as a private attribute:
`_analyzers`, whose value should be a list of analyzers.
"""
def _precalc(self, calc):
for a in self._analyzers:
a._precalc(calc)
def _dyncalc(self, calc):
for a in self._analyzers:
a._dyncalc(calc)
def _postcalc(self, *args):
for a in self._analyzers:
a._postcalc(*args)
def _update_fsys_ct_in_cms(data, pbc, fr, *_):
"""
Update atom coordinates of the full system CT inside the cms object using
the trajectory frame `fr`. Note that the component CTs' coordinates are not
updated.
"""
for key in data:
cms_model = key
# FIXME: With inactive-free API, set_nactive_gids() will cause _sync(),
# which is costly.
# At this time, update_ct() only updates coordinates, but does
# not add/delete atoms in fsys_ctsolvent CT. When the trajectory
# viewer is able to use inactive-free API, we can swap the orders
# of the following two statements.
# An alternative is to have the cms.Cms.active_total setter not
# call _sync(), but add/delete atoms as in topo.update_fsys_ct_from_frame_GF()
cms_model.set_nactive_gids(fr.nactive, fr.natoms)
topo.update_ct(cms_model.fsys_ct, cms_model, fr)
return data
[docs]class UpdatedCmsFsysCtAnalysis(GeomAnalyzerBase):
"""
This class updates the full system CT inside the Cms object to each
trajectory frame.
N.B.: Typically all analyzers share the same Cms object, thus the existence of
an instance of this class has side effect on other analyzers (This side
effect is probably wanted).
"""
[docs] def __init__(self, msys_model, cms_model, *arg, **kwarg):
super().__init__(*arg, **kwarg)
self._key = cms_model
def _precalc(self, calc):
calc.addCustom(_update_fsys_ct_in_cms, self._key)
[docs]class DynamicPositerAnalyzer(UpdatedCmsFsysCtAnalysis):
"""
This analyzer uses (dynamic) atom selection to create `Positer` for each
frame.
"""
[docs] def __init__(self, msys_model, cms_model, asl, initializer):
"""
:type msys_model: `msys.System`
:type cms_model: `structure.Structure`
:type asl: `str`
:type initializer: callable
:param initializer: It takes three input arguments (`msys_model`,
`cms_model`, and a `list` of AIDs) and returns a
`Positer` instance.
"""
super().__init__(msys_model,
cms_model,
is_dynamic=topo.is_dynamic_asl(cms_model, asl))
self._msys_model = msys_model
self._cms_model = cms_model
self._asl = asl
self._aids = cms_model.select_atom(asl)
self._initializer = initializer
self._dyninit()
def _dyninit(self):
self._positer = self._initializer(self._msys_model, self._cms_model,
self._aids)
def _precalc(self, calc):
super()._precalc(calc)
if not self._is_dynamic:
calc.addPosition(self._positer, self._positer.numPos())
def _dyncalc(self, calc):
self._aids = self._cms_model.select_atom(self._asl)
self._dyninit()
calc.addPosition(self._positer, self._positer.numPos(), is_dynamic=True)
[docs]class DynamicAslAnalyzer(UpdatedCmsFsysCtAnalysis):
"""
A base class for all analyzers that support dynamic ASL expressions.
It guarantees to update the atom IDs from the given ASL expression for each
frame and store the atom IDs into a private attribute `_aids`.
This class defines the private method `_dyncalc` to be called by the
`GeomCalc` object. This class expects its subclass to define a
`_dyninit(self)' method to be automatically called by the `_dyncalc` method
after the `_aids` has been updated. The subclass' `_dyninit` method should
redefine the geometry calculations based on the updated atom IDs and will
be called by both the `__init__` and the `_dyncalc` methods of this class.
"""
[docs] def __init__(self, msys_model, cms_model, asl):
"""
:type msys_model: `msys.System`
:type cms_model: `structure.Structure`
:type asl: `str`
"""
super().__init__(msys_model,
cms_model,
is_dynamic=topo.is_dynamic_asl(cms_model, asl))
self._msys_model = msys_model
self._cms_model = cms_model
self._asl = asl
self._aids = cms_model.select_atom(asl)
self._dyninit()
def _dyninit(self):
pass
def _dyncalc(self, calc):
self._aids = self._cms_model.select_atom(self._asl)
self._dyninit()
[docs]class CompositeDynamicAslAnalyzer(DynamicAslAnalyzer, CompositeAnalyzer):
"""
A base class for analyzers whose ALL subanalyzers are redefined for each
frame based on the results of the dynamic ASL expression. The redefinition
of subanalyzers are done by the analyzer's `_dyninit` method (see docstring of
`DynamicAslAnalyzer`), which will be automatically called by the
`DynamicAslAnalyzer._dyncalc` method (see docstring of `DynamicAslAnalyzer` for more detail).
"""
def _precalc(self, calc):
DynamicAslAnalyzer._precalc(self, calc)
CompositeAnalyzer._precalc(self, calc)
def _dyncalc(self, calc):
DynamicAslAnalyzer._dyncalc(self, calc)
CompositeAnalyzer._dyncalc(self, calc)
[docs]class Positer(object):
"""
This class appends analyzers' results to trajectory frames. The prominent
application is to treat `Com`, `Coc`, or `Centroid` analyzers as composite
"atoms" with GIDs, which enables geometric calculations among both regular
atoms and composite "atoms".
"""
# FIXME: Here num_pos is redundant. Each analyzer's result has to be a 1x3
# Numpy array (see __call__ below).
[docs] def __init__(self, analyzers, num_pos):
"""
:type analyzers: `list`. Each element can be an analyzer, for example,
a `Com`, or `Coc`, or `Centroid` object.
:param analyzers: A list of analyzers. Each analyzer will return a new
position.
"""
self._analyzers = analyzers
self._num_pos = num_pos
self._calc = GeomCalc()
self._gid_offset = 0
self._first_gid = 0
for a in self._analyzers:
self._calc.addAnalyzer(a)
def __call__(self, pbc, fr):
"""
:type pbc: `Pbc`
:type fr: `traj.Frame`
"""
self._calc(pbc, fr)
self._first_gid = fr.natoms + self._gid_offset
for i, a in enumerate(self._analyzers, start=self._first_gid):
fr.pos(i)[:] = a()
[docs] def setGidOffset(self, gid_offset):
"""
:type gid_offset: `int`
:param gid_offset: The GID of the first position added by this positer
will be natoms + gid_offset, where natoms is the
number of interaction sites in the original model
system.
"""
self._gid_offset = gid_offset
[docs] def numPos(self):
"""
:rtype: `int`
:return: The number of new positions to be added into the position array
of the given frame.
"""
return self._num_pos
[docs] def gids(self):
"""
:rtype: `list` of `int` objects
:return: The GIDs of the new positions to be added.
"""
return list(range(self._first_gid, self._first_gid + self._num_pos))
class _CustomCalc(dict):
"""
An instance of this class will store custom calculation requests (keys) and
results (values). Each request is represented as a calculation ID (or CID
for short). If the same calculation is requested by different analyzers,
using the same CID avoids duplicate calculation. A CID must be hashable,
whereas the results (values) can be of any type.
A CID can be either a callable or noncallable object. If it's callable, the
`GeomCalc` object will call it with the following arguments:
callable_object(data, pbc, fr, self)
where `data` is the current value corresponding to this CID, `pbc` is a
`Pbc` object, `fr` is a trajectory frame (`traj.Frame`), and `self` is
this `_CustomCalc` object. This function should return a new value to
replace the old one.
If a CID is noncallable, the analyzer(s) is
responsible to update the data from the current trajectory frame. Note that
if you have more than one analyzers reading the data, you have to somehow
ensure that the data is updated only once for each trajectory frame.
"""
def calc(self, pbc, fr, keys=None):
"""
:type keys: Iterable of (some) keys in this `_CustomCalc` object,
or `None`.
:param keys: If `keys` is not `None`, this function will calculate for
the given keys, or it will calculate for all keys in this
`_CustomCalc` object.
"""
for cid in (keys or self):
if callable(cid):
func = cid
data = self[cid]
self[cid] = func(data, pbc, fr, self)
class _Registry(object):
"""
A facility class used by `GeomCalc` (see below) to save the calculations
as currently registered into the `GeomCalc` instance.
The use case is to accommodate dynamic analyzers (dynalyzers), which may
register different geometric calculations for each frame. If we do NOT
remove the extra calculations added by the dynalyzers for the previous
frames, we will end up wastefully doing a lot of unused calculations. One
way to remove the additional calculations is to use this class to take a
snapshot of the currently registered calculations (we call it "registry")
before calling the dynalyzers, and after we finish doing the calculations,
we restore the registry.
"""
def __init__(self, calc):
self.vec = set(calc._vec)
self.dist = set(calc._dist)
self.angle = set(calc._angle)
self.torsion = set(calc._torsion)
self.custom = list(calc._custom)
self.static_positers = copy.copy(calc._static_positers)
self.dynamic_positers = copy.copy(calc._dynamic_positers)
self.staalyzer = copy.copy(calc._staalyzer)
self.dynalyzer = copy.copy(calc._dynalyzer)
self.num_extra_pos = calc._num_extra_pos
def delta(self, calc):
"""
Return the difference between calculations recorded in this `_Registry`
object and those registered into the given `calc`.
This function assumes that the calculations in `calc` is always a super
set of those in this object.
@type calc: `GeomCalc`
@rtype: `_Registry`
"""
reg = _Registry(calc)
reg.vec = reg.vec - self.vec
reg.dist = reg.dist - self.dist
reg.angle = reg.angle - self.angle
reg.torsion = reg.torsion - self.torsion
reg.custom = reg.custom[len(self.custom):]
reg.static_positers = reg.static_positers[len(self.static_positers):]
reg.dynamic_positers = reg.dynamic_positers[len(self.dynamic_positers):]
reg.staalyzer = reg.staalyzer[len(self.staalyzer):]
reg.dynalyzer = reg.dynalyzer[len(self.dynalyzer):]
return reg
def _dictmap(func, dictionary):
"""
Map each key-value pair in a dictionary with a function `func`. This will
create a new `dict` object with the same keys as in `dictionary`, but
the corresponding value of each key will be "mapped" by calling
`func(key)`.
:type func: Any callable object that takes a key from `dictionary` as the
input argument and returns a new value.
:type dictionary: `dict`
:rtype: `dict`
:return: A new `dict` object. `dictionary` will NOT be mutated.
"""
return {key: func(key) for key in dictionary}
def _extend_frame(fr, num_extra_pos):
"""
Hijack the given frame's position array with a copy that is extended by the
requested number of positions.
:type fr: `traj.Frame`
:type num_extra_pos: positive `int`
"""
curr_pos = fr.pos()
curr_size = curr_pos.shape[0]
impos = numpy.ndarray((curr_size + num_extra_pos, 3), dtype=numpy.float32)
impos[:curr_size, :] = curr_pos
fr.hijackPos(impos)
[docs]class GeomCalc(object):
"""
We use this class to batch geometry calculations and avoid duplications.
For example, you want to calculate the bond distance between atoms 1 and 2,
and also an dihedral angle involving these two atoms. Both calculations
require to calculate the vector between the minimum images of the two atoms,
but we don't want to do the calculation twice. With this class, we avoid
such duplications.
All geometry calculations take into account the periodic boundary condition.
Basic usage:
calc = GeomCalc()
# Loads geometry-calculation requests.
calc.addVector(...)
calc.addDistance(...)
calc.addAngle(...)
calc.addTorsion(...)
calc.addPlanarAngle(...)
# Does calculations.
calc(pbc, frame)
# Gets results.
vec = calc.getVector(...)
dis = calc.getDistance(...)
ang = calc.getAngle(...)
dih = calc.getTorsion(...)
pla = calc.getPlanarAngle(...)
"""
[docs] def __init__(self):
self._vec = {}
self._dist = {}
self._angle = {}
self._torsion = {}
self._planarangle = {}
self._custom = _CustomCalc()
self._static_positers = []
self._dynamic_positers = []
self._staalyzer = [] # Static analyzers
self._dynalyzer = [] # Dynamic analyzers
self._num_extra_pos = 0
def _snapshot_registry(self):
return _Registry(self)
def _restore_registry(self, registry):
self._vec = {k: self._vec[k] for k in registry.vec}
self._dist = {k: self._dist[k] for k in registry.dist}
self._angle = {k: self._angle[k] for k in registry.angle}
self._torsion = {k: self._torsion[k] for k in registry.torsion}
custom = self._custom
self._custom = _CustomCalc()
for k in registry.custom:
self._custom[k] = custom[k]
self._static_positers = copy.copy(registry.static_positers)
self._dynamic_positers = copy.copy(registry.dynamic_positers)
self._staalyzer = copy.copy(registry.staalyzer)
self._dynalyzer = copy.copy(registry.dynalyzer)
self._num_extra_pos = registry.num_extra_pos
def __call__(self, pbc, fr):
"""
Perform all registered geometry calculations.
:type pbc: `Pbc`
:type fr: `traj.Frame`
"""
def gids2vec(two_gids):
return pbc.calcMinimumDiff(*fr.pos(two_gids))
def gids2dist(two_gids):
return numpy.linalg.norm(self._vec[two_gids])
def gids2angle(three_gids):
i, j, k = three_gids
return msys._msys.calc_vec_angle(-self._vec[(i, j)],
+self._vec[(j, k)])
def gids2torsion(four_gids):
i, j, k, m = four_gids
return msys._msys.calc_vec_dihedral(self._vec[(i, j)],
self._vec[(j, k)],
self._vec[(k, m)])
def gids2planarangle(six_gids):
i, j, k, m, n, o = six_gids
r21 = self._vec[(j, i)]
r23 = self._vec[(j, k)]
n123 = numpy.cross(r21, r23)
r54 = self._vec[(n, m)]
r56 = self._vec[(n, o)]
n456 = numpy.cross(r54, r56)
# angle between the two plane normals
mag123 = numpy.linalg.norm(n123)
mag456 = numpy.linalg.norm(n456)
if numpy.isclose(mag123, 0) or numpy.isclose(mag456, 0):
return numpy.nan
n123 /= mag123
n456 /= mag456
angle = numpy.arccos(abs(numpy.clip(numpy.dot(n123, n456), -1.,
1.))) # acute angle
return angle
# Perform static positer calculations.
if self._num_extra_pos:
_extend_frame(fr, self._num_extra_pos)
for e in self._static_positers:
e(pbc, fr)
# Performs geometric calculations for all staalyzers.
self._vec = _dictmap(gids2vec, self._vec)
self._dist = _dictmap(gids2dist, self._dist)
self._angle = _dictmap(gids2angle, self._angle)
self._torsion = _dictmap(gids2torsion, self._torsion)
self._planarangle = _dictmap(gids2planarangle, self._planarangle)
self._custom.calc(pbc, fr)
# By this point, the position array of the current frame should be
# extended with the extra positions from the static positers (if any).
# Now performs geometric calculations for all dynalyzers.
registry = None
num_static_extra_pos = self._num_extra_pos
if self._dynalyzer:
registry = self._snapshot_registry()
for e in self._dynalyzer:
# This call registers the dynamic positers (if any).
e._dyncalc(self)
# Performs dynamic positer calculations.
num_dynamic_extra_pos = self._num_extra_pos - num_static_extra_pos
if num_dynamic_extra_pos:
_extend_frame(fr, num_dynamic_extra_pos)
for dp in self._dynamic_positers:
dp(pbc, fr)
reg = registry.delta(self)
self._vec.update(_dictmap(gids2vec, reg.vec))
self._dist.update(_dictmap(gids2dist, reg.dist))
self._angle.update(_dictmap(gids2angle, reg.angle))
self._torsion.update(_dictmap(gids2torsion, reg.torsion))
self._custom.calc(pbc, fr, reg.custom)
for e in self._staalyzer + self._dynalyzer:
e._postcalc(self, pbc, fr)
if registry:
# This removes dynamic positer counts as well.
self._restore_registry(registry)
if self._num_extra_pos:
fr.hijackPos(None) # recover the original positions
[docs] def addPosition(self, positer, num_pos, is_dynamic=False):
"""
Add extra position into the position array.
:type positer: Callable, will be called as: `positer(pbc, fr)`, where
`pbc` is a `Pbc` object, and `fr` is a `traj.Frame`
object.
:param positer: Function (or callable object) to append new positions
into the position array of the given frame.
:type num_pos: `int`, must be a nonnegative number.
:param num_pos: The number of new positions to be added by `positer`
:rtype: `int`
:return: The gid offset of the first new position that will be generated
by this `positer`.
"""
gid_offset = self._num_extra_pos
self._num_extra_pos += num_pos
if is_dynamic:
self._dynamic_positers.append(positer)
else:
self._static_positers.append(positer)
if isinstance(positer, Positer):
positer.setGidOffset(gid_offset)
return gid_offset
[docs] def addVector(self, from_gid, to_gid):
"""
Add a vector calculation request.
:type from_gid: `int`
:type to_gid: `int`
"""
key = (from_gid, to_gid)
self._vec[key] = None
[docs] def addDistance(self, i_gid, j_gid):
"""
Add a distance calculation request.
:type i_gid: `int`
:type j_gid: `int`
"""
key = (i_gid, j_gid)
self._dist[key] = None
self._vec[key] = None
[docs] def addAngle(self, i_gid, j_gid, k_gid):
"""
Add an angle calculation request.
:type i_gid: `int`
:type j_gid: `int`
:type k_gid: `int`
The angle is defined by the two vectors: j==>i and j==>k.
"""
key = (i_gid, j_gid, k_gid)
self._angle[key] = None
self._vec[(i_gid, j_gid)] = None
self._vec[(j_gid, k_gid)] = None
[docs] def addTorsion(self, i_gid, j_gid, k_gid, l_gid):
r"""
Add a torsion calculation request.
:type i_gid: `int`
:type j_gid: `int`
:type k_gid: `int`
:type l_gid: `int`
The torsion is defined by the four position vectors::
i o o l
\ /
\ /
j o-----o k
In other words, it's the dihedral angle between the two planes: i-j-k
and j-k-l.
"""
key = (i_gid, j_gid, k_gid, l_gid)
self._torsion[key] = None
self._vec[(i_gid, j_gid)] = None
self._vec[(j_gid, k_gid)] = None
self._vec[(k_gid, l_gid)] = None
[docs] def addPlanarAngle(self, i_gid, j_gid, k_gid, l_gid, m_gid, n_gid):
"""
Add a planar angle calculation request. The first three gids define the
first plane and the second gids define the second plane. The acute
angle between these two planes is returned.
:type ._gid: `int`
"""
key = (i_gid, j_gid, k_gid, l_gid, m_gid, n_gid)
self._planarangle[key] = None
self._vec[(j_gid, i_gid)] = None
self._vec[(j_gid, k_gid)] = None
self._vec[(m_gid, l_gid)] = None
self._vec[(m_gid, n_gid)] = None
[docs] def addCustom(self, cid, key=None, default=None):
"""
Add a custom calculation item.
:type cid: Any hashable object
:param cid: Specify the type of the calculation. The results of this
type of calculation can be obtained by calling
`.getCustom(c)`.
:type key: Any hashable object
:param key: A particular calculation item of the type `c`. The result
of this item can be obtained by this: `.getCustom(c)[key]`.
:param default: The default result of the calculation item `key`.
"""
self._custom.setdefault(cid, {})[key] = default
[docs] def addAnalyzer(self, analyzer):
"""
Add a custom analyzer that must define the following interface:
1. _precalc(self, calc)
where `calc` is a `GeomCalc` object.
This method should call `calc.addCustom` to add an calculation
item of a custom calculation type.
2. _postcalc(self, calc, pbc, fr)
where `calc` is a `GeomCalc` object, `pbc` is a `Pbc` object,
and `fr` is a `traj.Frame` object.
This method can get the calculation result by calling
`calc.getCustom` and do further calculations as necessary to get
the final analytic results.
"""
if analyzer.isDynamic():
self._dynalyzer.append(analyzer)
else:
self._staalyzer.append(analyzer)
analyzer._precalc(self)
[docs] def getVector(self, from_gid, to_gid):
"""
Get the vector between the atoms: `from_gid` and `to_gid`.
"""
return self._vec[(from_gid, to_gid)]
[docs] def getDistance(self, i_gid, j_gid):
"""
Get the distance (in Angstroms) between the atoms: `i_gid` and `j_gid`.
"""
return self._dist[(i_gid, j_gid)]
[docs] def getAngle(self, i_gid, j_gid, k_gid):
"""
Get the angle (in radians) between the two vectors: j==>i and j==>k.
"""
return self._angle[(i_gid, j_gid, k_gid)]
[docs] def getTorsion(self, i_gid, j_gid, k_gid, l_gid):
"""
Get the torsion (in radians) as defined by the four atoms: `i_gid`,
`j_gid`, `k_gid`, and `l_gid`. See the docstring of `addTorsion` for
more detail.
"""
return self._torsion[(i_gid, j_gid, k_gid, l_gid)]
[docs] def getPlanarAngle(self, i_gid, j_gid, k_gid, l_gid, m_gid, n_gid):
"""
Get the planar angle (in radians) as defined by the six atoms: `i_gid`,
`j_gid`, `k_gid`, `l_gid`, `m_gid`, and `n_gid`. See the docstring of
`addPlanarAngle` for more detail.
"""
return self._planarangle[(i_gid, j_gid, k_gid, l_gid, m_gid, n_gid)]
[docs] def getCustom(self, cid):
"""
Return all results of the custom calculation type `c`
:type cid: Any hashable object
"""
return self._custom[cid]