Visualize modulated structures

from pathlib import Path

import numpy as np
import phonopy
from phonopy import Phonopy
from phonopy.interface.vasp import write_vasp
from pymatgen.core import Structure
from pymatgen.io.phonopy import get_pmg_structure, get_phonopy_structure

from spgrep_modulation.modulation import Modulation
from spgrep_modulation.visualize import viewer, get_local_and_ghost_sites

path = Path().resolve().parent.parent / "tests" / "phonopy_mp-2998.yaml.xz"
ph = phonopy.load(path)

qpoint = [0.5, 0, 0]  # X point
dimension = [2, 2, 2]

md = Modulation.with_supercell_and_symmetry_search(
    dynamical_matrix=ph.dynamical_matrix,
    supercell_matrix=dimension,
    qpoint=qpoint,
    factor=ph.unit_conversion_factor,
)

# Degenerated imaginary mode
frequency_index = 0
dim = md.eigenspaces[frequency_index][2].shape[1]
print(f"Frequency (THz): {md.eigvals_to_frequencies(md.eigenspaces[frequency_index][0]):.2f}")
print(f"Irrep shape: {md.eigenspaces[frequency_index][2].shape}")

Frequency (THz): -4.88
Irrep shape: (16, 2, 2)

def get_modulated_cells(md: Modulation, frequency_index, idx, max_amplitude, num):
    argument = 0
    symmetry_adapted_eigvec = md.eigenspaces[frequency_index][1][idx]
    amplitudes = np.linspace(-max_amplitude, max_amplitude, num=num, endpoint=True)
    modulated_cells = []
    for amplitude in amplitudes:
        modulation = md._get_displacements(symmetry_adapted_eigvec, amplitude, argument=0)
        scaled_modulation = amplitude / np.max(np.abs(modulation)) * modulation
        modulated = md.apply_modulation_to_supercell(scaled_modulation)
        modulated_cells.append(modulated)
    return amplitudes, modulated_cells


def save_modulated_cells(modulated_cells, root_dir: str):
    root = Path(root_dir)
    root.mkdir(exist_ok=True)

    num = len(modulated_cells)
    # shift list to start with non-modulated cell
    shifted = modulated_cells[-num // 2 :] + modulated_cells[: num // 2]
    for i, cell in enumerate(shifted):
        structure = get_pmg_structure(cell)
        structure_with_ghosts = Structure.from_sites(get_local_and_ghost_sites(structure))
        cell_with_ghosts = get_phonopy_structure(structure_with_ghosts)

        write_vasp(root / f"POSCAR-{i:03d}", cell_with_ghosts)
        write_vasp(root / f"POSCAR-{2 * num - 1 - i:03d}", cell_with_ghosts)

amplitude = 1.0
num = 16
for idx in range(dim):
    _, modulated_cells = get_modulated_cells(md, frequency_index, idx, amplitude, num=num)
    save_modulated_cells(modulated_cells, f"disp-{idx}")