1"""
  2# Description
  3
  4This submodule contains tools to rotate molecular structures.
  5Works with Quantum ESPRESSO input files.
  6
  7
  8# Index
  9
 10| | |
 11| --- | --- |
 12| `rotate_qe()`     | Rotate specific atoms from a Quantum ESPRESSO input file |
 13| `rotate_coords()` | Rotate a specific list of coordinates |
 14
 15---
 16"""
 17
 18
 19import numpy as np
 20import os
 21import shutil
 22from scipy.spatial.transform import Rotation
 23from .constants import *
 24import aton.api as api
 25import aton.txt.extract as extract
 26import aton.txt.edit as edit
 27
 28
 29def rotate_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 pw.x 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.pwx.get_atom(filepath=filepath, position=position, precision=precision, literal=True)
 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.pwx.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.pwx.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
103def _save_qe(
104        filename,
105        output:str,
106        lines:list,
107        positions:list
108    ) -> str:
109    """Copies `filename` to `output`, updating the old `lines` with the new `positions`.
110    
111    The angle will be appended at the end of the input prefix to avoid overlapping calculations.
112    """
113    shutil.copy(filename, output)
114    for i, line in enumerate(lines):
115        strings = line.split()
116        atom = strings[0]
117        new_line = f"  {atom}   {positions[i][0]:.15f}   {positions[i][1]:.15f}   {positions[i][2]:.15f}"
118        #print(f'OLD LINE: {line}')  # DEBUG
119        #print(f'NEW_LINE: {new_line}')  # DEBUG
120        edit.replace_line(output, line, new_line, raise_errors=True)
121    if len(lines) + 2 == len(positions):  # In case show_axis=True
122        additional_positions = positions[-2:]
123        for pos in additional_positions:
124            pos.insert(0, 'He')
125            api.pwx.add_atom(output, pos)
126    elif len(lines) != len(positions):
127        raise ValueError(f"What?!  len(lines)={len(lines)} and len(positions)={len(positions)}")
128    # Add angle to calculation prefix
129    output_name = os.path.basename(output)
130    splits = output_name.split('_')
131    angle_str = splits[-1].replace('.in', '')
132    prefix = ''
133    content = api.pwx.read_in(output)
134    if 'prefix' in content.keys():
135        prefix = content['prefix']
136        prefix = prefix.strip("'")
137    prefix = "'" + prefix + angle_str + "'"
138    api.pwx.set_value(output, 'prefix', prefix)
139    return output
140
141
142def rotate_coords(
143        positions:list,
144        angle:float,
145        use_centroid:bool=True,
146        show_axis:bool=False,
147    ) -> list:
148    """Rotates geometrical coordinates.
149
150    Takes a list of atomic `positions` in cartesian coordinates, as
151    `[[x1,y1,z1], [x2,y2,z2], [x3,y3,z3], [etc]`.
152    Then rotates said coordinates by a given `angle` in degrees.
153    Returns a list with the updated positions.
154
155    By default, the rotation vector is defined by the perpendicular
156    passing through the geometrical center of the first three points.
157    To override this and use the vector between the first two atoms
158    as the rotation axis, set `use_centroid = False`.
159
160    **WARNING: The `positions` list is order-sensitive**.
161    If you rotate more than one chemical group in a structure,
162    be sure to follow the same direction for each group (e.g. all clockwise)
163    to ensure that all rotation vectors point in the same direction.
164
165    If `show_axis = True` it returns two additional coordinates at the end of the list,
166    with the centroid and the rotation vector. Only works with `use_centroid = True`.
167
168    The rotation uses Rodrigues' rotation formula,
169    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).
170    """
171    if len(positions) < 3:
172        raise ValueError("At least three atoms must be rotated.")
173    if not isinstance(positions[0], list):
174        raise ValueError(f"Atomic positions must have the form: [[x1,y1,z1], [x2,y2,z2], [x3,y3,z3], etc]. Yours were:\n{positions}")
175    positions = np.array(positions)
176    #print(f'POSITIONS: {positions}')  # DEBUG
177    # Define the geometrical center
178    center_atoms = positions[:2]
179    if use_centroid:
180        center_atoms = positions[:3]
181    center = np.mean(center_atoms, axis=0)
182    # Ensure the axis passes through the geometrical center
183    centered_positions = positions - center
184    # Define the perpendicular axis (normal to the plane formed by the first three points)
185    v1 = centered_positions[0] - centered_positions[1]
186    v2 = centered_positions[0] - centered_positions[2]
187    axis = v1  # Axis defined by the first two points
188    if use_centroid:  # Axis defined by the cross product of the first three points
189        axis = np.cross(v2, v1)
190    axis_length = np.linalg.norm(axis)
191    axis = axis / axis_length
192    # Create the rotation object using scipy
193    rotation = Rotation.from_rotvec(np.radians(angle) * axis)
194    # Rotate all coordinates around the geometrical center
195    rotated_centered_positions = rotation.apply(centered_positions)
196    rotated_positions = (rotated_centered_positions + center).tolist()
197    #print(f'ROTATED_POSITIONS: {rotated_positions}')  # DEBUG
198    if show_axis and use_centroid:
199        rotated_positions.append(center.tolist())
200        rotated_positions.append((center + axis).tolist())
201    return rotated_positions