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)`. In meV."""
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, gridsize:int=None, B: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 gridsize:
126            self.gridsize = gridsize
127        if B:
128            self.B = B
129        return energies(self)
130
131    def solve_potential(self, gridsize:int=None):
132        """Default user method to quickly solve the potential of the quantum system.
133
134        This method does not solve the energies of the system,
135        it just computes the potential and sets `System.potential_max`,
136        `System.potential_min` and `System.potential_offset` accordingly.
137        To solve the potential AND the energies, check `System.solve()`.
138
139        The potential can be interpolated to a new `gridsize`.
140
141        Same as running `qrotor.solve.potential(System)`
142        with an optional new gridsize.
143        """
144        from .solve import potential
145        if gridsize:
146            self.gridsize = gridsize
147        return potential(self)
148
149    def change_phase(self, phase:float, calculate:bool=True):
150        """Apply a phase shift to the grid and potential values.
151
152        The `phase` should be a multiple of $\\pi$ (e.g., 3/2 for $3\\pi/2$).
153        The resulting grid will be expressed between $-2\\pi$ and $2\\pi$.
154
155        The System is solved immediately after the phase change.
156        This last step ensures that all eigenvalues and wavefunctions are correct.
157        You can override this step with `calculate = False`,
158        but remember to solve the System later!
159        """
160        if not any(self.potential_values) or not any(self.grid):
161            raise ValueError("System.potential_values and System.grid must be set before applying a phase shift.")
162        # Normalise the phase between 0 and 2
163        if abs(phase) >= 2:
164            phase = phase % 2
165        while phase < 0:
166            phase = phase + 2
167        # Shift the grid, between -2pi and 2pi
168        self.grid = (self.grid + (phase * np.pi))
169        # Apply the phase shift to potential values
170        phase_points = int((phase / 2) * self.gridsize)
171        self.potential_values = np.roll(self.potential_values, phase_points)
172        # Check that the grid is still within -2pi and 2pi, otherwise normalise it for a final time
173        while self.grid[0] <= (-2 * np.pi + 0.1):  # With a small tolerance
174            self.grid = self.grid + 2 * np.pi
175        while self.grid[-1] >= 2.5 * np.pi:  # It was not a problem until reaching 5/2 pi
176            self.grid = self.grid -2 * np.pi
177        print(f'Potential shifted by {phase}π')
178        if calculate:
179            self.solve()
180        return self
181
182    def set_grid(self, gridsize:int=None):
183        """Sets the `System.grid` to the specified `gridsize` from 0 to $2\\pi$.
184
185        If the system had a previous grid and potential values,
186        it will interpolate those values to the new gridsize,
187        using `qrotor.potential.interpolate()`.
188        """
189        if gridsize == self.gridsize:
190            return self  # Nothing to do here
191        if gridsize:
192            self.gridsize = gridsize
193        # Should we interpolate?
194        if any(self.potential_values) and any(self.grid) and self.gridsize:
195            from .potential import interpolate
196            self = interpolate(self)
197        # Should we create the values from zero?
198        elif self.gridsize:
199                self.grid = np.linspace(0, 2*np.pi, self.gridsize)
200        else:
201            raise ValueError('gridsize must be provided if there is no System.gridsize')
202        return self
203
204    def reduce_size(self):
205        """Discard data that takes too much space,
206        like eigenvectors, potential values and grids."""
207        self.eigenvectors = []
208        self.potential_values = []
209        self.grid = []
210        return self
211
212    def summary(self):
213        """Returns a dict with a summary of the System data."""
214        return {
215            'version': self.version,
216            'comment': self.comment,
217            'tags': self.tags,
218            'searched_E': self.searched_E,
219            'correct_potential_offset': self.correct_potential_offset,
220            'save_eigenvectors': self.save_eigenvectors,
221            'B': self.B,
222            'gridsize': self.gridsize,
223            'potential_name': self.potential_name,
224            'potential_constants': self.potential_constants.tolist() if isinstance(self.potential_constants, np.ndarray) else self.potential_constants,
225            'potential_offset': self.potential_offset,
226            'potential_min': self.potential_min,
227            'potential_max': self.potential_max,
228            'eigenvalues': self.eigenvalues.tolist() if isinstance(self.eigenvalues, np.ndarray) else self.eigenvalues,
229            'E_levels': self.E_levels,
230            'deg': self.deg,
231            'excitations': self.excitations,
232            'splittings': self.splittings,
233            'E_activation': self.E_activation,
234            'runtime': self.runtime,
235        }