1"""
  2# Description
  3
  4This module provides straightforward functions to plot QRotor data.
  5
  6
  7# Index
  8
  9| | |
 10| --- | --- |
 11| `potential()`        | Potential values as a function of the angle |
 12| `energies()`         | Calculated eigenvalues |
 13| `reduced_energies()` | Reduced energies E/B as a function of the reduced potential V/B |
 14| `wavefunction()`     | Selected wavefunctions or squared wavefunctions of a system |
 15| `splittings()`       | Tunnel splitting energies of a list of systems |
 16| `convergence()`      | Energy convergence of a list of systems calculated with different parameters |
 17
 18---
 19"""
 20
 21
 22from .system import System
 23from . import systems
 24from . import constants
 25import matplotlib.pyplot as plt
 26import numpy as np
 27from copy import deepcopy
 28import aton.alias as alias
 29
 30
 31def potential(
 32        data:System|list,
 33        title:str=None,
 34        marker:str|list='',
 35        linestyle:str|list='-',
 36        cm:str=None,
 37        cm_range:tuple|list=(0.0,0.8),
 38        color:list|tuple|str=None,
 39        normalize:bool=False,
 40        ylim:tuple=None,
 41        rc:dict={},
 42        ) -> None:
 43    """Plot the potential values of `data` (System object, or list of systems).
 44
 45    Title can be customized with `title`.
 46    If empty, system[0].comment will be used as title if no more comments are present.
 47
 48    `marker` and `linestyle` can be a Matplotlib string or list of strings.
 49    Custom colormaps can be set through `cm`, in a custom `cm_range` from 0 to 1.
 50    Custom colors can also be set through `color`.
 51    Alternatively, custom color lists can be set through the `color` variable.
 52
 53    Set `normalize = True` to normalize by their respective `qrotor.system.System.potential_max`.
 54    This can be useful if you have performed subtractions or similar operations.
 55    In this case, you might also want to play with `ylim` to adjust the y-axis limits.
 56
 57    Additional matplotlib runtime configuration
 58    [rcParams](https://matplotlib.org/stable/api/matplotlib_configuration_api.html#matplotlib.RcParams)
 59    can be set with the `rc` dict.
 60    """
 61    data_copy = deepcopy(data)
 62    system = systems.as_list(data_copy)
 63    title_str = title if title else (system[0].comment if (system[0].comment and (len(system) == 1 or not system[-1].comment)) else 'Rotational potential energy')
 64    # Marker as a list
 65    if isinstance(marker, list):
 66        if len(marker) < len(system):
 67            marker.extend([''] * (len(system) - len(marker)))
 68    else:
 69        marker = [marker] * len(system)
 70    # Linestyle as a list
 71    if isinstance(linestyle, list):
 72        if len(linestyle) < len(system):
 73            linestyle.extend(['-'] * (len(system) - len(linestyle)))
 74    else:
 75        linestyle = [linestyle] * len(system)
 76    # Custom colormap or colors
 77    colors = None
 78    if isinstance(cm, str):
 79        cmap = plt.get_cmap(cm)
 80        colors = cmap(np.linspace(cm_range[0], cm_range[1], len(system)))
 81    elif color:
 82        colors = color
 83        if isinstance(colors, list):
 84            if len(colors) < len(system):
 85                colors.extend(['grey'] * (len(system) - len(colors)))
 86        else:
 87            colors = [colors] * len(system)
 88
 89    with plt.rc_context(rc):
 90        plt.figure(layout='constrained')
 91        plt.title(title_str)
 92        plt.xlabel('Angle / rad')
 93        plt.ylabel('Potential energy / meV')
 94        if normalize:
 95            plt.ylabel('Energy / V$_{3}$')
 96            for s in system:
 97                s.potential_values = s.potential_values / s.potential_max
 98        if ylim:
 99            plt.ylim(ylim)
