ase_get_ir_activity

#! /usr/bin/env python3

from pathlib import Path

import typer
from rich import print as echo
import numpy as np


def solve(F: np.ndarray, U: np.ndarray, masses: np.ndarray):
    """Solve for frequencies and mode vectors"""
    masses_scale = 1 / np.sqrt(masses.repeat(3))

    H = -(F.T @ np.linalg.pinv(U.T)).T

    D = H * masses_scale[:, None] * masses_scale[None, :]

    eigenvalues, eigenvectors = np.linalg.eigh(D)

    frequencies = np.sign(eigenvalues) * np.sqrt(np.abs(eigenvalues))

    return frequencies, eigenvalues, eigenvectors


app = typer.Typer(pretty_exceptions_show_locals=False)


@app.command()
def main(
    file_reference: Path = "geometry.in",
    infile_forces: str = "outfile_forces.dat",
    infile_dipoles: str = "outfile_dipoles.dat",
    infile_displacements: str = "outfile_displacements.dat",
    outfile_vibrations: str = "vibrations.extxyz",
    outfile_activity: str = "outfile.mode_activity.csv",
    decimals_activity: int = 5,
    amplitude_vibration: float = 0.2,
    format_in: str = None,
):
    """Get IR activities

    Args:
        file_reference: reference structure
        infile_forces: input file with forces
        infile_dipoles: input file with dipoles
        infile_displacements: input file with displacements
    """
    from ase.io import read, write
    from ase import units
    import numpy as np
    import pandas as pd

    to_thz = 1 / (units._hbar * 1e10 / np.sqrt(units._e * units._amu))
    to_wn = to_thz * 33.356

    echo(f'Reading reference structure from "{file_reference}"')
    atoms_reference = read(file_reference, format=format_in)
    echo(f"... {atoms_reference}")

    n_atoms = len(atoms_reference)
    masses = atoms_reference.get_masses()

    # Hessian = d^2 E / dx_i dx_j =  - dF_i / dx_j
    forces = np.loadtxt(infile_forces)
    displacements = np.loadtxt(infile_displacements)

    n_displacements = forces.shape[0] // n_atoms

    assert n_displacements == 6 * n_atoms, (n_displacements, 6 * n_atoms)

    echo(f"... number of displacements:          {n_displacements}")
    echo(f"--> number of displacements per atom: {n_displacements//n_atoms}")

    _F = forces.reshape(n_atoms, 3, 2, n_atoms, 3)
    _U = displacements.reshape(n_atoms, 3, 2, n_atoms, 3)

    # take positive displacements only:
    # echo()
    # echo("POSITIVE DISPLACEMENTS")
    F = _F[:, :, 0, :, :].reshape(3 * n_atoms, 3 * n_atoms)
    U = _U[:, :, 0, :, :].reshape(3 * n_atoms, 3 * n_atoms)

    frequencies_positive, _, eigenvectors = solve(F, U, masses)

    # take negatiev displacements only:
    # echo()
    # echo("NEGATIVE DISPLACEMENTS")

    # take positive displacements only:
    F = _F[:, :, 1, :, :].reshape(3 * n_atoms, 3 * n_atoms)
    U = _U[:, :, 1, :, :].reshape(3 * n_atoms, 3 * n_atoms)

    frequencies_negative, _, eigenvectors = solve(F, U, masses)

    # take all
    echo()
    echo("FREQUENCIES")

    F = np.rollaxis(_F, 2).reshape(-1, 3 * n_atoms)
    U = np.rollaxis(_U, 2).reshape(-1, 3 * n_atoms)

    frequencies, eigenvalues, eigenvectors = solve(F, U, masses)

    M = np.sqrt(masses.repeat(3))
    # mode displacements
    dus = (eigenvectors.T / M[None, :]).reshape(-1, n_atoms, 3)

    list_atoms = []
    for ii, du in enumerate(dus):
        _f = to_wn * frequencies[ii]
        _fp = to_wn * frequencies_positive[ii]
        _fn = to_wn * frequencies_negative[ii]
        _s = f"(pos: {_fp:10.2f}, neg: {_fn:10.2f})"
        echo(f"Mode {ii:3d}, frequency = {_f:10.2f} cm^-1 {_s}")
        _atoms = atoms_reference.copy()
        _atoms.set_velocities(amplitude_vibration * du)

        list_atoms.append(_atoms)

    echo(f'... write vibrations to "{outfile_vibrations}"')
    write(outfile_vibrations, list_atoms, format="extxyz")

    # intensities if applicable
    if Path(infile_dipoles).exists():
        echo()
        echo("ACTIVITIES")
        dipoles = np.loadtxt(infile_dipoles).reshape(n_atoms, 3, 2, 3)

        # only every second one
        P1 = dipoles[:, :, 0]
        P2 = dipoles[:, :, 1]
        dP = (P2 - P1) / abs(displacements).max() / 2

        # mode dipoles
        dps = (dus[:, :, :, None] * dP[None, :, :, :]).sum(axis=(1, 2))

        conv = (1.0 / units.Debye) ** 2
        activities = np.linalg.norm(dps, axis=1) ** 2 * conv
        activities = activities.round(decimals_activity)

        for ii, II in enumerate(activities):

            _f = to_wn * frequencies[ii]
            echo(f"Mode {ii+1:3d}, frequency = {_f:10.2f} cm^-1, intensity = {II:.3f}")

        _data = {
            "frequency": to_thz * frequencies,
            "frequency (cm^-1)": to_wn * frequencies,
            "activity": activities,
        }
        df = pd.DataFrame(_data)
        df.index = df.index + 1
        df.index.name = "mode"

        echo(f'... write activities to "{outfile_activity}"')
        df.to_csv(outfile_activity, index=True, float_format="%20.10f")


if __name__ == "__main__":
    app()