## ...........................................................................
##
## PIMMS (Polymer Interactions in Multicomponent Mixtures)
## Alex Holehouse, Pappu Lab, Holehouse Lab
## Copyright 2015 - 2026
## ...........................................................................
##
## restart
##
## The RestartObject implemements a way to read and write restart files. This allows PIMMS
## to restart from previous simulations. Other than chain position, no other state is saved.
##
import pickle
import copy
from . import CONFIG
from .latticeExceptions import RestartException
from . import pimmslogger
[docs]
class RestartObject:
"""
Object used to read and write restart files. Restart information ONLY contains information on
chain position, sequence, and type, and grid dimenisons, but does NOT include any information
Note that the self.chains object in a RestartObject has the following structure:
1. Is a dictionary
2. Keys are chainID (i.e. each seperate chain has it's own entry)
3. values is a list with three elements
[0] : bead positions (N->C)
[1] : chain sequence (which will be referenced against the parameter file)
[2] : chainType : a single value that defines the type of chain
"""
#-----------------------------------------------------------------
#
def __init__(self):
"""
Initialize an empty RestartObject.
Sets up the internal state with zero energy, empty dimensions, a
non-hardwall flag, and empty chain / sequence-to-chaintype / extra-chain
containers. These are subsequently populated via one of the
``build_from_*`` methods or by adding extra chains.
Returns
-------
None
No return value; the new object's attributes are initialised in place.
"""
self.energy = 0
self.dimensions = []
self.hardwall = False
self.chains = {}
self.seq2chainType = {}
self.extra_chains = {}
#-----------------------------------------------------------------
#
def __apply_position_offset(self, position_offset):
"""
Function that allows position of each residue to be offset by some fixed amount.
This is not relevant for traditional restart operations, but is useful when using
a Restart object to initialize a new (resized) lattice. This requires that the
restart object dimensions are big enough to contain the newly offset positions.
Parameters
--------------
position_offset : list
List of integers with length equal to the number of dimensions in the lattice. Each element
in the list defines the amount by which the position of each residue should be offset.
Returns
-------------
None
No return type, but the internal self.chains object will be appropriately
updated.
Raises
-------------
RestartException
If the dimensions of the restart object are not big enough to contain the
newly offset positions.
"""
# <><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>
# Internal function that tests if a position (pos) in dimension (dim) is valid given the
# restart lattice' dimensions
def valid_pos(pos, dim):
pos = pos+position_offset[dim]
if (pos < 0) or (pos >= self.dimensions[dim]):
return False
else:
return True
# <><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>
# check offset dimensions match chain dimensions
if len(position_offset) != len(self.dimensions):
raise RestartException('Trying to apply position offset to restart object, but dimensions do not match.')
n_dim = len(self.dimensions)
# First pass: validate all proposed positions so this update is atomic.
for chainID in self.chains:
for position in self.chains[chainID][0]:
for dim in range(0, n_dim):
if not valid_pos(position[dim], dim):
raise RestartException(f'Trying to offet a position on chain {chainID} from {position[dim]} to {position[dim] + position_offset[dim]} (dim={dim}) but lattice dimensions are {self.dimensions}')
# Second pass: apply position updates once we know every position is valid.
for chainID in self.chains:
for position in self.chains[chainID][0]:
for dim in range(0, n_dim):
position[dim] = position[dim] + position_offset[dim]
#-----------------------------------------------------------------
#
def __update_seq2chainType(self, local_chainType, local_seq, log):
"""
Internal function called by both build_from_lattice() and build_from_file()
which ensures an updated and dynamically constructed self.seq2chainType dictionary
exists which enables mapping of protein sequence to a chainType.
Note we don't not allow two identical chain sequences to have different chainTypes - there
are some circumstances where this might be preferable, so, the seq2chainType mapping
is a one-to-many mapping, although IN GENERAL we probably expect this mostly to be
a 1-to-1 mapping.
If a one-to-many mapping is found and log=True then this is written via the pimmslogger as a
warning
Parameters
--------------
local_chainType : int
The chainType associated with the passed chain
local_seq : str
The amino acid sequence of the passed chain.
log : bool
Flag which, if set to true, means if this seq already has a chainType defined but the
passed chainType is a DIFFERENT value it'll warn the user about this.
Returns
-------------
None
No return type, but the internal self.seq2chainType dictionary will be appropriately
updated
"""
# if we've seen this sequence before
if local_seq in self.seq2chainType:
# If the chainType assigned here is associated with that previous record
# move on...
if local_chainType in self.seq2chainType[local_seq]:
pass
else:
self.seq2chainType[local_seq].append(local_chainType)
# note this is not strictly a problem, just might be good to know about...
if log:
pimmslogger.log_warning(f'When building RestartObject from Lattice found two identical chains [{local_seq}] with different chainType indices. This is not a bug or problem, but may be undesired...')
# if we've never seen this sequence before this is easy...
else:
self.seq2chainType[local_seq] = [local_chainType]
#-----------------------------------------------------------------
#
#-----------------------------------------------------------------
#
[docs]
def set_energy(self, energy):
"""
Set the RestartObject's stored energy value.
Parameters
----------
energy : float
The system energy to record in the restart object (written out when
the restart file is saved).
Returns
-------
None
No return value, but ``self.energy`` is updated in place.
"""
self.energy = energy
#-----------------------------------------------------------------
#
[docs]
def build_from_lattice(self, LATTICE, hardwall=False, log=False):
"""
Construct a restart object using a lattice object to set the chain
positions.
Parameter
------------
LATTICE : pimms.lattice.Lattice
A standard PIMMS latticd object
hardwall : bool (default = False)
Flag which sets of the current system defines a hardwall or, if false
PBC.
log : bool (default = False)
Flag which if set to True means warnings are written to the standard PIMMS
logfile
Returns
----------
None
No return type, but the internal self.chains dictionary will be appropriately
updated.
"""
self.dimensions = LATTICE.dimensions
self.hardwall = hardwall
# reset chain info...
self.chains = {}
self.seq2chainType = {}
self.extra_chains = {}
for chainID in LATTICE.chains:
local_chainType = LATTICE.chains[chainID].chainType
local_seq = LATTICE.chains[chainID].sequence
# add the chain to the restart object
self.chains[chainID] = [copy.deepcopy(LATTICE.chains[chainID].positions), local_seq, local_chainType]
# udpate the self.seq2chainType dictionary
self.__update_seq2chainType(local_chainType, local_seq, log)
#-----------------------------------------------------------------
#
[docs]
def update_lattice_dimensions(self, new_dimensions, manual_offset=None):
"""
Resize the lattice and reposition the chains within the new lattice.
Updates the restart object's dimensions and shifts every chain position
by a per-dimension offset. By default the offset is computed so the
existing chains end up centred in the larger lattice; alternatively an
explicit ``manual_offset`` can be supplied. If the offset would move any
bead outside the new lattice, the original dimensions are restored and
the underlying :class:`RestartException` is re-raised.
Parameters
----------
new_dimensions : list of int
The new lattice dimensions (length 2 or 3). Should be greater than or
equal to the current dimensions for centring to make sense.
manual_offset : list of int, optional
Explicit per-dimension offset to apply to every chain position. If
``None`` (the default), a centring offset is computed automatically
from the difference between ``new_dimensions`` and the current
dimensions.
Returns
-------
None
No return value, but ``self.dimensions`` and the stored chain
positions are updated in place.
Raises
------
RestartException
If applying the offset would place a bead outside the new lattice
(in which case the prior dimensions are restored before re-raising).
"""
## -----------
if manual_offset is None:
# calculate offset so the chains are placed in the center of the new lattice
x_off = int((new_dimensions[0] - self.dimensions[0])/2)
y_off = int((new_dimensions[1] - self.dimensions[1])/2)
if len(new_dimensions) == 3:
z_off = int((new_dimensions[2] - self.dimensions[2])/2)
position_offset=[x_off, y_off, z_off]
else:
position_offset=[x_off, y_off]
## -----------
# Manually provide the offsets to convert from old dimensions -> new dimensions
# TODO - add check in keyfile parser that manual_offset is reasonable
else:
position_offset = manual_offset
# Next construct and instantiate a new restart object which has the new dimensions
# including applying the possition offset we calculated above
# finaly, apply the offset on this 'new' lattice (order matters, as __apply_position_offset
# assesses if, given self.dimensions, the offset is valid or not)
# If this fails, restore prior dimensions.
old_dimensions = list(self.dimensions)
self.dimensions = new_dimensions
try:
self.__apply_position_offset(position_offset)
except RestartException:
self.dimensions = old_dimensions
raise
#-----------------------------------------------------------------
#
[docs]
def build_from_file(self, filename, log=False):
"""
Function that constructs a restart object from a passed filename. Performs some sanity check
in reading in the file but doesn't actually check that the chain positions make sense on the
lattice. We can and should probably make this better going forwards...
Parameters
--------------
filename : str
Name of the file to be read
log : bool (default = False)
Flag which if set to True means warnings are written to the standard PIMMS
logfile
Returns
-------------
None but updates the current object to contain self.dimensions, self.energy, self.hardwall
and self.chains[] info.
"""
# if IO issue (not IndexError often thrown if a valid file is found
# but its not actually a pickle file!
try:
with open(filename, "rb") as fh:
input_dict = pickle.load(fh)
except (OSError, EOFError, pickle.UnpicklingError, IndexError, ValueError) as e:
raise RestartException("Error reading restart file. Error:\n\n%s" %(str(e)))
# extract out key info (throw exception if missing)
try:
self.dimensions = input_dict['DIMENSIONS']
self.energy = input_dict['ENERGY']
self.hardwall = input_dict['HARDWALL']
# local chains is a dictionary where keys are chainIDs and values are lists with three elements
# [0] : bead positions (N->C)
# [1] : chain sequence (which will be referenced against the parameter file)
# [2] : chainType : a single value that defines the type of chain (many chains can have the same chainType,
# but each chain has a unique chainID)
local_chains = input_dict['CHAINS']
except KeyError as e:
raise RestartException("Invalid restart file - missing entry for %s" % (e.args[0]))
if not isinstance(local_chains, dict):
raise RestartException("Invalid restart file - CHAINS entry must be a dictionary")
# reset chain info...
self.chains = {}
self.seq2chainType = {}
self.extra_chains = {}
# one entry PER chain (not per chain type)
for chainID in local_chains:
# extract info for each chain
try:
local_pos = local_chains[chainID][0]
local_seq = local_chains[chainID][1]
local_chainType = local_chains[chainID][2]
except (TypeError, IndexError):
raise RestartException(f"Invalid restart file - malformed chain entry for chainID={chainID}")
# assign local info to the self.chains dictionary
self.chains[chainID] = copy.deepcopy(local_chains[chainID])
# check sequence and number of positions match
if len(local_seq) != len(local_pos):
raise RestartException("Invalid restart file - sequence length does not match number of positions")
for position in local_pos:
if len(position) != len(self.dimensions):
raise RestartException("Invalid restart file - chain position dimensionality does not match DIMENSIONS")
# update the self.seq2chainType dictionary
self.__update_seq2chainType(local_chainType, local_seq, log)
#-----------------------------------------------------------------
#
[docs]
def write_to_file(self):
"""
Serialize the restart object to disk as a pickle file.
Writes a dictionary containing the chain information (``CHAINS``), lattice
dimensions (``DIMENSIONS``), recorded energy (``ENERGY``) and hardwall
flag (``HARDWALL``) to ``CONFIG.RESTART_FILENAME`` using :mod:`pickle`.
Note that ``extra_chains`` are not written; only the materialised
``self.chains`` are saved.
Returns
-------
None
No return value; the restart data is written to
``CONFIG.RESTART_FILENAME``.
"""
output={}
output['CHAINS'] = {}
for chainID in self.chains:
output['CHAINS'][chainID] = self.chains[chainID]
output['DIMENSIONS'] = self.dimensions
output['ENERGY'] = self.energy
output['HARDWALL'] = self.hardwall
with open(CONFIG.RESTART_FILENAME, "wb") as fh:
pickle.dump(output, fh)