digichem-core 6.0.0rc1__py3-none-any.whl

digichem/image/excited_states.py

@@ -0,0 +1,335 @@
+from matplotlib import ticker
+from matplotlib.ticker import FuncFormatter
+from itertools import chain
+import math
+
+from digichem.image.graph import OneD_graph_image_maker
+#from digichem.result.excited_state import Excited_state, Excited_state_list
+import digichem.result.excited_state
+
+class Excited_states_diagram_maker(OneD_graph_image_maker):
+    """
+    Class for making orbital energy diagrams.
+    """
+    
+    def __init__(self, output, excited_states, ground_state, y_limits = "auto", show_dest = True, **kwargs):
+        super().__init__(output, **kwargs)
+        self.excited_states = excited_states
+        self.ground_state = ground_state
+                        
+        # Each multiplicity will be plotted as a separate column, so first we need to organise our data.
+        states = chain(self.excited_states, [self.ground_state]) if self.ground_state is not None else self.excited_states
+        self.grouped_states = (type(self.excited_states)(sorted(states, key = lambda state: state.energy))).group()
+        # We'll also keep hold of just our excited states (so no ground).
+        self.grouped_excited_states = self.excited_states.group()
+        
+        # Whether to display our dEst label, this is forced off if we don't have only singlets and triplets.
+        self.show_dest = show_dest if 1 in self.grouped_excited_states and 3 in self.grouped_excited_states and len(self.grouped_excited_states) == 2 else False
+        self.left_padding = 0
+        self.right_padding = 0    
+            
+        # Set our output dpi (higher values = bigger image/quality).
+        self.output_dpi = 140
+        
+        # Initial image dimensions (we'll change these to suit).
+        self.init_image_width = 5.6
+        self.init_image_height = 5.6
+        
+        # Amount of space to allocate per unit of each axis.
+        self.inch_per_x = 1.66
+        #self.inch_per_y = 0.94
+        self.inch_per_y = 1.1
+        
+        # Y axis label (we don't have an x axis label).
+        self.y_label = "Energy /eV"
+        
+        self.plot_options = {
+            "orientation": "vertical",
+            "lineoffsets": 1,
+            "linelengths": 0.9,
+            "linewidths": None,
+            "colors": None
+        }
+        
+        # Save our scaling method.
+        self.limits_method = y_limits
+    
+    @classmethod
+    def from_options(self, output, *, excited_states, ground_state, options, **kwargs):
+        """
+        Constructor that takes a dictionary of config like options.
+        """    
+        return self(
+            output,
+            excited_states = excited_states,
+            ground_state = ground_state,
+            show_dest = options['excited_states_diagram']['show_dest'],
+            y_limits = options['excited_states_diagram']['y_limits'],
+            enable_rendering = options['excited_states_diagram']['enable_rendering'],
+            **kwargs
+        )
+            
+    def plot_lines(self):
+        """
+        Plot the lines that make up the body of the graph.
+        
+        This method is called automatically as part of the make() method.
+        :return: Nothing.
+        """
+        # Turn our dictionary of grouped excited states into an ordered list for matplotlib.
+        data = [
+            [excited_state.energy for excited_state in self.grouped_states[grouped_key]]
+            #self.grouped_states[grouped_key] 
+            for grouped_key 
+            in sorted(self.grouped_states)
+        ]
+            
+        # We can use an event plot which is good for 1D graphs like this.
+        self.axes.eventplot(data, linestyles = "solid", **self.plot_options)
+    
+    def get_x_pos_from_multiplicity(self, multiplicity):
+        """
+        Get the X coordinate of an energy state.
+        
+        :raises ValueError: If the given multiplicity is not plotted. 
+        :param multiplicity: The multiplicity of the energy state, which determines its x position.
+        :return: The X coord.
+        """
+        # Get all the multiplicities that we're using.
+        mult = sorted(self.grouped_states)
+        
+        # Our X position depends on the multiplicity of our GS and how many total multiplicities there are (as well as how spaced out each column on the graph is).
+        if len(self.grouped_states) > 1:
+            return mult.index(multiplicity) * self.plot_options['lineoffsets']
+        else:
+            # It appears matplotlib uses a different numbering if we only have 1 column, where x == 1 (normally the first column is x == 0). This makes me sad. Might be bug...
+            return 1
+    
+    def annotate_ground_state(self):
+        """
+        Add a label to our ground state.
+        
+        This method is called automatically as part of the make() method.
+        :return: Nothing.
+        """
+        # First work out where our GS is on the graph.
+        x_pos = self.get_x_pos_from_multiplicity(self.ground_state.multiplicity)
+        
+        # Now plot our text.
+        self.axes.annotate(
+             r"$\bf{{{}}}_{{{}}}$".format(self.ground_state.multiplicity_symbol, self.ground_state.multiplicity_level),
+             (x_pos, self.ground_state.energy),
+             textcoords = "offset pixels",
+             xytext = (0, 7),
+             horizontalalignment = "center",
+             fontsize = 12
+         )
+        
+    def annotate_excited_state(self, state):
+        """
+        Add a label to an excited state
+        
+        This method is called automatically as part of the make() method.
+        :param multiplicity_symbol: The symbol (a string) to annotate (eg, S1, T1).
+        :return: Nothing.
+        """
+        # First work out where our state is on the graph.
+        #state = self.excited_states.get_state(multiplicity_symbol)
+        x_pos = self.get_x_pos_from_multiplicity(state.multiplicity)
+        
+        # Assemble our label.
+        label_string =     r"$\bf{{{}}}_{{{}}}$: ".format(state.multiplicity_symbol, state.multiplicity_level) + "{:0.2f} eV".format(state.energy)
+                        
+        # Add oscillator strength (if we have it).
+        if state.oscillator_strength is not None:
+            label_string +=    "\n" + "f: {:0.2f}".format(state.oscillator_strength)
+        
+        # Now plot our text.
+        self.axes.annotate(
+             label_string,
+             (x_pos, state.energy),
+             textcoords = "offset pixels",
+             xytext = (0, -7),
+             verticalalignment = "top",
+             horizontalalignment = "center",
+             fontsize = 12
+         )
+    
+    def annotate_dest(self, x, y):
+        """
+        """
+        # First the text label.
+        self.axes.annotate(
+                r"$\bf{{ΔE_{{ST}}}}$: {:0.2f} eV".format(self.excited_states.singlet_triplet_energy),
+                (x, y),
+                textcoords = "offset pixels",
+                xytext = (0,0),
+                verticalalignment = "center",
+                horizontalalignment = "left",
+                fontsize = 12
+            )
+    
+    def plot_labels(self):
+        """
+        Plot the labels on our graph that identify various excited states and the singlet/triplet splitting energy (dEst) if available.
+        
+        This method is called automatically as part of the make() method.
+        :return: Nothing.
+        """
+        # First add a label for our ground state, unless there isn't room (this hack again).
+        lowest_state = self.excited_states[0]
+        # We need about 0.6 eV space.
+        if (lowest_state.energy - 0.6) >= 0:
+            self.annotate_ground_state()
+        
+        # Now we'll add a label for the lowest state of each multiplicity.
+        for multiplicity in self.grouped_excited_states:
+            self.annotate_excited_state(self.grouped_excited_states[multiplicity][0])
+            
+            
+        # We'll also try and add a dEst label (obviously this is only possible if there is an S1 and T1 available).
+        if self.show_dest:
+            # First get our two states.
+            S1 = self.excited_states.get_state("S(1)")
+            T1 = self.excited_states.get_state("T(1)")
+            
+            # Find the midpoint between them.
+            y_pos = min([S1.energy, T1.energy]) +  (abs(S1.energy - T1.energy) /2)
+            
+            # The label will be to the right of S1/T1.
+            x_pos = max(self.get_x_pos_from_multiplicity(S1.multiplicity), self.get_x_pos_from_multiplicity(T1.multiplicity)) + (self.plot_options['linelengths'] /2) + 0.2
+            # Add the label
+            self.annotate_dest(x_pos, y_pos)
+            
+            # Now draw our vertical arrow.
+            self.axes.annotate(
+                "",
+                xy = (x_pos-0.05, S1.energy),
+                xytext = (x_pos-0.05, T1.energy),
+                arrowprops=dict(arrowstyle="<->")
+            )
+            
+            # Now extend a line from S1 and T1 to meet the vertical arrow.
+            for state in (S1, T1):
+                self.axes.annotate(
+                    "",
+                    xy = (self.get_x_pos_from_multiplicity(state.multiplicity) + (self.plot_options['linelengths'] /2), state.energy),
+                    xytext = (x_pos-0.05, state.energy),
+                    arrowprops=dict(arrowstyle="-", linestyle="--")
+                )
+            
+            # Pad to the right so our dest label is chopped off.
+            self.right_padding = 0.7            
+
+        
+    def x_axis_formatter(self, x, pos):
+        """
+        This method tells matplotlib how to label our X axis (with multiplicity labels).
+        """
+        # The multiplicities that we're using.
+        mult = sorted(self.grouped_states)
+        
+        if len(self.grouped_states) > 1:
+            # Check to see if our x is in our list of multiplicities.
+            if x % self.plot_options['lineoffsets'] == 0 and x >= 0:
+                try:
+                    return digichem.result.excited_state.Excited_state.multiplicity_number_to_string((mult[int(x / self.plot_options['lineoffsets'])]))
+                    
+                except IndexError:
+                    # Out of range.
+                    return ""
+        else:
+            # Matplotlib uses different counting when there's only 1 column :(.
+            if x == 1:
+                return digichem.result.excited_state.Excited_state.multiplicity_number_to_string(mult[0])
+            else:
+                return ""
+            
+    def all_limits(self):
+        """
+        Limit the Y axis so all states are visible.
+        """
+        # Use auto scaling.
+        self.axes.autoscale(enable = True, axis = 'y')
+        
+        # We'll go a small amount below 0 so our S0 is off the bottom of the graph.
+        min_y = -0.1
+        # Use a sneaky hack so that if the max energy is close to an integer, we'll use the next highest integer.
+        #max_y = math.ceil(self.excited_states[-1].energy +0.1)
+        max_y = self.excited_states[-1].energy + 0.1
+        self.axes.set_ylim(min_y, max_y)
+        
+        # Adjusty hack.
+        self._adjust_limits_for_zero()
+        
+    def auto_limits(self):
+        """
+        Limit the Y axis so S1, T1 ... N1 etc are all visible.
+        """
+        self.all_limits()
+        
+        # We'll show the lowest excited state of each mult.
+        highest_state = max([self.grouped_excited_states[multiplicity][0].energy for multiplicity in self.grouped_excited_states])
+        # Use a sneaky hack so that if the max energy is close to an integer, we'll use the next highest integer.
+        max_y = math.ceil(highest_state +0.1)
+        self.axes.set_ylim(None, max_y)
+        
+        
+    def _adjust_limits_for_zero(self):
+        """
+        This method is a hack which lowers the y axis limit in cases where S1/T1 is very close to S0
+        
+        This is necessary because otherwise the state annotations will overflow the bottom of the diagram.
+        """
+        # First get our lowest E excited state.
+        lowest_state = self.excited_states[0]
+        # We need about 0.7 eV space.
+        if (lowest_state.energy - 0.5) < 0:
+            # Add on a bit extra.
+            self.axes.set_ylim(bottom = lowest_state.energy - (0.5))
+        
+        
+    def simple_limits(self, y_min = -0.1, y_max = 5):
+        """
+        Limit the Y axis between between two points.
+        
+        :param y_min: The start of the y axis.
+        :param y_max: The end of the y axis.
+        """
+        self.axes.set_ylim(y_min, y_max)
+        
+    def adjust_axes(self):
+        """
+        Adjust the axes of our graph.
+        
+        This method is called automatically as part of the make() method.
+        :return: Nothing.
+        """
+        super().adjust_axes()
+                
+        # Adjust axis
+        # Add a bit o' padding (often we actually reduce padding from the default).
+        xlim = self.axes.get_xlim()
+        self.axes.set_xlim(xlim[0] - (self.left_padding -0.25), xlim[1] + (self.right_padding -0.25))
+        #self.axes.set_xlim(xlim[0] +0.25, xlim[1] -0.25)
+        
+        # Now decide how big our y axis will be.
+        if self.limits_method == "all":
+            self.all_limits()
+        elif self.limits_method == "auto":
+            self.auto_limits()
+        elif isinstance(self.limits_method, tuple) or isinstance(self.limits_method, list):
+            self.simple_limits(self.limits_method[0], self.limits_method[1])
+        else:
+            raise ValueError("Unknown limits method '{}'".format(self.limits_method))
+        
+        # Set custom text on our x axis.
+        self.axes.xaxis.set_major_formatter(FuncFormatter(self.x_axis_formatter))
+        #self.axes.xaxis.set_major_locator(ticker.MultipleLocator(1))
+        self.axes.xaxis.set_major_locator(ticker.FixedLocator([self.get_x_pos_from_multiplicity(multiplicity) for multiplicity in self.grouped_excited_states]))
+        
+        # Both our dimensions can scale to fit more data.
+        self.figure.tight_layout()
+        self.constant_scale(0, self.inch_per_x)
+        self.constant_scale(1, self.inch_per_y)
+

