<commit_before>"""
Widget for multi-model constraints.
"""
import os

# numpy methods required for the validator! Don't remove.
# pylint: disable=unused-import,unused-wildcard-import,redefined-builtin
from numpy import *

from PyQt5 import QtCore
from PyQt5 import QtGui
from PyQt5 import QtWidgets
import webbrowser

from sas.qtgui.Perspectives.Fitting import FittingUtilities
import sas.qtgui.Utilities.GuiUtils as GuiUtils
from sas.qtgui.Perspectives.Fitting.Constraint import Constraint

#ALLOWED_OPERATORS = ['=','<','>','>=','<=']
ALLOWED_OPERATORS = ['=']

# Local UI
from sas.qtgui.Perspectives.Fitting.UI.ComplexConstraintUI import Ui_ComplexConstraintUI

class ComplexConstraint(QtWidgets.QDialog, Ui_ComplexConstraintUI):
    constraintReadySignal = QtCore.pyqtSignal(tuple)
    def __init__(self, parent=None, tabs=None):
        super(ComplexConstraint, self).__init__(parent)

        self.setupUi(self)
        self.setModal(True)

        # disable the context help icon
        windowFlags = self.windowFlags()
        self.setWindowFlags(windowFlags & ~QtCore.Qt.WindowContextHelpButtonHint)

        # Useful globals
        self.tabs = tabs
        self.params = None
        self.tab_names = None
        self.operator = '='
        self._constraint = Constraint()
        self.all_menu = None
        self.parent = parent
        self.redefining_warning = ""

        self.warning = self.lblWarning.text()
        self.setupData()
        self.setupSignals()
        self.setupWidgets()
        self.setupTooltip()

        # Default focus is on OK
        self.cmdOK.setFocus()

    def setupData(self):
        """
        Digs into self.tabs and pulls out relevant info
        """
        self.tab_names = [tab.kernel_module.name for tab in self.tabs]
        self.params = [tab.getParamNames() for tab in self.tabs]

    def setupSignals(self):
        """
        Signals from various elements
        """
        self.cmdOK.clicked.connect(self.onApply)
        self.cmdHelp.clicked.connect(self.onHelp)
        self.cmdAddAll.clicked.connect(self.onSetAll)

        self.txtConstraint.editingFinished.connect(self.validateFormula)
        self.cbModel1.currentIndexChanged.connect(self.onModelIndexChange)
        self.cbModel2.currentIndexChanged.connect(self.onModelIndexChange)

        self.cbParam1.currentIndexChanged.connect(self.onParamIndexChange)
        self.cbParam2.currentIndexChanged.connect(self.onParamIndexChange)
        self.cbOperator.currentIndexChanged.connect(self.onOperatorChange)

    def setupWidgets(self):
        """
        Setup widgets based on current parameters
        """
        self.cbModel1.insertItems(0, self.tab_names)
        # add an `All` option in the lhs if there are more than 3 tabs
        if len(self.tab_names) > 2:
            self.cbModel1.addItem("All")
        self.cbModel2.insertItems(0, self.tab_names)

        self.setupParamWidgets()

        self.setupMenu()

    def setupMenu(self):
        # Show Add All button, if necessary
        if self.cbModel1.currentText() ==self.cbModel2.currentText():
            self.cmdAddAll.setVisible(False)
        else:
            self.cmdAddAll.setVisible(True)
        return

    def setupParamWidgets(self):
        """
        Fill out comboboxes and set labels with non-constrained parameters
        """
        # Store previously select parameter
        previous_param1 = self.cbParam1.currentText()
        # Clear the combobox
        self.cbParam1.clear()
        # Populate the left combobox parameter arbitrarily with the parameters
        # from the first tab if `All` option is selected
        if self.cbModel1.currentText() == "All":
            items1 = self.tabs[1].main_params_to_fit + self.tabs[1].poly_params_to_fit
        else:
            tab_index1 = self.cbModel1.currentIndex()
            items1 = self.tabs[tab_index1].main_params_to_fit + self.tabs[tab_index1].poly_params_to_fit
        self.cbParam1.addItems(items1)
        # Show the previously selected parameter if available
        if previous_param1 in items1:
            index1 = self.cbParam1.findText(previous_param1)
            self.cbParam1.setCurrentIndex(index1)

        # Store previously select parameter
        previous_param2 = self.cbParam2.currentText()
        # M2 has to be non-constrained
        self.cbParam2.clear()
        tab_index2 = self.cbModel2.currentIndex()
        items2 = [param for param in self.params[tab_index2] if not self.tabs[tab_index2].paramHasConstraint(param)]

        self.cbParam2.addItems(items2)
        # Show the previously selected parameter if available
        if previous_param2 in items2:
            index2 = self.cbParam2.findText(previous_param2)
            self.cbParam2.setCurrentIndex(index2)

        self.txtParam.setText(self.cbModel1.currentText() + ":"
                              + self.cbParam1.currentText())

        self.cbOperator.clear()
        self.cbOperator.addItems(ALLOWED_OPERATORS)
        self.txtOperator.setText(self.cbOperator.currentText())

        self.txtConstraint.setText(self.tab_names[tab_index2]+"."+self.cbParam2.currentText())

        # disable Apply if no parameters available
        if len(items1)==0:
            self.cmdOK.setEnabled(False)
            self.cmdAddAll.setEnabled(False)
            txt = "No parameters in model "+self.tab_names[0] +\
                " are available for constraining.\n"+\
                "Please select at least one parameter for fitting when adding a constraint."
        else:
            txt = self.redefining_warning
            self.cmdOK.setEnabled(True)
            self.cmdAddAll.setEnabled(True)
        self.lblWarning.setText(txt)

        # disable Aplly all if `All` option on lhs has been selected
        if self.cbModel1.currentText() == "All":
            self.cmdAddAll.setEnabled(False)
        else:
            self.cmdAddAll.setEnabled(True)

    def setupTooltip(self):
        """
        Tooltip for txtConstraint
        """
        tooltip = "E.g. M1:scale = 2.0 * M2.scale\n"
        tooltip += "M1:scale = sqrt(M2.scale) + 5"
        self.txtConstraint.setToolTip(tooltip)

    def onParamIndexChange(self, index):
        """
        Respond to parameter combo box changes
        """
        # Find out the signal source
        source = self.sender().objectName()
        param1 = self.cbParam1.currentText()
        param2 = self.cbParam2.currentText()
        if source == "cbParam1":
            self.txtParam.setText(self.cbModel1.currentText() + ":" + param1)
        else:
            self.txtConstraint.setText(self.cbModel2.currentText() + "." + param2)
        # Check if any of the parameters are polydisperse
        params_list = [param1, param2]
        all_pars = [tab.model_parameters for tab in self.tabs]
        is2Ds = [tab.is2D for tab in self.tabs]
        txt = self.redefining_warning
        for pars, is2D in zip(all_pars, is2Ds):
            if any([FittingUtilities.isParamPolydisperse(p, pars, is2D) for p in params_list]):
                # no parameters are pd - reset the text to not show the warning
                txt = self.warning
        self.lblWarning.setText(txt)

    def onModelIndexChange(self, index):
        """
        Respond to mode combo box changes
        """
        # disable/enable Add All
        self.setupMenu()
        # Reload parameters
        self.setupParamWidgets()

    def onOperatorChange(self, index):
        """
        Respond to operator combo box changes
        """
        self.txtOperator.setText(self.cbOperator.currentText())

    def validateFormula(self):
        """
        Add visual cues when formula is incorrect
        """
        # temporarily disable validation
        return

        formula_is_valid = self.validateConstraint(self.txtConstraint.text())
        if not formula_is_valid:
            self.cmdOK.setEnabled(False)
            self.cmdAddAll.setEnabled(False)
            self.txtConstraint.setStyleSheet("QLineEdit {background-color: red;}")
        else:
            self.cmdOK.setEnabled(True)
            self.cmdAddAll.setEnabled(True)
            self.txtConstraint.setStyleSheet("QLineEdit {background-color: white;}")

    def validateConstraint(self, constraint_text):
        """
        Ensure the constraint has proper form
        """
        # 0. none or empty
        if not constraint_text or not isinstance(constraint_text, str):
            return False

        # M1.scale --> model_str='M1', constraint_text='scale'
        param_str = self.cbParam2.currentText()
        constraint_text = constraint_text.strip()
        model_str = self.cbModel2.currentText()

        # 0. Has to contain the model name
        if model_str != model_str:
            return False

        # Remove model name from constraint
        constraint_text = constraint_text.replace(model_str+".",'')

        # 1. just the parameter
        if param_str == constraint_text:
            return True

        # 2. ensure the text contains parameter name
        parameter_string_start = constraint_text.find(param_str)
        if parameter_string_start < 0:
            return False

        # 3. replace parameter name with "1" and try to evaluate the expression
        try:
            expression_to_evaluate = constraint_text.replace(param_str, "1.0")
            eval(expression_to_evaluate)
        except Exception:
            # Too many cases to cover individually, just a blanket
            # Exception should be sufficient
            # Note that in current numpy things like sqrt(-1) don't
            # raise but just return warnings
            return False

        return True

    def constraint(self):
        """
        Return the generated constraint
        """
        param = self.cbParam1.currentText()
        value = self.cbParam2.currentText()
        func = self.txtConstraint.text()
        value_ex = (self.cbModel1.currentText() + "."
                    + self.cbParam1.currentText())
        model1 = self.cbModel1.currentText()
        operator = self.cbOperator.currentText()

        con = Constraint(self,
                         param=param,
                         value=value,
                         func=func,
                         value_ex=value_ex,
                         operator=operator)

        return (model1, con)

    def constraintIsRedefined(self, cons_tuple):
        """
        Warn the user when a constraint is being redefined
        """
        # get the the parameter that is being redefined
        param = cons_tuple[1].param
        # get a list of all constrained parameters
        tab_index1 = self.cbModel1.currentIndex()
        items = [param for param in self.params[tab_index1] if self.tabs[tab_index1].paramHasConstraint(param)]
        # loop over the list of constrained parameters to check for redefinition
        for item in items:
            if item == param:
                return True
        return False

    def onApply(self):
        """
        Respond to Add constraint action.
        Send a signal that the constraint is ready to be applied
        """
        # if the combobox is set to `All` just call `applyAcrossTabs` and
        # return
        if self.cbModel1.currentText() == "All":
            # exclude the tab on the lhs
            tabs = [tab for tab in self.tabs if
                    tab.kernel_module.name != self.cbModel2.currentText()]
            self.applyAcrossTabs(tabs, self.cbParam1.currentText(),
                                 self.txtConstraint.text())
            self.setupParamWidgets()
            return

        cons_tuple = self.constraint()
        #check if constraint has been redefined
        if self.constraintIsRedefined(cons_tuple):
            txt = "Warning: parameter " + \
                  cons_tuple[0] + "." + cons_tuple[1].param +\
                  " has been redefined."
            self.redefining_warning = txt
        else:
            self.redefining_warning = ""
        self.constraintReadySignal.emit(cons_tuple)
        # reload the comboboxes
        if self.parent.constraint_accepted:
            self.setupParamWidgets()

    def applyAcrossTabs(self, tabs, param, expr):
        """
        Apply constraints across tabs, e.g. all `scale` parameters
        constrained to an expression. *tabs* is a list of active fit tabs
        for which the parameter string *param* will be constrained to the
        *expr* string.
        """
        for tab in tabs:
            if hasattr(tab, "kernel_module"):
                if param in tab.kernel_module.params:
                    value_ex = tab.kernel_module.name + "." +param
                    constraint = Constraint(param=param,
                                            value=param,
                                            func=expr,
                                            value_ex=value_ex,
                                            operator="=")
                    self.constraintReadySignal.emit((tab.kernel_module.name,
                                                     constraint))

    def onSetAll(self):
        """
        Set constraints on all identically named parameters between two fitpages
        """
        # loop over parameters in constrained model
        index1 = self.cbModel1.currentIndex()
        index2 = self.cbModel2.currentIndex()
        items1 = self.tabs[index1].kernel_module.params
        items2 = self.params[index2]
        # create an empty list to store redefined constraints
        redefined_constraints = []
        for item in items1:
            if item not in items2: continue
            param = item
            value = item
            func = self.cbModel2.currentText() + "." + param
            value_ex = self.cbModel1.currentText() + "." + param
            model1 = self.cbModel1.currentText()
            operator = self.cbOperator.currentText()

            con = Constraint(self,
                             param=param,
                             value=value,
                             func=func,
                             value_ex=value_ex,
                             operator=operator)

            # check for redefined constraints and add them to the list
            if self.constraintIsRedefined((model1, con)):
                redefined_constraints.append(model1 + "." + param)
            # warn the user if constraints have been redefined
            if redefined_constraints:
                constraint_txt = ""
                for redefined_constraint in redefined_constraints:
                    constraint_txt += redefined_constraint + ", "
                txt = "Warning: parameters " +\
                    constraint_txt[:-2] +\
                    " have been redefined."
                if len(redefined_constraints) == 1:
                    txt = txt.replace("parameters", "parameter")
                    txt = txt.replace("has", "been")
                self.redefining_warning = txt
            else:
                self.redefining_warning = ""
            self.constraintReadySignal.emit((model1, con))

        # reload the comboboxes
        self.setupParamWidgets()

    def onHelp(self):
        """
        Display related help section
        """
        tree_location = "/user/qtgui/Perspectives/Fitting/"
        helpfile = "fitting_help.html#simultaneous-fits-with-constraints"
        help_location = tree_location + helpfile
        self.parent.parent.parent.showHelp(help_location)
