from collections import OrderedDict
from schrodinger.Qt import QtCore
from schrodinger.Qt.QtCore import Qt
from schrodinger.ui.qt import delegates as qt_delegates
from schrodinger.ui.qt import pop_up_widgets
from ... import basis_selector
from ...input_tab_widgets import Basis
from . import base_widgets
[docs]class BasisSetColumns(object):
HEADERS = ["Atom", "ID", "Entry Title", "Basis Set"]
NUM_COLS = len(HEADERS)
ATOM, ID, TITLE, BASIS = list(range(NUM_COLS))
[docs]class PerAtomBasisRow(base_widgets.SubTabRow):
"""
Data about a per-atom basis set setting
"""
[docs] def __init__(self, entry_id, title, atom_name, atom_num, basis=None):
super(PerAtomBasisRow, self).__init__(entry_id, title)
self.atom_name = atom_name
self.atom_num = atom_num
if isinstance(basis, str):
basis = Basis.fromText(basis)
self.basis = basis
[docs] def copy(self):
"""
Create a new row object that is a copy of this row
:rtype: `PerAtomBasisRow`
:return: The row item that is a copy of this row
"""
eid = self.entry_id
title = self.title
name = self.atom_name
num = self.atom_num
basis = self.basis
return PerAtomBasisRow(eid, title, name, num, basis=basis)
[docs] def hasValidBasis(self):
"""
Does this row have a valid basis set?
:return: True if the specified row has a valid basis. False if the
basis is invalid or blank.
:rtype: bool
"""
if not self.basis:
return False
struc = self.getStructure()
basis = str(self.basis)
num_funcs, is_ps = basis_selector.num_basis_functions(
basis, struc, atom_num=self.atom_num)
return bool(num_funcs)
[docs]class BasisSetTableView(base_widgets.SubTabTableView):
COLUMN = BasisSetColumns
[docs] def __init__(self, parent=None):
super(BasisSetTableView, self).__init__(parent)
basis_delegate = AtomBasisSetDelegate(self)
self.setItemDelegateForColumn(self.COLUMN.BASIS, basis_delegate, True)
# TODO: label atoms with their basis set
[docs]class BasisSetModel(base_widgets.SubTabModel):
COLUMN = BasisSetColumns
UNEDITABLE = (COLUMN.ATOM, COLUMN.ID, COLUMN.TITLE)
ROW_CLASS = PerAtomBasisRow
MARKER_SETTINGS = {"color": "orange", "alt_color": "yellow"}
basisChanged = QtCore.pyqtSignal(str)
def _displayAndSortData(self, col, basis_row, role):
# See parent class for documentation on the arguments and return value
if col == self.COLUMN.BASIS:
if role == Qt.DisplayRole:
if basis_row.basis:
return str(basis_row.basis)
else:
return ""
elif role == base_widgets.SORT_ROLE:
return basis_row.basis
else:
return super(BasisSetModel,
self)._displayAndSortData(col, basis_row, role)
def _backgroundData(self, col, row_data):
"""
Mark any invalid basis sets as errors. Don't color cells where the user
hasn't yet entered a basis, since there's already a "Double-click to
edit" message in those cells.
See parent class for documentation on the arguments and return value
"""
if col == self.COLUMN.BASIS:
if row_data.basis and not row_data.hasValidBasis():
return self.ERROR_BACKGROUND_BRUSH
return super(BasisSetModel, self)._backgroundData(col, row_data)
def _setData(self, col, row_data, value, role, row_num):
# See parent class for documentation
if col == self.COLUMN.BASIS:
row_data.basis = Basis.fromText(value)
self.basisChanged.emit(row_data.entry_id)
return True
else:
return False
[docs] def removeRows(self, row, count, parent=None):
# See Qt documentation. Note that removeRow (no "s") calls this
# function
rows_to_delete = self._rows[row:row + count]
eids = {row.entry_id for row in rows_to_delete}
super(BasisSetModel, self).removeRows(row, count, parent)
list(map(self.basisChanged.emit, eids))
return True
[docs] def perAtomBasisForEid(self, eid, name=False):
"""
Return the per-atom basis sets for the specified entry ID
:param eid: The entry id
:type eid: str
:param name: If True, the return dictionary keys will be atom names. If
False, the keys will be atom numbers.
:type name: bool
:return: A dictionary of {atom number: basis set} or {atom name: basis
set}
:rtype: dict
"""
rows = self.rowsForEid(eid)
if name:
return {row.atom_name: str(row.basis) for row in rows if row.basis}
else:
return {row.atom_num: str(row.basis) for row in rows if row.basis}
[docs]class BasisSetProxyModel(base_widgets.SubTabProxyModel):
[docs] def checkBasisSets(self):
"""
Check to see if all rows have valid basis sets
:return: An OrderedDict of {entry title: list of atom names with invalid
basis sets}. The OrderedDict and atom lists are in the same order as
the table itself. If all rows have valid basis sets, then an empty
OrderedDict is returned.
:rtype: `OrderedDict`
"""
# use an OrderedDict so the errors appear in the same order as the table
bad_data = OrderedDict()
for row_num in range(self.rowCount()):
row_data = self._getRowData(row_num)
if not row_data.hasValidBasis():
cur_struc = bad_data.setdefault(row_data.title, [])
cur_struc.append(row_data.atom_name)
return bad_data
def _getRowData(self, row_num):
"""
Get the row data for the specified row
:param row_num: The row number (using proxy model row numbering, not
source model)
:type row_num: int
:return: The requested row data
:rtype: `PerAtomBasisRow`
"""
proxy_index = self.index(row_num, 1)
model_index = self.mapToSource(proxy_index)
model_row_num = model_index.row()
row_data = self.sourceModel()._rows[model_row_num]
return row_data
[docs]class AtomBasisSetDelegate(pop_up_widgets.PopUpDelegate,
qt_delegates.DefaultMessageDelegate):
"""
A delegate for selecting per-atom basis sets
"""
[docs] def __init__(self, parent):
super(AtomBasisSetDelegate, self).__init__(parent, None, True)
def _createEditor(self, parent, option, index):
# See parent class and Qt documentation
struc = index.data(base_widgets.STRUCTURE_ROLE)
atom_nums = index.data(base_widgets.ATOM_NUMS_ROLE)
atom_num = atom_nums[0]
editor = basis_selector.FilterBasisSelectorReadOnlyLineEdit(parent)
editor.setStructure(struc)
editor.setAtomNum(atom_num)
return editor