100
101        plt.xticks([-2*np.pi, -3*np.pi/2, -np.pi, -np.pi/2, 0, np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], [r'$-2\pi$', r'$-\frac{3\pi}{2}$', r'$-\pi$', r'$-\frac{\pi}{2}$', '0', r'$\frac{\pi}{2}$', r'$\pi$', r'$\frac{3\pi}{2}$', r'$2\pi$'])
102
103        if colors is not None:
104            for i, s in enumerate(system):
105                plt.plot(s.grid, s.potential_values, marker=marker[i], linestyle=linestyle[i], label=s.comment, color=colors[i])
106        else:  # Regular plot
107            for i, s in enumerate(system):
108                plt.plot(s.grid, s.potential_values, marker=marker[i], linestyle=linestyle[i], label=s.comment)
109
110        if all(s.comment for s in system) and len(system) != 1:
111            leg = plt.legend()
112            leg.set_draggable(True)
113
114        plt.show()
115
116
117def energies(
118        data,
119        title:str=None,
120        rc:dict={},
121        ) -> None:
122    """Plot the eigenvalues of `data` (System or a list of System objects).
123
124    You can use up to 1 tag per system to differentiate between molecular groups.
125
126    Additional matplotlib runtime configuration
127    [rcParams](https://matplotlib.org/stable/api/matplotlib_configuration_api.html#matplotlib.RcParams)
128    can be set with the `rc` dict.
129    """
130    if isinstance(data, System):
131        var = [data]
132    else:  # Should be a list
133        systems.as_list(data)
134        var = data
135
136    V_colors = ['C0', 'C1', 'C2', 'C3', 'C4']
137    E_colors = ['lightblue', 'sandybrown', 'lightgrey', 'lightcoral', 'plum']
138    E_linestyles = ['--', ':', '-.']
139    edgecolors = E_colors
140
141    V_linestyle = '-'
142    title = title if title else (var[0].comment if var[0].comment else 'Energy eigenvalues')
143    ylabel_text = f'Energy / meV'
144    xlabel_text = 'Angle / radians'
145
146    with plt.rc_context(rc):
147        plt.figure(layout='constrained')
148        plt.xlabel(xlabel_text)
149        plt.ylabel(ylabel_text)
150        plt.title(title)
151        plt.xticks([-2*np.pi, -3*np.pi/2, -np.pi, -np.pi/2, 0, np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], [r'$-2\pi$', r'$-\frac{3\pi}{2}$', r'$-\pi$', r'$-\frac{\pi}{2}$', '0', r'$\frac{\pi}{2}$', r'$\pi$', r'$\frac{3\pi}{2}$', r'$2\pi$'])
152
153        unique_potentials = []
154        unique_groups = []
155        for i, system in enumerate(var):
156            V_color = V_colors[i % len(V_colors)]
157            E_color = E_colors[i % len(E_colors)]
158            E_linestyle = E_linestyles[i % len(E_linestyles)]
159            edgecolor = edgecolors[i % len(edgecolors)]
160
161            # Plot potential energy if it is unique
162            if not any(np.array_equal(system.potential_values, value) for value in unique_potentials):
163                unique_potentials.append(system.potential_values)
164                plt.plot(system.grid, system.potential_values, color=V_color, linestyle=V_linestyle)
165
166            # Plot eigenvalues
167            if any(system.eigenvalues):
168                text_offset = 3 * len(unique_groups)
169                if system.tags not in unique_groups:
170                    unique_groups.append(system.tags)
171                for j, energy in enumerate(system.eigenvalues):
172                    plt.axhline(y=energy, color=E_color, linestyle=E_linestyle)
173                    # Textbox positions are a bit weird when plotting more than 2 systems, but whatever...
174                    plt.text(j%3*1.0 + text_offset, energy, f'$E_{{{j}}}$ = {round(energy,4):.04f}', va='top', bbox=dict(edgecolor=edgecolor, boxstyle='round,pad=0.2', facecolor='white', alpha=0.8))
175                if len(systems.get_tags(var)) > 1:
176                    plt.plot([], [], color=E_color, label=f'{system.tags} Energies')  # Add to legend
177
178        if len(systems.get_tags(var)) > 1:
179            plt.subplots_adjust(right=0.85)
180            leg = plt.legend() #(bbox_to_anchor=(1.1, 0.5), loc='center', fontsize='small')
181            leg.set_draggable(True)
182
183        plt.show()
184
185
186def reduced_energies(
187        data:list,
188        title:str=None,
189        values:list=[],
190        legend:list=[],
191        rc:dict={},
192        ) -> None:
193    """Plots the reduced energy of the system E/B vs the reduced potential energy V/B.
194
195    Takes a `data` list of System objects as input.
196    An optional `title` can be specified.
197
198    Optional maximum reduced potential `values` are plotted
199    as vertical lines (floats or ints) or regions
200    (lists inside the values list, from min to max).
201    A `legend` of the same len as `values` can be included.
202    These values are assumed to be divided by B by the user.
203
204    Additional matplotlib runtime configuration
205    [rcParams](https://matplotlib.org/stable/api/matplotlib_configuration_api.html#matplotlib.RcParams)
206    can be set with the `rc` dict.
207    """
208    if values and (isinstance(values, float) or isinstance(values, int) or isinstance(values, np.float64)):
209        values = [values]
210    if values and len(values) <= len(legend):
211        plot_legend = True
212    else:
213        plot_legend = False
214        legend = [''] * len(values)
215    systems.as_list(data)
216    title = title if title else (data[0].comment if data[0].comment else 'Reduced energies')
217    number_of_levels = data[0].searched_E
218    x = []
219    for system in data:
220        potential_max_B = system.potential_max / system.B
221        x.append(potential_max_B)
222    colors = plt.cm.viridis(np.linspace(0, 1, number_of_levels+1))  # +1 to avoid the lighter tones
223
224    with plt.rc_context(rc):
225        plt.figure(layout='constrained')
226
227        for i in range(number_of_levels):
228            y = []
229            for system in data:
230                eigenvalues_B_i = system.eigenvalues[i] / system.B
231                y.append(eigenvalues_B_i)
232            plt.plot(x, y, marker='', linestyle='-', color=colors[i])
233        # Add vertical lines in the specified values
234        line_colors = plt.cm.tab10(np.linspace(0, 1, len(values)))
235        for i, value in enumerate(values):
236            if isinstance(value, list):
237                min_value = min(value)
238                max_value = max(value)
239                plt.axvspan(min_value, max_value, color=line_colors[i], alpha=0.2, linestyle='', label=legend[i])
240            else:
241                plt.axvline(x=value, color=line_colors[i], linestyle='--', label=legend[i], alpha=0.5)
242        plt.xlabel('V$_{B}$ / B')
243        plt.ylabel('E / B')
244        plt.title(title)
245        if plot_legend:
246            leg = plt.legend()
247            leg.set_draggable(True)
248        plt.show()
249
250
251def wavefunction(
252        system:System,
253        title:str=None,
254        square:bool=True,
255        levels:list|int=[0, 1, 2],
256        overlap:bool|int=False,
257        yticks:bool=False,
258        rc:dict={},
259        ) -> None:
260    """Plot the wavefunction of a `system` for the specified `levels`.
261
262    Wavefunctions are squared by default, showing the probabilities;
263    To show the actual wavefunctions, set `square = False`.
264
265    `levels` can be a list of indexes, or the number of levels to plot.
266
267    Specific wavefunctions can be overlapped with `overlap` as a list with the target indexes.
268    The `overlap` value can also be the max number of wavefunctions to add.
269    All found wavefunctions can be added together with `overlap = True`;
270    but note that this overlap is limited by the number of System.searched_E,
271    that must be specified before solving the system.
272    Setting `overlap` will ignore the `levels` argument.
273
274    Set `yticks = True` to plot the wavefunction yticks.
275
276    Additional matplotlib runtime configuration
277    [rcParams](https://matplotlib.org/stable/api/matplotlib_configuration_api.html#matplotlib.RcParams)
278    can be set with the `rc` dict.
279    """
280    data = deepcopy(system)
281    eigenvectors = data.eigenvectors
282    title = title if title else (data.comment if data.comment else 'System wavefunction')
283    with plt.rc_context(rc):
284        fig, ax1 = plt.subplots(layout='constrained')
285        plt.title(title)
286        ax1.set_xlabel('Angle / radians')
287        ax1.set_ylabel('Potential / meV')
288        ax1.set_xticks([-2*np.pi, -3*np.pi/2, -np.pi, -np.pi/2, 0, np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], [r'$-2\pi$', r'$-\frac{3\pi}{2}$', r'$-\pi$', r'$-\frac{\pi}{2}$', '0', r'$\frac{\pi}{2}$', r'$\pi$', r'$\frac{3\pi}{2}$', r'$2\pi$'])
289        ax1.plot(data.grid, data.potential_values, color='blue', linestyle='-')
290        ax2 = ax1.twinx()
291        if not yticks:
292            ax2.set_yticks([])
293        ax2.set_ylabel('Squared wavefunction' if square else 'Wavefunction')
294        # Set levels list
295        if isinstance(levels, int):
296            levels = [levels]
297        if not isinstance(levels, list):
298            raise ValueError('levels must be an int or a list of ints')
299        # Set overlap if requested
300        if overlap == True and isinstance(overlap, bool):
301            eigenvectors = [np.sum(eigenvectors, axis=0)]
302            levels = [0]
303            show_legend = False
304        elif overlap is not False and (isinstance(overlap, int) or isinstance(overlap, float)):
305            max_int = int(overlap)
306            eigenvectors = [np.sum(eigenvectors[:max_int], axis=0)]
307            levels = [0]
308            show_legend = False
309        elif isinstance(overlap, list):
310            eigenvectors = [np.sum([eigenvectors[i] for i in overlap], axis=0)]
311            levels = [0]
312            show_legend = False
313        if len(levels) > 1:
314            show_legend = True
315        else:
316            show_legend = False
317        # Square values if so
318        if square:
319            eigenvectors = [vec**2 for vec in eigenvectors]
320        # Plot the wavefunction
321        for i in levels:
322            ax2.plot(data.grid, eigenvectors[i], linestyle='--', label=f'{i}')
323        if show_legend:
324            leg = fig.legend()#(loc='upper right', bbox_to_anchor=(0.9, 0.88), fontsize='small', title='Index')
325            leg.set_draggable(True)
326        plt.show()
327
328
329def splittings(
330        data:list,
331        title:str=None,
332        units:str='ueV',
333        rc:dict={},
334        ) -> None:
335    """Plot the tunnel splitting energies of a `data` list of systems.
336
337    The different `System.comment` are shown in the horizontal axis.
338    An optional `title` can be specified.
339    Default units shown are $\\mu$eV (`'ueV'`).
340    Available units are: `'ueV'`, `'meV'`, `'Ry'`, or `'B'` (free rotor units).
341
342    Additional matplotlib runtime configuration
343    [rcParams](https://matplotlib.org/stable/api/matplotlib_configuration_api.html#matplotlib.RcParams)
344    can be set with the `rc` dict.
345    """
346    title = title if title != None else 'Tunnel splitting energies'
347    calcs = deepcopy(data)
348    calcs = systems.as_list(calcs)
349
350    with plt.rc_context(rc):
351        fig, ax = plt.subplots(layout='constrained')
352        ax.set_ylabel("Energy / meV")
353
354        y = [c.splittings[0] for c in calcs]
355        x = [c.comment for c in calcs]
356        # What units do we want?
357        if units.lower() in alias.units['ueV']:
358            y = [j * 1000 for j in y]  # Convert meV to micro eV
359            ax.set_ylabel("Energy / $\\mu$eV")
360        elif units.lower() in alias.units['Ry']:
361            y = [j * constants.meV_to_Ry for j in y]
362            ax.set_ylabel("Energy / Ry")
363        elif units.upper() == 'B':
364            y = [j / c.B for j, c in zip(y, calcs)]
365            ax.set_ylabel("Energy / B")
366        #else:  # It's okay let's use meV
367
368        ax.bar(range(len(y)), y)
369        for i, comment in enumerate(x):
370            ax.text(x=i, y=0, s=comment+' ', rotation=45, verticalalignment='top', horizontalalignment='right')
371        ax.set_xlabel("")
372        ax.set_title(title)
373        ax.set_xticks([])
374        fig.tight_layout()
375        plt.show()
376
377
378def convergence(data:list, rc:dict={}) -> None:
379    """Plot the energy convergence of a `data` list of Systems as a function of the gridsize.
380
381    Additional matplotlib runtime configuration
382    [rcParams](https://matplotlib.org/stable/api/matplotlib_configuration_api.html#matplotlib.RcParams)
383    can be set with the `rc` dict.
384    """
385    systems.as_list(data)
386    gridsizes = [system.gridsize for system in data]
387    runtimes = [system.runtime for system in data]
388    deviations = []  # List of lists, containing all eigenvalue deviations for every system
389    searched_E = data[0].searched_E
390    for system in data:
391        deviation_list = []
392        for i, eigenvalue in enumerate(system.eigenvalues):
393            ideal_E = systems.calculate_ideal_E(i)
394            deviation = abs(ideal_E - eigenvalue)
395            deviation_list.append(deviation)
396        deviation_list = deviation_list[1:]  # Remove ground state
397        deviations.append(deviation_list)
398    # Plotting
399    with plt.rc_context(rc):
400        fig, ax1 = plt.subplots(layout='constrained')
401        ax1.set_xlabel('Grid size')
402        ax1.set_ylabel('Error / meV')
403        ax1.set_xscale('log')
404        ax1.set_yscale('log')
405        ax2 = ax1.twinx()
406        ax2.set_ylabel('Runtime / s')
407        ax2.set_yscale('log')
408        ax2.plot(gridsizes, runtimes, color='tab:grey', label='Runtime', linestyle='--')
409        colors = plt.cm.viridis(np.linspace(0, 1, searched_E))  # Should be searched_E-1 but we want to avoid lighter colors
410        for i in range(searched_E-1):
411            if i % 2 == 0:  # Ignore even numbers, since those levels are degenerated.
412                continue
413            ax1.plot(gridsizes, [dev[i] for dev in deviations], label=f'$E_{{{int((i+1)/2)}}}$', color=colors[i])
414        leg = fig.legend()#(loc='upper right', bbox_to_anchor=(0.9, 0.88), fontsize='small')
415        leg.set_draggable(True)
416        plt.title(data[0].comment if data[0].comment else 'Energy convergence vs grid size')
417        plt.show()