Display Kinetic Energy Barrier Values

To achieve this, you can modify the add_barrier method to automatically determine a good spot for displaying the kinetic energy barrier values while avoiding overlap with other elements in the figure. Here's an updated version of the method:

# ...

class FED:
    # ... (existing code)

    def add_barrier(self, energy, start_level_id, end_level_id,
                    ls='-', color='k'):
        '''
        This is a method of the FED class that will take in a start/end level id 
        and an energy and create a barrier between those two levels. 

        Parameters
        ----------
        start_level_id : int
            The id (index of the level) of the level that you want the
            barrier to start from.
        end_level_id : int
            The id (index of the level) of the level that you want the
            barrier to end.
        energy : int
            The energy (saddle point) of the barrier
        ls : str
            The barrier style. (Default = '-')
        color : str
            The colour of the barrier. (Default = 'k')

        Returns
        -------
        Appends a tuple containing the information about the barrier to a
        list in the class barriers attribute. 
        '''

        self.auto_adjust()

        start = self.positions[start_level_id-1]*(self.dimension+self.space)
        x1 = start + self.dimension
        y1 = self.energies[start_level_id-1]
        x2 = self.positions[end_level_id-1]*(self.dimension+self.space)
        y2 = self.energies[end_level_id-1]
        vert_x = x1 + ((x2-x1)/2)
        vert_y = energy
        a1 = (y1 - vert_y)/((x1 - vert_x)**2)
        a2 = (y2 - vert_y)/((x2 - vert_x)**2)
        left_xspace = list(np.linspace(x1, vert_x, 500))
        right_xspace = list(np.linspace(vert_x, x2, 500))
        left_yspace = []
        right_yspace = []
        for y in left_xspace:
            left_y = (a1*((y - vert_x)**2)) + vert_y
            left_yspace.append(left_y)
        for x in right_xspace:
            right_y = (a2*((x - vert_x)**2)) + vert_y
            right_yspace.append(right_y)
        overall_xspace = left_xspace + right_xspace
        overall_yspace = left_yspace + right_yspace
        barrier_data = (overall_xspace, overall_yspace, ls, color, vert_x, vert_y)

        # Find a good spot for displaying the barrier value
        barrier_text_x, barrier_text_y = self.find_good_barrier_text_spot(barrier_data)

        self.barriers.append(barrier_data)
        self.display_barrier_text(barrier_text_x, barrier_text_y, energy)

    def find_good_barrier_text_spot(self, barrier_data):
        '''
        Finds a good spot for displaying the kinetic energy barrier value 
        next to (or above/below) the maxima of the barrier.

        Parameters
        ----------
        barrier_data : tuple
            Tuple containing information about the barrier.

        Returns
        -------
        Tuple (x, y) representing the coordinates for displaying the barrier value text.
        '''

        # Unpack barrier_data
        overall_xspace, overall_yspace, _, _, vert_x, vert_y = barrier_data

        # Find the maxima of the barrier
        maxima_index = np.argmax(overall_yspace)
        maxima_x = overall_xspace[maxima_index]
        maxima_y = overall_yspace[maxima_index]

        # Calculate text position relative to maxima
        text_x = maxima_x + 0.1 * self.dimension  # Adjust the position as needed
        text_y = maxima_y

        return text_x, text_y

    def display_barrier_text(self, x, y, energy):
        '''
        Displays the kinetic energy barrier value text.

        Parameters
        ----------
        x : float
            X-coordinate for displaying the text.
        y : float
            Y-coordinate for displaying the text.
        energy : int
            The energy (saddle point) of the barrier.
        '''

        # Check for overlap with other elements and adjust if necessary
        # You may need to further customize this based on your specific requirements

        # Example: check for overlap with existing barrier texts
        for existing_barrier_data in self.barriers:
            existing_x, existing_y = self.find_good_barrier_text_spot(existing_barrier_data)
            if abs(existing_x - x) < self.dimension * 0.2 and abs(existing_y - y) < self.dimension * 0.2:
                # Adjust the y-coordinate to avoid overlap
                y += self.dimension * 0.3

        # Display the barrier value text
        self.ax.text(x, y, f'{energy} eV', color='black', fontsize=8,
                     horizontalalignment='left', verticalalignment='center')

    # ...

In this modification:

I added the find_good_barrier_text_spot method, which calculates a good spot for displaying the kinetic energy barrier value.
The display_barrier_text method checks for overlap with existing barrier texts and adjusts the position if necessary.
The add_barrier method now calls these helper methods to find a good spot for the barrier text and displays the kinetic energy barrier value accordingly.

Note: Adjust the conditions for checking overlap based on your specific requirements, and further customization might be needed depending on your use case.

cia-ran / free_energy_diagrams

Display Kinetic Energy Barrier Values #2