digichem/image/graph.py

@@ -0,0 +1,293 @@
+from matplotlib import ticker
+from matplotlib.transforms import Bbox
+import matplotlib.figure
+import math
+
+from digichem.image import Image_maker
+
+class Graph_image_maker(Image_maker):
+    """
+    A class used for creating graph images.
+    
+    Most of the heavy lifting here is achieved by the matplotlib library.
+    """
+    
+    def __init__(self, *args, max_width = None, max_height = None, **kwargs):
+        """
+        Constructor for Graph_image_maker objects.
+        
+        :param max_width: The maximum image width in pixels, None for no limit.
+        :param max_height: The maximum image height in pixels, None for no limit.
+        """
+        super().__init__(*args, **kwargs)
+        
+        # Set our output dpi (higher values = bigger image/quality).
+        self.output_dpi = 130
+        
+        # The width and height (in inch) that we initially create our figure with. Note that the true figure size may change from this value.
+        self.init_image_width = 5
+        self.init_image_height = 5
+        
+        # If true, the label and tick labels will be hidden for that axis.
+        self.hide_x = False
+        self.hide_y = False
+        
+        # The label to use for the x and y axes.
+        self.x_label = ""
+        self.y_label = ""
+        
+        # The matplotlib figure object that represents our graph.
+        self.figure = None
+        
+        # The plotting area within our figure.
+        self.axes = None
+        
+        # Image clamping.
+        self.max_width = max_width
+        self.max_height = max_height
+        
+    def make_files(self):
+        """
+        Make the image described by this object.
+        """
+        # Change font and other defaults.
+        #pyplot.rcParams.update({'font.size': 14, 'font.family': 'DejaVu Sans', 'axes.linewidth': 2})
+        matplotlib.rcParams.update({'font.size': 14, 'font.family': 'DejaVu Sans', 'axes.linewidth': 2})
+                
+        # Make our (empty) plot.
+        # pyplot uses a nasty interface, will switch to OO at some point.
+        #pyplot.figure(figsize = (self.init_image_width, self.init_image_height), tight_layout = True)
+        #pyplot.figure(figsize = (self.init_image_width, self.init_image_height), tight_layout = False)
+        
+        # Plot our graph.
+        self.make_graph()
+        
+        # And save to file.
+        self.figure.savefig(self.output, dpi = self.output_dpi)
+        #pyplot.savefig(self.output, dpi = self.output_dpi)
+        
+    def make_graph(self):
+        """
+        Make the graph data described by this object.
+        
+        The matplotlib figure is available at self.figure (and is also returned by this method for convenience). The graph is not saved to file.
+        
+        :return: A reference to the created matplotlib Figure object.
+        """
+        # First make our figure object, which is the top-level matplotlib container.
+        self.figure = matplotlib.figure.Figure(figsize = (self.init_image_width, self.init_image_height))
+        
+        # Also get our main axes (all the graph types we support so far contain a single plot/axes).
+        self.axes = self.figure.add_subplot(111)
+        
+        # Call plot(), which is defined by sub-classes to actually draw the graph.
+        self.plot()
+        
+        # Call adjust_axes(), which does what it suggests (also adds axis labels etc.)
+        self.adjust_axes()
+        
+        # Finally, clamp our image size if we're going to be too big.
+        if self.max_width is not None:
+            self.clamp_dimension(0, self.max_width)
+        if self.max_height is not None:
+            self.clamp_dimension(1, self.max_height)
+        
+        # Return a reference to the figure (for convenience).
+        return self.figure
+        
+    def plot(self):
+        """
+        Main workhorse of the make_files method. Inheriting classes should write their own implementation.
+        
+        This is called automatically by make_graph().
+        """
+        pass
+    
+    def adjust_axes(self):
+        """
+        Adjust the axes of our graph.
+        
+        This method is called automatically as part of the make() method.
+        """ 
+        if self.hide_x:
+            self.x_label = "Arbitrary units"
+            self.axes.xaxis.set_ticks([])
+        if self.hide_y:
+            self.y_label = "Arbitrary units"
+            self.axes.yaxis.set_ticks([])
+        
+        # Add our axes labels.
+        if not self.x_label == "": 
+            self.axes.set_xlabel(self.x_label, labelpad = 8, fontsize = 16, weight = 'bold')
+        if not self.y_label == "":
+            self.axes.set_ylabel(self.y_label, labelpad = 8, fontsize = 16, weight = 'bold')
+    
+    def clamp_dimension(self, dim, maximum):
+        """
+        Ensure the image does not exceed a certain number of pixels in a given dimension.
+        
+        :param dim: The dimension to clamp; 0 for x, 1 for y.
+        :param maximum: The maximum allowed pixels.
+        """
+        # Check to see if we've exceeded our max size.
+        fig_size = self.figure.get_size_inches()
+        if (fig_size[dim] * self.output_dpi) > maximum:
+            # We're too big, reduce to max.
+            fig_size[dim] = maximum / self.output_dpi
+            self.figure.set_size_inches(fig_size)
+                        
+            # Update layout. Calling this twice is necessary for some reason loool
+            self.figure.tight_layout()
+            self.figure.tight_layout()
+    
+    def constant_scale(self, axis, inch_per_axis):
+        """
+        Adjust the size of our figure so a given axis has a constant scale.
+        
+        You should adjust the corresponding axis' limits before calling this method.
+        
+        :param axis: The axis to adjust, 0 for the x axis, 1 for the y axis.
+        :param inch_per_axis: The number of inches per unit of the given axis to adjust to. 
+        """
+        
+        # First determine how much space we need for our margins.
+        axes_position = self.axes.get_position(False).get_points()
+        figure_size = self.figure.get_size_inches()
+        # Axes_position is in fractions, so multiply by our current figsize.
+        start_margin = axes_position[0][axis] * figure_size[axis]
+        end_margin = (1 - axes_position[1][axis]) * figure_size[axis]
+        
+        # Next determine how much space we need for our actual data.
+        axes_lim = (self.axes.get_xlim(), self.axes.get_ylim())
+        axis_min = axes_lim[axis][0]
+        axis_max = axes_lim[axis][1]
+        #axis_min = pyplot.axis()[axis *2]
+        #axis_max = pyplot.axis()[axis *2 +1]
+        
+        # We take the absolute because min can actually be bigger than max if our axes are inverted.
+        axis_length = math.fabs(axis_max - axis_min)
+        axis_size = axis_length * inch_per_axis
+        
+        # Update size.
+        figure_size[axis] = start_margin + end_margin + axis_size
+        self.figure.set_size_inches(figure_size)
+        # And update the position of our axes.
+        axes_position[0][axis] = start_margin / figure_size[axis]
+        axes_position[1][axis] = 1 - (end_margin / figure_size[axis])
+        self.axes.set_position(Bbox(axes_position))
+        
+    
+class OneD_graph_image_maker(Graph_image_maker):
+    """
+    Classes for creating '1D' graphs using matplotlib's event plot.
+    """
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        
+        # Options that we'll pass to matplotlib's event plot function.
+        self.plot_options = {
+            "orientation": "vertical",
+            "lineoffsets": 0.75,
+            "linelengths": 1,
+            "linewidths": None,
+            "colors": None
+        }
+        
+        # Set our output width and height.
+        self.init_image_width = 4.1
+        self.init_image_height = 5.6
+        
+    def plot(self):
+        """
+        Make the image described by this object.
+        """
+        # Plot our 1D data.
+        self.plot_lines()
+        
+        # Annotate with data labels.
+        self.plot_labels()
+            
+    def plot_lines(self):
+        """
+        Plot the lines that make up our graph.
+        
+        This default implementation does nothing, inheriting classes should write their own.
+        """
+        pass
+    
+    def plot_labels(self):
+        """
+        Plot annotation labels that explain the data in out graph.
+        
+        This default implementation does nothing, inheriting classes should write their own.
+        """
+        pass
+        
+class Convergence_graph_maker(Graph_image_maker):
+    """
+    A class for creating graphs that show the convergence of energy in optimisation calculations.
+    """
+    
+    def __init__(self, output, energies, **kwargs):
+        """
+        Constructor for graph makers.
+        
+        :param energies: A list of energies to plot (probably an Energy_list object).
+        :param output: A path to an output file to write to. The extension of this path is used to determine the format of the file (eg, png, jpeg).
+        """
+        super().__init__(output, **kwargs)
+        self.energies = energies
+        
+        # Our output width and height (in inch).
+        self.init_image_width = 7
+        self.init_image_height = 5.265
+        
+        # Axis titles.
+        self.x_label = "Step number"
+        self.y_label = "Energy /eV"
+        
+    @classmethod
+    def from_options(self, output, *, energies, options, **kwargs):
+        """
+        Constructor that takes a dictionary of config like options.
+        """        
+        return self(
+            output,
+            energies = energies,
+            **kwargs
+        )
+    
+    def plot(self):
+        """
+        Plot the contents of our graph.
+        """
+        # Plot a simple 'scatter' style graph.
+        self.axes.plot(range(1, len(self.energies) +1), self.energies)    
+            
+    def adjust_axes(self):
+        """
+        Adjust the axes of our graph.
+        
+        This method is called automatically as part of the make() method.
+        """
+        super().adjust_axes()
+        
+        # Axes range, the format is [xmin, xmax, ymin, ymax].
+        # We'll add a bit on to the y axis so we can see better.
+        y_diff = max(self.energies) - min(self.energies)
+        y_fudge = y_diff * 0.05
+        # Also the x.
+        x_diff = len(self.energies) - 1
+        x_fudge = x_diff * 0.05
+        
+        self.axes.set_xlim(1 - x_fudge, len(self.energies) + x_fudge)
+        self.axes.set_ylim(min(self.energies) - y_fudge, max(self.energies) + y_fudge)
+        
+        # Change spacing.
+        self.axes.xaxis.set_major_locator(ticker.MaxNLocator(nbins = '12', steps = [1, 2, 2.5, 5, 10], integer = True))
+        
+        # Use matplotlib to automatically layout our graph.
+        self.figure.tight_layout()
+        
+