mechanism

from analysis_base import *
import numpy as np
from math import log
from catmap.functions import convert_formation_energies

class MechanismAnalysis(MechanismPlot,ReactionModelWrapper,MapPlot):
    """
      A simple tool for the generation of potential energy diagrams
      from a reaction network.
    """
    def __init__(self,reaction_model=None):
        """
        Class for generating potential energy diagrams.
        :param reaction_model: The ReactionModel object to load.
        """
        self._rxm = reaction_model
        defaults = {'pressure_correction':True,
                    'min_pressure':1e-12,
                    'energy_type':'free_energy',
                    'include_labels':False,
                    'subplots_adjust_kwargs':{},
                    'line_offset':0,
                    'kwarg_dict':{}}
        self._rxm.update(defaults)
        self.data_dict = {}
        MechanismPlot.__init__(self,[0])

    def plot(self,ax=None,plot_variants=None,mechanisms=None,
            labels=None,save=True):
        """
        Generates the potential energy diagram plot
        :param ax: Matplotlib Axes object to optionally plot into
        :param plot_variants: Which PEDs to plot. Defaults to all surfaces
                              or all applied voltages
        :param plot_variants: list of voltages (if electrochemical) or
                              descriptor tuples to plot
        :param mechanisms: Which reaction pathways to plot.  Each integer
                           corresponds to an elementary step. Elementary
                           steps are indexed in the order that they are
                           input with 1 being the first index. Negative
                           integers are used to designate reverse reactions.
                           Read in from model.rxn_mechanisms by default
        :type mechanisms: {string:[int]}
        :param labels: Labels for each state.  Can be generated automatically
        :type labels: [string]
        :param save: Whether to write plot to file
        :type save: bool
        """
        if not ax:
            fig = plt.figure()
            ax = fig.add_subplot(111)
        else:
            fig = ax.get_figure()
        if not mechanisms:
            mechanisms = self.rxn_mechanisms.values()
        if not plot_variants and self.descriptor_dict:
            plot_variants = self.surface_names
        elif not plot_variants and 'voltage' in self.descriptor_names:
            voltage_idx = self.descriptor_names.index('voltage')
            v_min, v_max = self.descriptor_ranges[voltage_idx]
            plot_variants = np.linspace(v_min, v_max, 5)
        if not self.plot_variant_colors:
            self.plot_variant_colors = get_colors(max(len(plot_variants),len(mechanisms)))

        self.kwarg_list = []
        for key in self.rxn_mechanisms.keys():
            self.kwarg_list.append(self.kwarg_dict.get(key,{}))

        for n,mech in enumerate(mechanisms):
            for i, variant in enumerate(plot_variants):
                if self.descriptor_dict:
                    xy = self.descriptor_dict[variant]
                elif 'voltage' in self.descriptor_names:
                    voltage_idx = self.descriptor_names.index('voltage')
                    xy = [0, 0]
                    xy[voltage_idx] = variant
                    xy[1-voltage_idx] = self.descriptor_ranges[1-voltage_idx][0]
                else:
                    xy = variant
                if '-' not in xy:
                    self.thermodynamics.current_state = None #force recalculation
                    self._rxm._descriptors = xy

                    if self.energy_type == 'free_energy':
                        energy_dict = self.scaler.get_free_energies(xy)
                    elif self.energy_type == 'potential_energy':
                        energy_dict = self.scaler.get_free_energies(xy)
                        energy_dict.update(self.scaler.get_electronic_en