<commit_msg>Make All button to work also for PD constraints<commit_after>"""
Widget for multi-model constraints.
"""
import os

# numpy methods required for the validator! Don't remove.
# pylint: disable=unused-import,unused-wildcard-import,redefined-builtin
from numpy import *

from PyQt5 import QtCore
from PyQt5 import QtGui
from PyQt5 import QtWidgets
import webbrowser

from sas.qtgui.Perspectives.Fitting import FittingUtilities
import sas.qtgui.Utilities.GuiUtils as GuiUtils
from sas.qtgui.Perspectives.Fitting.Constraint import Constraint

#ALLOWED_OPERATORS = ['=','<','>','>=','<=']
ALLOWED_OPERATORS = ['=']

# Local UI
from sas.qtgui.Perspectives.Fitting.UI.ComplexConstraintUI import Ui_ComplexConstraintUI

class ComplexConstraint(QtWidgets.QDialog, Ui_ComplexConstraintUI):
    constraintReadySignal = QtCore.pyqtSignal(tuple)
    def __init__(self, parent=None, tabs=None):
        super(ComplexConstraint, self).__init__(parent)

        self.setupUi(self)
        self.setModal(True)

        # disable the context help icon
        windowFlags = self.windowFlags()
        self.setWindowFlags(windowFlags & ~QtCore.Qt.WindowContextHelpButtonHint)

        # Useful globals
        self.tabs = tabs
        self.params = None
        self.tab_names = None
        self.operator = '='
        self._constraint = Constraint()
        self.all_menu = None
        self.parent = parent
        self.redefining_warning = ""

        self.warning = self.lblWarning.text()
        self.setupData()
        self.setupSignals()
        self.setupWidgets()
        self.setupTooltip()

        # Default focus is on OK
        self.cmdOK.setFocus()

    def setupData(self):
        """
        Digs into self.tabs and pulls out relevant info
        """
        self.tab_names = [tab.kernel_module.name for tab in self.tabs]
        self.params = [tab.getParamNames() for tab in self.tabs]

    def setupSignals(self):
        """
        Signals from various elements
        """
        self.cmdOK.clicked.connect(self.onApply)
        self.cmdHelp.clicked.connect(self.onHelp)
        self.cmdAddAll.clicked.connect(self.onSetAll)

        self.txtConstraint.editingFinished.connect(self.validateFormula)
        self.cbModel1.currentIndexChanged.connect(self.onModelIndexChange)
        self.cbModel2.currentIndexChanged.connect(self.onModelIndexChange)

        self.cbParam1.currentIndexChanged.connect(self.onParamIndexChange)
        self.cbParam2.currentIndexChanged.connect(self.onParamIndexChange)
        self.cbOperator.currentIndexChanged.connect(self.onOperatorChange)

    def setupWidgets(self):
        """
        Setup widgets based on current parameters
        """
        self.cbModel1.insertItems(0, self.tab_names)
        # add an `All` option in the lhs if there are more than 3 tabs
        if len(self.tab_names) > 2:
            self.cbModel1.addItem("All")
        self.cbModel2.insertItems(0, self.tab_names)

        self.setupParamWidgets()

        self.setupMenu()

    def setupMenu(self):
        # Show Add All button, if necessary
        if self.cbModel1.currentText() ==self.cbModel2.currentText():
            self.cmdAddAll.setVisible(False)
        else:
            self.cmdAddAll.setVisible(True)
        return

    def setupParamWidgets(self):
        """
        Fill out comboboxes and set labels with non-constrained parameters
        """
        # Store previously select parameter
        previous_param1 = self.cbParam1.currentText()
        # Clear the combobox
        self.cbParam1.clear()
        # Populate the left combobox parameter arbitrarily with the parameters
        # from the first tab if `All` option is selected
        if self.cbModel1.currentText() == "All":
            items1 = self.tabs[1].main_params_to_fit + self.tabs[1].poly_params_to_fit
        else:
            tab_index1 = self.cbModel1.currentIndex()
            items1 = self.tabs[tab_index1].main_params_to_fit + self.tabs[tab_index1].poly_params_to_fit
        self.cbParam1.addItems(items1)
        # Show the previously selected parameter if available
        if previous_param1 in items1:
            index1 = self.cbParam1.findText(previous_param1)
            self.cbParam1.setCurrentIndex(index1)

        # Store previously select parameter
        previous_param2 = self.cbParam2.currentText()
        # M2 has to be non-constrained
        self.cbParam2.clear()
        tab_index2 = self.cbModel2.currentIndex()
        items2 = [param for param in self.params[tab_index2] if not self.tabs[tab_index2].paramHasConstraint(param)]

        self.cbParam2.addItems(items2)
        # Show the previously selected parameter if available
        if previous_param2 in items2:
            index2 = self.cbParam2.findText(previous_param2)
            self.cbParam2.setCurrentIndex(index2)

        self.txtParam.setText(self.cbModel1.currentText() + ":"
                              + self.cbParam1.currentText())

        self.cbOperator.clear()
        self.cbOperator.addItems(ALLOWED_OPERATORS)
        self.txtOperator.setText(self.cbOperator.currentText())

        self.txtConstraint.setText(self.tab_names[tab_index2]+"."+self.cbParam2.currentText())

        # disable Apply if no parameters available
        if len(items1)==0:
            self.cmdOK.setEnabled(False)
            self.cmdAddAll.setEnabled(False)
            txt = "No parameters in model "+self.tab_names[0] +\
                " are available for constraining.\n"+\
                "Please select at least one parameter for fitting when adding a constraint."
        else:
            txt = self.redefining_warning
            self.cmdOK.setEnabled(True)
            self.cmdAddAll.setEnabled(True)
        self.lblWarning.setText(txt)

        # disable Aplly all if `All` option on lhs has been selected
        if self.cbModel1.currentText() == "All":
            self.cmdAddAll.setEnabled(False)
        else:
            self.cmdAddAll.setEnabled(True)

    def setupTooltip(self):
        """
        Tooltip for txtConstraint
        """
        tooltip = "E.g. M1:scale = 2.0 * M2.scale\n"
        tooltip += "M1:scale = sqrt(M2.scale) + 5"
        self.txtConstraint.setToolTip(tooltip)

    def onParamIndexChange(self, index):
        """
        Respond to parameter combo box changes
        """
        # Find out the signal source
        source = self.sender().objectName()
        param1 = self.cbParam1.currentText()
        param2 = self.cbParam2.currentText()
        if source == "cbParam1":
            self.txtParam.setText(self.cbModel1.currentText() + ":" + param1)
        else:
            self.txtConstraint.setText(self.cbModel2.currentText() + "." + param2)
        # Check if any of the parameters are polydisperse
        params_list = [param1, param2]
        all_pars = [tab.model_parameters for tab in self.tabs]
        is2Ds = [tab.is2D for tab in self.tabs]
        txt = self.redefining_warning
        for pars, is2D in zip(all_pars, is2Ds):
            if any([FittingUtilities.isParamPolydisperse(p, pars, is2D) for p in params_list]):
                # no parameters are pd - reset the text to not show the warning
                txt = self.warning
        self.lblWarning.setText(txt)

    def onModelIndexChange(self, index):
        """
        Respond to mode combo box changes
        """
        # disable/enable Add All
        self.setupMenu()
        # Reload parameters
        self.setupParamWidgets()

    def onOperatorChange(self, index):
        """
        Respond to operator combo box changes
        """
        self.txtOperator.setText(self.cbOperator.currentText())

    def validateFormula(self):
        """
        Add visual cues when formula is incorrect
        """
        # temporarily disable validation, as not yet fully operational
        return

        formula_is_valid = self.validateConstraint(self.txtConstraint.text())
        if not formula_is_valid:
            self.cmdOK.setEnabled(False)
            self.cmdAddAll.setEnabled(False)
            self.txtConstraint.setStyleSheet("QLineEdit {background-color: red;}")
        else:
            self.cmdOK.setEnabled(True)
            self.cmdAddAll.setEnabled(True)
            self.txtConstraint.setStyleSheet("QLineEdit {background-color: white;}")

    def validateConstraint(self, constraint_text):
        """
        Ensure the constraint has proper form
        """
        # 0. none or empty
        if not constraint_text or not isinstance(constraint_text, str):
            return False

        # M1.scale --> model_str='M1', constraint_text='scale'
        param_str = self.cbParam2.currentText()
        constraint_text = constraint_text.strip()
        model_str = self.cbModel2.currentText()

        # 0. Has to contain the model name
        if model_str != model_str:
            return False

        # Remove model name from constraint
        constraint_text = constraint_text.replace(model_str+".",'')

        # 1. just the parameter
        if param_str == constraint_text:
            return True

        # 2. ensure the text contains parameter name
        parameter_string_start = constraint_text.find(param_str)
        if parameter_string_start < 0:
            return False

        # 3. replace parameter name with "1" and try to evaluate the expression
        try:
            expression_to_evaluate = constraint_text.replace(param_str, "1.0")
            eval(expression_to_evaluate)
        except Exception:
            # Too many cases to cover individually, just a blanket
            # Exception should be sufficient
            # Note that in current numpy things like sqrt(-1) don't
            # raise but just return warnings
            return False

        return True

    def constraint(self):
        """
        Return the generated constraint
        """
        param = self.cbParam1.currentText()
        value = self.cbParam2.currentText()
        func = self.txtConstraint.text()
        value_ex = (self.cbModel1.currentText() + "."
                    + self.cbParam1.currentText())
        model1 = self.cbModel1.currentText()
        operator = self.cbOperator.currentText()

        con = Constraint(self,
                         param=param,
                         value=value,
                         func=func,
                         value_ex=value_ex,
                         operator=operator)

        return (model1, con)

    def constraintIsRedefined(self, cons_tuple):
        """
        Warn the user when a constraint is being redefined
        """
        # get the the parameter that is being redefined
        param = cons_tuple[1].param
        # get a list of all constrained parameters
        tab_index1 = self.cbModel1.currentIndex()
        items = [param for param in self.params[tab_index1] if self.tabs[tab_index1].paramHasConstraint(param)]
        # loop over the list of constrained parameters to check for redefinition
        for item in items:
            if item == param:
                return True
        return False

    def onApply(self):
        """
        Respond to Add constraint action.
        Send a signal that the constraint is ready to be applied
        """
        # if the combobox is set to `All` just call `applyAcrossTabs` and
        # return
        if self.cbModel1.currentText() == "All":
            # exclude the tab on the lhs
            tabs = [tab for tab in self.tabs if
                    tab.kernel_module.name != self.cbModel2.currentText()]
            self.applyAcrossTabs(tabs, self.cbParam1.currentText(),
                                 self.txtConstraint.text())
            self.setupParamWidgets()
            return

        cons_tuple = self.constraint()
        #check if constraint has been redefined
        if self.constraintIsRedefined(cons_tuple):
            txt = "Warning: parameter " + \
                  cons_tuple[0] + "." + cons_tuple[1].param +\
                  " has been redefined."
            self.redefining_warning = txt
        else:
            self.redefining_warning = ""
        self.constraintReadySignal.emit(cons_tuple)
        # reload the comboboxes
        if self.parent.constraint_accepted:
            self.setupParamWidgets()

    def applyAcrossTabs(self, tabs, param, expr):
        """
        Apply constraints across tabs, e.g. all `scale` parameters
        constrained to an expression. *tabs* is a list of active fit tabs
        for which the parameter string *param* will be constrained to the
        *expr* string.
        """
        for tab in tabs:
            if hasattr(tab, "kernel_module"):
                if (param in tab.kernel_module.params or
                    param in tab.poly_params or
                    param in tab.magnet_params):
                    value_ex = tab.kernel_module.name + "." +param
                    constraint = Constraint(param=param,
                                            value=param,
                                            func=expr,
                                            value_ex=value_ex,
                                            operator="=")
                    self.constraintReadySignal.emit((tab.kernel_module.name,
                                                     constraint))

    def onSetAll(self):
        """
        Set constraints on all identically named parameters between two fitpages
        """
        # loop over parameters in constrained model
        index1 = self.cbModel1.currentIndex()
        index2 = self.cbModel2.currentIndex()
        items1 = self.tabs[index1].kernel_module.params
        items2 = self.params[index2]
        # create an empty list to store redefined constraints
        redefined_constraints = []
        for item in items1:
            if item not in items2: continue
            param = item
            value = item
            func = self.cbModel2.currentText() + "." + param
            value_ex = self.cbModel1.currentText() + "." + param
            model1 = self.cbModel1.currentText()
            operator = self.cbOperator.currentText()

            con = Constraint(self,
                             param=param,
                             value=value,
                             func=func,
                             value_ex=value_ex,
                             operator=operator)

            # check for redefined constraints and add them to the list
            if self.constraintIsRedefined((model1, con)):
                redefined_constraints.append(model1 + "." + param)
            # warn the user if constraints have been redefined
            if redefined_constraints:
                constraint_txt = ""
                for redefined_constraint in redefined_constraints:
                    constraint_txt += redefined_constraint + ", "
                txt = "Warning: parameters " +\
                    constraint_txt[:-2] +\
                    " have been redefined."
                if len(redefined_constraints) == 1:
                    txt = txt.replace("parameters", "parameter")
                    txt = txt.replace("has", "been")
                self.redefining_warning = txt
            else:
                self.redefining_warning = ""
            self.constraintReadySignal.emit((model1, con))

        # reload the comboboxes
        self.setupParamWidgets()

    def onHelp(self):
        """
        Display related help section
        """
        tree_location = "/user/qtgui/Perspectives/Fitting/"
        helpfile = "fitting_help.html#simultaneous-fits-with-constraints"
        help_location = tree_location + helpfile
        self.parent.parent.parent.showHelp(help_location)
