1"""
  2# Description
  3
  4This module is used to solve any given quantum `qrotor.system.System()`.
  5
  6Although the functions of this module can be used independently,
  7it is highly recommended to use the methods `System.solve()` or
  8`System.solve_potential()` instead to solve the whole quantum system
  9or just the potential values.
 10These user methods perform all calculations automatically,
 11see `qrotor.system.System.solve()` and
 12`qrotor.system.System.solve_potential()` respectively for more details.
 13
 14This documentation page is left for reference and advanced users only.
 15
 16
 17# Index
 18
 19| | |
 20| --- | --- |
 21| `energies()`              | Solve the quantum system, including eigenvalues and eigenvectors |
 22| `potential()`             | Solve the potential values of the system |
 23| `schrodinger()`           | Solve the Schrödiger equation for the system |
 24| `hamiltonian_matrix()`    | Calculate the hamiltonian matrix of the system |
 25| `laplacian_matrix()`      | Calculate the second derivative matrix for a given grid |
 26| `excitations()`           | Get excitation levels and tunnel splitting energies |
 27| `E_levels`                | Group a list of degenerated eigenvalues by energy levels |
 28
 29---
 30"""
 31
 32
 33from .system import System
 34from .potential import solve as solve_potential
 35from .potential import interpolate
 36import time
 37import numpy as np
 38from scipy import sparse
 39import aton
 40from ._version import __version__
 41
 42
 43def energies(system:System, filename:str=None) -> System:
 44    """Solves the quantum `system`.
 45
 46    This includes solving the potential, the eigenvalues and the eigenvectors.
 47
 48    The resulting System object is saved with pickle to `filename` if specified.
 49    """
 50    system = potential(system)
 51    system = schrodinger(system)
 52    if filename:
 53        aton.file.save(system, filename)
 54    return system
 55
 56
 57def potential(system:System, gridsize:int=None) -> System:
 58    """Solves the potential values of the `system`.
 59
 60    Creates a grid if not yet present.
 61    It also interpolates the potential if `system.gridsize` is larger than the current grid;
 62    optionally, an alternative `gridsize` can be specified.
 63
 64    It then solves the potential according to the potential name.
 65    Then it applies extra operations, such as removing the potential offset
 66    if `system.correct_potential_offset = True`.
 67    """
 68    if gridsize:
 69        system.gridsize = gridsize
 70    if not any(system.grid):
 71        system.set_grid()
 72    if system.gridsize and any(system.grid):
 73        if system.gridsize > len(system.grid):
 74            system = interpolate(system)
 75    V = solve_potential(system)
 76    if system.correct_potential_offset is True:
 77        offset = min(V)
 78        V = V - offset
 79        system.potential_offset = offset
 80    system.potential_max = max(V)
 81    system.potential_min = min(V)
 82    system.potential_values = V
 83    return system
 84
 85
 86def schrodinger(system:System) -> System:
 87    """Solves the Schrödinger equation for a given `system`.
 88    
 89    Uses ARPACK in shift-inverse mode to solve the hamiltonian sparse matrix.
 90    """
 91    time_start = time.time()
 92    V = system.potential_values
 93    H = hamiltonian_matrix(system)
 94    print('Solving Schrodinger equation...')
 95    # Solve eigenvalues with ARPACK in shift-inverse mode, with a sparse matrix
 96    eigenvalues, eigenvectors = sparse.linalg.eigsh(H, system.searched_E, which='LM', sigma=0, maxiter=10000)
 97    if any(eigenvalues) is None:
 98        print('WARNING:  Not all eigenvalues were found.\n')
 99    else: print('Done.')
