1"""
  2# Description
  3
  4This module contains functions to calculate the actual `potential_values` of the system.
  5
  6
  7# Index
  8
  9User functions:
 10
 11| | |
 12| --- | --- |
 13| `save()`        | Save the potential from a System to a data file |
 14| `load()`        | Load a System with a custom potential from a potential data file |
 15| `from_qe()`     | Creates a potential data file from Quantum ESPRESSO outputs |
 16| `merge()`       | Add and subtract potentials from systems |
 17| `scale()`       | Scale potential values by a given factor |
 18
 19To solve the system, optionally interpolating to a new gridsize, use the `System.solve(gridsize)` method.  
 20However, if you just want to quickly solve or interpolate the potential, check the `System.solve_potential(gridsize)` method.
 21This will run several checks before applying the following functions automatically:
 22
 23| | |
 24| --- | --- |
 25| `interpolate()` | Interpolates the current `System.potential_values` to a new `System.gridsize` |
 26| `solve()`       | Solve the potential values based on the potential name |
 27
 28A synthetic potential can be created by specifying its name in `System.potential_name`,
 29along with the corresponding `System.potential_constants` if required.
 30Available potentials are:
 31
 32| | |
 33| --- | --- |
 34| `zero()`        | Zero potential |
 35| `sine()`        | Sine potential |
 36| `cosine()`      | Cosine potential |
 37| `titov2023()`   | Potential of the hindered methyl rotor, as in titov2023. |
 38
 39---
 40"""
 41
 42
 43from .system import System
 44from . import constants
 45from . import systems
 46import numpy as np
 47import os
 48from copy import deepcopy
 49from scipy.interpolate import CubicSpline
 50import aton.alias as alias
 51import aton.file as file
 52import aton.api.pwx as pwx
 53from ._version import __version__
 54
 55
 56def save(
 57        system:System,
 58        comment:str='',
 59        filepath:str='potential.csv',
 60        angle:str='deg',
 61        energy:str='meV',
 62        ) -> None:
 63    """Save the rotational potential from a `system` to a CSV file.
 64
 65    The output `filepath` contains angle and energy columns,
 66    in degrees and meVs by default.
 67    The units can be changed with `angle` and `energy`,
 68    but only change these defaults if you know what you are doing.
 69    An optional `comment` can be included in the header of the file.
 70    """
 71    print('Saving potential data file...')
 72    # Check if a previous potential.dat file exists, and ask to overwrite it
 73    previous_potential_file = file.get(filepath, return_anyway=True)
 74    if previous_potential_file:
 75        print(f"WARNING: Previous '{filepath}' file will be overwritten, proceed anyway?")
 76        answer = input("(y/n): ")
 77        if not answer.lower() in alias.boolean[True]:
 78            print("Aborted.")
 79            return None
 80    # Set header
 81    potential_data = f'## {comment}\n' if comment else f'## {system.comment}\n' if system.comment else ''
 82    potential_data += '# Rotational potential dataset\n'
 83    potential_data += f'# Saved with QRotor {__version__}\n'
 84    potential_data += '# https://pablogila.github.io/qrotor\n'
 85    potential_data += '#\n'
 86    # Check that grid and potential values are the same size
 87    if len(system.grid) != len(system.potential_values):
 88        raise ValueError('len(system.grid) != len(system.potential_values)')
 89    grid = system.grid
 90    potential_values = system.potential_values
 91    # Convert angle units
 92    if angle.lower() in alias.units['rad']:
 93        potential_data += '# Angle/rad,    '
 94    else:
 95        grid = np.degrees(grid)
 96        potential_data += '# Angle/deg,    '
 97        if not angle.lower() in alias.units['deg']:
 98            print(f"WARNING: Unrecognised '{angle}' angle units, using degrees instead")
 99    # Convert energy units
