1"""
  2# Description
  3
  4The `System` object contains all the information needed for a single QRotor calculation.
  5This class can be loaded directly as `qrotor.System()`.
  6
  7---
  8"""
  9
 10
 11import numpy as np
 12from .constants import *
 13from aton import alias
 14from ._version import __version__
 15
 16
 17class System:
 18    """Quantum system.
 19
 20    Contains all the data for a single QRotor calculation, with both inputs and outputs.
 21
 22    Energy units are in meV and angles are in radians, unless stated otherwise.
 23    """
 24    def __init__(
 25            self,
 26            comment: str = None,
 27            B: float = B_CH3,
 28            gridsize: int = 200000,
 29            searched_E: int = 21,
 30            correct_potential_offset: bool = True,
 31            save_eigenvectors: bool = True,
 32            potential_name: str = '',
 33            potential_constants: list = None,
 34            tags: str = '',
 35            ):
 36        """A new quantum system can be instantiated as `system = qrotor.System()`.
 37        This new system will contain the default values listed above.
 38        """
 39        ## Technical
 40        self.version = __version__
 41        """Version of the package used to generate the data."""
 42        self.comment: str = comment
 43        """Custom comment for the dataset."""
 44        self.searched_E: int = searched_E
 45        """Number of energy eigenvalues to be searched."""
 46        self.correct_potential_offset: bool = correct_potential_offset
 47        """Correct the potential offset as `V - min(V)` or not."""
 48        self.save_eigenvectors: bool = save_eigenvectors
 49        """Save or not the eigenvectors. Final file size will be bigger."""
 50        self.tags: str = tags
 51        """Custom tags separated by spaces, such as the molecular group, etc.
 52
 53        Can be used to filter between datasets.
 54        """
 55        ## Potential
 56        self.B: float = B
 57        """Kinetic rotational energy, as in $B=\\frac{\\hbar^2}{2I}$.
 58
 59        Defaults to the value for a methyl group.
 60        """
 61        self.gridsize: int = gridsize
 62        """Number of points in the grid."""
 63        self.grid = []
 64        """The grid with the points to be used in the calculation.
 65
 66        Can be set automatically over $2 \\pi$ with `System.set_grid()`.
 67        Units must be in radians.
 68        """
 69        self.potential_name: str = potential_name
 70        """Name of the desired potential: `'zero'`, `'titov2023'`, `'test'`...
 71
 72        If empty or unrecognised, the custom potential values inside `System.potential_values` will be used. 
 73        """
 74        self.potential_constants: list = potential_constants
 75        """List of constants to be used in the calculation of the potential energy, in the `qrotor.potential` module."""
 76        self.potential_values = []
 77        """Numpy [ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html) with the potential values for each point in the grid.
 78
 79        Can be calculated with a function available in the `qrotor.potential` module,
 80        or loaded externally with the `qrotor.potential.load()` function.
 81        Potential energy units must be in meV.
 82        """
 83        # Potential values determined upon solving
 84        self.potential_offset: float = None
 85        """`min(V)` before offset correction when `correct_potential_offset = True`"""
 86        self.potential_min: float = None
 87        """`min(V)`"""
 88        self.potential_max: float = None
 89        """`max(V)`"""
 90        # Energies determined upon solving
 91        self.eigenvectors = []
 92        """Eigenvectors, if `save_eigenvectors` is True. Beware of the file size."""
 93        self.eigenvalues = []
 94        """Calculated eigenvalues of the system. In meV."""
 95        self.E_levels: list = []
 96        """List of `eigenvalues` grouped by energy levels, found below `potential_max`."""
 97        self.deg: float = None
 98        """Estimated degeneracy of the `E_levels` found below `potential_max`."""
 99        self.E_activation: float = None