100    system.version = __version__
101    system.runtime = time.time() - time_start
102    system.eigenvalues = eigenvalues
103    system.E_activation = max(V) - min(eigenvalues)
104    # Solve excitations and tunnel splittings, assuming triplet degeneracy
105    system = excitations(system)
106    # Do we really need to save eigenvectors?
107    if system.save_eigenvectors == True:
108        system.eigenvectors = np.transpose(eigenvectors)
109    # Save potential max and min, in case these are not already saved
110    system.potential_max = max(V)
111    system.potential_min = min(V)
112    return system
113
114
115def hamiltonian_matrix(system:System):
116    """Calculates the Hamiltonian sparse matrix for a given `system`."""
117    print(f'Creating Hamiltonian sparse matrix of size {system.gridsize}...')
118    V = system.potential_values.tolist()
119    potential = sparse.diags(V, format='lil')
120    B = system.B
121    x = system.grid
122    H = -B * laplacian_matrix(x) + potential
123    return H
124
125
126def laplacian_matrix(grid):
127    """Calculates the Laplacian (second derivative) matrix for a given `grid`."""
128    x = grid
129    n = len(x)
130    diagonals = [-2*np.ones(n), np.ones(n), np.ones(n)]
131    laplacian_matrix = sparse.spdiags(diagonals, [0, -1, 1], m=n, n=n, format='lil')
132    # Periodic boundary conditions
133    laplacian_matrix[0, -1] = 1
134    laplacian_matrix[-1, 0] = 1
135    dx = x[1] - x[0]
136    laplacian_matrix /= dx**2
137    return laplacian_matrix
138
139
140def excitations(system: System) -> System:
141    """Calculate the excitation levels and the tunnel splitting energies of a system.
142
143    Automatically detects degenerated energy levels by looking at significant jumps
144    between consecutive eigenvalues. Within each level, finds two subgroups
145    to calculate tunnel splittings. Stops when energies reach the maximum potential.
146
147    Excitations are calculated as the energy difference between the mean energy of the
148    ground state level and the mean energy of each excited level.
149
150    Tunnel splittings are calculated as the difference between the mean values of
151    the two subgroups within each degenerate level.
152    """
153    # Get eigenvalues, stop before any possible None value
154    eigenvalues = system.eigenvalues
155    if not isinstance(eigenvalues, (list, np.ndarray)) or len(eigenvalues) == 0:
156        return system
157    if None in eigenvalues:
158        none_index = eigenvalues.tolist().index(None)
159        eigenvalues = eigenvalues[:none_index]
160    if len(eigenvalues) < 3:
161        return system
162    # Group degenerated eigenvalues into energy levels
163    levels, degeneracy = E_levels(eigenvalues, system.potential_max)
164    system.E_levels = levels
165    system.deg = degeneracy
166    # Calculate excitations and splittings
167    ground_energy = np.mean(levels[0])  # Mean of ground state level
168    excitations = []
169    tunnel_splittings = []
170    for level in levels:
171        level_mean = np.mean(level)
172        excitations.append(level_mean - ground_energy)
173        # Get the tunnel splitting within the level
174        if len(level) > 1:
175            # Find the largest gap within the level to split into two subgroups
176            internal_gaps = np.diff(level)
177            split_idx = np.argmax(internal_gaps) + 1
178            # Split into two subgroups
179            subgroup1 = level[:split_idx]
180            subgroup2 = level[split_idx:]
181            # Medians of subgroups
182            median1 = np.median(subgroup1)
183            median2 = np.median(subgroup2)
184            # Tunnel splitting is the difference between medians
185            tunnel_splittings.append(abs(median2 - median1))
186        else:
187            tunnel_splittings.append(0)
188    system.excitations = excitations[1:]  # Exclude ground state
189    system.splittings = tunnel_splittings
190    return system
191
192
193def E_levels(eigenvalues, vmax:float=None) -> list:
194    """Group a list of degenerated eigenvalues by energy levels.
195
196    Automatically detects degenerated energy levels by
197    looking at significant jumps between consecutive eigenvalues.
198
199    An optional `vmax` can be specified,
200    to avoid including too many eigenvalues
201    above a certain potential maximum.
202    Only two more eigenvalues are considered after `vmax`,
203    to properly detect energy levels around the maximum.
204
205    Example:
206    ```python
207    levels, deg = qr.solve.E_levels(array([1.1, 1.2, 1.3, 5.4, 5.5, 5.6]))
208    levels  # [array([1.1, 1.2, 1.3]), array([5.4, 5.5, 5.6])]
209    deg  # 3
210    ```
211    """
212    if vmax:  # Include all eigenvalues below Vmax plus 3 more eigenvalues
213        # Check if any values are above vmax
214        eigenvalues_above_vmax = eigenvalues > vmax
215        if np.any(eigenvalues_above_vmax):
216            index_first_above_vmax = np.where(eigenvalues_above_vmax)[0][0]
217            eigenvalues = eigenvalues[:(index_first_above_vmax + 2)]
218    # Group degenerated eigenvalues into energy levels
219    for scale in np.arange(2, 4, 0.25):  # First search going to bigger scales
220        levels, degeneracy = _get_E_levels_by_gap(eigenvalues, scale)
221        if (degeneracy > 1) and (degeneracy % 1 == 0):
222            break
223        else:
224            levels, degeneracy = None, None
225    if not degeneracy:  # If it didn't work, search with tighter values
226        for scale in np.arange(0.75, 2, 0.25):
227            levels, degeneracy = _get_E_levels_by_gap(eigenvalues, scale)
228            if (degeneracy > 1) and (degeneracy % 1 == 0):
229                break
230    if not (degeneracy > 1) or not (degeneracy % 1) == 0:
231        return levels, degeneracy  # I give up
232    # Correct the last two levels
233    if len(levels) >= 2 and len(levels[-2]) != degeneracy:
234        levels[-2] = np.concatenate((levels[-2], levels[-1]))
235        levels.pop(-1)
236    # Split last level into groups of size = degeneracy
237    last_level = levels[-1]
238    additional_levels = len(last_level) // degeneracy
239    if additional_levels > 0:
240        # Replace last level with list of complete degeneracy groups
241        complete_groups = [last_level[i:i+degeneracy] for i in range(0, additional_levels*degeneracy, degeneracy)]
242        levels.pop(-1)  # Remove original last level
243        levels.extend(complete_groups)  # Add all complete groups
244    else:
245        levels.pop(-1)  # Remove incomplete last level
246    return levels, degeneracy
247
248
249def _get_E_levels_by_gap(eigenvalues, scale:float=2) -> tuple:
250    """Split a list of eigenvalues into energy levels by looking at gaps.
251    
252    If the gap is bigger than the average gap times `scale`, it is considered a new level.
253
254    Returns a tuple with the estimated levels and the average degeneracy.
255    The last two levels are not taken into account to estimate the degeneracy.
256    """
257    # Find gaps between consecutive eigenvalues
258    gaps = np.diff(eigenvalues)
259    # Use mean gap times scale as threshold to distinguish energy levels
260    med_gap = np.mean(gaps)
261    level_breaks = np.where(gaps > scale * med_gap)[0] + 1
262    levels = np.split(eigenvalues, level_breaks)
263    # Calculate average degeneracy excluding last two levels if possible
264    if len(levels) > 2:
265        avg_degeneracy = float(np.mean([len(level) for level in levels[:-2]]))
266    else:
267        avg_degeneracy = float(len(levels[0]))
268    if avg_degeneracy % 1 == 0:
269        avg_degeneracy = int(avg_degeneracy)
270    return levels, avg_degeneracy