1"""
  2# Description
  3
  4Functions to work with [Phonopy](https://phonopy.github.io/phonopy/) calculations,
  5along with [Quantum ESPRESSO](https://www.quantum-espresso.org/).  
  6
  7
  8# Index
  9
 10| | |
 11| --- | --- |
 12| `make_supercells()` | Build supercell SCF inputs for phonon calculations |
 13
 14
 15# Examples
 16
 17To create the 2x2x2 supercells and run the phonon calculations
 18from a folder with `relax.in` and `relax.out` files,
 19using a `template.slurm` file,
 20```python
 21from aton import api
 22api.phonopy.make_supercells('relax.in', 'relax.out')
 23api.slurm.sbatch('supercell-', 'template.slurm')
 24```
 25
 26---
 27"""
 28
 29
 30import os
 31from aton._version import __version__
 32import aton.file as file
 33import aton.call as call
 34import aton.txt.find as find
 35import aton.txt.edit as edit
 36import aton.txt.extract as extract
 37import aton.api.pwx as pwx
 38import aton.api.slurm as slurm
 39import shutil
 40import scipy.constants as const
 41
 42
 43def make_supercells(
 44        relax_in:str='relax.in',
 45        relax_out:str='relax.out',
 46        scf:str=None,
 47        slurm_template:str='template.slurm',
 48        folder:str=None,
 49        dimension:str='2 2 2',
 50        amplitude:float=None,
 51        update:dict=None,
 52        update_E:bool=True,
 53    ) -> None:
 54    """
 55    Creates and prepares the supercell inputs of Phonopy,
 56
 57    These supercells are created from the `relax_in` and `relax_out` files in the `folder`
 58    ('relax.in', 'relax.out' and CWD by default, respectively),
 59    needed for the Phonopy calculations with Quantum ESPRESSO.
 60    Alternatively, a previously relaxed `scf` input file can be provided,
 61    which will override the creation of a new scf file
 62    from the `relax_in` and `relax_out` files.
 63    The cells have a given `dimension` ('2 2 2' by default),
 64    and the displacement `amplitude` is set to phonopy's default unless specified.
 65
 66    Extensive convergence values for the energy (`etot_conv_thr` and `conv_thr`)
 67    are updated automatically according to the supercell size.
 68    This can be disabled with `update_E=False`.
 69    Any input value can be updated with an `update` dict,
 70    overriding automatic values.
 71
 72    By default, at the end of the execution it will check
 73    that an `slurm_template` ('template.slurm') is present and valid;
 74    this is, containing the keywords `JOBNAME`, `INPUT` and `OUTPUT`.
 75    If not, an example with instructions will be provided.
 76    This check can be skipped with `slurm_template=''`.
 77    The template will allow to easily run the Phonopy calculations with the one-line command
 78    `aton.api.slurm.sbatch('supercell-', 'template.slurm')`.
 79    """
 80    print(f'\nWelcome to aton.api.phonopy {__version__}\n'
 81          'Creating all supercell inputs with Phonopy for Quantum ESPRESSO...\n')
 82    if not scf:
 83        pwx.scf_from_relax(folder, relax_in, relax_out, update=update)
 84        scf = 'scf.in'
 85    _check_dims = extract.coords(dimension)
 86    if len(_check_dims) != 3:
 87        raise ValueError('Supercell dimension must be given as "nx ny nz"!')
 88    _supercells_from_scf(folder, scf, dimension, amplitude)
 89    _copy_scf_header_to_supercells(folder, scf, update, update_E)
 90    print('\n------------------------------------------------------\n'
 91          'PLEASE CHECH BELOW THE CONTENT OF supercell-001.in\n'
 92          '------------------------------------------------------\n')
 93    call.bash('head -n 100 supercell-001.in')
 94    print('\n------------------------------------------------------\n'
 95          'PLEASE CHECH THE CONTENT OF supercell-001.in\n'
 96          'The first 100 lines of the input were printed above!\n'
 97          '------------------------------------------------------\n\n'
 98          'If it seems correct, run the calculations with:\n'
 99          f"aton.api.slurm.sbatch('supercell-', '{slurm_template}')\n")