100        """Activation energy or energy barrier, from the ground torsional state to the top of the potential barrier, `max(V) - min(eigenvalues)`"""
101        self.excitations: list = []
102        """Torsional excitations, as the difference between each energy level with respect to the ground state.
103
104        Considers the means between degenerated eigenvalues for all energy levels below `potential_max`.
105        """
106        self.splittings: list = []
107        """Tunnel splitting energies, for every degenerated energy level.
108        
109        Calculated for all `E_levels` as the difference between
110        the mean of the eigenvalues from A and the mean of the eigenvalues from E,
111        see [R. M. Dimeo, American Journal of Physics 71, 885–893 (2003)](https://doi.org/10.1119/1.1538575).
112        """
113        self.runtime: float = None
114        """Time taken to solve the eigenvalues."""
115
116    def solve(self, new_gridsize:int=None):
117        """Default user method to solve the quantum system.
118
119        The potential can be interpolated to a `new_gridsize`.
120
121        Same as running `qrotor.solve.energies(System)`
122        with an optional new gridsize.
123        """
124        from .solve import energies
125        if new_gridsize:
126            self.gridsize = new_gridsize
127        return energies(self)
128
129    def solve_potential(self, new_gridsize:int=None):
130        """Default user method to quickly solve the potential of the quantum system.
131
132        This method does not solve the energies of the system,
133        it just computes the potential and sets `System.potential_max`,
134        `System.potential_min` and `System.potential_offset` accordingly.
135        To solve the potential AND the energies, check `System.solve()`.
136
137        The potential can be interpolated to a `new_gridsize`.
138
139        Same as running `qrotor.solve.potential(System)`
140        with an optional new gridsize.
141        """
142        from .solve import potential
143        if new_gridsize:
144            self.gridsize = new_gridsize
145        return potential(self)
146
147    def change_phase(self, phase:float, calculate:bool=True):
148        """Apply a phase shift to the grid and potential values.
149
150        The `phase` should be a multiple of $\\pi$ (e.g., 3/2 for $3\\pi/2$).
151        The resulting grid will be expressed between $-2\\pi$ and $2\\pi$.
152
153        The System is solved immediately after the phase change.
154        This last step ensures that all eigenvalues and wavefunctions are correct.
155        You can override this step with `calculate = False`,
156        but remember to solve the System later!
157        """
158        if not any(self.potential_values) or not any(self.grid):
159            raise ValueError("System.potential_values and System.grid must be set before applying a phase shift.")
160        # Normalise the phase between 0 and 2
161        if abs(phase) >= 2:
162            phase = phase % 2
163        while phase < 0:
164            phase = phase + 2
165        # Shift the grid, between -2pi and 2pi
166        self.grid = (self.grid + (phase * np.pi))
167        # Apply the phase shift to potential values
168        phase_points = int((phase / 2) * self.gridsize)
169        self.potential_values = np.roll(self.potential_values, phase_points)
170        # Check that the grid is still within -2pi and 2pi, otherwise normalise it for a final time
171        while self.grid[0] <= (-2 * np.pi + 0.1):  # With a small tolerance
172            self.grid = self.grid + 2 * np.pi
173        while self.grid[-1] >= 2.5 * np.pi:  # It was not a problem until reaching 5/2 pi
174            self.grid = self.grid -2 * np.pi
175        print(f'Potential shifted by {phase}π')
176        if calculate:
177            self.solve()
178        return self
179
180    def set_grid(self, gridsize:int=None):
181        """Sets the `System.grid` to the specified `gridsize` from 0 to $2\\pi$.
182
183        If the system had a previous grid and potential values,
184        it will interpolate those values to the new gridsize,
185        using `qrotor.potential.interpolate()`.
186        """
187        if gridsize == self.gridsize:
188            return self  # Nothing to do here
189        if gridsize:
190            self.gridsize = gridsize
191        # Should we interpolate?
192        if any(self.potential_values) and any(self.grid) and self.gridsize:
193            from .potential import interpolate
194            self = interpolate(self)
195        # Should we create the values from zero?
196        elif self.gridsize:
197                self.grid = np.linspace(0, 2*np.pi, self.gridsize)
198        else:
199            raise ValueError('gridsize must be provided if there is no System.gridsize')
200        return self
201
202    def reduce_size(self):
203        """Discard data that takes too much space,
204        like eigenvectors, potential values and grids."""
205        self.eigenvectors = []
206        self.potential_values = []
207        self.grid = []
208        return self
209
210    def summary(self):
211        """Returns a dict with a summary of the System data."""
212        return {
213            'version': self.version,
214            'comment': self.comment,
215            'tags': self.tags,
216            'searched_E': self.searched_E,
217            'correct_potential_offset': self.correct_potential_offset,
218            'save_eigenvectors': self.save_eigenvectors,
219            'B': self.B,
220            'gridsize': self.gridsize,
221            'potential_name': self.potential_name,
222            'potential_constants': self.potential_constants.tolist() if isinstance(self.potential_constants, np.ndarray) else self.potential_constants,
223            'potential_offset': self.potential_offset,
224            'potential_min': self.potential_min,
225            'potential_max': self.potential_max,
226            'eigenvalues': self.eigenvalues.tolist() if isinstance(self.eigenvalues, np.ndarray) else self.eigenvalues,
227            'E_levels': self.E_levels,
228            'deg': self.deg,
229            'excitations': self.excitations,
230            'splittings': self.splittings,
231            'E_activation': self.E_activation,
232            'runtime': self.runtime,
233        }