qrotor 4.0.0a1__py3-none-any.whl

qrotor/solve.py

@@ -0,0 +1,271 @@
+"""
+# Description
+
+This module is used to solve any given quantum system.
+
+Although the functions of this module can be used independently,
+it is highly recommended to use the methods `System.solve()` or
+`System.solve_potential()` instead to solve the whole quantum system
+or just the potential values.
+These user methods perform all calculations automatically,
+see `qrotor.system.System.solve()` and
+`qrotor.system.System.solve_potential()` respectively for more details.
+
+This documentation page is left for reference and advanced users only.
+
+
+# Index
+
+| | |
+| --- | --- |
+| `energies()`              | Solve the quantum system, including eigenvalues and eigenvectors |
+| `potential()`             | Solve the potential values of the system |
+| `schrodinger()`           | Solve the Schrödiger equation for the system |
+| `hamiltonian_matrix()`    | Calculate the hamiltonian matrix of the system |
+| `laplacian_matrix()`      | Calculate the second derivative matrix for a given grid |
+| `excitations()`           | Get excitation levels and tunnel splitting energies |
+| `E_levels`                | Group a list of degenerated eigenvalues by energy levels |
+
+---
+"""
+
+
+from .system import System
+from .potential import solve as solve_potential
+from .potential import interpolate
+import time
+import numpy as np
+from scipy import sparse
+import aton
+from ._version import __version__
+
+
+def energies(system:System, filename:str=None) -> System:
+    """Solves the quantum `system`.
+
+    This includes solving the potential, the eigenvalues and the eigenvectors.
+
+    The resulting System object is saved with pickle to `filename` if specified.
+    """
+    system = potential(system)
+    system = schrodinger(system)
+    if filename:
+        aton.st.file.save(system, filename)
+    return system
+
+
+def potential(system:System, gridsize:int=None) -> System:
+    """Solves the potential values of the `system`.
+
+    Creates a grid if not yet present.
+    It also interpolates the potential if `system.gridsize` is larger than the current grid;
+    optionally, an alternative `gridsize` can be specified.
+
+    It then solves the potential according to the potential name.
+    Then it applies extra operations, such as removing the potential offset
+    if `system.correct_potential_offset = True`.
+    """
+    if gridsize:
+        system.gridsize = gridsize
+    if not any(system.grid):
+        system.set_grid()
+    if system.gridsize and any(system.grid):
+        if system.gridsize > len(system.grid):
+            system = interpolate(system)
+    V = solve_potential(system)
+    if system.correct_potential_offset is True:
+        offset = min(V)
+        V = V - offset
+        system.potential_offset = offset
+    system.potential_max = max(V)
+    system.potential_min = min(V)
+    system.potential_values = V
+    return system
+
+
+def schrodinger(system:System) -> System:
+    """Solves the Schrödinger equation for a given `system`.
+    
+    Uses ARPACK in shift-inverse mode to solve the hamiltonian sparse matrix.
+    """
+    time_start = time.time()
+    V = system.potential_values
+    H = hamiltonian_matrix(system)
+    print('Solving Schrodinger equation...')
+    # Solve eigenvalues with ARPACK in shift-inverse mode, with a sparse matrix
+    eigenvalues, eigenvectors = sparse.linalg.eigsh(H, system.searched_E, which='LM', sigma=0, maxiter=10000)
+    if any(eigenvalues) is None:
+        print('WARNING:  Not all eigenvalues were found.\n')
+    else: print('Done.')
+    system.version = __version__
+    system.runtime = time.time() - time_start
+    system.eigenvalues = eigenvalues
+    system.energy_barrier = max(V) - min(eigenvalues)
+    # Solve excitations and tunnel splittings, assuming triplet degeneracy
+    system = excitations(system)
+    # Do we really need to save eigenvectors?
+    if system.save_eigenvectors == True:
+        system.eigenvectors = np.transpose(eigenvectors)
+    # Save potential max and min, in case these are not already saved
+    system.potential_max = max(V)
+    system.potential_min = min(V)
+    return system
+
+
+def hamiltonian_matrix(system:System):
+    """Calculates the Hamiltonian sparse matrix for a given `system`."""
+    print(f'Creating Hamiltonian sparse matrix of size {system.gridsize}...')
+    V = system.potential_values.tolist()
+    potential = sparse.diags(V, format='lil')
+    B = system.B
+    x = system.grid
+    H = -B * laplacian_matrix(x) + potential
+    return H
+
+
+def laplacian_matrix(grid):
+    """Calculates the Laplacian (second derivative) matrix for a given `grid`."""
+    x = grid
+    n = len(x)
+    diagonals = [-2*np.ones(n), np.ones(n), np.ones(n)]
+    laplacian_matrix = sparse.spdiags(diagonals, [0, -1, 1], m=n, n=n, format='lil')
+    # Periodic boundary conditions
+    laplacian_matrix[0, -1] = 1
+    laplacian_matrix[-1, 0] = 1
+    dx = x[1] - x[0]
+    laplacian_matrix /= dx**2
+    return laplacian_matrix
+
+
+def excitations(system: System) -> System:
+    """Calculate the excitation levels and the tunnel splitting energies of a system.
+
+    Automatically detects degenerated energy levels by looking at significant jumps
+    between consecutive eigenvalues. Within each level, finds two subgroups
+    to calculate tunnel splittings. Stops when energies reach the maximum potential.
+
+    Excitations are calculated as the energy difference between the mean energy of the
+    ground state level and the mean energy of each excited level.
+
+    Tunnel splittings are calculated as the difference between the mean values of
+    the two subgroups within each degenerate level.
+    """
+    # Get eigenvalues, stop before any possible None value
+    eigenvalues = system.eigenvalues
+    if not isinstance(eigenvalues, (list, np.ndarray)) or len(eigenvalues) == 0:
+        return system
+    if None in eigenvalues:
+        none_index = eigenvalues.tolist().index(None)
+        eigenvalues = eigenvalues[:none_index]
+    if len(eigenvalues) < 3:
+        return system
+    # Group degenerated eigenvalues into energy levels
+    levels, degeneracy = E_levels(eigenvalues, system.potential_max)
+    system.E_levels = levels
+    system.deg = degeneracy
+    # Calculate excitations and splittings
+    ground_energy = np.mean(levels[0])  # Mean of ground state level
+    excitations = []
+    tunnel_splittings = []
+    for level in levels:
+        level_mean = np.mean(level)
+        excitations.append(level_mean - ground_energy)
+        # Get the tunnel splitting within the level
+        if len(level) > 1:
+            # Find the largest gap within the level to split into two subgroups
+            internal_gaps = np.diff(level)
+            split_idx = np.argmax(internal_gaps) + 1
+            # Split into two subgroups
+            subgroup1 = level[:split_idx]
+            subgroup2 = level[split_idx:]
+            # Medians of subgroups
+            median1 = np.median(subgroup1)
+            median2 = np.median(subgroup2)
+            # Tunnel splitting is the difference between medians
+            tunnel_splittings.append(abs(median2 - median1))
+        else:
+            tunnel_splittings.append(0)
+    system.excitations = excitations[1:]  # Exclude ground state
+    system.splittings = tunnel_splittings
+    return system
+
+
+def E_levels(eigenvalues, vmax:float=None) -> list:
+    """Group a list of degenerated eigenvalues by energy levels.
+
+    Automatically detects degenerated energy levels by
+    looking at significant jumps between consecutive eigenvalues.
+
+    An optional `vmax` can be specified,
+    to avoid including too many eigenvalues
+    above a certain potential maximum.
+    Only two more eigenvalues are considered after `vmax`,
+    to properly detect energy levels around the maximum.
+
+    Example:
+    ```python
+    levels, deg = qr.solve.E_levels(array([1.1, 1.2, 1.3, 5.4, 5.5, 5.6]))
+    levels  # [array([1.1, 1.2, 1.3]), array([5.4, 5.5, 5.6])]
+    deg  # 3
+    ```
+    """
+    if vmax:  # Include all eigenvalues below Vmax plus 3 more eigenvalues
+        # Check if any values are above vmax
+        eigenvalues_above_vmax = eigenvalues > vmax
+        if np.any(eigenvalues_above_vmax):
+            index_first_above_vmax = np.where(eigenvalues_above_vmax)[0][0]
+            eigenvalues = eigenvalues[:(index_first_above_vmax + 2)]
+    # Group degenerated eigenvalues into energy levels
+    for scale in np.arange(2, 4, 0.25):  # First search going to bigger scales
+        levels, degeneracy = _get_E_levels_by_gap(eigenvalues, scale)
+        if (degeneracy > 1) and (degeneracy % 1 == 0):
+            break
+        else:
+            levels, degeneracy = None, None
+    if not degeneracy:  # If it didn't work, search with tighter values
+        for scale in np.arange(0.75, 2, 0.25):
+            levels, degeneracy = _get_E_levels_by_gap(eigenvalues, scale)
+            if (degeneracy > 1) and (degeneracy % 1 == 0):
+                break
+    if not (degeneracy > 1) and not (degeneracy % 1) == 0:
+        return levels, degeneracy  # I give up
+    # Correct the last two levels
+    if len(levels) >= 2 and len(levels[-2]) != degeneracy:
+        levels[-2] = np.concatenate((levels[-2], levels[-1]))
+        levels.pop(-1)
+    # Split last level into groups of size = degeneracy
+    last_level = levels[-1]
+    additional_levels = len(last_level) // degeneracy
+    if additional_levels > 0:
+        # Replace last level with list of complete degeneracy groups
+        complete_groups = [last_level[i:i+degeneracy] for i in range(0, additional_levels*degeneracy, degeneracy)]
+        levels.pop(-1)  # Remove original last level
+        levels.extend(complete_groups)  # Add all complete groups
+    else:
+        levels.pop(-1)  # Remove incomplete last level
+    return levels, degeneracy
+
+
+def _get_E_levels_by_gap(eigenvalues, scale:float=2) -> tuple:
+    """Split a list of eigenvalues into energy levels by looking at gaps.
+    
+    If the gap is bigger than the average gap times `scale`, it is considered a new level.
+
+    Returns a tuple with the estimated levels and the average degeneracy.
+    The last two levels are not taken into account to estimate the degeneracy.
+    """
+    # Find gaps between consecutive eigenvalues
+    gaps = np.diff(eigenvalues)
+    # Use mean gap times scale as threshold to distinguish energy levels
+    med_gap = np.mean(gaps)
+    level_breaks = np.where(gaps > scale * med_gap)[0] + 1
+    levels = np.split(eigenvalues, level_breaks)
+    # Calculate average degeneracy excluding last two levels if possible
+    if len(levels) > 2:
+        avg_degeneracy = float(np.mean([len(level) for level in levels[:-2]]))
+    else:
+        avg_degeneracy = float(len(levels[0]))
+    if avg_degeneracy % 1 == 0:
+        avg_degeneracy = int(avg_degeneracy)
+    return levels, avg_degeneracy
+

