Merge pull request #3 from phonopy/plot-transport

Add velph-transport-plot

src/phelel/velph/cli/transport/__init__.py

@@ -63,3 +63,6 @@ def cmd_generate(toml_filename: str, hdf5_filename: str, dry_run: bool):
 def cmd_check_fft(toml_filename: str):
     """Show [NGX, NGY, NGZ] in vasprun.xml."""
     check_fft(toml_filename, "transport")
+
+
+from phelel.velph.cli.transport.plot import cmd_plot  # noqa F401

src/phelel/velph/cli/transport/plot/__init__.py

@@ -0,0 +1,74 @@
+"""velph command line tool / velph-transport."""
+
+from __future__ import annotations
+
+import pathlib
+from typing import Optional
+
+import click
+import h5py
+
+from phelel.velph.cli.transport import cmd_transport
+from phelel.velph.cli.transport.plot.plot_selfenergy import plot_selfenergy
+from phelel.velph.cli.transport.plot.plot_transport import plot_transport
+
+
+@cmd_transport.group("plot")
+@click.help_option("-h", "--help")
+def cmd_plot():
+    """Choose plot options."""
+    pass
+
+
+@cmd_plot.command("selfenergy")
+@click.argument(
+    "vaspout_filename",
+    nargs=1,
+    type=click.Path(),
+    default="transport/vaspout.h5",
+)
+@click.help_option("-h", "--help")
+def cmd_plot_selfenergy(vaspout_filename: str):
+    """Plot self-energy in transports."""
+    args = _get_f_h5py_and_plot_filename(
+        "selfenergy", vaspout_filename=pathlib.Path(vaspout_filename)
+    )
+    if args[0] is not None:
+        plot_selfenergy(*args)
+
+
+@cmd_plot.command("transport")
+@click.argument(
+    "vaspout_filename",
+    nargs=1,
+    type=click.Path(),
+    default="transport/vaspout.h5",
+)
+@click.help_option("-h", "--help")
+def cmd_plot_transport(vaspout_filename: str):
+    """Plot transport in transports."""
+    args = _get_f_h5py_and_plot_filename(
+        "transport", vaspout_filename=pathlib.Path(vaspout_filename)
+    )
+    if args[0] is not None:
+        plot_transport(*args)
+
+
+def _get_f_h5py_and_plot_filename(
+    property_name: str,
+    vaspout_filename: pathlib.Path = pathlib.Path("transport/vaspout.h5"),
+    plot_filename: Optional[pathlib.Path] = None,
+) -> tuple[h5py._hl.files.File, pathlib.Path]:
+    if not vaspout_filename.exists():
+        click.echo(f'"{vaspout_filename}" (default path) not found.')
+        click.echo("Please specify vaspout.h5 file path.")
+        return None, None
+
+    if plot_filename is None:
+        _plot_filename = vaspout_filename.parent / f"{property_name}.pdf"
+    else:
+        _plot_filename = plot_filename
+
+    f_h5py = h5py.File(vaspout_filename)
+
+    return f_h5py, _plot_filename

src/phelel/velph/cli/transport/plot/plot_selfenergy.py

@@ -0,0 +1,86 @@
+"""Implementation of velph-el_bands-plot."""
+
+from __future__ import annotations
+
+import click
+import h5py
+import matplotlib.pyplot as plt
+
+
+def plot_selfenergy(f_h5py: h5py._hl.files.File, plot_filename: str, show: bool = True):
+    """Plot imaginary part of self-energies."""
+    selfens = {}
+    for key in f_h5py["results"]["electron_phonon"]["electrons"]:
+        if "self_energy_" in key:
+            selfens[int(key.split("_")[2])] = f_h5py["results"]["electron_phonon"][
+                "electrons"
+            ][key]
+
+    if len(selfens) == 1:
+        fig, axs = plt.subplots(1, 1, figsize=(4, 4))
+    else:
+        nrows = len(selfens) // 2
+        fig, axs = plt.subplots(nrows, 2, figsize=(8, 4 * nrows), squeeze=True)
+
+    for i in range(len(selfens)):
+        selfen = selfens[i + 1]
+        _show(selfen, i + 1)
+        _plot(axs[i], selfen)
+
+    plt.tight_layout()
+    if show:
+        plt.show()
+    else:
+        plt.rcParams["pdf.fonttype"] = 42
+        plt.savefig(plot_filename)
+    plt.close()
+
+    click.echo(f'Transport plot was saved in "{plot_filename}".')
+
+
+def _plot(ax, selfen):
+    for i_nw in range(selfen["nw"][()]):
+        for i_temp, _ in enumerate(selfen["temps"]):
+            ax.plot(
+                selfen["energies"][:, i_nw],
+                selfen["selfen_fan"][:, i_nw, i_temp, 1],
+                ".",
+            )
+
+
+def _show(selfen: h5py._hl.group.Group, index: int):
+    """Show self-energy properties.
+
+    ['band_start', 'band_stop', 'bks_idx', 'carrier_per_cell',
+    'carrier_per_cell0', 'delta', 'efermi', 'energies', 'enwin', 'nbands',
+    'nbands_sum', 'nw', 'scattering_approximation', 'select_energy_window',
+    'selfen_dw', 'selfen_fan', 'static', ' ', 'tetrahedron']
+
+    """
+    print(f"- parameters:  # {index}")
+    print(
+        "    scattering_approximation:",
+        selfen["scattering_approximation"][()].decode("utf-8"),
+    )
+    print(f"    static_approximation: {bool(selfen['static'][()])}")
+    print(f"    use_tetrahedron_method: {bool(selfen["tetrahedron"][()])}")
+    if not selfen["tetrahedron"][()]:
+        print(f"    smearing_width: {selfen['delta'][()]}")
+    print(
+        f"    band_start_stop: [{selfen['band_start'][()]}, {selfen['band_stop'][()]}]"
+    )
+    print(f"    nbands: {selfen['nbands'][()]}")
+    print(f"    nbands_sum: {selfen['nbands_sum'][()]}")
+    print(f"    nw: {selfen['nw'][()]}")
+    print("    temperatures:")
+    for i, t in enumerate(selfen["temps"]):
+        print(f"    - {t}  # {i + 1}")
+
+    print("  data_scalar:")
+    print(f"    carrier_per_cell0: {selfen["carrier_per_cell0"][()]}")
+
+    print("  data_array_shapes:")
+    print(f"    carrier_per_cell: {list(selfen["carrier_per_cell"].shape)}")
+    print(f"    Fan_self_energy: {list(selfen["selfen_fan"].shape)}")
+    print(f"    sampling_energy_points: {list(selfen["energies"].shape)}")
+    print(f"    Fermi_energies: {list(selfen["efermi"].shape)}")

