view_phonons.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext aiida\n",
    "%aiida"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.parse as urlparse\n",
    "\n",
    "import ase\n",
    "import ipywidgets as ipw\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from ase.visualize import view\n",
    "from IPython.display import HTML, clear_output, display\n",
    "from matplotlib.ticker import AutoMinorLocator\n",
    "\n",
    "mpl.rcParams[\"figure.dpi\"] = 200\n",
    "mpl.rcParams[\"axes.linewidth\"] = 1.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "IPython.OutputArea.prototype._should_scroll = function(lines) {\n",
    "    return false;\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read CP2K vibrational analysis data from Molden format\n",
    "\n",
    "\n",
    "def read_cp2k_molden(pk):\n",
    "    node = load_node(pk)\n",
    "    with node.outputs.retrieved.open(\"aiida-VIBRATIONS-1.mol\", mode=\"rb\") as fhandle:\n",
    "        data = fhandle.readlines()\n",
    "\n",
    "    freq = []\n",
    "    eq_coords = []  # element, x, y, z\n",
    "    vibr_displacements = []  # [vibration_nr][coord]\n",
    "    inten = []\n",
    "\n",
    "    section = \"\"\n",
    "    b2A = 0.52917721067\n",
    "    # Parse the datafile\n",
    "    for line in data:\n",
    "        line = line.strip().decode()\n",
    "        # Are we entering into a new section?\n",
    "        if line[0] == \"[\":\n",
    "            section = line.strip(\"[]\").lower()\n",
    "            continue\n",
    "\n",
    "        if section == \"freq\":\n",
    "            freq.append(float(line))\n",
    "\n",
    "        if section == \"fr-coord\":\n",
    "            el, x, y, z = line.split()\n",
    "            eq_coords.append([el, float(x) * b2A, float(y) * b2A, float(z) * b2A])\n",
    "\n",
    "        if section == \"fr-norm-coord\":\n",
    "            if line.startswith(\"vibration\"):\n",
    "                vibr_displacements.append([])\n",
    "                continue\n",
    "            coords = [float(x) for x in line.split()]\n",
    "            vibr_displacements[-1].append(coords)\n",
    "\n",
    "        if section == \"int\":\n",
    "            inten.append(float(line))\n",
    "\n",
    "    return freq, eq_coords, vibr_displacements, inten\n",
    "\n",
    "\n",
    "# READ QE spectrum from dynmat.x output\n",
    "\n",
    "\n",
    "def read_dynmat_output(filename):\n",
    "    with open(filename) as f:\n",
    "        data = f.readlines()\n",
    "    freq = []\n",
    "    inten = []\n",
    "    for line in data:\n",
    "        line = line.strip()\n",
    "        parts = line.split()\n",
    "        if len(parts) == 4 and parts[0].isdigit():\n",
    "            freq.append(float(parts[1]))\n",
    "            inten.append(float(parts[3]))\n",
    "    return freq, inten\n",
    "\n",
    "\n",
    "datasets = {}  # {label} = [0 - label, 1 - file, 2 - freq, 3 - intensity]\n",
    "vibr_disp_datasets = (\n",
    "    {}\n",
    ")  # {label} = [0 - freq, 1 - eq_coords, 2 - vibr_displacements, 3 - intensity]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate spectra\n",
    "\n",
    "# broadening\n",
    "fwhm = 5.0\n",
    "\n",
    "\n",
    "def lorentzian(x):\n",
    "    return 0.5 * fwhm / (np.pi * (x**2 + (0.5 * fwhm) ** 2))\n",
    "\n",
    "\n",
    "def gaussian(x):\n",
    "    sigma = fwhm / 2.3548\n",
    "    return np.exp(-(x**2) / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi))\n",
    "\n",
    "\n",
    "def shape(x):\n",
    "    return lorentzian(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lighten_color(color, amount=1.3):\n",
    "    \"\"\"\n",
    "    Lightens the given color by multiplying (1-luminosity) by the given amount.\n",
    "    Input can be matplotlib color string, hex string, or RGB tuple.\n",
    "\n",
    "    Examples:\n",
    "    >> lighten_color('g', 0.3)\n",
    "    >> lighten_color('#F034A3', 0.6)\n",
    "    >> lighten_color((.3,.55,.1), 0.5)\n",
    "    \"\"\"\n",
    "    import colorsys\n",
    "\n",
    "    import matplotlib.colors as mc\n",
    "\n",
    "    if color in mc.cnames:\n",
    "        c = mc.cnames[color]\n",
    "    elif isinstance(color, str):\n",
    "        h = color.lstrip(\"#\")\n",
    "        c = tuple(int(h[i : i + 2], 16) / 256.0 for i in (0, 2, 4))\n",
    "    else:\n",
    "        c = color\n",
    "    c = colorsys.rgb_to_hls(*mc.to_rgb(c))\n",
    "    new_c = np.array(colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2]))\n",
    "    new_c[new_c < 0.0] = 0.0\n",
    "    new_c[new_c > 1.0] = 1.0\n",
    "    return new_c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vis_mode(mode):\n",
    "    ase_atoms = ase.Atoms(chem_symbols, eq_coords)\n",
    "    plt.figure(figsize=(16, 5))\n",
    "\n",
    "    ax1 = plt.subplot(2, 1, 1)\n",
    "    ax2 = plt.subplot(2, 1, 2)\n",
    "\n",
    "    enhance = 1.0\n",
    "\n",
    "    for i_at, displ in enumerate(vibr_displacements[mode]):\n",
    "        pos = ase_atoms[i_at].position\n",
    "        dis = enhance * np.array(displ)\n",
    "\n",
    "        if ase_atoms[i_at].symbol == \"C\":\n",
    "            col = \"gray\"\n",
    "            rad = 0.4\n",
    "        else:\n",
    "            col = \"white\"\n",
    "            rad = 0.2\n",
    "\n",
    "        p = (pos[0], pos[1])\n",
    "        d = (dis[0], dis[1])\n",
    "        cir_border = plt.Circle(p, rad + 0.05, color=\"black\")\n",
    "        circ = plt.Circle(p, rad, color=col)\n",
    "        ax1.add_artist(cir_border)\n",
    "        ax1.add_artist(circ)\n",
    "        ax1.arrow(\n",
    "            p[0], p[1], d[0], d[1], head_width=0.5, head_length=0.1, fc=\"r\", ec=\"r\"\n",
    "        )\n",
    "\n",
    "        p = (pos[0], pos[2])\n",
    "        d = (dis[0], dis[2])\n",
    "        cir_border = plt.Circle(p, rad + 0.05, color=\"black\")\n",
    "        circ = plt.Circle(p, rad, color=col)\n",
    "        ax2.add_artist(cir_border)\n",
    "        ax2.add_artist(circ)\n",
    "        ax2.arrow(\n",
    "            p[0], p[1], d[0], d[1], head_width=0.5, head_length=0.1, fc=\"r\", ec=\"r\"\n",
    "        )\n",
    "\n",
    "    all_pos = ase_atoms.positions\n",
    "\n",
    "    ax1.axis(\"off\")\n",
    "    ax1.set_xlim([np.min(all_pos[:, 0]) - 1.0, np.max(all_pos[:, 0]) + 1.0])\n",
    "    ax1.set_ylim([np.min(all_pos[:, 1]) - 1.0, np.max(all_pos[:, 1]) + 1.0])\n",
    "    ax1.set_aspect(\"equal\")\n",
    "\n",
    "    ax2.axis(\"off\")\n",
    "    ax2.set_xlim([np.min(all_pos[:, 0]) - 1.0, np.max(all_pos[:, 0]) + 1.0])\n",
    "    ax2.set_ylim([np.min(all_pos[:, 2]) - 4.0, np.max(all_pos[:, 2]) + 4.0])\n",
    "    ax2.set_aspect(\"equal\")\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_pk(_=None):\n",
    "    global datasets, vibr_disp_datasets, chem_symbols, eq_coords, vibr_displacements\n",
    "\n",
    "    try:\n",
    "        freq, eq_coords, vibr_displacements, inten = read_cp2k_molden(pk_select.value)\n",
    "        datasets = {\"IR\": [\"label\", \"file\", freq, inten]}\n",
    "        vibr_disp_datasets = {\"IR\": [freq, eq_coords, vibr_displacements, inten]}\n",
    "        with ir_out:\n",
    "            clear_output()\n",
    "            xmin = 600  # -15\n",
    "            xmax = 1650  #  3400\n",
    "            x = np.arange(xmin, xmax, 0.01)\n",
    "\n",
    "            spectra = {}  # {label} = [x, spectrum]\n",
    "\n",
    "            for label, ds in datasets.items():\n",
    "                spectrum = np.zeros(len(x))\n",
    "\n",
    "                freq = ds[2]\n",
    "                inten = np.array(ds[3]) ** 2\n",
    "\n",
    "                for i, f in enumerate(freq):\n",
    "                    spectrum += inten[i] * shape(x - f)\n",
    "\n",
    "                spectra[label] = [x, spectrum]\n",
    "            label = \"IR\"\n",
    "            freq, eq_coords, vibr_displacements, inten = vibr_disp_datasets[label]\n",
    "            chem_symbols = np.array(eq_coords)[:, 0]\n",
    "            eq_coords = np.array(eq_coords)[:, 1:].astype(np.float32)\n",
    "            vibr_displacements = np.array(vibr_displacements)\n",
    "            [x, spectrum] = spectra[label]\n",
    "\n",
    "            plot_data = [\n",
    "                [\"IR\", [\"#d62728\", \"-\", 2, 10], \"c2h2\"],\n",
    "            ]\n",
    "\n",
    "            title = \"main\"\n",
    "\n",
    "            scaling = 1.0\n",
    "\n",
    "            # Preprocess before plottings\n",
    "            spectra_max = 0.0\n",
    "\n",
    "            plot_spectra = []\n",
    "\n",
    "            for i, pd in enumerate(plot_data):\n",
    "                label = pd[0]\n",
    "                [x, spectrum] = spectra[label]\n",
    "\n",
    "                plot_spectrum = np.copy(spectrum)\n",
    "\n",
    "                # Scale highest peak to 1.0\n",
    "                plot_spectrum *= 1.0 / np.max(plot_spectrum)\n",
    "                # flip spectrum\n",
    "                plot_spectrum = 1.0 - plot_spectrum\n",
    "\n",
    "                if np.max(plot_spectrum) > spectra_max:\n",
    "                    spectra_max = np.max(plot_spectrum)\n",
    "\n",
    "                plot_spectra.append(plot_spectrum)\n",
    "\n",
    "            plt.figure(figsize=(2, 2))\n",
    "\n",
    "            for i, pd in enumerate(plot_data):\n",
    "                label = pd[0]\n",
    "                pl = pd[1]\n",
    "\n",
    "                [x, spectrum] = spectra[label]\n",
    "                plot_spectrum = plot_spectra[i]\n",
    "\n",
    "                if len(pd) > 2:\n",
    "                    label = pd[2]\n",
    "\n",
    "                # Plot the SQUARE as main\n",
    "                p = plt.plot(\n",
    "                    scaling * x,\n",
    "                    plot_spectrum + (len(plot_data) - 1 - i) * spectra_max * 1.05,\n",
    "                    label=label,\n",
    "                    color=\"k\",\n",
    "                    linestyle=pl[1],\n",
    "                    linewidth=1,\n",
    "                    zorder=pl[3],\n",
    "                )\n",
    "\n",
    "                plt.text(\n",
    "                    680,\n",
    "                    (len(plot_data) - i) * spectra_max * 1.05 - 0.25,\n",
    "                    label,\n",
    "                    va=\"center\",\n",
    "                    ha=\"center\",\n",
    "                    fontdict={\"size\": 8},\n",
    "                )\n",
    "\n",
    "            plt.xlim([scaling * xmin, scaling * xmax])\n",
    "            plt.xlim([640, 1060])\n",
    "            plt.ylim([-0.1, 0.1 + 1.05 * spectra_max * (len(plot_data))])\n",
    "\n",
    "            plt.xlabel(\"Energy [cm$^{-1}$]\")\n",
    "            plt.ylabel(\"Relative intensity\")\n",
    "\n",
    "            ax = plt.gca()\n",
    "            ax.axes.yaxis.set_ticks([])\n",
    "\n",
    "            ax.xaxis.set_minor_locator(AutoMinorLocator(5))\n",
    "            ax.tick_params(axis=\"x\", which=\"both\", direction=\"in\")\n",
    "            plt.show()\n",
    "\n",
    "            plt.figure(figsize=(4, 2))\n",
    "\n",
    "            max_intensity = np.max(inten)\n",
    "            min_f = 1e6\n",
    "            max_f = 0.0\n",
    "            y_shift = 0\n",
    "            f_last = 0.0\n",
    "\n",
    "            for i in range(len(freq)):\n",
    "                f = freq[i]\n",
    "                if np.min(x) < f < np.max(x) and inten[i] > 0.10 * max_intensity:\n",
    "                    plt.axvline(f, linestyle=\"--\", color=\"b\")\n",
    "                    if f - f_last > 50:\n",
    "                        y_shift = 0\n",
    "                    f_last = f\n",
    "                    y_pos = (\n",
    "                        np.max(spectrum) - 1e-6 - y_shift * np.max(spectrum) / 20\n",
    "                    ) % np.max(spectrum)\n",
    "                    y_shift += 1\n",
    "                    t = plt.text(f, y_pos, i)\n",
    "                    t.set_bbox(dict(facecolor=\"white\", edgecolor=\"black\"))\n",
    "                    if f < min_f:\n",
    "                        min_f = f\n",
    "                    if f > max_f:\n",
    "                        max_f = f\n",
    "            plt.plot(x, spectrum, \"r\", lw=2.0)\n",
    "            plt.xlim(min_f - 1, max_f + 1)\n",
    "            # plt.savefig(\"./fig/zoom.png\", dpi=300, bbox_inches='tight')\n",
    "            plt.show()\n",
    "\n",
    "    except AssertionError:\n",
    "        print(\"Incorrect pk.\")\n",
    "        return\n",
    "\n",
    "\n",
    "style = {\"description_width\": \"50px\"}\n",
    "layout = {\"width\": \"70%\"}\n",
    "\n",
    "pk_select = ipw.IntText(value=0, description=\"pk\", style=style, layout=layout)\n",
    "\n",
    "load_pk_btn = ipw.Button(description=\"Load pk\", style=style, layout=layout)\n",
    "load_pk_btn.on_click(load_pk)\n",
    "\n",
    "ir_out = ipw.Output()\n",
    "display(pk_select, load_pk_btn, ir_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_animate_click(c):\n",
    "    with animation_out:\n",
    "        clear_output()\n",
    "        trajectory = []\n",
    "        time_arr = np.linspace(0.0, 2 * np.pi, 20)\n",
    "        print(\"Mode %d\" % mode_to_display.value)\n",
    "        vis_mode(mode_to_display.value)\n",
    "        for time in time_arr:\n",
    "            vibr_atoms = ase.Atoms(\n",
    "                chem_symbols,\n",
    "                eq_coords\n",
    "                + enhance.value\n",
    "                * np.sin(time)\n",
    "                * vibr_displacements[mode_to_display.value],\n",
    "            )\n",
    "            trajectory.append(vibr_atoms)\n",
    "        nglv = view(trajectory, viewer=\"ngl\")\n",
    "        display(nglv)\n",
    "\n",
    "\n",
    "mode_to_display = ipw.IntText(description=\"Mode\", value=1)\n",
    "enhance = ipw.FloatText(description=\"Amplitude\", value=2)\n",
    "animation_out = ipw.Output()\n",
    "animate = ipw.Button(description=\"Animate\")\n",
    "animate.on_click(on_animate_click)\n",
    "\n",
    "display(ipw.VBox([mode_to_display, enhance, animate, animation_out]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the URL after everything is set up.\n",
    "try:\n",
    "    url = urlparse.urlsplit(jupyter_notebook_url)\n",
    "    pk_select.value = urlparse.parse_qs(url.query)[\"pk\"][0]\n",
    "    load_pk(0)\n",
    "except Exception as exception:\n",
    "    print(exception)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}