Source code for schrodinger.application.jaguar.gui.tabs.sub_tab_widgets.charge_selector

from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.ui.qt import pop_up_widgets

from ... import utils as gui_utils


class ChargeSelectorLineEdit(pop_up_widgets.LineEditWithPopUp):
    """
    A line edit that can be used to select constraint's charge and
    individual atom weights. A weight selector pop up will appear
    whenever this line edit has focus.
    """

    def __init__(self, parent):
        super(ChargeSelectorLineEdit, self).__init__(parent,
                                                     _ChargeSelectorPopUp)
        validator = QtGui.QDoubleValidator(parent)
        self.setValidator(validator)

    def setWeights(self, weights):
        """
        Set charge constraint weights that will be shown in the popup.

        :param weights: charge constraint weights
        :type weights: dict
        """

        self._pop_up.populateWeights(weights)

    def getWeights(self):
        """
        This function returns constraint weights that were set in the popup.

        :return: constraint weights
        :rtype: dict
        """

        return self._pop_up.getWeights()


class _ChargeSelectorTableWidget(QtWidgets.QTableWidget):
    """
    The table used to specify weights for a charge constraint
    """

    def __init__(self, parent):
        super(_ChargeSelectorTableWidget, self).__init__(parent)
        self.setSizePolicy(QtWidgets.QSizePolicy.Preferred,
                           QtWidgets.QSizePolicy.Preferred)
        self.setSelectionMode(self.NoSelection)
        self.verticalHeader().hide()
        self.setFrameShape(self.NoFrame)

        self.setColumnCount(2)
        self.setRowCount(0)
        atom_item = QtWidgets.QTableWidgetItem(" Atom ")
        weight_item = QtWidgets.QTableWidgetItem(" Weight ")
        self.setHorizontalHeaderItem(0, atom_item)
        self.setHorizontalHeaderItem(1, weight_item)

    def sizeHint(self):
        """
        This table should be as small as possible.  We don't want space
        allocated for rows and columns that don't exist.
        """

        height = self.horizontalHeader().height()
        height += self.verticalHeader().length()
        width = self.horizontalHeader().length()
        return QtCore.QSize(width, height)


class _ChargeSelectorPopUp(pop_up_widgets.PopUp):
    """
    The pop up window that is displayed adjacent to L(ChargeSelectorLineEdit}
    """

    def setup(self):
        self.table_widget = _ChargeSelectorTableWidget(self)
        layout = QtWidgets.QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.addWidget(self.table_widget)

    def populateWeights(self, weights):
        """
        Populate atom weights table.

        :parame weights: atom weights
        :type weights: dict
        """

        self.table_widget.setRowCount(len(weights))
        sorted_weights = sorted(
            weights.items(), key=lambda x: gui_utils.atom_name_sort_key(x[0]))
        for i, (atom, weight) in enumerate(sorted_weights):
            atom_item = QtWidgets.QTableWidgetItem(atom)
            atom_item.setFlags(QtCore.Qt.ItemIsEnabled)
            weight_item = QtWidgets.QDoubleSpinBox(self.table_widget)
            weight_item.setFrame(False)
            weight_item.setButtonSymbols(weight_item.NoButtons)
            weight_item.setValue(weight)
            self.table_widget.setItem(i, 0, atom_item)
            self.table_widget.setCellWidget(i, 1, weight_item)
        self.table_widget.resizeColumnsToContents()
        self.table_widget.resizeRowsToContents()
        self.table_widget.updateGeometry()

    def getWeights(self):
        """
        This function returns dictionary of individual atom weights keyed
        on the atom names.

        :return: dictionary of atom weights
        :rtype: dict
        """

        weights = {}
        for row in range(self.table_widget.rowCount()):
            atom = self.table_widget.item(row, 0).text()
            cur_weight = self.table_widget.cellWidget(row, 1).value()
            weights[atom] = cur_weight
        return weights