Source code for p4util.optproc

#
#@BEGIN LICENSE
#
# PSI4: an ab initio quantum chemistry software package
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
#@END LICENSE
#

r"""Module to provide mechanism to store and restore option states in driver.

"""
import sys
import psi4


class OptionState(object):
    """Class to store the state of a single *option*. If *module* given, the *option*
    value and has_changed value is stored for global, local to *module*, and used by
    *module* scopes; otherwise (used for BASIS keywords), only global scope is stored.
    Class can store, print, and restore option values. ::

        >>> OptionState('SCF_TYPE', 'SCF')

        >>> print(OptionState('DF_BASIS_MP2'))

    """
    def __init__(self, option, module=None):
        self.option = option.upper()
        if module:
            self.module = module.upper()
        else:
            self.module = None

        self.value_global = psi4.get_global_option(option)
        self.haschanged_global = psi4.has_global_option_changed(option)
        if self.module:
            self.value_local = psi4.get_local_option(self.module, option)
            self.haschanged_local = psi4.has_local_option_changed(self.module, option)
            self.value_used = psi4.get_option(self.module, option)
            self.haschanged_used = psi4.has_option_changed(self.module, option)
        else:
            self.value_local = None
            self.haschanged_local = None
            self.value_used = None
            self.haschanged_used = None

    def __str__(self):
        text = ''
        if self.module:
            text += """  ==> %s Option in Module %s <==\n\n""" % (self.option, self.module)
            text += """  Global (has changed?) value: %7s %s\n""" % ('(' + str(self.haschanged_global) + ')', self.value_global)
            text += """  Local (has changed?) value:  %7s %s\n""" % ('(' + str(self.haschanged_local) + ')', self.value_local)
            text += """  Used (has changed?) value:   %7s %s\n""" % ('(' + str(self.haschanged_used) + ')', self.value_used)
        else:
            text += """  ==> %s Option in Global Scope <==\n\n""" % (self.option)
            text += """  Global (has changed?) value: %7s %s\n""" % ('(' + str(self.haschanged_global) + ')', self.value_global)
        text += """\n"""
        return text

    def restore(self):
        psi4.set_global_option(self.option, self.value_global)
        if not self.haschanged_global:
            psi4.revoke_global_option_changed(self.option)
        if self.module:
            psi4.set_local_option(self.module, self.option, self.value_local)
            if not self.haschanged_local:
                psi4.revoke_local_option_changed(self.module, self.option)


class OptionsState(object):
    """Class to contain multiple :py:func:`~optproc.OptionsState` objects.
    Used in python driver functions to collect several options before altering
    them, then restoring before function return. ::

        >>> optstash = OptionsState(
                ['SCF', 'DFT_FUNCTIONAL'],
                ['DF_BASIS_SCF'],
                ['SCF', 'SCF_TYPE'],
                ['SCF', 'REFERENCE'])

        >>> print(optstash)

        >>> optstash.restore()

    """
    def __init__(self, *largs):
        self.data = []
        for item in largs:
            if len(item) == 2:
                self.data.append(OptionState(item[1], item[0]))
            elif len(item) == 1:
                self.data.append(OptionState(item[0]))
            else:
                print('ERROR: Each argument to OptionsState should be an array, the first element')
                print('       of which is the module scope and the second element of which is the')
                print('       module name. Bad argument: %s' % (item))
                sys.exit()

    def __str__(self):
        text = ''
        for item in self.data:
            text += str(item)
        return text

    def restore(self):
        for item in self.data:
            item.restore()