def __init__(self, cms_file, trj_dir, queue, frames=None, params=None):
"""
Docstring
:param cms_file:
:param trj_dir:
:param frames:
:param params:
"""
multiprocessing.Process.__init__(self)
self.queue = queue
# load cms_model
self.msys_model, self.cms_model = topo.read_cms(str(cms_file))
# Get framelist
if frames is not None:
self.frame_list = [
frame for (i, frame) in enumerate(traj.read_traj(str(trj_dir)))
if i in frames
]
else:
self.frame_list = traj.read_traj(str(trj_dir))
self.total_frames = len(self.frame_list)
self.align = Superimposer()
# Calculation parameters
self.calculation_parameters = []
if params is not None:
for param in params:
self.add_param(**param)
def extract_st(cms_file, trj_dir, asl='all', frames=1):
"""
Extract N frames from the trajectory and save them in structure.Structure
:param cms_file:
:param trj_dir:
:param asl:
:param frames: if int: number of frames uniformly distributed, if iterable: list fo frames
:return structures: list of schrodinger.structure.Structure
:type structures: list
"""
structures = []
msys_model, cms_model = topo.read_cms(cms_file)
frame_arch = traj.read_traj(trj_dir)
if type(frames) == int:
for f in np.linspace(0, len(frame_arch) - 1, frames).astype(int):
st = topo.update_cms(cms_model, frame_arch[f])
st = st.extract(evaluate_asl(cms_model, asl))
st.title = 'Frame {}'.format(f)
structures.append(st)
else:
try:
for f in frames:
st = topo.update_cms(cms_model, frame_arch[f])
st = st.extract(asl)
st.title = 'Frame {}'.format(f)
structures.append(st)
except Exception as e:
raise RuntimeError(e)
return structures
def from_files(cls, cms_fname, traj_fname):
from schrodinger.application.desmond.packages import topo, traj
_, model = topo.read_cms(cms_fname)
traj = traj.read_traj(traj_fname)
return cls(model, traj)
def _process(structure_dict):
"""
:param structure_dict:
:return:
"""
t = time.time()
fork = None
# Check if transformers is called as part of a pipeline
if 'pipeline' in structure_dict['custom']:
pipeline = structure_dict['custom']['pipeline']
fork = [pipeline[0], ]
if len(pipeline) == 1:
del (structure_dict['custom']['pipeline'])
else:
structure_dict['custom']['pipeline'] = pipeline[1:]
structure_code = structure_dict['structure']['code']
logger.info('Load simulation files')
cms_file = structure_dict['files']['desmond_cms']
msys_model, cms_model = topo.read_cms(str(cms_file))
trjtar = structure_dict['files']['desmond_trjtar']
# If run from command line it does not make sense to provide a tarfile
if os.path.isdir(trjtar):
trj_dir = trjtar
elif tarfile.is_tarfile(trjtar):
with tarfile.open(name=trjtar, mode='r:gz') as tfile:
tfile.extractall()
trj_dir = tfile.getnames()[0]
else:
raise RuntimeError('trjtar is neither a directory nor a tarfile')
frame_list = traj.read_traj(str(trj_dir))
frame_list = [frame_list[i] for i in range(0, len(frame_list), STEP)]
logger.info('Calculating torsion angles')
cms_model = set_original_atom_index(cms_model)
ligand_ct = cms_model.extract(evaluate_asl(cms_model, 'ligand'), copy_props=True)
torsion_list = get_hetero_torsion_atoms(ligand_ct, element_priority=ELEMENT_PRIORITY)
torsion_list.extend(get_protein_torsion_atoms(cms_model))
analyzers = []
torsion_ids = pd.DataFrame(columns=['index', 'aid1', 'aid2', 'aid3', 'aid4'])
torsion_ids.set_index('index', inplace=True)
for i, atom_set in enumerate(torsion_list):
atom_set = list(map(get_original_atom_index, atom_set))
analyzers.append(Torsion(msys_model, cms_model, *atom_set))
torsion_ids.loc[i, ['aid1', 'aid2', 'aid3', 'aid4']] = atom_set
results = analyze(frame_list, *analyzers,
**{"progress_feedback": functools.partial(print_iframe, logger=logger)})
out_arch = '{}_torsion.tar.gz'.format(structure_code)
with tarfile.open(out_arch, 'w:gz') as tar:
torsion_ids.to_csv('torsion_ids.csv', sep=',')
tar.add('torsion_ids.csv')
for i, timeseries in enumerate(results):
fname = 'torsion_{}.csv'.format(i)
np.savetxt(fname, timeseries, delimiter=',')
tar.add(fname)
logger.info('Calculated torsion angles in {:.0f} seconds'.format(time.time() - t))
# Return structure dict
transformer_dict = {
'structure': {
'parent_structure_id':
structure_dict['structure']['structure_id'],
'searchable': False
},
'files': {'trj_torsion': out_arch},
'custom': structure_dict['custom']
}
if fork is not None:
logger.info('Forking pipeline: ' + ' '.join(fork))
transformer_dict['control'] = {'forks': fork}
yield transformer_dict