<|endoftext|>


<commit_before>from jinja2 import Template
from collections import OrderedDict
import yaml

f = open("osmaxx_schema.yaml", 'r')
h = open("header.json", 'r')
inp = f.read()
header = h.read()
out = open("osmaxx.json", 'w')

data = yaml.load(inp)
table = data['layers'].keys()
out.write(header)
out.write("    \"tags\": [")
for i in range(0, len(table)):
    dicts = data['layers'][table[i]]
    attr = dicts['attributes'].keys()
    for j in range(0, len(attr)):
        if attr[j] == 'type':

            type_dict = dicts['attributes'][attr[j]]['values']
            type_dict = OrderedDict(sorted(type_dict.items()))
            type_attr = type_dict.keys()

            for k in range(0, len(type_attr)):
                type_tags = dicts['attributes'][attr[j]]['values'][type_attr[k]]['osm_tags']
                type_tags_description = dicts['attributes'][attr[j]]['values'][type_attr[k]]['description']
                type_len = len(type_tags)

                for l in range(0, type_len):
                    if type_tags[l] == None:
                        continue
                    else:

                        s = "{{ '\n{'|indent(16, true) }}"
                        s += "{{ '\n\"key\": '|indent(16, true)}}"
                        s += "\"{{ type_tags[m].keys()[0] }}\",\n"
                        s += "{{ '\"value\": '|indent(16, true) }}"
                        s += "\"{{ type_tags[m].values()[0] }}\",\n"
                        s += "{{ '\"description\": '|indent(16, true) }}"
                        s += "\"{{ type_tags_description }}\""
                        s += "{{ '\n},'|indent(16, true) }}"

                        t = Template(s)
                        out.write(t.render(type_tags=type_tags, type_tags_description=type_tags_description, m=l))

        if 'osm_tags' in dicts['attributes'][attr[j]].keys() or \
                        table[i] == 'road_ground_l' or table[i] == 'road_bridge_l' or table[i] == 'road_tunnel_l' \
                or table[i] == 'railway_ground_l' or table[i] == 'railway_bridge_l' or table[i] == 'railway_tunnel_l':
            if attr[j] == 'bridge' or attr[j] == 'tunnel':

                attr_tags_keys = dicts['attributes'][attr[j]]['values'][dicts['attributes'][attr[j]]['values'].keys()[0]] \
                    ['osm_tags'][0].keys()
                attr_tags_keys.sort()
                attr_tags_values = dicts['attributes'][attr[j]]['values'][dicts['attributes'][attr[j]]['values'].keys()[0]] \
                    ['osm_tags'][0].values()
                attr_tags_description = dicts['attributes'][attr[j]]['description']

                s = "{{ '\n{'|indent(16, true) }}"
                s += "{{ '\n\"key\": '|indent(16, true)}}"
                s += "\"{{ attr_tags_keys[0] }}\",\n"
                s += "{{ '\"value\": '|indent(16, true) }}"
                s += "\"{{ attr_tags_values[0] }}\",\n"
                s += "{{ '\"description\": '|indent(16, true) }}"
                s += "\"{{ attr_tags_description }}\""
                s += "{{ '\n},'|indent(16, true) }}"

                t = Template(s)
                out.write(t.render(attr_tags_keys=attr_tags_keys, attr_tags_description=attr_tags_description, \
                                   attr_tags_values=attr_tags_values))
            elif attr[j] == 'aggtype' or attr[j] == 'type':
                continue

            else:

                attr_tags_keys = dicts['attributes'][attr[j]]['osm_tags'][0].keys()
                attr_tags_keys.sort()
                attr_tags_description = dicts['attributes'][attr[j]]['description']

                s = "{{ '\n{'|indent(16, true) }}"
                s += "{{ '\n\"key\": '|indent(16, true)}}"
                s += "\"{{ attr_tags_keys[0] }}\",\n"
                s += "{{ '\"description\": '|indent(16, true) }}"
                s += "\"{{ attr_tags_description }}\""
                s += "{{ '\n},'|indent(16, true) }}"

                t = Template(s)
                out.write(t.render(attr_tags_keys=attr_tags_keys, attr_tags_description=attr_tags_description))

        elif table[i] == 'nonop_l':
            continue

        elif attr[j] == 'aggtype':
            continue