100    if energy.lower() in alias.units['meV']:
101        potential_data += 'Potential/meV\n'
102    elif energy.lower() in alias.units['eV']:
103        potential_values = potential_values * 1e-3
104        potential_data += 'Potential/eV\n'
105    elif energy.lower() in alias.units['Ry']:
106        potential_values = potential_values * constants.meV_to_Ry
107        potential_data += 'Potential/Ry\n'
108    else:
109        print(f"WARNING:  Unrecognised '{energy}' energy units, using meV instead")
110        potential_data += 'Potential/meV\n'
111    potential_data += '#\n'
112    # Save all values
113    for angle_value, energy_value in zip(grid, potential_values):
114        potential_data += f'{angle_value},    {energy_value}\n'
115    with open(filepath, 'w') as f:
116        f.write(potential_data)
117    print(f'Saved to {filepath}')
118    # Warn the user if not in default units
119    if angle.lower() not in alias.units['deg']:
120        print(f"WARNING: You saved the potential in '{angle}' angle units! Remember that QRotor works in degrees!")
121    if energy.lower() not in alias.units['meV']:
122        print(f"WARNING: You saved the potential in '{energy}' energy units! Remember that QRotor works in meVs!")
123
124
125def load(
126        filepath:str='potential.csv',
127        comment:str=None,
128        tags:str='',
129        system:System=None,
130        angle:str='deg',
131        energy:str='meV',
132        ) -> System:
133    """Read a rotational potential energy datafile.
134
135    The input file in `filepath` should contain two columns with angle and potential energy values.
136    Degrees and meV are assumed as default units unless stated in `angle` and `energy`.
137    Units will be converted automatically to radians and meV.
138
139    An optional `comment` can be included in the output System.
140    Set to the parent folder name by default.
141
142    A previous System object can be provided through `system` to update its potential values.
143    """
144    file_path = file.get(filepath)
145    system = System() if system is None else system
146    with open(file_path, 'r') as f:
147        lines = f.readlines()
148    # Read the comment
149    loaded_comment = ''
150    if lines[0].startswith('## '):
151        loaded_comment = lines[0][3:].strip()
152    # Read data
153    positions = []
154    potentials = []
155    for line in lines:
156        if line.startswith('#'):
157            continue
158        position, potential = line.split()
159        positions.append(float(position.strip().strip(',').strip()))
160        potentials.append(float(potential.strip()))
161    # Save angles to numpy arrays
162    if angle.lower() in alias.units['deg']:
163        positions = np.radians(positions)
164    elif angle.lower() in alias.units['rad']:
165        positions = np.array(positions)
166    else:
167        raise ValueError(f"Angle unit '{angle}' not recognized.")
168    # Save energies to numpy arrays
169    if energy.lower() in alias.units['eV']:
170        potentials = np.array(potentials) * 1000
171    elif energy.lower() in alias.units['meV']:
172        potentials = np.array(potentials)
173    elif energy.lower() in alias.units['Ry']:
174        potentials = np.array(potentials) * constants.Ry_to_meV
175    else:
176        raise ValueError(f"Energy unit '{energy}' not recognized.")
177    # Set the system
178    system.grid = np.array(positions)
179    system.gridsize = len(positions)
180    system.potential_values = np.array(potentials)
181    # System comment as the loaded comment or the parent folder name
182    if comment:
183        system.comment = comment
184    elif loaded_comment:
185        system.comment = loaded_comment
186    else:
187        system.comment = os.path.basename(os.path.dirname(file_path))
188    if tags:
189        system.tags = tags
190    print(f"Loaded {filepath}")
191    return system
192
193
194def from_qe(
195        folder=None,
196        filepath:str='potential.csv',
197        include:list=['.out'],
198        exclude:list=['slurm-'],
199        energy:str='meV',
200        comment:str=None,
201        ) -> System:
202    """Compiles a rotational potential CSV file from Quantum ESPRESSO pw.x outputs,
203    created with `qrotor.rotation.rotate_qe()`.
204    Returns a `System` object with the new potential values.
205
206    The angle in degrees is extracted from the output filenames,
207    which must follow `whatever_ANGLE.out`.
208
209    Outputs from SCF calculations must be located in the provided `folder` (CWD if None).
210    Files can be filtered by those containing the specified `include` filters,
211    excluding those containing any string from the `exclude` list. 
212    The output `filepath` name is `'potential.dat'` by default.
213
214    Energy values are saved to meV by dafault, unless specified in `energy`.
215    Only change the energy units if you know what you are doing;
216    remember that default energy units in QRotor are meV!
217    """
218    folder = file.get_dir(folder)
219    # Check if a previous potential.dat file exists, and ask to overwrite it
220    previous_potential_file = file.get(filepath, return_anyway=True)
221    if previous_potential_file:
222        print(f"WARNING: Previous '{filepath}' file will be overwritten, proceed anyway?")
223        answer = input("(y/n): ")
224        if not answer.lower() in alias.boolean[True]:
225            print("Aborted.")
226            return None
227    # Get the files to read
228    files = file.get_list(folder=folder, include=include, exclude=exclude, abspath=True)
229    folder_name = os.path.basename(folder)
230    # Set header
231    potential_data = f'## {comment}\n' if comment else f'## {folder_name}\n'
232    potential_data += '# Rotational potential dataset\n'
233    potential_data += f'# Calculated with QE pw.x using QRotor {__version__}\n'
234    potential_data += '# https://pablogila.github.io/qrotor\n'
235    potential_data += '#\n'
236    if energy.lower() in alias.units['eV']:
237        potential_data += '# Angle/deg,    Potential/eV\n'
238    elif energy.lower() in alias.units['meV']:
239        potential_data += '# Angle/deg,    Potential/meV\n'
240    elif energy.lower() in alias.units['Ry']:
241        potential_data += '# Angle/deg,    Potential/Ry\n'
242    else:
243        potential_data += '# Angle/deg,    Potential/meV\n'
244    potential_data += '#\n'
245    potential_data_list = []
246    print('Extracting the potential as a function of the angle...')
247    print('----------------------------------')
248    counter_success = 0
249    counter_errors = 0
250    for file_path in files:
251        filename = os.path.basename(file_path)
252        file_path = file.get(filepath=file_path, include='.out', return_anyway=True)
253        if not file_path:  # Not an output file, skip it
254            continue
255        content = pwx.read_out(file_path)
256        if not content['Success']:  # Ignore unsuccessful calculations
257            print(f'x   {filename}')
258            counter_errors += 1
259            continue
260        if energy.lower() in alias.units['eV']:
261            energy_value = content['Energy'] * constants.Ry_to_eV
262        elif energy.lower() in alias.units['meV']:
263            energy_value = content['Energy'] * constants.Ry_to_meV
264        elif energy.lower() in alias.units['Ry']:
265            energy_value = content['Energy']
266        else:
267            print(f"WARNING: Energy unit '{energy}' not recognized, using meV instead.")
268            energy = 'meV'
269            energy_value = content['Energy'] * constants.Ry_to_meV
270        splits = filename.split('_')
271        angle_value = splits[-1].replace('.out', '')
272        angle_value = float(angle_value)
273        potential_data_list.append((angle_value, energy_value))
274        print(f'OK  {filename}')
275        counter_success += 1
276    # Sort by angle
277    potential_data_list_sorted = sorted(potential_data_list, key=lambda x: x[0])
278    # Append the sorted values as a string
279    for angle_value, energy_value in potential_data_list_sorted:
280        potential_data += f'{angle_value},    {energy_value}\n'
281    with open(filepath, 'w') as f:
282        f.write(potential_data)
283    print('----------------------------------')
284    print(f'Succesful calculations (OK): {counter_success}')
285    print(f'Faulty calculations     (x): {counter_errors}')
286    print('----------------------------------')
287    print(f'Saved angles and potential values at {filepath}')
288    # Warn the user if not in default units
289    if energy.lower() not in alias.units['meV']:
290        print(f"WARNING: You saved the potential in '{energy}' units! Remember that QRotor works in meVs!")
291    new_system = None
292    try:
293        new_system = load(filepath=filepath, comment=comment, energy=energy)
294    except:
295        pass
296    return new_system
297
298
299def merge(
300        add=[],
301        subtract=[],
302        comment:str=None
303        ) -> System:
304    """Add or subtract potentials from different systems.
305
306    Adds the potential values from the systems in `add`,
307    removes the ones from `subtract`.
308    All systems will be interpolated to the bigger gridsize if needed.
309
310    A copy of the first System will be returned with the resulting potential values,
311    with an optional `comment` if indicated.
312    """
313    add = systems.as_list(add)
314    subtract = systems.as_list(subtract)
315    gridsizes = systems.get_gridsizes(add)
316    gridsizes.extend(systems.get_gridsizes(subtract))
317    max_gridsize = max(gridsizes)
318    # All gridsizes should be max_gridsize
319    for s in add:
320        if s.gridsize != max_gridsize:
321            s.gridsize = max_gridsize
322            s = interpolate(s)
323    for s in subtract:
324        if s.gridsize != max_gridsize:
325            s.gridsize = max_gridsize
326            s = interpolate(s)
327
328    if len(add) == 0:
329        if len(subtract) == 0:
330            raise ValueError('No systems were provided!')
331        result = deepcopy(subtract[0])
332        result.potential_values = -result.potential_values
333        subtract.pop(0)
334    else:
335        result = deepcopy(add[0])
336        add.pop(0)
337
338    for system in add:
339        result.potential_values = np.sum([result.potential_values, system.potential_values], axis=0)
340    for system in subtract:
341        result.potential_values = np.sum([result.potential_values, -system.potential_values], axis=0)
342    if comment != None:
343        result.comment = comment
344    return result
345
346
347def scale(
348        system:System,
349        factor:float,
350        comment:str=None
351        ) -> System:
352    """Returns a copy of `system` with potential values scaled by a `factor`.
353
354    An optional `comment` can be included.
355    """
356    result = deepcopy(system)
357    if factor != 0:
358        result.potential_values = system.potential_values * factor
359    else:
360        result.potential_values = np.zeros(system.gridsize)
361    if comment != None:
362        result.comment = comment
363    return result
364
365
366def interpolate(system:System) -> System:
367    """Interpolates the current `System.potential_values`
368    to a new grid of size `System.gridsize`.
369
370    This basic function is called by `qrotor.solve.potential()`,
371    which is the recommended way to interpolate potentials.
372    """
373    print(f"Interpolating potential to a grid of size {system.gridsize}...")
374    V = system.potential_values
375    grid = system.grid
376    gridsize = system.gridsize
377    new_grid = np.linspace(0, 2*np.pi, gridsize)
378    # Impose periodic boundary conditions
379    grid_periodic = np.append(grid, grid[0] + 2*np.pi)
380    V_periodic = np.append(V, V[0])
381    cubic_spline = CubicSpline(grid_periodic, V_periodic, bc_type='periodic')
382    new_V = cubic_spline(new_grid)
383    system.grid = new_grid
384    system.potential_values = new_V
385    return system
386
387
388def solve(system:System):
389    """Solves `System.potential_values`
390    according to the `System.potential_name`,
391    returning the new `potential_values`.
392    Avaliable potential names are `zero`, `sine` and `titov2023`.
393
394    If `System.potential_name` is not present or not recognised,
395    the current `System.potential_values` are used.
396
397    This basic function is called by `qrotor.solve.potential()`,
398    which is the recommended way to solve potentials.
399    """
400    data = deepcopy(system)
401    # Is there a potential_name?
402    if not data.potential_name:
403        if data.potential_values is None or len(data.potential_values) == 0:
404            raise ValueError(f'No potential_name and no potential_values found in the system!')
405    elif data.potential_name.lower() == 'titov2023':
406        data.potential_values = titov2023(data)
407    elif data.potential_name.lower() in alias.math['0']:
408        data.potential_values = zero(data)
409    elif data.potential_name.lower() in alias.math['sin']:
410        data.potential_values = sine(data)
411    elif data.potential_name.lower() in alias.math['cos']:
412        data.potential_values = cosine(data)
413    # At least there should be potential_values
414    #elif not any(data.potential_values):
415    elif data.potential_values is None or len(data.potential_values) == 0:
416        raise ValueError("Unrecognised potential_name '{data.potential_name}' and no potential_values found")
417    return data.potential_values
418
419
420def zero(system:System):
421    """Zero potential.
422
423    $V(x) = 0$
424    """
425    x = system.grid
426    return 0 * np.array(x)
427
428
429def sine(system:System):
430    """Sine potential.
431
432    $V(x) = C_0 + \\frac{C_1}{2} sin(x C_2 + C_3)$  
433    With $C_0$ as the potential offset,
434    $C_1$ as the max potential value (without considering the offset),
435    $C_2$ as the frequency, and $C_3$ as the phase.
436    If no `System.potential_constants` are provided, defaults to $sin(3x)$  
437    """
438    x = system.grid
439    C = system.potential_constants
440    C0 = 0
441    C1 = 1
442    C2 = 3
443    C3 = 0
444    if C:
445        if len(C) > 0:
446            C0 = C[0]
447        if len(C) > 1:
448            C1 = C[1]
449        if len(C) > 2:
450            C2 = C[2]
451        if len(C) > 3:
452            C3 = C[3]
453    return C0 + (C1 / 2) * np.sin(np.array(x) * C2 + C3)
454
455
456def cosine(system:System):
457    """Cosine potential.
458
459    $V(x) = C_0 + \\frac{C_1}{2} cos(x C_2 + C_3)$  
460    With $C_0$ as the potential offset,
461    $C_1$ as the max potential value (without considering the offset),
462    $C_2$ as the frequency, and $C_3$ as the phase.
463    If no `System.potential_constants` are provided, defaults to $cos(3x)$  
464    """
465    x = system.grid
466    C = system.potential_constants
467    C0 = 0
468    C1 = 1
469    C2 = 3
470    C3 = 0
471    if C:
472        if len(C) > 0:
473            C0 = C[0]
474        if len(C) > 1:
475            C1 = C[1]
476        if len(C) > 2:
477            C2 = C[2]
478        if len(C) > 3:
479            C3 = C[3]
480    return C0 + (C1 / 2) * np.cos(np.array(x) * C2 + C3)
481
482
483def titov2023(system:System):
484    """Potential energy function of the hindered methyl rotor, from
485    [K. Titov et al., Phys. Rev. Mater. 7, 073402 (2023)](https://link.aps.org/doi/10.1103/PhysRevMaterials.7.073402).  
486
487    $V(x) = C_0 + C_1 sin(3x) + C_2 cos(3x) + C_3 sin(6x) + C_4 cos(6x)$  
488    Default constants are `qrotor.constants.constants_titov2023`[0].  
489    """
490    x = system.grid
491    C = system.potential_constants
492    if C is None:
493        C = constants.constants_titov2023[0]
494    return C[0] + C[1] * np.sin(3*x) + C[2] * np.cos(3*x) + C[3] * np.sin(6*x) + C[4] * np.cos(6*x)