qrotor 4.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of qrotor might be problematic. Click here for more details.

qrotor/potential.py ADDED
@@ -0,0 +1,472 @@
1
+ """
2
+ # Description
3
+
4
+ This module contains functions to calculate the actual `potential_values` of the system.
5
+
6
+
7
+ # Index
8
+
9
+ User 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
+
19
+ To solve the system, optionally interpolating to a new gridsize, use the `System.solve(gridsize)` method.
20
+ However, if you just want to quickly solve or interpolate the potential, check the `System.solve_potential(gridsize)` method.
21
+ This 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
+
28
+ A synthetic potential can be created by specifying its name in `System.potential_name`,
29
+ along with the corresponding `System.potential_constants` if required.
30
+ Available 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
+
43
+ from .system import System
44
+ from . import constants
45
+ from . import systems
46
+ import numpy as np
47
+ import os
48
+ from copy import deepcopy
49
+ from scipy.interpolate import CubicSpline
50
+ import aton.alias as alias
51
+ import aton.file as file
52
+ import aton.api.qe as qe
53
+ from ._version import __version__
54
+
55
+
56
+ def 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
+ """
70
+ print('Saving potential data file...')
71
+ # Check if a previous potential.dat file exists, and ask to overwrite it
72
+ previous_potential_file = file.get(filepath, return_anyway=True)
73
+ if previous_potential_file:
74
+ print(f"WARNING: Previous '{filepath}' file will be overwritten, proceed anyway?")
75
+ answer = input("(y/n): ")
76
+ if not answer.lower() in alias.boolean[True]:
77
+ print("Aborted.")
78
+ return None
79
+ # Set header
80
+ potential_data = f'# {comment}\n' if comment else f'# {system.comment}\n' if system.comment else ''
81
+ potential_data += '# Rotational potential dataset\n'
82
+ potential_data += f'# Saved with QRotor {__version__}\n'
83
+ potential_data += '# https://pablogila.github.io/qrotor\n'
84
+ potential_data += '#\n'
85
+ # Check that grid and potential values are the same size
86
+ if len(system.grid) != len(system.potential_values):
87
+ raise ValueError('len(system.grid) != len(system.potential_values)')
88
+ grid = system.grid
89
+ potential_values = system.potential_values
90
+ # Convert angle units
91
+ if angle.lower() in alias.units['rad']:
92
+ potential_data += '# Angle/rad, '
93
+ else:
94
+ grid = np.degrees(grid)
95
+ potential_data += '# Angle/deg, '
96
+ if not angle.lower() in alias.units['deg']:
97
+ print(f"WARNING: Unrecognised '{angle}' angle units, using degrees instead")
98
+ # Convert energy units
99
+ if energy.lower() in alias.units['meV']:
100
+ potential_data += 'Potential/meV\n'
101
+ elif energy.lower() in alias.units['eV']:
102
+ potential_values = potential_values * 1e-3
103
+ potential_data += 'Potential/eV\n'
104
+ elif energy.lower() in alias.units['Ry']:
105
+ potential_values = potential_values * constants.meV_to_Ry
106
+ potential_data += 'Potential/Ry\n'
107
+ else:
108
+ print(f"WARNING: Unrecognised '{energy}' energy units, using meV instead")
109
+ potential_data += 'Potential/meV\n'
110
+ potential_data += '#\n'
111
+ # Save all values
112
+ for angle_value, energy_value in zip(grid, potential_values):
113
+ potential_data += f'{angle_value}, {energy_value}\n'
114
+ with open(filepath, 'w') as f:
115
+ f.write(potential_data)
116
+ print(f'Saved to {filepath}')
117
+ # Warn the user if not in default units
118
+ if angle.lower() not in alias.units['deg']:
119
+ print(f"WARNING: You saved the potential in '{angle}' angle units! Remember that QRotor works in degrees!")
120
+ if energy.lower() not in alias.units['meV']:
121
+ print(f"WARNING: You saved the potential in '{energy}' energy units! Remember that QRotor works in meVs!")
122
+
123
+
124
+ def load(
125
+ filepath:str='potential.csv',
126
+ comment:str=None,
127
+ system:System=None,
128
+ angle:str='deg',
129
+ energy:str='meV',
130
+ ) -> System:
131
+ """Read a rotational potential energy datafile.
132
+
133
+ The input file in `filepath` should contain two columns with angle and potential energy values.
134
+ Degrees and meV are assumed as default units unless stated in `angle` and `energy`.
135
+ Units will be converted automatically to radians and meV.
136
+
137
+ An optional `comment` can be included in the output System.
138
+ Set to the parent folder name by default.
139
+
140
+ A previous System object can be provided through `system` to update its potential values.
141
+ """
142
+ file_path = file.get(filepath)
143
+ system = System() if system is None else system
144
+ with open(file_path, 'r') as f:
145
+ lines = f.readlines()
146
+ positions = []
147
+ potentials = []
148
+ for line in lines:
149
+ if line.startswith('#'):
150
+ continue
151
+ position, potential = line.split()
152
+ positions.append(float(position.strip().strip(',').strip()))
153
+ potentials.append(float(potential.strip()))
154
+ # Save angles to numpy arrays
155
+ if angle.lower() in alias.units['deg']:
156
+ positions = np.radians(positions)
157
+ elif angle.lower() in alias.units['rad']:
158
+ positions = np.array(positions)
159
+ else:
160
+ raise ValueError(f"Angle unit '{angle}' not recognized.")
161
+ # Save energies to numpy arrays
162
+ if energy.lower() in alias.units['eV']:
163
+ potentials = np.array(potentials) * 1000
164
+ elif energy.lower() in alias.units['meV']:
165
+ potentials = np.array(potentials)
166
+ elif energy.lower() in alias.units['Ry']:
167
+ potentials = np.array(potentials) * constants.Ry_to_meV
168
+ else:
169
+ raise ValueError(f"Energy unit '{energy}' not recognized.")
170
+ # Set the system
171
+ system.grid = np.array(positions)
172
+ system.gridsize = len(positions)
173
+ system.potential_values = np.array(potentials)
174
+ # System comment as the parent folder name
175
+ system.comment = os.path.basename(os.path.dirname(file_path)) if comment==None else comment
176
+ print(f"Loaded {filepath}")
177
+ return system
178
+
179
+
180
+ def from_qe(
181
+ folder=None,
182
+ filepath:str='potential.csv',
183
+ include:list=['.out'],
184
+ exclude:list=['slurm-'],
185
+ energy:str='meV',
186
+ comment:str=None,
187
+ ) -> None:
188
+ """Compiles a rotational potential CSV file from Quantum ESPRESSO outputs,
189
+ created with `qrotor.rotate.structure_qe()`.
190
+
191
+ The angle in degrees is extracted from the output filenames,
192
+ which must follow `whatever_ANGLE.out`.
193
+
194
+ Outputs from SCF calculations must be located in the provided `folder` (CWD if None).
195
+ Files can be filtered by those containing the specified `include` filters,
196
+ excluding those containing any string from the `exclude` list.
197
+ The output `filepath` name is `'potential.dat'` by default.
198
+
199
+ Energy values are saved to meV by dafault, unless specified in `energy`.
200
+ Only change the energy units if you know what you are doing;
201
+ remember that default energy units in QRotor are meV!
202
+ """
203
+ folder = file.get_dir(folder)
204
+ # Check if a previous potential.dat file exists, and ask to overwrite it
205
+ previous_potential_file = file.get(filepath, return_anyway=True)
206
+ if previous_potential_file:
207
+ print(f"WARNING: Previous '{filepath}' file will be overwritten, proceed anyway?")
208
+ answer = input("(y/n): ")
209
+ if not answer.lower() in alias.boolean[True]:
210
+ print("Aborted.")
211
+ return None
212
+ # Get the files to read
213
+ files = file.get_list(folder=folder, include=include, exclude=exclude, abspath=True)
214
+ folder_name = os.path.basename(folder)
215
+ # Set header
216
+ potential_data = f'# {comment}\n' if comment else f'# {folder_name}\n'
217
+ potential_data += '# Rotational potential dataset\n'
218
+ potential_data += f'# Calculated with QE using QRotor {__version__}\n'
219
+ potential_data += '# https://pablogila.github.io/qrotor\n'
220
+ potential_data += '#\n'
221
+ if energy.lower() in alias.units['eV']:
222
+ potential_data += '# Angle/deg, Potential/eV\n'
223
+ elif energy.lower() in alias.units['meV']:
224
+ potential_data += '# Angle/deg, Potential/meV\n'
225
+ elif energy.lower() in alias.units['Ry']:
226
+ potential_data += '# Angle/deg, Potential/Ry\n'
227
+ else:
228
+ potential_data += '# Angle/deg, Potential/meV\n'
229
+ potential_data += '#\n'
230
+ potential_data_list = []
231
+ print('Extracting the potential as a function of the angle...')
232
+ print('----------------------------------')
233
+ counter_success = 0
234
+ counter_errors = 0
235
+ for file_path in files:
236
+ filename = os.path.basename(file_path)
237
+ file_path = file.get(filepath=file_path, include='.out', return_anyway=True)
238
+ if not file_path: # Not an output file, skip it
239
+ continue
240
+ content = qe.read_out(file_path)
241
+ if not content['Success']: # Ignore unsuccessful calculations
242
+ print(f'x {filename}')
243
+ counter_errors += 1
244
+ continue
245
+ if energy.lower() in alias.units['eV']:
246
+ energy_value = content['Energy'] * constants.Ry_to_eV
247
+ elif energy.lower() in alias.units['meV']:
248
+ energy_value = content['Energy'] * constants.Ry_to_meV
249
+ elif energy.lower() in alias.units['Ry']:
250
+ energy_value = content['Energy']
251
+ else:
252
+ print(f"WARNING: Energy unit '{energy}' not recognized, using meV instead.")
253
+ energy = 'meV'
254
+ energy_value = content['Energy'] * constants.Ry_to_meV
255
+ splits = filename.split('_')
256
+ angle_value = splits[-1].replace('.out', '')
257
+ angle_value = float(angle_value)
258
+ potential_data_list.append((angle_value, energy_value))
259
+ print(f'OK {filename}')
260
+ counter_success += 1
261
+ # Sort by angle
262
+ potential_data_list_sorted = sorted(potential_data_list, key=lambda x: x[0])
263
+ # Append the sorted values as a string
264
+ for angle_value, energy_value in potential_data_list_sorted:
265
+ potential_data += f'{angle_value}, {energy_value}\n'
266
+ with open(filepath, 'w') as f:
267
+ f.write(potential_data)
268
+ print('----------------------------------')
269
+ print(f'Succesful calculations (OK): {counter_success}')
270
+ print(f'Faulty calculations (x): {counter_errors}')
271
+ print('----------------------------------')
272
+ print(f'Saved angles and potential values at {filepath}')
273
+ # Warn the user if not in default units
274
+ if energy.lower() not in alias.units['meV']:
275
+ print(f"WARNING: You saved the potential in '{energy}' units! Remember that QRotor works in meVs!")
276
+ return None
277
+
278
+
279
+ def merge(
280
+ add=[],
281
+ subtract=[],
282
+ comment:str=None
283
+ ) -> System:
284
+ """Add or subtract potentials from different systems.
285
+
286
+ Adds the potential values from the systems in `add`,
287
+ removes the ones from `subtract`.
288
+ All systems will be interpolated to the bigger gridsize if needed.
289
+
290
+ A copy of the first System will be returned with the resulting potential values,
291
+ with an optional `comment` if indicated.
292
+ """
293
+ add = systems.as_list(add)
294
+ subtract = systems.as_list(subtract)
295
+ gridsizes = systems.get_gridsizes(add)
296
+ gridsizes.extend(systems.get_gridsizes(subtract))
297
+ max_gridsize = max(gridsizes)
298
+ # All gridsizes should be max_gridsize
299
+ for s in add:
300
+ if s.gridsize != max_gridsize:
301
+ s.gridsize = max_gridsize
302
+ s = interpolate(s)
303
+ for s in subtract:
304
+ if s.gridsize != max_gridsize:
305
+ s.gridsize = max_gridsize
306
+ s = interpolate(s)
307
+
308
+ if len(add) == 0:
309
+ if len(subtract) == 0:
310
+ raise ValueError('No systems were provided!')
311
+ result = deepcopy(subtract[0])
312
+ result.potential_values = -result.potential_values
313
+ subtract.pop(0)
314
+ else:
315
+ result = deepcopy(add[0])
316
+ add.pop(0)
317
+
318
+ for system in add:
319
+ result.potential_values = np.sum([result.potential_values, system.potential_values], axis=0)
320
+ for system in subtract:
321
+ result.potential_values = np.sum([result.potential_values, -system.potential_values], axis=0)
322
+ if comment != None:
323
+ result.comment = comment
324
+ return result
325
+
326
+
327
+ def scale(
328
+ system:System,
329
+ factor:float,
330
+ comment:str=None
331
+ ) -> System:
332
+ """Returns a copy of `system` with potential values scaled by a `factor`.
333
+
334
+ An optional `comment` can be included.
335
+ """
336
+ result = deepcopy(system)
337
+ if factor != 0:
338
+ result.potential_values = system.potential_values * factor
339
+ else:
340
+ result.potential_values = np.zeros(system.gridsize)
341
+ if comment != None:
342
+ result.comment = comment
343
+ return result
344
+
345
+
346
+ def interpolate(system:System) -> System:
347
+ """Interpolates the current `System.potential_values`
348
+ to a new grid of size `System.gridsize`.
349
+
350
+ This basic function is called by `qrotor.solve.potential()`,
351
+ which is the recommended way to interpolate potentials.
352
+ """
353
+ print(f"Interpolating potential to a grid of size {system.gridsize}...")
354
+ V = system.potential_values
355
+ grid = system.grid
356
+ gridsize = system.gridsize
357
+ new_grid = np.linspace(0, 2*np.pi, gridsize)
358
+ cubic_spline = CubicSpline(grid, V)
359
+ new_V = cubic_spline(new_grid)
360
+ system.grid = new_grid
361
+ system.potential_values = new_V
362
+ return system
363
+
364
+
365
+ def solve(system:System):
366
+ """Solves `System.potential_values`
367
+ according to the `System.potential_name`,
368
+ returning the new `potential_values`.
369
+ Avaliable potential names are `zero`, `sine` and `titov2023`.
370
+
371
+ If `System.potential_name` is not present or not recognised,
372
+ the current `System.potential_values` are used.
373
+
374
+ This basic function is called by `qrotor.solve.potential()`,
375
+ which is the recommended way to solve potentials.
376
+ """
377
+ data = deepcopy(system)
378
+ # Is there a potential_name?
379
+ if not data.potential_name:
380
+ if data.potential_values is None or len(data.potential_values) == 0:
381
+ raise ValueError(f'No potential_name and no potential_values found in the system!')
382
+ elif data.potential_name.lower() == 'titov2023':
383
+ data.potential_values = titov2023(data)
384
+ elif data.potential_name.lower() in alias.math['0']:
385
+ data.potential_values = zero(data)
386
+ elif data.potential_name.lower() in alias.math['sin']:
387
+ data.potential_values = sine(data)
388
+ elif data.potential_name.lower() in alias.math['cos']:
389
+ data.potential_values = cosine(data)
390
+ # At least there should be potential_values
391
+ #elif not any(data.potential_values):
392
+ elif data.potential_values is None or len(data.potential_values) == 0:
393
+ raise ValueError("Unrecognised potential_name '{data.potential_name}' and no potential_values found")
394
+ return data.potential_values
395
+
396
+
397
+ def zero(system:System):
398
+ """Zero potential.
399
+
400
+ $V(x) = 0$
401
+ """
402
+ x = system.grid
403
+ return 0 * np.array(x)
404
+
405
+
406
+ def sine(system:System):
407
+ """Sine potential.
408
+
409
+ $V(x) = C_0 + \\frac{C_1}{2} sin(x C_2 + C_3)$
410
+ With $C_0$ as the potential offset,
411
+ $C_1$ as the max potential value (without considering the offset),
412
+ $C_2$ as the frequency, and $C_3$ as the phase.
413
+ If no `System.potential_constants` are provided, defaults to $sin(3x)$
414
+ """
415
+ x = system.grid
416
+ C = system.potential_constants
417
+ C0 = 0
418
+ C1 = 1
419
+ C2 = 3
420
+ C3 = 0
421
+ if C:
422
+ if len(C) > 0:
423
+ C0 = C[0]
424
+ if len(C) > 1:
425
+ C1 = C[1]
426
+ if len(C) > 2:
427
+ C2 = C[2]
428
+ if len(C) > 3:
429
+ C3 = C[3]
430
+ return C0 + (C1 / 2) * np.sin(np.array(x) * C2 + C3)
431
+
432
+
433
+ def cosine(system:System):
434
+ """Cosine potential.
435
+
436
+ $V(x) = C_0 + \\frac{C_1}{2} cos(x C_2 + C_3)$
437
+ With $C_0$ as the potential offset,
438
+ $C_1$ as the max potential value (without considering the offset),
439
+ $C_2$ as the frequency, and $C_3$ as the phase.
440
+ If no `System.potential_constants` are provided, defaults to $cos(3x)$
441
+ """
442
+ x = system.grid
443
+ C = system.potential_constants
444
+ C0 = 0
445
+ C1 = 1
446
+ C2 = 3
447
+ C3 = 0
448
+ if C:
449
+ if len(C) > 0:
450
+ C0 = C[0]
451
+ if len(C) > 1:
452
+ C1 = C[1]
453
+ if len(C) > 2:
454
+ C2 = C[2]
455
+ if len(C) > 3:
456
+ C3 = C[3]
457
+ return C0 + (C1 / 2) * np.cos(np.array(x) * C2 + C3)
458
+
459
+
460
+ def titov2023(system:System):
461
+ """Potential energy function of the hindered methyl rotor, from
462
+ [K. Titov et al., Phys. Rev. Mater. 7, 073402 (2023)](https://link.aps.org/doi/10.1103/PhysRevMaterials.7.073402).
463
+
464
+ $V(x) = C_0 + C_1 sin(3x) + C_2 cos(3x) + C_3 sin(6x) + C_4 cos(6x)$
465
+ Default constants are `qrotor.constants.constants_titov2023`[0].
466
+ """
467
+ x = system.grid
468
+ C = system.potential_constants
469
+ if C is None:
470
+ C = constants.constants_titov2023[0]
471
+ 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)
472
+
qrotor/rotate.py ADDED
@@ -0,0 +1,202 @@
1
+ """
2
+ # Description
3
+
4
+ This submodule contains tools to rotate molecular structures.
5
+ Works with Quantum ESPRESSO input files.
6
+
7
+
8
+ # Index
9
+
10
+ | | |
11
+ | --- | --- |
12
+ | `structure_qe()` | Rotate specific atoms from a Quantum ESPRESSO input file |
13
+ | `rotate_coords()` | Rotate a specific list of coordinates |
14
+
15
+ ---
16
+ """
17
+
18
+
19
+ import numpy as np
20
+ import os
21
+ import shutil
22
+ from scipy.spatial.transform import Rotation
23
+ from .constants import *
24
+ import aton.api as api
25
+ import aton.txt.extract as extract
26
+ import aton.txt.edit as edit
27
+
28
+
29
+ def structure_qe(
30
+ filepath:str,
31
+ positions:list,
32
+ angle:float,
33
+ repeat:bool=False,
34
+ precision:int=3,
35
+ use_centroid:bool=True,
36
+ show_axis:bool=False,
37
+ ) -> list:
38
+ """Rotates atoms from a Quantum ESPRESSO input file.
39
+
40
+ Takes a `filepath` with a molecular structure, and three or more atomic `positions` (list).
41
+ These input positions can be approximate, and are used to identify the target atoms.
42
+ The decimal precision in the search for these positions is controlled by `precision`.
43
+
44
+ It rotates the atoms by a specific `angle` in degrees.
45
+ Additionally, if `repeat = True` it repeats the same rotation over the whole circunference.
46
+ Finally, it writes the rotated structure(s) to a new structural file(s).
47
+ Returns a list with the output filename(s).
48
+
49
+ By default, the rotation axis is defined by the perpendicular vector
50
+ passing through the geometrical center of the first three points.
51
+ To override this and instead use the vector between the first two atoms
52
+ as the rotation axis, set `use_centroid = False`.
53
+
54
+ **WARNING: The `positions` list is order-sensitive**.
55
+ If you rotate more than one chemical group in a structure,
56
+ be sure to follow the same direction for each group (e.g. all clockwise)
57
+ to ensure that all axes of rotation point in the same direction.
58
+
59
+ To debug, `show_axis = True` adds two additional helium atoms as the rotation vector.
60
+
61
+ The resulting rotational potential can be compiled to a CSV file with `qrotor.potential.from_qe()`.
62
+ """
63
+ print('Rotating Quantum ESPRESSO input structure with QRotor...')
64
+ if len(positions) < 3:
65
+ raise ValueError("At least three positions are required to define the rotation axis.")
66
+ lines = []
67
+ full_positions = []
68
+ for position in positions:
69
+ line = api.qe.get_atom(filepath, position, precision)
70
+ lines.append(line)
71
+ pos = extract.coords(line)
72
+ if len(pos) > 3: # Keep only the first three coordinates
73
+ pos = pos[:3]
74
+ # Convert to cartesian
75
+ pos_cartesian = api.qe.to_cartesian(filepath, pos)
76
+ full_positions.append(pos_cartesian)
77
+ print(f'Found atom: "{line}"')
78
+ # Set the angles to rotate
79
+ if not repeat:
80
+ angles = [angle]
81
+ else:
82
+ angles = range(0, 360, angle)
83
+ # Rotate and save the structure
84
+ outputs = []
85
+ path = os.path.dirname(filepath)
86
+ basename = os.path.basename(filepath)
87
+ name, ext = os.path.splitext(basename)
88
+ print('Rotating the structure...')
89
+ for angle in angles:
90
+ output_name = name + f'_{angle}' + ext
91
+ output = os.path.join(path, output_name)
92
+ rotated_positions_cartesian = rotate_coords(full_positions, angle, use_centroid, show_axis)
93
+ rotated_positions = []
94
+ for coord in rotated_positions_cartesian:
95
+ pos = api.qe.from_cartesian(filepath, coord)
96
+ rotated_positions.append(pos)
97
+ _save_qe(filepath, output, lines, rotated_positions)
98
+ outputs.append(output)
99
+ print(output)
100
+ return outputs
101
+
102
+
103
+ def rotate_coords(
104
+ positions:list,
105
+ angle:float,
106
+ use_centroid:bool=True,
107
+ show_axis:bool=False,
108
+ ) -> list:
109
+ """Rotates geometrical coordinates.
110
+
111
+ Takes a list of atomic `positions` in cartesian coordinates, as
112
+ `[[x1,y1,z1], [x2,y2,z2], [x3,y3,z3], [etc]`.
113
+ Then rotates said coordinates by a given `angle` in degrees.
114
+ Returns a list with the updated positions.
115
+
116
+ By default, the rotation vector is defined by the perpendicular
117
+ passing through the geometrical center of the first three points.
118
+ To override this and use the vector between the first two atoms
119
+ as the rotation axis, set `use_centroid = False`.
120
+
121
+ **WARNING: The `positions` list is order-sensitive**.
122
+ If you rotate more than one chemical group in a structure,
123
+ be sure to follow the same direction for each group (e.g. all clockwise)
124
+ to ensure that all rotation vectors point in the same direction.
125
+
126
+ If `show_axis = True` it returns two additional coordinates at the end of the list,
127
+ with the centroid and the rotation vector. Only works with `use_centroid = True`.
128
+
129
+ The rotation uses Rodrigues' rotation formula,
130
+ powered by [`scipy.spatial.transform.Rotation.from_rotvec`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.from_rotvec.html#scipy.spatial.transform.Rotation.from_rotvec).
131
+ """
132
+ if len(positions) < 3:
133
+ raise ValueError("At least three atoms must be rotated.")
134
+ if not isinstance(positions[0], list):
135
+ raise ValueError(f"Atomic positions must have the form: [[x1,y1,z1], [x2,y2,z2], [x3,y3,z3], etc]. Yours were:\n{positions}")
136
+ positions = np.array(positions)
137
+ #print(f'POSITIONS: {positions}') # DEBUG
138
+ # Define the geometrical center
139
+ center_atoms = positions[:2]
140
+ if use_centroid:
141
+ center_atoms = positions[:3]
142
+ center = np.mean(center_atoms, axis=0)
143
+ # Ensure the axis passes through the geometrical center
144
+ centered_positions = positions - center
145
+ # Define the perpendicular axis (normal to the plane formed by the first three points)
146
+ v1 = centered_positions[0] - centered_positions[1]
147
+ v2 = centered_positions[0] - centered_positions[2]
148
+ axis = v1 # Axis defined by the first two points
149
+ if use_centroid: # Axis defined by the cross product of the first three points
150
+ axis = np.cross(v2, v1)
151
+ axis_length = np.linalg.norm(axis)
152
+ axis = axis / axis_length
153
+ # Create the rotation object using scipy
154
+ rotation = Rotation.from_rotvec(np.radians(angle) * axis)
155
+ # Rotate all coordinates around the geometrical center
156
+ rotated_centered_positions = rotation.apply(centered_positions)
157
+ rotated_positions = (rotated_centered_positions + center).tolist()
158
+ #print(f'ROTATED_POSITIONS: {rotated_positions}') # DEBUG
159
+ if show_axis and use_centroid:
160
+ rotated_positions.append(center.tolist())
161
+ rotated_positions.append((center + axis).tolist())
162
+ return rotated_positions
163
+
164
+
165
+ def _save_qe(
166
+ filename,
167
+ output:str,
168
+ lines:list,
169
+ positions:list
170
+ ) -> str:
171
+ """Copies `filename` to `output`, updating the old `lines` with the new `positions`.
172
+
173
+ The angle will be appended at the end of the input prefix to avoid overlapping calculations.
174
+ """
175
+ shutil.copy(filename, output)
176
+ for i, line in enumerate(lines):
177
+ strings = line.split()
178
+ atom = strings[0]
179
+ new_line = f" {atom} {positions[i][0]:.15f} {positions[i][1]:.15f} {positions[i][2]:.15f}"
180
+ #print(f'OLD LINE: {line}') # DEBUG
181
+ #print(f'NEW_LINE: {new_line}') # DEBUG
182
+ edit.replace_line(output, line, new_line, raise_errors=True)
183
+ if len(lines) + 2 == len(positions): # In case show_axis=True
184
+ additional_positions = positions[-2:]
185
+ for pos in additional_positions:
186
+ pos.insert(0, 'He')
187
+ api.qe.add_atom(output, pos)
188
+ elif len(lines) != len(positions):
189
+ raise ValueError(f"What?! len(lines)={len(lines)} and len(positions)={len(positions)}")
190
+ # Add angle to calculation prefix
191
+ output_name = os.path.basename(output)
192
+ splits = output_name.split('_')
193
+ angle_str = splits[-1].replace('.in', '')
194
+ prefix = ''
195
+ content = api.qe.read_in(output)
196
+ if 'prefix' in content.keys():
197
+ prefix = content['prefix']
198
+ prefix = prefix.strip("'")
199
+ prefix = "'" + prefix + angle_str + "'"
200
+ api.qe.set_value(output, 'prefix', prefix)
201
+ return output
202
+