Source code for schrodinger.trajectory.trajectoryviewer

from schrodinger import get_maestro
from schrodinger.trajectory import utils
from schrodinger.trajectory.trajectory_gui_dir import snapshot_panel
from schrodinger.ui import maestro_ui

# Global instance of trajectory player toolbar.
# This toolbar is docked in the Maestro main window.
player_toolbar = None

# Global instance of trajectory snapshot panel.
frame_snapshot_panel = None
maestro = get_maestro()


def has_player_toolbar():
    """
    :rtype: bool
    :return: True if player toolbar instance is already created.
    """

    return player_toolbar is not None


def has_snapshot_panel():
    """
    :rtype: bool
    :return: True if snapshot panel instance is already created.
    """

    return frame_snapshot_panel is not None


def get_player_toolbar():
    """
    Return trajectory player toolbar or None if desmond is not available.
    """

    if not utils.is_desmond_product_installed():
        return None
    global player_toolbar
    if player_toolbar is None:
        from schrodinger.trajectory.trajectory_gui_dir import playertoolbar
        maestro_hub = maestro_ui.MaestroHub.instance()
        trj_hub = maestro_hub.getTrajectoryViewerHub()
        player_toolbar = playertoolbar.TrajectoryPlayer(
            trj_hub, maestro_hub.isMaterialsScienceProfileMode())
        # It is set only via maestro.
        trj_layout = trj_hub.getTrajectoryPlayerLayout()
        if trj_layout is not None:
            # Add this widget in Maestro MainWindow trajectory player layout
            trj_layout.addWidget(player_toolbar)
        ws_hub = maestro_ui.WorkspaceHub.instance()
        ws_hub.inclusionChanged.connect(player_toolbar.inclusionUpdated)
        ws_hub.atomDisplayChanged.connect(player_toolbar.visibilityChanged)
        ws_hub.workspaceChanged.connect(
            player_toolbar.workspaceEntryStructureChanged)
        player_toolbar.showTrajectorySnapshotPanel.connect(
            lambda: setup_snapshot_panel(True))
        trj_hub.stopTrajectoryAndExitSnapshotMode.connect(
            stop_player_and_exit_snapshot_mode)
    return player_toolbar


def include_only_entry(eid: int, unlock_locked_entries: bool):
    """
    Include only given entry and exclude all other entries
    depending on unlocked locked entries option value.

    :param eid: Entry id to be included only.

    :param unlock_locked_entries: True if locked entries should
        be unlocked before including specified entry.

    """

    if unlock_locked_entries:
        maestro.command('entrywsincludeunlock all')
    maestro.command(f'entrywsincludeonly entry {eid}')


def include_and_show_trajectory(eid: int, include_only: bool,
                                unlock_locked_entries: bool):
    """
    Include only given trajectory entry and display trajectory player.

    :param eid: Trajectory entry id.

    :param include_only: Exclude all other entries and include only given
        trajectory entry.

    :param unlock_locked_entries: Unlock locked entries if it is True.

    Note caller must pass valid trajectory entry id present in the project.
    """

    if include_only:
        include_only_entry(eid, unlock_locked_entries)
    setup_player_toolbar(eid, True)
    # If snapshot panel is visible, we should update it as well.
    if has_snapshot_panel() and get_snapshot_panel().isVisible():
        setup_snapshot_panel(True)


def include_and_show_trajectory_snapshot(eid: int, include_only: bool):
    """
    Include only given trajectory entry and display trajectory snapshot panel.

    :param eid: Trajectory entry id.

    :param include_only: Exclude all other entries and include only given
        trajectory entry.

    Note caller must pass valid trajectory entry id present in the project.
    """

    if include_only:
        include_only_entry(eid, True)
    snapshot_panel = get_snapshot_panel()
    # Maybe desmond is not installed.
    if snapshot_panel is None:
        return
    visible = get_player_toolbar().isVisible()
    setup_player_toolbar(eid, visible)
    setup_snapshot_panel(visible)


def setup_snapshot_panel(player_visible: bool):
    """
    Setup snapshot panel and show it.

    :param player_visible: True if player is visible.
    """

    snapshot_panel = get_snapshot_panel()
    snapshot_panel.run()
    snapshot_panel.updatePanel()
    snapshot_panel.setEnabledCurrentFrameOptions(player_visible)
    snapshot_panel.setVisible(True)


def setup_player_toolbar(eid, visible: bool):
    """
    Setup player toolbar using given trajectory entry id.

    :param eid: Trajectory entry id.

    Note caller must pass valid trajectory entry id present in the project.
    """

    player = get_player_toolbar()
    if player:
        player.setupTrajectory([eid], visible)


def get_snapshot_panel():
    """
    Construct trajectory snapshot panel instance.

    :rtype: snapshot_pane.SnapshotPanel
    :return: Snapshot panel instance or None if Desmond product is not
        installed.
    """

    global frame_snapshot_panel
    if frame_snapshot_panel is None:
        player = get_player_toolbar()
        if player:
            frame_snapshot_panel = snapshot_panel.SnapshotPanel(player)
            player.currentFrameChanged.connect(
                frame_snapshot_panel.updateCurrentFrameLabel)
            player.trajectoryLoaded.connect(frame_snapshot_panel.updatePanel)
            player.trajectoryUnloaded.connect(
                lambda: frame_snapshot_panel.setEnabled(False))
            player.closed.connect(lambda: frame_snapshot_panel.
                                  setEnabledCurrentFrameOptions(False))
            frame_snapshot_panel.modeChanged.connect(player.snapshotModeChanged)
            frame_snapshot_panel.closed.connect(player.snapshotPanelClosed)
            frame_snapshot_panel.visibilityChanged.connect(
                player.setSnapshotPanelVisibility)
            frame_snapshot_panel.atomsColorChanged.connect(
                player.trj_hub.allAtomsChanged.emit)
            maestro_hub = maestro_ui.MaestroHub.instance()
            maestro_hub.aboutToQuitMaestro.connect(
                frame_snapshot_panel.aboutToQuitMaestro)
    return frame_snapshot_panel


def stop_player_and_exit_snapshot_mode():
    """
    If user includes non-trajectory entry, then stop player if it is running,
    and exit from snapshot mode if it is active.
    """

    player = None
    if has_player_toolbar():
        player = get_player_toolbar()
    snapshot = None
    if has_snapshot_panel():
        snapshot = get_snapshot_panel()
    if player and player.isVisible() and player.player_active:
        player.frame_timer.stop()
    if snapshot and snapshot.isVisible() and player.is_snapshot_mode_active:
        player.snapshotModeChanged(False)
        snapshot.ui.display_snapshots_toolbutton.setChecked(False)