out.write('\n]}')
<commit_msg>end JSON file with EOL<commit_after>from jinja2 import Template
from collections import OrderedDict
import yaml

f = open("osmaxx_schema.yaml", 'r')
h = open("header.json", 'r')
inp = f.read()
header = h.read()
out = open("osmaxx.json", 'w')

data = yaml.load(inp)
table = data['layers'].keys()
out.write(header)
out.write("    \"tags\": [")
for i in range(0, len(table)):
    dicts = data['layers'][table[i]]
    attr = dicts['attributes'].keys()
    for j in range(0, len(attr)):
        if attr[j] == 'type':

            type_dict = dicts['attributes'][attr[j]]['values']
            type_dict = OrderedDict(sorted(type_dict.items()))
            type_attr = type_dict.keys()

            for k in range(0, len(type_attr)):
                type_tags = dicts['attributes'][attr[j]]['values'][type_attr[k]]['osm_tags']
                type_tags_description = dicts['attributes'][attr[j]]['values'][type_attr[k]]['description']
                type_len = len(type_tags)

                for l in range(0, type_len):
                    if type_tags[l] == None:
                        continue
                    else:

                        s = "{{ '\n{'|indent(16, true) }}"
                        s += "{{ '\n\"key\": '|indent(16, true)}}"
                        s += "\"{{ type_tags[m].keys()[0] }}\",\n"
                        s += "{{ '\"value\": '|indent(16, true) }}"
                        s += "\"{{ type_tags[m].values()[0] }}\",\n"
                        s += "{{ '\"description\": '|indent(16, true) }}"
                        s += "\"{{ type_tags_description }}\""
                        s += "{{ '\n},'|indent(16, true) }}"

                        t = Template(s)
                        out.write(t.render(type_tags=type_tags, type_tags_description=type_tags_description, m=l))

        if 'osm_tags' in dicts['attributes'][attr[j]].keys() or \
                        table[i] == 'road_ground_l' or table[i] == 'road_bridge_l' or table[i] == 'road_tunnel_l' \
                or table[i] == 'railway_ground_l' or table[i] == 'railway_bridge_l' or table[i] == 'railway_tunnel_l':
            if attr[j] == 'bridge' or attr[j] == 'tunnel':

                attr_tags_keys = dicts['attributes'][attr[j]]['values'][dicts['attributes'][attr[j]]['values'].keys()[0]] \
                    ['osm_tags'][0].keys()
                attr_tags_keys.sort()
                attr_tags_values = dicts['attributes'][attr[j]]['values'][dicts['attributes'][attr[j]]['values'].keys()[0]] \
                    ['osm_tags'][0].values()
                attr_tags_description = dicts['attributes'][attr[j]]['description']

                s = "{{ '\n{'|indent(16, true) }}"
                s += "{{ '\n\"key\": '|indent(16, true)}}"
                s += "\"{{ attr_tags_keys[0] }}\",\n"
                s += "{{ '\"value\": '|indent(16, true) }}"
                s += "\"{{ attr_tags_values[0] }}\",\n"
                s += "{{ '\"description\": '|indent(16, true) }}"
                s += "\"{{ attr_tags_description }}\""
                s += "{{ '\n},'|indent(16, true) }}"

                t = Template(s)
                out.write(t.render(attr_tags_keys=attr_tags_keys, attr_tags_description=attr_tags_description, \
                                   attr_tags_values=attr_tags_values))
            elif attr[j] == 'aggtype' or attr[j] == 'type':
                continue

            else:

                attr_tags_keys = dicts['attributes'][attr[j]]['osm_tags'][0].keys()
                attr_tags_keys.sort()
                attr_tags_description = dicts['attributes'][attr[j]]['description']

                s = "{{ '\n{'|indent(16, true) }}"
                s += "{{ '\n\"key\": '|indent(16, true)}}"
                s += "\"{{ attr_tags_keys[0] }}\",\n"
                s += "{{ '\"description\": '|indent(16, true) }}"
                s += "\"{{ attr_tags_description }}\""
                s += "{{ '\n},'|indent(16, true) }}"

                t = Template(s)
                out.write(t.render(attr_tags_keys=attr_tags_keys, attr_tags_description=attr_tags_description))

        elif table[i] == 'nonop_l':
            continue

        elif attr[j] == 'aggtype':
            continue

out.write('\n]}\n')
<|endoftext|>


<commit_before># -*- coding: utf-8  -*-
"""
Basic HTTP access interface.

This module handles communication between the bot and the HTTP threads.

This module is responsible for
    - Setting up a connection pool
    - Providing a (blocking) interface for HTTP requests
    - Translate site objects with query strings into urls
    - Urlencoding all data
    - Basic HTTP error handling
"""

#
# (C) Pywikipedia bot team, 2007
#
# Distributed under the terms of the MIT license.
#

__version__ = '$Id: $'
__docformat__ = 'epytext'

import Queue
import urllib
import urlparse
import logging
import atexit

import config2
import cookielib
import threadedhttp

logger = logging.getLogger("comm")

# global variables

useragent = 'Pywikipediabot/2.0' # This should include some global version string
numthreads = 1
threads = []

connection_pool = threadedhttp.ConnectionPool()
http_queue = Queue.Queue()

cookie_jar = threadedhttp.LockableCookieJar(
                 config2.datafilepath("%s-%s-%s.lwp"
                        % (config2.family,
                           config2.mylang,
                           config2.usernames[config2.family][config2.mylang])))
try:
    cookie_jar.load()
except (IOError, cookielib.LoadError):
    logger.debug("Loading cookies failed.")
else:
    logger.debug("Loaded cookies from file.")

# Build up HttpProcessors
logger.info('Starting %(numthreads)i threads...' % locals())
for i in range(numthreads):
    proc = threadedhttp.HttpProcessor(http_queue, cookie_jar, connection_pool)
    proc.setDaemon(True)
    threads.append(proc)
    proc.start()

# Prepare flush on quit
def _flush():
    for i in threads:
        http_queue.put(None)
    logger.info('Waiting for threads to finish... ')
    for i in threads:
        i.join()
    logger.debug('All threads finished.')
atexit.register(_flush)

# export cookie_jar to global namespace
import pywikibot
pywikibot.cookie_jar = cookie_jar