100    if slurm_template:
101        slurm.check_template(slurm_template, folder)
102    return None
103
104
105def _supercells_from_scf(
106        folder:str=None,
107        scf:str='scf.in',
108        dimension:str='2 2 2',
109        amplitude:float=None,
110    ) -> None:
111    """
112    Creates supercells of a given `dimension` (`2 2 2` by default) inside a `folder`,
113    from a Quantum ESPRESSO `scf` input (`scf.in` by default).
114    """
115    folder = call.here(folder)
116    scf_in = file.get(folder, scf, True)
117    scf_temp1 = _ensure_bohr_units(folder, scf_in)
118    if scf_temp1 is None:
119        raise FileNotFoundError('No SCF input found in path!')
120    if amplitude:
121        call.bash(f'phonopy --qe -d --dim="{dimension}" --amplitude="{amplitude}" -c {scf_temp1}')
122    else:
123        call.bash(f'phonopy --qe -d --dim="{dimension}" -c {scf_temp1}')
124    os.remove(scf_temp1)
125    return None
126
127
128def _ensure_bohr_units(folder:str=None, scf:str='scf.in') -> None:
129    """Check that the lattice units are bohr instead of angstrom in the `scf` input file."""
130    folder = call.here(folder)
131    scf_in = file.get(folder, scf, True)
132    scf_temp1 = '_temp_scf_in_bohr_units.in'
133    shutil.copy(scf_in, scf_temp1)
134    input_values = pwx.read_in(scf_in)
135    if 'A' in input_values:  # Convert angstrom to bohr
136        celldm = input_values['A'] / (const.physical_constants['Bohr radius'][0] * 1e10)
137        pwx.set_value(scf_temp1, 'celldm(1)', celldm)
138        print(f'Updated celldm(1) from A:  {input_values["A"]} AA  ->  {celldm} bohr')
139    return scf_temp1
140
141
142def _copy_scf_header_to_supercells(
143        folder:str=None,
144        scf:str='scf.in',
145        update:dict=None,
146        update_E:bool=True,
147    ) -> None:
148    """Paste the header from the `scf` file in `folder` to the supercells created by Phonopy."""
149    print(f'Creating header...\n')
150    folder = call.here(folder)
151    # Check if the header file, the scf.in, exists
152    scf_file = file.get(folder, scf, True)
153    if scf_file is None:
154        raise FileNotFoundError('No header file found in path!')
155    # Check if the supercells exist
156    supercells = file.get_list(folder, include='supercell-')
157    if supercells is None:
158        raise FileNotFoundError('No supercells found in path!')
159    # Check if the supercells contains '&CONTROL' and abort if so
160    supercell_sample = supercells[0]
161    is_control = find.lines(supercell_sample, r'(&CONTROL|&control)', 1, 0, False, True)
162    if is_control:
163        raise RuntimeError('Supercells already contain &CONTROL! Did you do this already?')
164    # Check if the keyword is in the scf file
165    is_header = find.lines(scf_file, r'ATOMIC_SPECIES', 1, 0, False, False)
166    if not is_header:
167        raise RuntimeError('No ATOMIC_SPECIES found in header!')
168    # Copy the scf to a temp file
169    scf_temp2 = '_temp_scf_with_updated_supercell_values.in'
170    shutil.copy(scf_file, scf_temp2)
171    values = pwx.read_in(scf_temp2)
172    # Find the new number of atoms and replace the line
173    updated_values = find.lines(supercell_sample, 'ibrav', 1)  # !    ibrav = 0, nat = 384, ntyp = 5
174    if not updated_values:
175        print("!!! Okay listen, this is weird. This line of code should never be running, "
176              "but for some reson I couldn't find the updated values in the supercells. "
177              "Please, introduce the NEW NUMBER OF ATOMS in the supercells manually (int):")
178        nat = int(input('nat = '))
179    else:
180        nat = int(extract.number(updated_values[0], 'nat'))
181    pwx.set_value(scf_temp2, 'nat', nat)
182    print(f'Updated nat:  {values['nat']}  ->  {nat}')
183    # Remove the lattice parameters, since Phonopy already indicates units
184    pwx.set_value(scf_temp2, 'celldm(1)', '')
185    pwx.set_value(scf_temp2, 'A', '')
186    pwx.set_value(scf_temp2, 'B', '')
187    pwx.set_value(scf_temp2, 'C', '')
188    pwx.set_value(scf_temp2, 'cosAB', '')
189    pwx.set_value(scf_temp2, 'cosAC', '')
190    pwx.set_value(scf_temp2, 'cosBC', '')
191    print('Updated lattice parameters')
192    # Remove the top content from the temp file
193    edit.delete_under(scf_temp2, 'K_POINTS', -1, 2, False)
194    # Update extensive energy values
195    if update_E:
196        old_nat = values['nat']
197        new_atoms_factor = nat / old_nat
198        etot_conv_thr = values['etot_conv_thr']
199        conv_thr = values['conv_thr']
200        new_etot_conv_thr = etot_conv_thr * new_atoms_factor
201        new_conv_thr      = conv_thr      * new_atoms_factor
202        pwx.set_value(scf_temp2, 'etot_conv_thr', new_etot_conv_thr)
203        pwx.set_value(scf_temp2, 'conv_thr', new_conv_thr)
204        print(f'Updated etot_conv_thr:  {etot_conv_thr}  ->  {new_etot_conv_thr}')
205        print(f'Updated conv_thr:  {conv_thr}  ->  {new_conv_thr}')
206    # Update any other user-defined values
207    if update:
208        pwx.set_values(scf_temp2, update)
209        for key, value in update.items():
210            print(f'Updated {key}:  {values[key]}  ->  {value}')
211    # Add the header to the supercells
212    print(f'\nAdding header to the supercells...')
213    with open(scf_temp2, 'r') as f:
214        header = f.read()
215    for supercell in supercells:
216        edit.insert_at(supercell, header, 0)
217    # Remove the temp file
218    os.remove(scf_temp2)
219    print('Done.')
220    return None