MsSpec-DFM/src/python/msspec_dfm/plotting.py

#!/usr/bin/env python
#
# Copyright © 2020-2020 - Rennes Institute of Physics.
#
# This file is part of plotdf.
#
# plotdf is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# plotdf is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with plotdf.  If not, see <http://www.gnu.org/licenses/>.
#
# Source file  : src/old_plotdf.py
# Last modified: Tue, 27 Apr 2021 11:00:13 +0200
# Committed by : Sylvain Tricot <sylvain.tricot@univ-rennes1.fr>

import sys

try:
    from matplotlib import pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm
    from matplotlib.colorbar import Colorbar
except ImportError:
    print('You need matplotlib to run this script!')
    exit(-1)

try:
    import numpy as np
except ImportError:
    print('You need NumPy to run this script!')
    exit(-1)

try:
    import argparse
except ImportError as err:
    print('Please, install the package \"{}\" to run that script.'.format(err.name))
    exit(-1)

try:
    import logging
except ImportError as err:
    print('Please, install the package \"{}\" to run that script.'.format(err.name))
    exit(-1)





# Choose diverging colormap
# (PiYG, PRGn, BrBG, PuOr, RdGy, RdBu, RdYlBu, RdYlGn, Spectral, coolwarm, bwr, seismic)

# custom cmap
cmap = cm.bwr.from_list("custom_cmap", [
                        (0.00, (.6, 0.8, 1.0)),
                        (0.48, (0.2, 0.2, .8)),
                        (0.50, (1, 1, 1)),
                        (0.52, (1., 0., 0.)),
                        (1.00, (1.0, 0.8, 0)),
                        ])

#cmap = cm.bwr

def load_data(df='diel_func.dat', pd='plas_disp.dat', eh='elec_hole.dat'):
    logging.info('loading data, please wait...')
    # load data
    logging.debug('loading {:s} file'.format(df))
    data = np.loadtxt(df)
    logging.debug('loading {:s} file'.format(pd))
    pd = np.loadtxt(pd)
    logging.debug('loading {:s} file'.format(eh))
    eh = np.loadtxt(eh)

    # convert to grid
    logging.debug('Creating the grid')
    x = np.unique(data[:,0])
    y = np.unique(data[:,1])

    X, Y = np.meshgrid(x, y)
    size = (len(x), len(y))
    ReZ = data[:,2].reshape(size).T
    ImZ = data[:,3].reshape(size).T

    logging.debug('loading data done.')

    return X, Y, ReZ, ImZ, pd, eh

def plot_2d(X, Y, Z, cmap, fig=None, ax=None, vmin=None, vmax=None):
    # plotting
    if fig is None:
        fig = plt.figure()
    if ax is None:
        ax = fig.add_subplot(111)

    xmin, xmax, ymin, ymax = [np.min(X), np.max(X), np.min(Y), np.max(Y)]
    extent = [xmin, xmax, ymin, ymax]
    im = ax.imshow(Z, cmap=cmap, aspect='auto', origin='lower', vmin=vmin, vmax=vmax, extent=extent,
               interpolation='bilinear')
    ax.set_xlim(None, xmax)
    ax.set_ylim(None, ymax)

    # colorbar
    cb_val = np.linspace(vmin, vmax, 9, endpoint=True)
    #cb_lbl = [f"{_:.1f}" for _ in cb_val]
    cb_lbl = ["{:.1f}".format(_) for _ in cb_val]
    cb_lbl[0] = "$\leq$ {:.1f}".format(cb_val[0])
    cb_lbl[-1] = "$\geq$ {:.1f}".format(cb_val[-1])
    cb = fig.colorbar(im, ax=ax, ticks=cb_val)
    cb.ax.set_yticklabels(cb_lbl)

    return fig, ax


def plot_contour(X, Y, Z, colors='black', cmap=None, fig=None, ax=None, levels=8,
                 vmin=None, vmax=None):
    # plotting
    if fig is None:
        fig = plt.figure()
    if ax is None:
        ax = fig.add_subplot(111)

    xmin, xmax, ymin, ymax = [np.min(X), np.max(X), np.min(Y), np.max(Y)]
    extent = [xmin, xmax, ymin, ymax]
    cset = ax.contour(X, Y, Z, zdir='z', levels=np.linspace(vmin,vmax,levels), offset=np.min(Z),
                      colors=colors, cmap=cmap)
    ax.clabel(cset, cset.levels, inline=True, fontsize=10)

    ax.set_xlim(None, xmax)
    ax.set_ylim(None, ymax)

    return fig, ax


# Plot the 3D surface
def plot_3d(X, Y, Z, cmap=None, fig=None, ax=None, vmin=None, vmax=None, strides=64):
    # plotting
    if fig is None:
        fig = plt.figure()
    if ax is None:
        ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(X, Y, Z, vmin=vmin, vmax=vmax, rstride=strides, cstride=strides, alpha=1., cmap=cmap)

    return fig, ax


def plot_pdeh(pd, eh, Ef=1, color='black', fig=None, ax=None):
    # plotting
    if fig is None:
        fig = plt.figure()
    if ax is None:
        ax = fig.add_subplot(111)

    ax.plot(pd[:,0], pd[:,1], color=color, linestyle="--", label="Plasmon energy")
    ax.plot(eh[:,0], eh[:,1]/Ef, color=color, linestyle=":", label=r"e-h pair continuum $\omega_+$")
    ax.plot(eh[:,0], np.abs(eh[:,2]/Ef), color=color, linestyle="-.", label=r"e-h pair continuum $\omega_-$")
    ax.legend()

    return fig, ax


def entitle(fig, ax, img=False):
    title = r"$\Im(\epsilon)$" if img else r"$\Re(\epsilon)$"
    ax.set_title(title)
    ax.set_xlabel(r"$q/k_F$")
    ax.set_ylabel(r"$E/E_F$")