def request(site, uri, *args, **kwargs):
    """Queue a request to be submitted to Site.

    All parameters not listed below are the same as
    L{httplib2.Http.request}, but the uri is relative

    @param site: The Site to connect to
    @return: The received data (a unicode string).
    """
    baseuri = "%s://%s/" % (site.protocol(), site.hostname())
    uri = urlparse.urljoin(baseuri, uri)

    # set default user-agent string
    kwargs.setdefault("headers", {})
    kwargs["headers"].setdefault("user-agent", useragent)
    request = threadedhttp.HttpRequest(uri, *args, **kwargs)
    http_queue.put(request)
    request.lock.acquire()

    #TODO: do some error correcting stuff

    #if all else fails
    if isinstance(request.data, Exception):
        raise request.data

    if request.data[0].status != 200:
        logger.warning("Http response status %(status)s"
                       % {'status': request.data[0].status})

    return request.data[1]
<commit_msg>fix import<commit_after># -*- coding: utf-8  -*-
"""
Basic HTTP access interface.

This module handles communication between the bot and the HTTP threads.

This module is responsible for
    - Setting up a connection pool
    - Providing a (blocking) interface for HTTP requests
    - Translate site objects with query strings into urls
    - Urlencoding all data
    - Basic HTTP error handling
"""

#
# (C) Pywikipedia bot team, 2007
#
# Distributed under the terms of the MIT license.
#

__version__ = '$Id: $'
__docformat__ = 'epytext'

import Queue
import urllib
import urlparse
import logging
import atexit

from pywikibot import config2
import cookielib
import threadedhttp

logger = logging.getLogger("comm")

# global variables

useragent = 'Pywikipediabot/2.0' # This should include some global version string
numthreads = 1
threads = []

connection_pool = threadedhttp.ConnectionPool()
http_queue = Queue.Queue()

cookie_jar = threadedhttp.LockableCookieJar(
                 config2.datafilepath("%s-%s-%s.lwp"
                        % (config2.family,
                           config2.mylang,
                           config2.usernames[config2.family][config2.mylang])))
try:
    cookie_jar.load()
except (IOError, cookielib.LoadError):
    logger.debug("Loading cookies failed.")
else:
    logger.debug("Loaded cookies from file.")

# Build up HttpProcessors
logger.info('Starting %(numthreads)i threads...' % locals())
for i in range(numthreads):
    proc = threadedhttp.HttpProcessor(http_queue, cookie_jar, connection_pool)
    proc.setDaemon(True)
    threads.append(proc)
    proc.start()

# Prepare flush on quit
def _flush():
    for i in threads:
        http_queue.put(None)
    logger.info('Waiting for threads to finish... ')
    for i in threads:
        i.join()
    logger.debug('All threads finished.')
atexit.register(_flush)

# export cookie_jar to global namespace
import pywikibot
pywikibot.cookie_jar = cookie_jar

def request(site, uri, *args, **kwargs):
    """Queue a request to be submitted to Site.

    All parameters not listed below are the same as
    L{httplib2.Http.request}, but the uri is relative

    @param site: The Site to connect to
    @return: The received data (a unicode string).
    """
    baseuri = "%s://%s/" % (site.protocol(), site.hostname())
    uri = urlparse.urljoin(baseuri, uri)

    # set default user-agent string
    kwargs.setdefault("headers", {})
    kwargs["headers"].setdefault("user-agent", useragent)
    request = threadedhttp.HttpRequest(uri, *args, **kwargs)
    http_queue.put(request)
    request.lock.acquire()

    #TODO: do some error correcting stuff

    #if all else fails
    if isinstance(request.data, Exception):
        raise request.data

    if request.data[0].status != 200:
        logger.warning("Http response status %(status)s"
                       % {'status': request.data[0].status})

    return request.data[1]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
Created by 'bens3' on 2013-06-21.
Copyright (c) 2013 'bens3'. All rights reserved.
"""

from pylons import config
from paste.deploy.converters import asbool

TEST_PREFIX = '10.5072'

ENDPOINT = 'https://test.datacite.org/mds'
TEST_ENDPOINT = 'https://test.datacite.org/mds'

def get_test_mode():
    """
    Get test mode as boolean
    @return:
    """

    return asbool(config.get("ckanext.doi.test_mode"))

def get_prefix():
    """
    Get the prefix to use for DOIs
    @return: test prefix if we're in test mode, otherwise config prefix setting
    """

    return TEST_PREFIX if get_test_mode() else config.get("ckanext.doi.prefix")

def get_endpoint():

    """
    Get the DataCite endpoint
    @return: test endpoint if we're in test mode
    """
    return TEST_ENDPOINT if get_test_mode() else ENDPOINT

def get_site_url():
    """
    Get the site URL
    Try and use ckanext.doi.site_url but if that's not set use ckan.site_url
    @return:
    """
    site_url = config.get("ckanext.doi.site_url")

    if not site_url:
        site_url = config.get('ckan.site_url')

    return site_url.rstrip('/')<commit_msg>Updating endpoint<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
Created by 'bens3' on 2013-06-21.
Copyright (c) 2013 'bens3'. All rights reserved.
"""

from pylons import config
from paste.deploy.converters import asbool

TEST_PREFIX = '10.5072'

ENDPOINT = 'https://mds.datacite.org'
TEST_ENDPOINT = 'https://test.datacite.org/mds'

def get_test_mode():
    """
    Get test mode as boolean
    @return:
    """

    return asbool(config.get("ckanext.doi.test_mode"))

def get_prefix():
    """
    Get the prefix to use for DOIs
    @return: test prefix if we're in test mode, otherwise config prefix setting
    """

    return TEST_PREFIX if get_test_mode() else config.get("ckanext.doi.prefix")

def get_endpoint():

    """
    Get the DataCite endpoint
    @return: test endpoint if we're in test mode
    """
    return TEST_ENDPOINT if get_test_mode() else ENDPOINT

def get_site_url():
    """
    Get the site URL
    Try and use ckanext.doi.site_url but if that's not set use ckan.site_url
    @return:
    """
    site_url = config.get("ckanext.doi.site_url")

    if not site_url:
        site_url = config.get('ckan.site_url')

    return site_url.rstrip('/')<|endoftext|>


<commit_before>"""Class to gather and analse various metal line statistics"""

import numpy as np
import hdfsim
import convert_cloudy
import line_data
import spectra

class MetalLines:
    """Generate metal line spectra from simulation snapshot"""
    def __init__(self,num, base, los_table, cloudy_dir="/home/spb/codes/ArepoCoolingTables/tmp_spb/", nbins = 1024):
        f = hdfsim.get_file(num, base, 0)
        self.box = f["Header"].attrs["BoxSize"]
        self.hubble = f["Header"].attrs["HubbleParam"]
        self.atime = f["Header"].attrs["Time"]
        self.redshift = f["Header"].attrs["Redshift"]
        Omega0 = f["Header"].attrs["Omega0"]
        OmegaLambda = f["Header"].attrs["OmegaLambda"]
        f.close()
        self.Hz = 100.0*self.hubble * np.sqrt(Omega0/self.atime**3 + OmegaLambda)
        self.nbins = nbins
        self.xbins = np.arange(0,self.nbins)*self.box/self.nbins
        self.species = ['He', 'C', 'N', 'O', 'Ne', 'Mg', 'Si', 'Fe']
        self.NumLos = np.size(los_table.axis)
        #Line data
        self.lines = line_data.LineData()
        #generate metal and hydrogen spectral densities
        #Indexing is: rho_metals [ NSPECTRA, NBIN ]
        (self.rho_H, self.metals) = spectra.SPH_Interpolate_metals(num, base, los_table, nbins)
        #rescale H density
        self.rho_H = spectra.rescale_units_rho_H(self.rho_H, self.hubble, self.atime)
        #Rescale metals
        for (key, value) in self.metals.iteritems():
            mass = self.lines.get_mass(key)
            value.rescale_units(self.hubble, self.atime, mass)

        #Generate cloudy tables
        self.cloudy = convert_cloudy.CloudyTable(cloudy_dir)

    def get_lines(self, elem, ion):
        """Get the optical depth for a particular element out of:
           (He, C, N, O, Ne, Mg, Si, Fe)
           and some ion number
           NOTE: May wish to special-case SiIII at some point
        """
        species = self.metals[elem]
        line = self.lines.get_line(elem,ion)
        mass = self.lines.get_mass(elem)
        #To convert H density from kg/m^3 to atoms/cm^3
        conv = 1./(spectra.PROTONMASS*100**3)
        ion_density = np.array(species.rho)
        #Compute tau for this metal ion
        tau_metal=np.empty(np.shape(species.rho))
        for n in np.arange(0,self.NumLos):
            #For the density parameter use the hydrogen density at this pixel
            #For metallicity pass the metallicity of this species at this bin (rho_Z/ rho_H) and it will be converted to cloudy format
            ind = np.where(ion_density[n,:] > 0)
            for i in np.array(ind).ravel():
                ion_density[n,i] *= self.cloudy.ion(elem, ion, self.redshift, ion_density[n,i]/self.rho_H[n,i], self.rho_H[n,i]/conv)
            tau_metal[n] = spectra.compute_absorption(self.xbins, ion_density[n], species.vel[n], species.temp[n],line,self.Hz,self.hubble, self.box, self.atime,mass)
        return tau_metal

    def vel_width(self, tau):
        """Find the velocity width of a line"""
        #  Size of a single velocity bin
        dvbin = self.box / (1.*self.nbins) * self.Hz *self.atime /self.hubble / 1000 # velocity bin size (kms^-1)

        tot_tau = np.sum(tau)
        cent = np.where(tau == np.max(tau))
        tot = tau[cent]
        if tau[cent] > 0.9 *tot_tau:
            return dvbin
        i = 0
        #Extend the region of interest until we have enough tau
        while tot < 0.9*tot_tau:
            i+=1
            tot += tau[cent-i] + tau[cent+i]
        #Return the width
        return dvbin*i*2

