Reorganization Energy¶

from collections import OrderedDict
from ..core.functions import add_to_instance
from ..core.basejob import MultiJob
from ..core.results import Results
from ..core.settings import Settings
from ..mol.molecule import Molecule
from ..interfaces.adfsuite.ams import AMSJob

__all__ = ['ReorganizationEnergyJob', 'ReorganizationEnergyResults']


# using this function to pass a molecule inside a MultiJob ensures proper parallel execution
def pass_molecule(source, target):
    @add_to_instance(target)
    def prerun(self):
        self.molecule = source.results.get_main_molecule()


class ReorganizationEnergyResults(Results):
    """Results class for reorganization energy.
    """
    def reorganization_energy(self, unit='au'):
        energies = self.get_all_energies(unit)

        reorganization_energy = (energies['state B geo A'] - energies['state A geo A'] + 
                                 energies['state A geo B'] - energies['state B geo B'])
        return reorganization_energy

    def get_all_energies(self, unit='au'):
        energies = {}
        energies['state A geo A'] = self.job.children['go_A'].results.get_energy(unit=unit)
        energies['state B geo B'] = self.job.children['go_B'].results.get_energy(unit=unit)
        energies['state A geo B'] = self.job.children['sp_A_for_geo_B'].results.get_energy(unit=unit)
        energies['state B geo A'] = self.job.children['sp_B_for_geo_A'].results.get_energy(unit=unit)

        return energies


class ReorganizationEnergyJob(MultiJob):
    """A class for calculating the reorganization energy using AMS. 
    Given two states, A and B, the reorganization energy is defined as follows: 

    reorganization energy = 
    E(state B at optimal geometry for state A) - 
    E(state A at optimal geometry for state A) + 
    E(state A at optimal geometry for state B) - 
    E(state B at optimal geometry for state B)
    
    This job will run two geometry optimizations and two single point calculations.
    """

    _result_type = ReorganizationEnergyResults

    def __init__(self, molecule, common_settings, settings_state_A, settings_state_B, **kwargs):
        """
        molecule: the molecule 
        common_settings: a setting object for an AMSJob containing the shared settings for all the calculations
        settings_state_A: Setting object for an AMSJob containing exclusivelt the options defining the state A (e.g. charge and spin)
        settings_state_B: Setting object for an AMSJob containing exclusivelt the options defining the state B (e.g. charge and spin)
        kwargs: other options to be passed to the MultiJob constructor
        """
        MultiJob.__init__(self, children=OrderedDict(), **kwargs)

        go_settings = common_settings.copy()
        go_settings.input.ams.task = 'GeometryOptimization'

        sp_settings = common_settings.copy()
        sp_settings.input.ams.task = 'SinglePoint'

        # copy the settings so that we wont modify the ones provided as input by the user
        settings_state_A = settings_state_A.copy()
        settings_state_B = settings_state_B.copy()
        # In case the charge key is not specified, excplicitely set the value to 0.
        # This is to prevent the charge in molecule.properties.charge (set by get_main_molecule())
        # to be used in case of neutral systems
        for s in [settings_state_A, settings_state_B]:
            if not 'charge' in s.input.ams.system:
                s.input.ams.system.charge = 0

        self.children['go_A'] = AMSJob(molecule=molecule, settings=go_settings+settings_state_A, name='go_A')
        self.children['go_B'] = AMSJob(molecule=molecule, settings=go_settings+settings_state_B, name='go_B')
        self.children['sp_A_for_geo_B'] = AMSJob(settings=sp_settings+settings_state_A, name='sp_A_geo_B')
        self.children['sp_B_for_geo_A'] = AMSJob(settings=sp_settings+settings_state_B, name='sp_B_geo_A')
        pass_molecule(self.children['go_A'], self.children['sp_B_for_geo_A'])
        pass_molecule(self.children['go_B'], self.children['sp_A_for_geo_B'])