import os
from schrodinger import get_maestro
from schrodinger import structure
from schrodinger.infra import mm
from schrodinger.Qt import QtWidgets
from schrodinger.structutils import analyze
from schrodinger.trajectory import utils
from schrodinger.trajectory.trajectory_gui_dir import export_snapshot_structures
from schrodinger.trajectory.trajectory_gui_dir import export_structures
from schrodinger.trajectory.trajectory_gui_dir.display_settings_manager import \
DisplaySettingsManager
from schrodinger.trajectory.trajectory_gui_dir.export_structure_enums import \
ExportFrameOption
from schrodinger.trajectory.trajectory_gui_dir.export_structure_enums import \
ExportMode
from schrodinger.trajectory.trajectory_gui_dir.export_structure_enums import \
ExportToOption
from schrodinger.trajectory.trajectory_gui_dir.playback_settings_data import \
CURRENT_FRAME
from schrodinger.trajectory.trajectory_gui_dir.secondary_structure_data import \
SecondaryStructureData
from schrodinger.ui import maestro_ui
from schrodinger.ui.qt import filedialog
from schrodinger.ui.qt.decorators import wait_cursor
from schrodinger.utils import fileutils
try:
from schrodinger.application.desmond.packages import topo
from schrodinger.application.desmond.packages import traj
except ImportError as err:
topo = None # Test may not have desmond.
maestro = get_maestro()
[docs]class FrameStructureExporter:
"""
This class provides ability to export structures in the file.
Client can export trajectory structures in three ways
1. Export current frame.
2. Export current frame full system ct and all its components.
3. Export set of frames specified in the range.
"""
EXPORT_TRJ_FILE_PREF = 'TRJ_PLAYER_EXPORT_TRJ_FILE'
[docs] def __init__(self, player_obj, export_mode: ExportMode,
parent: QtWidgets.QWidget):
"""
:type player_obj: TrajectoryPlayer
:param player_obj: Trajectory player toolbar object.
:param export_mode: Indicates whether exports need to be done in
snapshot mode or player mode.
:param parent: Parent widget (used to set export structures dialog
parent)
"""
self.player_obj = player_obj
self._export_mode = export_mode
self._parent = parent
self._dlg = None
if self._export_mode == ExportMode.TRAJECTORY_VIEWER:
self._dlg = export_structures.ExportStructures(parent)
else:
self._dlg = export_snapshot_structures.ExportStructures(parent)
self._dlg.exportButtonClicked.connect(self._handleExportStrcutres)
[docs] def clearData(self):
"""
Clean all temporary data members.
"""
self._working_ct = None
self.ss_data = None
self._saved_frame_ct = None
self._saved_ss_data = None
[docs] def allInRange(self, start_frame, end_frame, step_size, atoms_asl, writer):
"""
Export all structures specified in the range including start and end
frame if step size is 1.
Otherwise, export frames in the range based on step size, but last frame
is need not to be included.
:type start_frame: int
:param start_frame: Starting frame in the range.
:type end_frame: int
:param end_frame: End frame in the range.
:type step_size: int
:param step_size: Step size to be used when stepping in the range.
:type atoms_asl: str
:param atoms_asl: Set of atoms to be exported from a given frame.
:type writer: structure.StructureWriter
:param writer: Write object which writes structures in the file.
:rtype: bool
:return: Whether frame is exported successfully.
"""
for frame in range(start_frame, end_frame + 1, step_size):
if not self.currentFrameOnly(frame, atoms_asl, writer):
return False
return True
def _updateVelocity(self, ct, frame):
"""
Update velocity in the ct for all atoms of frame.
"""
tpt = self.player_obj
fr = tpt.entry_traj.getFrame(frame)
velocity = fr.vel()
if velocity is not None:
allaid_vel = velocity[tpt.entry_traj.cms_model.allaid_gids]
allaid_vel_len = len(allaid_vel)
# It is possible that working ct has lesser or greater atoms than
# cms_model due to varying atoms frame.
for at in self._working_ct.atom:
atom_index = at.index
if atom_index <= allaid_vel_len:
v = allaid_vel[atom_index - 1]
topo.set_atom_velocity(at, v)
[docs] def currentFrameOnly(self, frame, atoms_asl, writer):
"""
Export only frame number associated structure.
:type frame: int
:param frame: Frame number of structure to be exported.
:type atoms_asl: str
:param atoms_asl: Set of atoms to be exported from a given frame.
:type writer: structure.StructureWriter
:param writer: Write object which writes structures in the file.
:rtype: bool
:return: Whether frame is exported successfully.
"""
tpt = self.player_obj
self._working_ct = self.getUpdatedFrameStructure(frame)
# Update velocity
self._updateVelocity(self._working_ct, frame)
# Extract atoms based on asl.
if atoms_asl is not None and atoms_asl != 'all':
atoms = analyze.evaluate_asl(self._working_ct, atoms_asl)
if not atoms:
maestro.warning(f'ASL {atoms_asl} does not result any atom.')
return False
new_ct = self._working_ct.extract(atoms, copy_props=True)
else:
new_ct = self._working_ct.copy()
new_ct.property[CURRENT_FRAME] = frame
if mm.M2IO_DATA_TRAJECTORY_FILE in new_ct.property:
del new_ct.property[mm.M2IO_DATA_TRAJECTORY_FILE]
new_ct.title = tpt.entry_traj.cms_model.fsys_ct.title + " - Frame " + str(
frame)
writer.append(new_ct)
return True
[docs] def getUpdatedFrameStructure(self, frame: int):
"""
:param frame: Frame number of structure to be exported.
:return: Update frame structure.
:rtype: structure.Structure
"""
tpt = self.player_obj
# Update structure from frame.
if self._export_mode == ExportMode.TRAJECTORY_VIEWER:
self._working_ct, _, _, _, _ = tpt.getUpdatedCT(tpt.entry_traj,
frame,
self._working_ct,
honor_replica=False)
else:
self._working_ct = tpt.getUpdatedFrameSnapshot(
frame, self._working_ct, True, self._use_trajectory_settings)
# If display settings manager is provided, it implies that we need
# to apply custom display settings.
if self._display_settings_manager:
varying_atoms_asl = tpt.isVaryingAtomsFrame()
ref_ct = tpt.getOriginalEntryCT(copy=False)
if self._display_settings_manager.shouldRestoreVisibility(
varying_atoms_asl):
self._display_settings_manager.restoreDefaultVisibility(
ref_ct, self._working_ct)
self._display_settings_manager.applyShowHideAtomsAsl(
self._working_ct, varying_atoms_asl)
# If cooltowarmcolor_map is provided, it implies that we need to
# apply color gradient on exported frames.
if self._cooltowarmcolor_map:
self._cooltowarmcolor_map.applyColors(self._working_ct.handle,
[frame])
return self._working_ct
[docs] def getUpdatedCMSModelCopy(self, frame):
"""
"""
tpt = self.player_obj
# Take a copy of cms model.
cms_model_copy = tpt.entry_traj.cms_model.copy()
frame_obj = tpt.entry_traj.getFrame(frame)
topo.update_cms(cms_model_copy, frame_obj)
cts = [cms_model_copy.fsys_ct] + cms_model_copy.comp_ct
for st in cts:
if mm.M2IO_DATA_TRAJECTORY_FILE in st.property:
del st.property[mm.M2IO_DATA_TRAJECTORY_FILE]
return cms_model_copy
[docs] def currentFrameWithComponentStructures(self, frame, writer):
"""
Export current frame along with its component structures.
:type frame: int
:param frame: Frame number of structure to be exported.
:type writer: structure.StructureWriter
:param writer: Write object which writes structures in the file.
:rtype: bool
:return: Whether frame is exported successfully.
"""
cms_model = self.getUpdatedCMSModelCopy(frame)
# Write component cts.
cts = [cms_model.fsys_ct] + cms_model.comp_ct
for st in cts:
writer.append(st)
return True
def _getFilePath(self, dlg):
"""
Return file path provided by user if export to external file.
Return pre-defined structure file path if export to project table.
:type dlg: export_structures.ExportStructures
:param dlg: The instance of export structure dialog.
:rtype: str
:return: File full path.
"""
frame_option = dlg.export_frame_option
pt_export = dlg.export_to_option == ExportToOption.PROJECT_TABLE
if pt_export:
# Export to project table is achieved first export into an external
# file and later import manually.
file_name = {
ExportFrameOption.ONE_PER_STEP_IN_RANGE: "per_step_frames.mae",
ExportFrameOption.ALL_IN_RANGE: "All Frames in Range.mae",
ExportFrameOption.ONLY_CURRENT_FRAME: "currentframe.mae",
ExportFrameOption.CURRENT_FRAME_WITH_COMPONENT_CTS: "Component Structures.mae",
ExportFrameOption.SNAPSHOT_FRAMES: "Snapshot Frames.mae",
ExportFrameOption.ALL_IN_LIMITED_RANGE: "All Frames (limited range).mae",
}[frame_option]
file_path = os.path.join(maestro_ui.mm_get_maestro_temp_location(),
file_name)
else:
file_path = dlg.export_file_path
return file_path
[docs] def showExportStructuresDlg(self):
"""
Show export structure dialog and collect necessary input data.
"""
tpt = self.player_obj
self._dlg.showDlg(tpt.getStartFrame(), tpt.getEndFrame(),
tpt.step_size, tpt.total_frame,
len(tpt.entry_traj.cms_model.comp_ct))
[docs] def showExportSnapshotStructuresDlg(
self, all_frames: list, selected_frames: list,
use_trajectory_settings: bool,
display_settings_manager: DisplaySettingsManager,
cooltowarmcolor_map):
"""
Show export snapshot structure dialog and collect necessary input data.
:param all_frames: List of all frames listed in the snapshot panel.
:param selected_frames: List of all selected frames in the snapshot
panel.
:param use_trajectory_settings: True if trajectory settings should be
used exporting a frame.
:param display_settings_manager: Valid display setting manager to apply
the display settings, otherwise None
:param cooltowarmcolor_map: Instance of cool to warm color map
object. It is used to apply color on exported frames. It will be
None if cool to warm color gradient should not be applied.
:type cooltowarmcolor_map: maestro_ui.CoolToWarmColorMap
"""
self._use_trajectory_settings = use_trajectory_settings
self._display_settings_manager = display_settings_manager
self._cooltowarmcolor_map = cooltowarmcolor_map
self._dlg.showDlg(all_frames, selected_frames)
def _handleExportStrcutres(self):
"""
Handle 'Export Structures' operation when user clicks on 'Export'
button in 'Export Structures' dialog
"""
tpt = self.player_obj
# Collect user settings to export structure.
file_path = self._getFilePath(self._dlg)
if file_path is None:
maestro.warning("File path must be specified.")
return
# Export structures in the file.
self._exportStructures(file_path, self._dlg)
# Import structures in the project and remove temporary file.
pt_export = self._dlg.export_to_option == ExportToOption.PROJECT_TABLE
if pt_export and os.path.exists(file_path):
tpt.proj.importStructureFile(file_path,
wsreplace=False,
wsinclude='none')
os.remove(file_path)
def _playerToolbarSaveState(self):
"""
Save playertoolbar current state, so that export structure
process does not change workspace.
"""
tpt = self.player_obj
self._saved_ss_data = tpt.entry_traj.ss_data
self._saved_frame_ct = tpt.entry_traj.frame_ct
self.ss_data = SecondaryStructureData()
tpt.entry_traj.ss_data = self.ss_data
self._working_ct = tpt.getOriginalEntryCT(copy=True)
utils.clone_workspace_selection(self._working_ct)
tpt.entry_traj.updateSavedFrameStructureForExport(self._working_ct)
tpt.setToolbarEnabled(False)
def _playerToolbarRestoreState(self):
"""
Restore playertoolbar current state.
"""
tpt = self.player_obj
tpt.entry_traj.ss_data = self._saved_ss_data
tpt.entry_traj.frame_ct = self._saved_frame_ct
self.clearData()
tpt.setToolbarEnabled(True)
@wait_cursor
def _exportStructures(self, file_path, dlg):
"""
Export structures in the given file based on frame option setting.
Export to project table if pt_export is True. It is achieved by
exporting structures into an external file and later import manually.
:type file_path: str
:param file_path: File path in which exported structures are added.
:type dlg: export_structures.ExportStructures
:param dlg: The instance of export structure dialog.
"""
tpt = self.player_obj
self._playerToolbarSaveState()
writer = structure.StructureWriter(file_path)
asl = dlg.export_asl
frame_option = dlg.export_frame_option
step_size = tpt.step_size if (
frame_option == ExportFrameOption.ONE_PER_STEP_IN_RANGE) else 1
if frame_option in (ExportFrameOption.ONE_PER_STEP_IN_RANGE,
ExportFrameOption.ALL_IN_RANGE,
ExportFrameOption.ALL_IN_LIMITED_RANGE):
ok_status = self.allInRange(dlg.start, dlg.end, step_size, asl,
writer)
elif frame_option == ExportFrameOption.ONLY_CURRENT_FRAME:
if (fileutils.is_cms_file(file_path)):
new_cms_model = self.getUpdatedCMSModelCopy(tpt.current_value)
new_cms_model.write(file_path)
ok_status = True
else:
ok_status = self.currentFrameOnly(tpt.current_value, asl,
writer)
elif frame_option == ExportFrameOption.SNAPSHOT_FRAMES:
for frame in dlg.frames_list:
ok_status = self.currentFrameOnly(frame, asl, writer)
if not ok_status:
break
else:
ok_status = self.currentFrameWithComponentStructures(
tpt.current_value, writer)
# If export failed for some reason, delete partially exported structure
# file.
if not ok_status and os.path.exists(file_path):
os.remove(file_path)
self._playerToolbarRestoreState()
[docs] def exportTrajectory(self):
"""
Export trajectory to file.
"""
tpt = self.player_obj
filter_string = filedialog.filter_string_from_formats([fileutils.CMS])
abs_cms_fn = filedialog.get_save_file_name(
parent=self._parent,
caption='Choose File for Export',
id=self.EXPORT_TRJ_FILE_PREF,
filter=filter_string)
if not abs_cms_fn:
return
path_dir, out_cms_fn = os.path.split(abs_cms_fn)
root, ext = fileutils.splitext(out_cms_fn)
out_trj_fn = '%s_trj' % root
# write_traj doesn't like Windows absolute paths (MATSCI-9050)
with fileutils.chdir(path_dir):
traj.write_traj(tpt.entry_traj.trajectory, out_trj_fn)
cms_model_copy = self.getUpdatedCMSModelCopy(
len(tpt.entry_traj.trajectory))
cms_model_copy.fix_filenames(out_cms_fn, out_trj_fn)
cms_model_copy.write(out_cms_fn)