<commit_msg>Better way of computing velocity widths<commit_after>"""Class to gather and analse various metal line statistics"""

import numpy as np
import hdfsim
import convert_cloudy
import line_data
import spectra

class MetalLines:
    """Generate metal line spectra from simulation snapshot"""
    def __init__(self,num, base, los_table, cloudy_dir="/home/spb/codes/ArepoCoolingTables/tmp_spb/", nbins = 1024):
        f = hdfsim.get_file(num, base, 0)
        self.box = f["Header"].attrs["BoxSize"]
        self.hubble = f["Header"].attrs["HubbleParam"]
        self.atime = f["Header"].attrs["Time"]
        self.redshift = f["Header"].attrs["Redshift"]
        Omega0 = f["Header"].attrs["Omega0"]
        OmegaLambda = f["Header"].attrs["OmegaLambda"]
        f.close()
        self.Hz = 100.0*self.hubble * np.sqrt(Omega0/self.atime**3 + OmegaLambda)
        self.nbins = nbins
        self.xbins = np.arange(0,self.nbins)*self.box/self.nbins
        self.species = ['He', 'C', 'N', 'O', 'Ne', 'Mg', 'Si', 'Fe']
        self.NumLos = np.size(los_table.axis)
        #Line data
        self.lines = line_data.LineData()
        #generate metal and hydrogen spectral densities
        #Indexing is: rho_metals [ NSPECTRA, NBIN ]
        (self.rho_H, self.metals) = spectra.SPH_Interpolate_metals(num, base, los_table, nbins)
        #rescale H density
        self.rho_H = spectra.rescale_units_rho_H(self.rho_H, self.hubble, self.atime)
        #Rescale metals
        for (key, value) in self.metals.iteritems():
            mass = self.lines.get_mass(key)
            value.rescale_units(self.hubble, self.atime, mass)

        #Generate cloudy tables
        self.cloudy = convert_cloudy.CloudyTable(cloudy_dir)

    def get_lines(self, elem, ion):
        """Get the optical depth for a particular element out of:
           (He, C, N, O, Ne, Mg, Si, Fe)
           and some ion number
           NOTE: May wish to special-case SiIII at some point
        """
        species = self.metals[elem]
        line = self.lines.get_line(elem,ion)
        mass = self.lines.get_mass(elem)
        #To convert H density from kg/m^3 to atoms/cm^3
        conv = 1./(spectra.PROTONMASS*100**3)
        ion_density = np.array(species.rho)
        #Compute tau for this metal ion
        tau_metal=np.empty(np.shape(species.rho))
        for n in np.arange(0,self.NumLos):
            #For the density parameter use the hydrogen density at this pixel
            #For metallicity pass the metallicity of this species at this bin (rho_Z/ rho_H) and it will be converted to cloudy format
            ind = np.where(ion_density[n,:] > 0)
            for i in np.array(ind).ravel():
                ion_density[n,i] *= self.cloudy.ion(elem, ion, self.redshift, ion_density[n,i]/self.rho_H[n,i], self.rho_H[n,i]/conv)
            tau_metal[n] = spectra.compute_absorption(self.xbins, ion_density[n], species.vel[n], species.temp[n],line,self.Hz,self.hubble, self.box, self.atime,mass)
        return tau_metal

    def vel_width(self, tau):
        """Find the velocity width of a line"""
        #  Size of a single velocity bin
        dvbin = self.box / (1.*self.nbins) * self.Hz *self.atime /self.hubble / 1000 # velocity bin size (kms^-1)

        tot_tau = np.sum(tau)
        cum_tau = np.cumsum(tau)
        ind_low = np.where(cum_tau > 0.05 * tot_tau)
        low = np.ravel(ind_low)[0]
        ind_high = np.where(cum_tau > 0.95 * tot_tau)
        high = np.ravel(ind_high)[0]
        #Return the width
        return dvbin*(high - low)

<|endoftext|>


<commit_before>import logging

from PyQt5.QtWidgets import QDialog, QPushButton, QWidget
from PyQt5.QtCore import Qt
from typing import Optional, Callable, Dict, Any

from sas.system import config
from sas.qtgui.Utilities.Preferences.UI.PreferencesUI import Ui_preferencesUI
from sas.qtgui.Utilities.Preferences.PreferencesWidget import PreferencesWidget

# The PreferencesPanel object will instantiate all widgets during its instantiation.
#  e.g:
#  `from foo.bar import BarWidget  # BarWidget is a child of PreferencesWidget`
#  `BASE_PANELS = {"Bar Widget Options": BarWidget}`
# PreferenceWidget Imports go here and then are added to the BASE_PANELS, but not instantiated.
from .DisplayPreferencesWidget import DisplayPreferencesWidget
from .PlottingPreferencesWidget import PlottingPreferencesWidget
# Pre-made option widgets

BASE_PANELS = {"Plotting Settings": PlottingPreferencesWidget,
               "Display Settings": DisplayPreferencesWidget,
               }  # Type: Dict[str, Union[Type[PreferencesWidget], Callable[[],QWidget]]

logger = logging.getLogger(__name__)

class PreferencesPanel(QDialog, Ui_preferencesUI):
    """A preferences panel to house all SasView related settings. The left side of the window is a listWidget with a
    options menus available. The right side of the window is a stackedWidget object that houses the options
    associated with each listWidget item.
    """

    def __init__(self, parent: Optional[Any] = None):
        super(PreferencesPanel, self).__init__(parent)
        self.setupUi(self)
        self._staged_changes = {}
        self.parent = parent
        self.setWindowTitle("Preferences")
        # Add predefined widgets to window
        self.addWidgets(BASE_PANELS)
        # Set defaults values for the list and stacked widgets
        self.stackedWidget.setCurrentIndex(0)
        self.listWidget.setCurrentRow(0)
        # Add window actions
        self.listWidget.currentItemChanged.connect(self.prefMenuChanged)
        self.buttonBox.clicked.connect(self.onClick)

    def addWidgets(self, widgets: Dict[str, Callable]):
        """Add a list of named widgets to the window"""
        for name, widget in widgets.items():
            if isinstance(widget, PreferencesWidget):
                self.addWidget(widget, name)
            elif callable(widget):
                self.addWidget(widget(), name)

    def prefMenuChanged(self):
        """When the preferences menu selection changes, change to the appropriate preferences widget """
        row = self.listWidget.currentRow()
        self.setWidgetIndex(self.listWidget.currentRow())

    def setMenuByName(self, name: str):
        """Set the index of the listWidget and stackedWidget, using the display name as the search term"""
        for item in self.listWidget.findItems(name, Qt.MatchContains):
            if item.text() == name:
                self.listWidget.setCurrentItem(item)
        self.setWidgetIndex(self.listWidget.currentRow())

    def setWidgetIndex(self, row: int):
        """Set the menu and options stack to a given index"""
        self.listWidget.setCurrentRow(row)
        self.stackedWidget.setCurrentIndex(row)

    def onClick(self, btn: QPushButton):
        """Handle button click events in one area"""
        # Reset to the default preferences
        if btn.text() == 'Restore Defaults':
            self.restoreDefaultPreferences()
        elif btn.text() == 'OK':
            self._okClicked()
        elif btn.text() == 'Help':
            self.help()

    def restoreDefaultPreferences(self):
        """Reset preferences for the active widget to the default values."""
        widget = self.stackedWidget.currentWidget()
        if hasattr(widget, 'restoreDefaults') and callable(widget.restoreDefaults):
            widget.restoreDefaults()

    def stageSingleChange(self, key, value):
        """ Preferences widgets should call this method when changing a variable to prevent direct configuration
        changes"""
        self._staged_changes[key] = value

    def _okClicked(self):
        """ Action triggered when the OK button is clicked"""
        self._saveStagedChanges()
        self.close()

    def _saveStagedChanges(self):
        """ When OK or Apply are clicked, all staged changes should be applied to the config. """
        for k, v in self._staged_changes.items():
            setattr(config, k, v)
        self._staged_changes = {}

    def _cancelStaging(self):
        """ When the Cancel button is clicked, throw away any staged changes and hide the window"""
        self._staged_changes = {}
        self.close()

    def close(self):
        """Save the configuration values when the preferences window is closed"""
        super(PreferencesPanel, self).close()

    def addWidget(self, widget: QWidget, name: Optional[str] = None):
        """Add a single widget to the panel"""
        # Set the parent of the new widget to the parent of this window
        widget.parent = self
        self.stackedWidget.addWidget(widget)
        # Set display name in the listWidget with the priority of
        #  name passed to method > widget.name > "Generic Preferences"
        name = name if name is not None else getattr(widget, 'name', None)
        name = name if name is not None else "Generic Preferences"
        # Add the widget default reset method to the global set
        self.listWidget.addItem(name)

    def help(self):
        """Open the help window associated with the preferences window"""
        tree_location = "/user/qtgui/MainWindow/preferences_help.html"
        self.parent.guiManager.showHelp(tree_location)
<commit_msg>Separate button box signals in preferences panel for better maintainability<commit_after>import logging

from PyQt5.QtWidgets import QDialog, QPushButton, QWidget, QDialogButtonBox
from PyQt5.QtCore import Qt
from typing import Optional, Callable, Dict, Any

