Source code for comet.snmr.misc.IO_pdf

"""
Part of comet/snmr/misc
"""
# COMET - COupled Magnetic resonance Electrical resistivity Tomography
# Copyright (C) 2019  Nico Skibbe

# This program 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.

# This program 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 this program. If not, see <https://www.gnu.org/licenses/>.

import matplotlib.pyplot as plt
import pygimli as pg
import numpy as np
import matplotlib as mpl


def returnFigure(ax):
    import matplotlib.figure
    if not isinstance(ax, matplotlib.figure.Figure):
        try:
            fig = ax.get_figure()
        except AttributeError:  # 1d axis
            try:
                fig = ax[0].get_figure()
            except AttributeError:  # 2d axis
                fig = ax[0, 0].get_figure()
            except TypeError as t:
                print('ax is None')
                raise t

        return fig
    else:
        return ax


def closeAxis(ax):
    plt.close(returnFigure(ax))


def robustPDFSave(fig_or_ax, name, **kwargs):
    import time
    from comet.pyhed.IO import checkDirectory
    checkDirectory(name, filename=True)
    saved = False
    fig = returnFigure(fig_or_ax)
    kwargs.setdefault('bbox_inches', 'tight')
    while saved is False:
        try:
            fig.savefig(name, **kwargs)
            saved = True
            print('saved "{}"'.format(name))
        except PermissionError:
            print('Permission denied, close pdf viewer. \
I will try again in 5 seconds.')
            time.sleep(5.0)


def exportKernelPDF(kern, fig=None, ax=None, savename=None, dpi=300,
                    noYLabel=False, index=0, xl_add='', rotate=False,
                    plotlims=None, colorbar=False, figsize=[10, 6]):
    import kernel as k
    if isinstance(kern, str):
        kern = k.Kernel(kern, verbose=True)

    if kern.dimension != 2:
        raise Exception('Only for 2D Kernel')

    cmin = -2
    cmax = 2
    if fig is None or ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)
    if rotate:
        kern.kernelMesh2D.rotate(pg.RVector3((0.0, 0.0, np.pi/2.)))
    im, cbar = kern.show(indices=[index], cMin=cmin, cMax=cmax, size=14, ax=ax,
                         colorBar=colorbar)
    ax.set_ylabel('X [m]')
    ax.set_xlabel('Z [m]')
    if plotlims is not None:
        ax.set_xlim(plotlims[0], plotlims[1])
        ax.set_ylim(plotlims[2], plotlims[3])
    if noYLabel is True:
        ax.set_ylabel('')
        ax.yaxis.set_major_locator(plt.NullLocator())
    if savename is not None:
        if savename.endswith('.png'):
            robustPDFSave(fig, savename, bbox_inches='tight', dpi=dpi)
        else:
            robustPDFSave(fig, savename, bbox_inches='tight')
        plt.close()


def exportColorBarPDF(name=None, cMap='viridis', cmin=0, cmax=1,
                      orientation='horizontal', label='colorbar label',
                      size=14, ax=None, dpi=300):
    # import matplotlib.pyplot as plt
    # import matplotlib as mpl
    if ax is None:
        if orientation == 'horizontal':
            fig = plt.figure(figsize=(8, 2))
            ax = fig.add_axes([0.05, 0.80, 0.9, 0.15])
        else:
            fig = plt.figure(figsize=(2, 8))
            ax = fig.add_axes([0.05, 0.05, 0.25, 0.9])

    if isinstance(cMap, str):
        cMap = plt.get_cmap(cMap)
    norm = mpl.colors.Normalize(vmin=cmin, vmax=cmax)
    cbar = mpl.colorbar.ColorbarBase(ax, cmap=cMap,
                                     norm=norm,
                                     orientation=orientation)
    cbar.ax.tick_params(labelsize=size)
    cbar.set_label(label, size=size)
    cbar.set_ticks([cmin, 0, cmax])
    if orientation == 'horizontal':
        cbar.ax.set_xticklabels(['< $10^{' + str(cmin) + '}$', '1',
                                 '> $10^{}$'.format(cmax)])
    else:
        cbar.ax.set_yticklabels(['< $10^{' + str(cmin) + '}$', '1',
                                 '> $10^{}$'.format(cmax)])
    if name is not None:
        if name.endswith('.png'):
            robustPDFSave(fig, name, bbox_inches='tight', dpi=dpi)
        else:
            robustPDFSave(fig, name, bbox_inches='tight')
        plt.close()
    return ax

# The End