ase_plot_ir_activity

#! /usr/bin/env python3

from pathlib import Path

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


def lorentzian(x, x0, gamma):
    """Simple normalized Lorentzian"""
    return gamma / np.pi / ((x - x0) ** 2 + gamma**2)


def lorentzian2(x, x0, Gamma):
    """normalized Lorentzian with FWHM"""
    gamma = Gamma / 2
    return gamma / np.pi / ((x - x0) ** 2 + gamma**2)


def gaussian2(x, x0, Gamma):
    """normalized Gaussian with FWHM parameter"""
    sigma = Gamma / np.sqrt(2 * np.log(16))
    return np.exp(-((x - x0) ** 2) / 2 / sigma**2) / sigma / np.sqrt(2 * np.pi)


app = typer.Typer(pretty_exceptions_show_locals=False)


@app.command()
def main(
    file: Path = "outfile.mode_activity.csv",
    width: float = 10.0,
    xmin: float = None,
    xmax: float = None,
    xpoints: int = 1001,
    gaussian: bool = False,
    outfile: Path = "outfile.mode_activity.pdf",
):
    """Plot IR activities"""
    from matplotlib import pyplot as plt
    import numpy as np
    import pandas as pd

    df = pd.read_csv(file)

    x = df["frequency (cm^-1)"]
    y = df["activity"]

    if xmin is None:
        xmin = x.min()

    if xmax is None:
        xmax = 1.2 * x.max()

    xs = np.linspace(xmin, xmax, xpoints)

    ys = []

    for _x, _y in zip(x, y):
        if gaussian:
            ys.append(_y * gaussian2(xs, _x, width))
        else:
            ys.append(_y * lorentzian2(xs, _x, width))

    ys = np.array(ys)

    fig, ax = plt.subplots()

    ax.plot(xs, ys.sum(axis=0), label="Total")

    ax.set_xlabel("Frequency (cm$^{-1}$)")
    ax.set_ylabel("Intensity")

    echo(f'... save to "{outfile}"')
    fig.savefig(outfile)


if __name__ == "__main__":
    app()