from sas.system import config
from sas.qtgui.Utilities.Preferences.UI.PreferencesUI import Ui_preferencesUI
from sas.qtgui.Utilities.Preferences.PreferencesWidget import PreferencesWidget

# The PreferencesPanel object will instantiate all widgets during its instantiation.
#  e.g:
#  `from foo.bar import BarWidget  # BarWidget is a child of PreferencesWidget`
#  `BASE_PANELS = {"Bar Widget Options": BarWidget}`
# PreferenceWidget Imports go here and then are added to the BASE_PANELS, but not instantiated.
from .DisplayPreferencesWidget import DisplayPreferencesWidget
from .PlottingPreferencesWidget import PlottingPreferencesWidget
# Pre-made option widgets

BASE_PANELS = {"Plotting Settings": PlottingPreferencesWidget,
               "Display Settings": DisplayPreferencesWidget,
               }  # Type: Dict[str, Union[Type[PreferencesWidget], Callable[[],QWidget]]

logger = logging.getLogger(__name__)

class PreferencesPanel(QDialog, Ui_preferencesUI):
    """A preferences panel to house all SasView related settings. The left side of the window is a listWidget with a
    options menus available. The right side of the window is a stackedWidget object that houses the options
    associated with each listWidget item.
    """

    def __init__(self, parent: Optional[Any] = None):
        super(PreferencesPanel, self).__init__(parent)
        self.setupUi(self)
        self._staged_changes = {}
        self.parent = parent
        self.setWindowTitle("Preferences")
        # Add predefined widgets to window
        self.addWidgets(BASE_PANELS)
        # Set defaults values for the list and stacked widgets
        self.stackedWidget.setCurrentIndex(0)
        self.listWidget.setCurrentRow(0)
        # Add window actions
        self.listWidget.currentItemChanged.connect(self.prefMenuChanged)
        self.buttonBox.button(QDialogButtonBox.RestoreDefaults).connect(self.restoreDefaultPreferences)
        self.buttonBox.button(QDialogButtonBox.Ok).connect(self._okClicked())
        self.buttonBox.button(QDialogButtonBox.Cancel).connect(self._cancelStaging)
        self.buttonBox.button(QDialogButtonBox.Apply).connect(self._saveStagedChanges)
        self.buttonBox.button(QDialogButtonBox.Help).connect(self.help)

    def addWidgets(self, widgets: Dict[str, Callable]):
        """Add a list of named widgets to the window"""
        for name, widget in widgets.items():
            if isinstance(widget, PreferencesWidget):
                self.addWidget(widget, name)
            elif callable(widget):
                self.addWidget(widget(), name)

    def prefMenuChanged(self):
        """When the preferences menu selection changes, change to the appropriate preferences widget """
        row = self.listWidget.currentRow()
        self.setWidgetIndex(self.listWidget.currentRow())

    def setMenuByName(self, name: str):
        """Set the index of the listWidget and stackedWidget, using the display name as the search term"""
        for item in self.listWidget.findItems(name, Qt.MatchContains):
            if item.text() == name:
                self.listWidget.setCurrentItem(item)
        self.setWidgetIndex(self.listWidget.currentRow())

    def setWidgetIndex(self, row: int):
        """Set the menu and options stack to a given index"""
        self.listWidget.setCurrentRow(row)
        self.stackedWidget.setCurrentIndex(row)

    def restoreDefaultPreferences(self):
        """Reset preferences for the active widget to the default values."""
        widget = self.stackedWidget.currentWidget()
        if hasattr(widget, 'restoreDefaults') and callable(widget.restoreDefaults):
            widget.restoreDefaults()

    def stageSingleChange(self, key, value):
        """ Preferences widgets should call this method when changing a variable to prevent direct configuration
        changes"""
        self._staged_changes[key] = value

    def _okClicked(self):
        """ Action triggered when the OK button is clicked"""
        self._saveStagedChanges()
        self.close()

    def _saveStagedChanges(self):
        """ When OK or Apply are clicked, all staged changes should be applied to the config. """
        for k, v in self._staged_changes.items():
            setattr(config, k, v)
        self._staged_changes = {}

    def _cancelStaging(self):
        """ When the Cancel button is clicked, throw away any staged changes and hide the window"""
        self._staged_changes = {}
        self.close()

    def close(self):
        """Save the configuration values when the preferences window is closed"""
        super(PreferencesPanel, self).close()

    def addWidget(self, widget: QWidget, name: Optional[str] = None):
        """Add a single widget to the panel"""
        # Set the parent of the new widget to the parent of this window
        widget.parent = self
        self.stackedWidget.addWidget(widget)
        # Set display name in the listWidget with the priority of
        #  name passed to method > widget.name > "Generic Preferences"
        name = name if name is not None else getattr(widget, 'name', None)
        name = name if name is not None else "Generic Preferences"
        # Add the widget default reset method to the global set
        self.listWidget.addItem(name)

    def help(self):
        """Open the help window associated with the preferences window"""
        tree_location = "/user/qtgui/MainWindow/preferences_help.html"
        self.parent.guiManager.showHelp(tree_location)
<|endoftext|>


<commit_before>##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2014 Stanford University and the Authors
#
# Authors: Kyle A. Beauchamp
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
#############################################################################

##############################################################################
# Imports
##############################################################################

import os
from distutils.spawn import find_executable as _find_executable
import numpy as np

from mdtraj.utils import import_, six
from mdtraj.utils import enter_temp_directory

##############################################################################
# Globals
##############################################################################

# Possible names for the external commands -- these are expected
# to be found in the PATH.
SHIFTX2 = ['shiftx2.py']
SPARTA_PLUS = ['sparta+', 'SPARTA+', 'SPARTA+.linux']
PPM = ['ppm_linux_64.exe']

__all__ = ['chemical_shifts_shiftx2', 'chemical_shifts_ppm', 'chemical_shifts_spartaplus', "reindex_dataframe_by_atoms"]

def find_executable(names):
    for possible in names:
        result = _find_executable(possible)
        if result is not None:
            return result
    return None

##############################################################################
# Code
##############################################################################

def compute_chemical_shifts(trj, model="shiftx2", **kwargs):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    model : str, optional, default="shiftx2"
        The program to use for calculating chemical shifts.  Must be one
        of shiftx2, ppm, or sparta+

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have the appropriate chemical soft programs installed
    and in your executable path.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference, see docstrings for chemical_shifts_*
    for the various possible models.
    """
    if model == "shiftx2":
        return chemical_shifts_shiftx2(trj, **kwargs)
    elif model == "ppm":
        return chemical_shifts_ppm(trj, **kwargs)
    elif model == "sparta+":
        return chemical_shifts_spartaplus(trj, **kwargs)
    else:
        raise(ValueError("model must be one of shiftx2, ppm, or sparta+"))

def chemical_shifts_shiftx2(trj):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ShiftX2 available on your path; see (http://www.shiftx2.ca/).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Beomsoo Han, Yifeng Liu, Simon Ginzinger, and David Wishart.
       "SHIFTX2: significantly improved protein chemical shift
       prediction." J. Biomol. NMR, 50, 1 43-57 (2011)
    """
    pd = import_('pandas')
    binary = find_executable(SHIFTX2)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_shiftx2` requires the external program SHIFTX2, available at http://www.shiftx2.ca/' % ', '.join(SHIFTX2))

    if six.PY3:
        raise(OSError("Cannot run shiftx2.py on Python 3.X"))

    results = []
    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)
        cmd = "%s -b 'trj*.pdb'" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your ShiftX2 installation or your input trajectory." % cmd))

        for i in range(trj.n_frames):
            d = pd.read_csv("./trj%d.pdb.cs" % i)
            d.rename(columns={"NUM": "resSeq", "RES": "resName", "ATOMNAME": "name"}, inplace=True)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")
    return results

def chemical_shifts_ppm(trj):
    """Predict chemical shifts of a trajectory using ppm.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ppm available on your path; see
    (http://spin.ccic.ohio-state.edu/index.php/download/index).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Li, DW, and Bruschweiler, R. "PPM: a side-chain and backbone chemical
       shift predictor for the assessment of protein conformational ensembles."
       J Biomol NMR. 2012 Nov;54(3):257-65.
    """
    pd = import_('pandas')
    binary = find_executable(PPM)

    first_resSeq = trj.top.residue(0).resSeq

    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_ppm` requires the external program PPM, available at http://spin.ccic.ohio-state.edu/index.php/download/index' % ', '.join(PPM))

    with enter_temp_directory():
        trj.save("./trj.pdb")
        cmd = "%s -pdb trj.pdb -mode detail" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your PPM installation or your input trajectory." % cmd))

        d = pd.read_csv("./bb_details.dat", delim_whitespace=True, header=None, index_col=False).drop([3], axis=1)
        d = d.rename(columns={0: "resSeq", 1: "resName", 2: "name"})
        d["resSeq"] += first_resSeq - 1  # Fix bug in PPM that reindexes to 1
        d = d.drop("resName", axis=1)
        d = d.set_index(["resSeq", "name"])
        d.columns = np.arange(trj.n_frames)
        d.columns.name = "frame"

    return d

def _get_lines_to_skip(filename):
    """Determine the number of comment lines in a SPARTA+ output file."""
    format_string = """FORMAT %4d %4s %4s %9.3f %9.3f %9.3f %9.3f %9.3f %9.3f"""
    handle = open(filename)
    for i, line in enumerate(handle):
        if line.find(format_string) != -1:
            return i + 2

    raise(Exception("No format string found in SPARTA+ file!"))

