def unitcell_vectors(self, vectors):
"""Set the three vectors that define the shape of the unit cell
Parameters
----------
vectors : tuple of three arrays, each of shape=(n_frames, 3)
The semantics of this array are that the shape of the unit cell
in frame ``i`` are given by the three vectors, ``value[i, 0, :]``,
``value[i, 1, :]``, and ``value[i, 2, :]``.
"""
if vectors is None:
self._unitcell_lengths = None
self._unitcell_angles = None
return
if not len(vectors) == len(self):
raise TypeError('unitcell_vectors must be the same length as '
'the trajectory. you provided %s' % vectors)
v1 = vectors[:, 0, :]
v2 = vectors[:, 1, :]
v3 = vectors[:, 2, :]
a, b, c, alpha, beta, gamma = unitcell.box_vectors_to_lengths_and_angles(v1, v2, v3)
self._unitcell_lengths = np.vstack((a, b, c)).T
self._unitcell_angles = np.vstack((alpha, beta, gamma)).T
def report(self, simulation, state):
"""Generate a report.
Parameters
----------
simulation : simtk.openmm.app.Simulation
The Simulation to generate a report for
state : simtk.openmm.State
The current state of the simulation
"""
if not self._is_intialized:
self._initialize(simulation)
self._is_intialized = True
self._checkForErrors(simulation, state)
args = ()
kwargs = {}
if self._coordinates:
coordinates = state.getPositions(asNumpy=True)[self._atomSlice]
coordinates = coordinates.value_in_unit(getattr(unit, self._traj_file.distance_unit))
args = (coordinates,)
if self._time:
kwargs['time'] = state.getTime()
if self._cell:
vectors = state.getPeriodicBoxVectors(asNumpy=True)
vectors = vectors.value_in_unit(getattr(unit, self._traj_file.distance_unit))
a, b, c, alpha, beta, gamma = unitcell.box_vectors_to_lengths_and_angles(*vectors)
kwargs['cell_lengths'] = np.array([a, b, c])
kwargs['cell_angles'] = np.array([alpha, beta, gamma])
if self._potentialEnergy:
kwargs['potentialEnergy'] = state.getPotentialEnergy()
if self._kineticEnergy:
kwargs['kineticEnergy'] = state.getKineticEnergy()
if self._temperature:
kwargs['temperature'] = 2*state.getKineticEnergy()/(self._dof*unit.MOLAR_GAS_CONSTANT_R)
if self._velocities:
kwargs['velocities'] = state.getVelocities(asNumpy=True)[self._atomSlice, :]
#add a portion like this to store things other than the protocol work
if self._protocolWork:
protocol_work = simulation.integrator.get_protocol_work(dimensionless=True)
kwargs['protocolWork'] = np.array([protocol_work])
if self._alchemicalLambda:
kwargs['alchemicalLambda'] = np.array([simulation.integrator.getGlobalVariableByName('lambda')])
if self._title:
kwargs['title'] = self._title
if self._parameters:
kwargs['parameters'] = self._parameters
if self._environment:
kwargs['environment'] = self._environment
self._traj_file.write(*args, **kwargs)
# flush the file to disk. it might not be necessary to do this every
# report, but this is the most proactive solution. We don't want to
# accumulate a lot of data in memory only to find out, at the very
# end of the run, that there wasn't enough space on disk to hold the
# data.
if hasattr(self._traj_file, 'flush'):
self._traj_file.flush()
def report(self, simulation, state):
"""Generate a report.
Parameters
----------
simulation : simtk.openmm.app.Simulation
The Simulation to generate a report for
state : simtk.openmm.State
The current state of the simulation
"""
if not self._is_intialized:
self._initialize(simulation)
self._is_intialized = True
self._checkForErrors(simulation, state)
args = ()
kwargs = {}
if self._coordinates:
coordinates = state.getPositions(asNumpy=True)[self._atomSlice]
coordinates = coordinates.value_in_unit(getattr(units, self._traj_file.distance_unit))
args = (coordinates,)
if self._time:
kwargs['time'] = state.getTime()
if self._cell:
vectors = state.getPeriodicBoxVectors(asNumpy=True)
vectors = vectors.value_in_unit(getattr(units, self._traj_file.distance_unit))
a, b, c, alpha, beta, gamma = unitcell.box_vectors_to_lengths_and_angles(*vectors)
kwargs['cell_lengths'] = np.array([a, b, c])
kwargs['cell_angles'] = np.array([alpha, beta, gamma])
if self._potentialEnergy:
kwargs['potentialEnergy'] = state.getPotentialEnergy()
if self._kineticEnergy:
kwargs['kineticEnergy'] = state.getKineticEnergy()
if self._temperature:
kwargs['temperature'] = 2*state.getKineticEnergy()/(self._dof*units.MOLAR_GAS_CONSTANT_R)
if self._velocities:
kwargs['velocities'] = state.getVelocities(asNumpy=True)[self._atomSlice, :]
self._traj_file.write(*args, **kwargs)
# flush the file to disk. it might not be necessary to do this every
# report, but this is the most proactive solution. We don't want to
# accumulate a lot of data in memory only to find out, at the very
# end of the run, that there wasn't enough space on disk to hold the
# data.
if hasattr(self._traj_file, 'flush'):
self._traj_file.flush()
def get_state_as_mdtraj(simulation):
"""Construct a length-1 trajectory with unitcells from the current simulation state"""
state = simulation.context.getState(getPositions=True)
xyz = state.getPositions(asNumpy=True).value_in_unit(unit.nanometer)
box_vectors = [
np.array(v.value_in_unit(unit.nanometer))
for v in state.getPeriodicBoxVectors()
]
a_length, b_length, c_length, alpha, beta, gamma = unitcell.box_vectors_to_lengths_and_angles(
*box_vectors)
return md.Trajectory([xyz],
topology=md.Topology().from_openmm(
simulation.topology),
unitcell_lengths=(a_length, b_length, c_length),
unitcell_angles=(alpha, beta, gamma))