qrotor/system.py

@@ -0,0 +1,275 @@
+"""
+# Description
+
+The `System` object contains all the information needed for a single QRotor calculation.
+This class can be loaded directly as `qrotor.System()`.
+
+---
+"""
+
+
+import numpy as np
+from .constants import *
+from aton import alias
+from ._version import __version__
+
+
+class System:
+    """Quantum system.
+
+    Contains all the data for a single QRotor calculation, with both inputs and outputs.
+
+    Energy units are in meV and angles are in radians, unless stated otherwise.
+    """
+    def __init__(
+            self,
+            comment: str = None,
+            searched_E: int = 21,
+            correct_potential_offset: bool = True,
+            save_eigenvectors: bool = True,
+            group: str = '',
+            B: float = B_CH3,
+            gridsize: int = 200000,
+            grid = [],
+            potential_name: str = '',
+            potential_constants: list = None,
+            potential_values = [],
+            ):
+        """A new quantum system can be instantiated as `system = qrotor.System()`.
+        This new system will contain the default values listed above.
+        """
+        ## Technical
+        self.version = __version__
+        """Version of the package used to generate the data."""
+        self.comment: str = comment
+        """Custom comment for the dataset."""
+        self.searched_E: int = searched_E
+        """Number of energy eigenvalues to be searched."""
+        self.correct_potential_offset: bool = correct_potential_offset
+        """Correct the potential offset as `V - min(V)` or not."""
+        self.save_eigenvectors: bool = save_eigenvectors
+        """Save or not the eigenvectors. Final file size will be bigger."""
+        self.group: str = group
+        """Chemical group, methyl or amine: `'CH3'`, `'CD3'`, `'NH3'`, `'ND3'`.
+
+        Can be used to set the value of `B` automatically at startup.
+        It can also be configured afterwards with `System.set_group()`.
+        This group can be used as metadata to analyse different datasets.
+        """
+        self.set_group(group)  # Normalise the group name, and set the value of B
+        ## Potential
+        if not B:
+            B = self.B
+        self.B: float = B
+        """Rotational inertia, as in $B=\\frac{\\hbar^2}{2I}$.
+
+        Defaults to the value for a methyl group.
+        """
+        self.gridsize: int = gridsize
+        """Number of points in the grid."""
+        self.grid = grid
+        """The grid with the points to be used in the calculation.
+
+        Can be set automatically over $2 \\pi$ with `System.set_grid()`.
+        Units must be in radians.
+        """
+        self.potential_name: str = potential_name
+        """Name of the desired potential: `'zero'`, `'titov2023'`, `'test'`...
+
+        If empty or unrecognised, the custom potential values inside `System.potential_values` will be used. 
+        """
+        self.potential_constants: list = potential_constants
+        """List of constants to be used in the calculation of the potential energy, in the `qrotor.potential` module."""
+        self.potential_values = potential_values
+        """Numpy [ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html) with the potential values for each point in the grid.
+
+        Can be calculated with a function available in the `qrotor.potential` module,
+        or loaded externally with the `qrotor.potential.load()` function.
+        Potential energy units must be in meV.
+        """
+        # Potential values determined upon solving
+        self.potential_offset: float = None
+        """`min(V)` before offset correction when `correct_potential_offset = True`"""
+        self.potential_min: float = None
+        """`min(V)`"""
+        self.potential_max: float = None
+        """`max(V)`"""
+        # Energies determined upon solving
+        self.eigenvectors = []
+        """Eigenvectors, if `save_eigenvectors` is True. Beware of the file size."""
+        self.eigenvalues = []
+        """Calculated eigenvalues of the system. In meV."""
+        self.E_levels: list = []
+        """List of `eigenvalues` grouped by energy levels, found below `potential_max`."""
+        self.deg: float = None
+        """Estimated degeneracy of the `E_levels` found below `potential_max`."""
+        self.excitations: list = []
+        """Torsional excitations, as the difference between each energy level with respect to the ground state.
+
+        Considers the means between degenerated eigenvalues for all energy levels below `potential_max`.
+        """
+        self.splittings: list = []
+        """Tunnel splitting energies, for every degenerated energy level.
+        
+        Calculated for all `E_levels` as the difference between
+        the mean of the eigenvalues from A and the mean of the eigenvalues from E,
+        see [R. M. Dimeo, American Journal of Physics 71, 885–893 (2003)](https://doi.org/10.1119/1.1538575).
+        """
+        self.energy_barrier: float = None
+        """Activation energy or energy barrier, from the ground torsional state to the top of the potential barrier, `max(V) - min(eigenvalues)`"""
+        self.runtime: float = None
+        """Time taken to solve the eigenvalues."""
+
+    def solve(self, new_gridsize:int=None):
+        """Default user method to solve the quantum system.
+
+        The potential can be interpolated to a `new_gridsize`.
+
+        Same as running `qrotor.solve.energies(System)`
+        with an optional new gridsize.
+        """
+        from .solve import energies
+        if new_gridsize:
+            self.gridsize = new_gridsize
+        return energies(self)
+
+    def solve_potential(self, new_gridsize:int=None):
+        """Default user method to quickly solve the potential of the quantum system.
+
+        This method does not solve the energies of the system,
+        it just computes the potential and sets `System.potential_max`,
+        `System.potential_min` and `System.potential_offset` accordingly.
+        To solve the potential AND the energies, check `System.solve()`.
+
+        The potential can be interpolated to a `new_gridsize`.
+
+        Same as running `qrotor.solve.potential(System)`
+        with an optional new gridsize.
+        """
+        from .solve import potential
+        if new_gridsize:
+            self.gridsize = new_gridsize
+        return potential(self)
+
+    def change_phase(self, phase:float, calculate:bool=True):
+        """Apply a phase shift to the grid and potential values.
+
+        The `phase` should be a multiple of $\\pi$ (e.g., 3/2 for $3\\pi/2$).
+        The resulting grid will be expressed between $-2\\pi$ and $2\\pi$.
+
+        The System is solved immediately after the phase change.
+        This last step ensures that all eigenvalues and wavefunctions are correct.
+        You can override this step with `calculate = False`,
+        but remember to solve the System later!
+        """
+        if not any(self.potential_values) or not any(self.grid):
+            raise ValueError("System.potential_values and System.grid must be set before applying a phase shift.")
+        # Normalise the phase between 0 and 2
+        if abs(phase) >= 2:
+            phase = phase % 2
+        while phase < 0:
+            phase = phase + 2
+        # Shift the grid, between -2pi and 2pi
+        self.grid = (self.grid + (phase * np.pi))
+        # Apply the phase shift to potential values
+        phase_points = int((phase / 2) * self.gridsize)
+        self.potential_values = np.roll(self.potential_values, phase_points)
+        # Check that the grid is still within -2pi and 2pi, otherwise normalise it for a final time
+        while self.grid[0] <= (-2 * np.pi + 0.1):  # With a small tolerance
+            self.grid = self.grid + 2 * np.pi
+        while self.grid[-1] >= 2.5 * np.pi:  # It was not a problem until reaching 5/2 pi
+            self.grid = self.grid -2 * np.pi
+        print(f'Potential shifted by {phase}π')
+        if calculate:
+            self.solve()
+        return self
+
+    def set_grid(self, gridsize:int=None):
+        """Sets the `System.grid` to the specified `gridsize` from 0 to $2\\pi$.
+
+        If the system had a previous grid and potential values,
+        it will interpolate those values to the new gridsize,
+        using `qrotor.potential.interpolate()`.
+        """
+        if gridsize == self.gridsize:
+            return self  # Nothing to do here
+        if gridsize:
+            self.gridsize = gridsize
+        # Should we interpolate?
+        if any(self.potential_values) and any(self.grid) and self.gridsize:
+            from .potential import interpolate
+            self = interpolate(self)
+        # Should we create the values from zero?
+        elif self.gridsize:
+                self.grid = np.linspace(0, 2*np.pi, self.gridsize)
+        else:
+            raise ValueError('gridsize must be provided if there is no System.gridsize')
+        return self
+    
+    def set_group(self, group:str=None, B:float=None):
+        """Normalise `System.group` name, and set `System.B` based on it."""
+        for name in alias.chemical['CH3']:
+            if group.lower() == name:
+                self.group = 'CH3'
+                if not B:
+                    B = B_CH3
+                self.B = B
+                return self
+        for name in alias.chemical['CD3']:
+            if group.lower() == name:
+                self.group = 'CD3'
+                if not B:
+                    B = B_CD3
+                self.B = B
+                return self
+        for name in alias.chemical['NH3']:
+            if group.lower() == name:
+                self.group = 'NH3'
+                if not B:
+                    B = B_NH3
+                self.B = B
+                return self
+        for name in alias.chemical['ND3']:
+            if group.lower() == name:
+                self.group = 'ND3'
+                if not B:
+                    B = B_ND3
+                self.B = B
+                return self
+        self.group = group  # No match was found
+        self.B = None
+        return self
+
+    def reduce_size(self):
+        """Discard data that takes too much space,
+        like eigenvectors, potential values and grids."""
+        self.eigenvectors = []
+        self.potential_values = []
+        self.grid = []
+        return self
+
+    def summary(self):
+        """Returns a dict with a summary of the System data."""
+        return {
+            'version': self.version,
+            'comment': self.comment,
+            'searched_E': self.searched_E,
+            'correct_potential_offset': self.correct_potential_offset,
+            'save_eigenvectors': self.save_eigenvectors,
+            'group': self.group,
+            'B': self.B,
+            'gridsize': self.gridsize,
+            'potential_name': self.potential_name,
+            'potential_constants': self.potential_constants.tolist() if isinstance(self.potential_constants, np.ndarray) else self.potential_constants,
+            'potential_offset': self.potential_offset,
+            'potential_min': self.potential_min,
+            'potential_max': self.potential_max,
+            'eigenvalues': self.eigenvalues.tolist() if isinstance(self.eigenvalues, np.ndarray) else self.eigenvalues,
+            'E_levels': self.E_levels,
+            'deg': self.deg,
+            'excitations': self.excitations,
+            'splittings': self.splittings,
+            'energy_barrier': self.energy_barrier,
+            'runtime': self.runtime,
+        }
+