src/phelel/velph/cli/transport/plot/plot_transport.py

@@ -0,0 +1,130 @@
+"""Implementation of velph-el_bands-plot."""
+
+from __future__ import annotations
+
+import click
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_transport(f_h5py: h5py._hl.files.File, plot_filename: str, show: bool = True):
+    """Plot transport properties."""
+    transports = {}
+    for key in f_h5py["results"]["electron_phonon"]["electrons"]:
+        if "transport_" in key:
+            transports[int(key.split("_")[1])] = f_h5py["results"]["electron_phonon"][
+                "electrons"
+            ][key]
+
+    for i in range(len(transports)):
+        transport = transports[i + 1]
+        _show(transport, i + 1)
+
+    transports_temps = {}
+    for i in range(len(transports)):
+        transport = transports[i + 1]
+        key = transport["scattering_approximation"][()].decode("utf-8")
+        if key in transports_temps:
+            transports_temps[key].append(transports[i + 1])
+        else:
+            transports_temps[key] = [
+                transports[i + 1],
+            ]
+
+    nrows = len(transports_temps)
+    fig, axs = plt.subplots(nrows, 5, figsize=(20, 4 * nrows))
+
+    for i, key in enumerate(transports_temps):
+        _plot(axs[i, :], transports_temps[key])
+
+    plt.tight_layout()
+    if show:
+        plt.show()
+    else:
+        plt.rcParams["pdf.fonttype"] = 42
+        plt.savefig(plot_filename)
+    plt.close()
+
+    click.echo(f'Transport plot was saved in "{plot_filename}".')
+
+
+def _plot(axs, tranports_temps):
+    property_names = (
+        "e_conductivity",
+        "mobility",
+        "e_t_conductivity",
+        "peltier",
+        "seebeck",
+    )
+    properties = [[] for _ in property_names]
+    temps = []
+    for trpt in tranports_temps:
+        temps.append(trpt["temperature"][()])
+        for i, property in enumerate(property_names):
+            properties[i].append(np.trace(trpt[property][:]) / 3)
+
+    for i, property in enumerate(property_names):
+        axs[i].plot(temps, properties[i], ".-")
+        axs[i].set_xlabel("temperature (K)")
+        axs[i].set_ylabel(property)
+
+
+def _show(transport: h5py._hl.group.Group, index: int):
+    """Show transport properties.
+
+    ['cell_volume', 'dfermi_tol', 'e_conductivity', 'e_t_conductivity', 'emax',
+    'emin', 'energy', 'lab', 'mobility', 'mu', 'n', 'n0', 'ne', 'nedos',
+    'nelect', 'nh', 'peltier', 'scattering_approximation', 'seebeck', 'static',
+    'tau_average', 'temperature', 'transport_function']
+
+    """
+    print(f"- parameters:  # {index}")
+    print(
+        "    scattering_approximation:",
+        transport["scattering_approximation"][()].decode("utf-8"),
+    )
+    print(f"    temperature: {transport["temperature"][()]}")
+    for key in (
+        "cell_volume",
+        "n",
+        ("n0", "number_of_electrons_gausslegendre"),
+        "nedos",
+        "nelect",
+        "cell_volume",
+        "dfermi_tol",
+    ):
+        if isinstance(key, tuple):
+            print(f"    {key[1]}: {transport[key[0]][()]}")
+        else:
+            print(f"    {key}: {transport[key][()]}")
+    print(f"    static_approximation: {bool(transport["static"][()])}")
+
+    print("  data_scalar:")
+    for key in (
+        ("emax", "emax_for_transport_function"),
+        ("emin", "emin_for_transport_function"),
+        ("mu", "chemical_potential"),
+        ("ne", "ne_in_conduction_band"),
+        ("nh", "nh_in_valence_band"),
+        "tau_average",
+    ):
+        if isinstance(key, tuple):
+            print(f"    {key[1]}: {transport[key[0]][()]}")
+        else:
+            print(f"    {key}: {transport[key][()]}")
+    print("  data_array_shapes:")
+    for key in (
+        "e_conductivity",
+        "e_t_conductivity",
+        "energy",
+        ("lab", "Onsager_coefficients"),
+        "mobility",
+        "peltier",
+        "seebeck",
+        "transport_function",
+    ):
+        if isinstance(key, tuple):
+            print(f"    {key[1]}: {list(transport[key[0]].shape)}")
+        else:
+            print(f"    {key}: {list(transport[key].shape)}")