import itertools
import os
from typing import Union
import numpy
from schrodinger import get_maestro
from schrodinger import structure
from schrodinger.application.desmond import cms
from schrodinger.application.matsci import cgforcefield
from schrodinger.application.matsci.nano import xtal
from schrodinger.infra import mm
from schrodinger.infra.mm import M2IO_DATA_ORIGINAL_CMS_FILE
from schrodinger.infra.mm import M2IO_DATA_TRAJECTORY_FILE
from schrodinger.infra.mmproj import M2IO_DATA_SOURCE_FILE
from schrodinger.project import utils as projutils
from schrodinger.Qt import QtCore
HIDDEN_GENERATED_CMS_FILE = 's_m_hidden_generated_cms_file'
maestro = get_maestro()
[docs]def is_desmond_product_installed():
"""
Check if desmond product is installed or not. We can not view trajectory
features if there is no Desmond product.
"""
try:
from schrodinger.application.desmond.packages import topo # noqa: F401
from schrodinger.application.desmond.packages import traj # noqa: F401
from schrodinger.application.desmond.packages import traj_util # noqa: F401
except ImportError:
maestro.warning("Desmond product is not available.")
return False
return True
[docs]def get_trajectory_path(proj, eid):
"""
Return trajectory file path if any, or None
Old trajectory was not storing trajectory directory path in the property,
but stores .idx file path. So, we need to do idx trajectory path treatment
to know the trajectory directory name.
:param proj: Project on which to operate.
:type proj: Project
:param eid: Entry id associated with the given project.
:type eid: int or str
"""
traj_path = projutils.get_trajectory_path(proj, eid)
if traj_path:
# cms.fix_idx_traj_path does necessary path treatment only if required,
# otherwise it returns same value.
return cms.fix_idx_traj_path(traj_path)
[docs]def get_cms_file_path(proj, eid):
"""
Return trajectory file path if there is any, or None
It attemtps to find cms file path in the same directory where trajectory
exists, otherwise retrieves from the project.
:param proj: Project on which to operate.
:type proj: Project
:param eid: Entry id associated with the given project.
:type eid: int or str
"""
trj_path = get_trajectory_path(proj, eid)
cms_file_path = proj[eid][M2IO_DATA_ORIGINAL_CMS_FILE]
if trj_path and cms_file_path:
cms_file_full_path = os.path.join(os.path.dirname(trj_path),
cms_file_path)
if os.path.exists(cms_file_full_path):
return cms_file_full_path
return proj[eid].cms_file
[docs]def get_hidden_cms_file_path(proj, eid):
"""
Return trajectory hidden cms file path, or None
:param proj: Project on which to operate.
:type proj: Project
:param eid: Entry id associated with the given project.
:type eid: int or str
"""
trj_path = get_trajectory_path(proj, eid)
hidden_cms_file_path = proj[eid][HIDDEN_GENERATED_CMS_FILE]
if trj_path and hidden_cms_file_path:
return os.path.join(os.path.dirname(trj_path), hidden_cms_file_path)
[docs]def get_unique_cms_file_name(directory_path, base_name):
"""
Gets unique cms file name in the given directory for the given base name. It
will prefix an increasing number in case the file already exists.
:param directory_path: Directory path in which unique cms file name is
requested.
:type directory_path: str
:param base_name: Base name for the unique cms file name.
:type base_name: str
"""
directory = QtCore.QDir(directory_path)
filter_list = ['*.cms']
cms_files = directory.entryList(filter_list, QtCore.QDir.Files)
unique_cms_name = base_name + ".cms"
prefix = 0
while unique_cms_name in cms_files:
prefix = prefix + 1
unique_cms_name = base_name + str(prefix) + '.cms'
return unique_cms_name
[docs]def generate_cms_file_from_entry(proj, eid):
"""
Generates cms file from entry and sets the generated cms file path to entry's
s_m_hidden_generated_cms_file property.
:param proj: Project on which to operate.
:type proj: Project
:param eid: Entry id associated with the given project.
:type eid: int or str
"""
# Get the trajectory path
trj_path = get_trajectory_path(proj, eid)
# Trajectory parent directory, this is where the newly generated cms file
# will be placed
trj_parent_directory = os.path.dirname(trj_path)
# Base file name for the cms file to be generated
source_file = proj[eid][M2IO_DATA_SOURCE_FILE]
# In case there is no 's_m_Source_File' property then we use 'generated' as
# the base file name
if source_file is None:
source_file = 'generated'
# Remove the extension from the base file name
base_name = os.path.splitext(source_file)[0]
# Get the unique name in the trajectory parent directory for the base file
# name
hidden_cms_file_name = get_unique_cms_file_name(trj_parent_directory,
base_name)
# Add the directory to make it absolute file path
hidden_cms_file_path = os.path.join(trj_parent_directory,
hidden_cms_file_name)
row = proj.getRow(eid)
st = row.getStructure()
# Set the trajectory file property to the entry
st.property[M2IO_DATA_TRAJECTORY_FILE] = get_trajectory_path(proj, eid)
cgffiost = cgforcefield.CGFFIOStructure(st)
cgffiost.writeCMS(st, hidden_cms_file_path)
# Set the hidden property of the entry so that it will not be generated
# again
proj[eid][HIDDEN_GENERATED_CMS_FILE] = hidden_cms_file_name
[docs]def set_trajectory_path(eid: Union[str, int], path: str):
"""
Update the s_chorus_trajectory_file property.
:param eid: The entry id being checked.
:param path: The new s_chorus_trajectory_file value to set
"""
cmd = (
f"entrysetprop property={M2IO_DATA_TRAJECTORY_FILE} value=\"{path}\" entry {eid}"
)
maestro.command(cmd)
[docs]def clone_workspace_selection(st: structure.Structure):
"""
Clones the selection from workspace structure to given structure.
:param st: Structure for which selection needs to be applied.
"""
main_ct = maestro.workspace_get(copy=False)
atom_total = min(st.atom_total, main_ct.atom_total)
for atom in range(1, atom_total + 1):
if mm.mmct_atom_is_selected(main_ct, atom):
mm.mmct_atom_select_atom(st, atom)
[docs]def wrap_trajectory(cms_model, frames, fract_offset=xtal.FRACT_OFFSET):
"""
Wraps the trajectory within a unit cell box.
:param cms_model: The trajectory's cms model.
:type cms_model: schrodinger.application.desmond.cms.Cms
:param frames: The list of trajectory frames.
:type frames: list(schrodinger.application.desmond.packages.traj.Frame)
:param float fract_offset: The threshold used to compare floating point
fractional coordinate values and in particular those that are on
the cell boundary
:return: Returns a tuple of modified cms_model and the list of trajectory
frames.
:rtype: tuple(schrodinger.application.desmond.cms.Cms,
list(schrodinger.application.desmond.packages.traj.Frame))
"""
# Lazy load it, so people who don't have desmond installed, dont
# get exception.
from schrodinger.application.desmond.packages import topo
velocity = None
for fr in frames:
pos = fr.pos()
pos[:] = xtal.translate_to_1st_cell(pos,
fr.box,
fract_offset=fract_offset)
topo.update_cms(cms_model, frames[-1])
return cms_model, frames
[docs]def unroll_pos(ct, rep_vec, xyz0):
"""
Set coordinates for all copies of the replicated CT.
:type ct: `structure.Structure`
:type rep_vec: `tuple` of 3 `int`
:type xyz0: Nx3 `numpy.array` where N is the number of atoms in the
un-replicated CT.
"""
box = numpy.reshape(cms.get_box(ct), [3, 3])
xyz = ct.getXYZ(copy=False)
n_copies = numpy.prod(rep_vec)
shifts = numpy.empty((n_copies, 1, 3))
for i, (a, b, c) in enumerate(itertools.product(*map(range, rep_vec))):
shifts[i][0] = numpy.dot((a, b, c), box)
xyz.shape = (n_copies, len(xyz0), 3) # Reshape for Numpy broadcasting
numpy.copyto(xyz, shifts + xyz0)