def chemical_shifts_spartaplus(trj, rename_HN=True):
    """Predict chemical shifts of a trajectory using SPARTA+.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    rename_HN : bool, optional, default=True
        SPARTA+ calls the amide proton "HN" instead of the standard "H".
        When True, this option renames the output as "H" to match the PDB
        and BMRB nomenclature.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have SPARTA+ available on your path; see
    (http://spin.niddk.nih.gov/bax/software/SPARTA+/). Also, the SPARTAP_DIR
    environment variable must be set so that SPARTA+ knows where to find
    its database files.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Shen, Y., and Bax, Ad. "SPARTA+: a modest improvement in empirical
       NMR chemical shift prediction by means of an artificial neural network."
       J. Biomol. NMR, 48, 13-22 (2010)
    """
    pd = import_('pandas')
    binary = find_executable(SPARTA_PLUS)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_spartaplus` requires the external program SPARTA+, available at http://spin.niddk.nih.gov/bax/software/SPARTA+/' % ', '.join(SPARTA_PLUS))

    names = ["resSeq", "resName", "name", "SS_SHIFT", "SHIFT", "RC_SHIFT", "HM_SHIFT", "EF_SHIFT", "SIGMA"]

    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)

        cmd = "%s -in %s" % (binary, ' '.join("trj%d.pdb" % i for i in range(trj.n_frames)))

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your SPARTA+ installation or your input trajectory." % cmd))

        lines_to_skip = _get_lines_to_skip("trj0_pred.tab")

        results = []
        for i in range(trj.n_frames):
            d = pd.read_table("./trj%d_pred.tab" % i, names=names, header=None, sep="\s*", skiprows=lines_to_skip)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)

    if rename_HN:
        results.name[results.name == "HN"] = "H"

    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")

    return results

def reindex_dataframe_by_atoms(trj, frame):
    """Reindex chemical shift output to use atom number (serial) indexing.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Returns
    -------
    new_frame : pandas.DataFrame
        Dataframe containing results, with index consisting of atom
        indices (AKA the 'serial' entry in a PDB).  Columns correspond to
        each frame in trj.

    Notes
    -----
    Be aware that this function may DROP predictions if the atom naming
    is different between the input trajectory and the output of various
    chemical shift prediction tools.
    """

    top, bonds = trj.top.to_dataframe()
    top["serial"] = top.index
    top = top.set_index(["resSeq", "name"])

    new_frame = frame.copy()

    new_frame["serial"] = top.ix[new_frame.index].serial
    new_frame = new_frame.dropna().reset_index().set_index("serial").drop(["resSeq", "name"], axis=1)

    return new_frame
<commit_msg>Remove py3 test from shiftx2 wrapper<commit_after>##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2014 Stanford University and the Authors
#
# Authors: Kyle A. Beauchamp
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
#############################################################################

##############################################################################
# Imports
##############################################################################

import os
from distutils.spawn import find_executable as _find_executable
import numpy as np

from mdtraj.utils import import_
from mdtraj.utils import enter_temp_directory

##############################################################################
# Globals
##############################################################################

# Possible names for the external commands -- these are expected
# to be found in the PATH.
SHIFTX2 = ['shiftx2.py']
SPARTA_PLUS = ['sparta+', 'SPARTA+', 'SPARTA+.linux']
PPM = ['ppm_linux_64.exe']

__all__ = ['chemical_shifts_shiftx2', 'chemical_shifts_ppm', 'chemical_shifts_spartaplus', "reindex_dataframe_by_atoms"]

def find_executable(names):
    for possible in names:
        result = _find_executable(possible)
        if result is not None:
            return result
    return None

##############################################################################
# Code
##############################################################################

def compute_chemical_shifts(trj, model="shiftx2", **kwargs):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    model : str, optional, default="shiftx2"
        The program to use for calculating chemical shifts.  Must be one
        of shiftx2, ppm, or sparta+

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have the appropriate chemical soft programs installed
    and in your executable path.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference, see docstrings for chemical_shifts_*
    for the various possible models.
    """
    if model == "shiftx2":
        return chemical_shifts_shiftx2(trj, **kwargs)
    elif model == "ppm":
        return chemical_shifts_ppm(trj, **kwargs)
    elif model == "sparta+":
        return chemical_shifts_spartaplus(trj, **kwargs)
    else:
        raise(ValueError("model must be one of shiftx2, ppm, or sparta+"))

def chemical_shifts_shiftx2(trj):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ShiftX2 available on your path; see (http://www.shiftx2.ca/).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Beomsoo Han, Yifeng Liu, Simon Ginzinger, and David Wishart.
       "SHIFTX2: significantly improved protein chemical shift
       prediction." J. Biomol. NMR, 50, 1 43-57 (2011)
    """
    pd = import_('pandas')
    binary = find_executable(SHIFTX2)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_shiftx2` requires the external program SHIFTX2, available at http://www.shiftx2.ca/' % ', '.join(SHIFTX2))

    results = []
    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)
        cmd = "%s -b 'trj*.pdb'" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your ShiftX2 installation or your input trajectory." % cmd))

        for i in range(trj.n_frames):
            d = pd.read_csv("./trj%d.pdb.cs" % i)
            d.rename(columns={"NUM": "resSeq", "RES": "resName", "ATOMNAME": "name"}, inplace=True)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")
    return results

def chemical_shifts_ppm(trj):
    """Predict chemical shifts of a trajectory using ppm.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ppm available on your path; see
    (http://spin.ccic.ohio-state.edu/index.php/download/index).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Li, DW, and Bruschweiler, R. "PPM: a side-chain and backbone chemical
       shift predictor for the assessment of protein conformational ensembles."
       J Biomol NMR. 2012 Nov;54(3):257-65.
    """
    pd = import_('pandas')
    binary = find_executable(PPM)

    first_resSeq = trj.top.residue(0).resSeq

    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_ppm` requires the external program PPM, available at http://spin.ccic.ohio-state.edu/index.php/download/index' % ', '.join(PPM))

    with enter_temp_directory():
        trj.save("./trj.pdb")
        cmd = "%s -pdb trj.pdb -mode detail" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your PPM installation or your input trajectory." % cmd))

        d = pd.read_csv("./bb_details.dat", delim_whitespace=True, header=None, index_col=False).drop([3], axis=1)
        d = d.rename(columns={0: "resSeq", 1: "resName", 2: "name"})
        d["resSeq"] += first_resSeq - 1  # Fix bug in PPM that reindexes to 1
        d = d.drop("resName", axis=1)
        d = d.set_index(["resSeq", "name"])
        d.columns = np.arange(trj.n_frames)
        d.columns.name = "frame"

    return d

def _get_lines_to_skip(filename):
    """Determine the number of comment lines in a SPARTA+ output file."""
    format_string = """FORMAT %4d %4s %4s %9.3f %9.3f %9.3f %9.3f %9.3f %9.3f"""
    handle = open(filename)
    for i, line in enumerate(handle):
        if line.find(format_string) != -1:
            return i + 2

    raise(Exception("No format string found in SPARTA+ file!"))

def chemical_shifts_spartaplus(trj, rename_HN=True):
    """Predict chemical shifts of a trajectory using SPARTA+.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    rename_HN : bool, optional, default=True
        SPARTA+ calls the amide proton "HN" instead of the standard "H".
        When True, this option renames the output as "H" to match the PDB
        and BMRB nomenclature.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have SPARTA+ available on your path; see
    (http://spin.niddk.nih.gov/bax/software/SPARTA+/). Also, the SPARTAP_DIR
    environment variable must be set so that SPARTA+ knows where to find
    its database files.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Shen, Y., and Bax, Ad. "SPARTA+: a modest improvement in empirical
       NMR chemical shift prediction by means of an artificial neural network."
       J. Biomol. NMR, 48, 13-22 (2010)
    """
    pd = import_('pandas')
    binary = find_executable(SPARTA_PLUS)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_spartaplus` requires the external program SPARTA+, available at http://spin.niddk.nih.gov/bax/software/SPARTA+/' % ', '.join(SPARTA_PLUS))

    names = ["resSeq", "resName", "name", "SS_SHIFT", "SHIFT", "RC_SHIFT", "HM_SHIFT", "EF_SHIFT", "SIGMA"]

    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)

        cmd = "%s -in %s" % (binary, ' '.join("trj%d.pdb" % i for i in range(trj.n_frames)))

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your SPARTA+ installation or your input trajectory." % cmd))

        lines_to_skip = _get_lines_to_skip("trj0_pred.tab")

        results = []
        for i in range(trj.n_frames):
            d = pd.read_table("./trj%d_pred.tab" % i, names=names, header=None, sep="\s*", skiprows=lines_to_skip)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)

    if rename_HN:
        results.name[results.name == "HN"] = "H"

    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")

    return results

def reindex_dataframe_by_atoms(trj, frame):
    """Reindex chemical shift output to use atom number (serial) indexing.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Returns
    -------
    new_frame : pandas.DataFrame
        Dataframe containing results, with index consisting of atom
        indices (AKA the 'serial' entry in a PDB).  Columns correspond to
        each frame in trj.

    Notes
    -----
    Be aware that this function may DROP predictions if the atom naming
    is different between the input trajectory and the output of various
    chemical shift prediction tools.
    """

    top, bonds = trj.top.to_dataframe()
    top["serial"] = top.index
    top = top.set_index(["resSeq", "name"])

    new_frame = frame.copy()

    new_frame["serial"] = top.ix[new_frame.index].serial
    new_frame = new_frame.dropna().reset_index().set_index("serial").drop(["resSeq", "name"], axis=1)

    return new_frame
<|endoftext|>