Source code for analysator.plot

#
# This file is part of Analysator.
# Copyright 2013-2016 Finnish Meteorological Institute
# Copyright 2017-2018 University of Helsinki
#
# For details of usage, see the COPYING file and read the "Rules of the Road"
# at http://www.physics.helsinki.fi/vlasiator/
#
# 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 2 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, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#

''' The plot module has all the functions related to plotting variables.

   .. code-block:: python

      # Example:
      import analysator as pt
      pt.pt.
      #press [tab] -> get the functions

'''


#for usage with "from (package) import *"
__all__ = ['plot_colormap', 'plot_vdf', 'plot_vdfdiff', 'plot_vdf_profiles', 'plot_colormap3dslice',
           'plot_threeslice', 'plot_ionosphere', 'plot_isosurface', 'plot_neutral_sheet',
           'plot_variables','colormaps', 'plot_helpers']

from .plot_variables import plot_variables, plot_multiple_variables

import analysator as pt

import logging
import matplotlib.pyplot as plt
import matplotlib
from . import colormaps

from . import plot_helpers
from .plot_colormap import plot_colormap
from .plot_vdf import plot_vdf
from .plot_vdfdiff import plot_vdfdiff
from .plot_vdf_profiles import plot_vdf_profiles
from .plot_colormap3dslice import plot_colormap3dslice
from .plot_threeslice import plot_threeslice
from .plot_ionosphere import plot_ionosphere
import platform
from packaging.version import Version
from numbers import Number

try:
    from .plot_isosurface import plot_isosurface, plot_neutral_sheet
    if Version(platform.python_version()) < Version("3.7"):
       raise ImportError("Python>=3.7 required for RBFInterpolator.")
except:
    logging.warning("plot_isosurface not imported. To access it, use Python version >=3.7 and install scikit-image.")


import numpy as np, os

if Version(matplotlib.__version__) < Version("3.3.0"):
    plt.register_cmap(name='viridis', cmap=colormaps.viridis)
    plt.register_cmap(name='viridis_r', cmap=colormaps.viridis_r)
    plt.register_cmap(name='plasma', cmap=colormaps.plasma)
    plt.register_cmap(name='plasma_r', cmap=colormaps.plasma_r)
    plt.register_cmap(name='inferno', cmap=colormaps.inferno)
    plt.register_cmap(name='inferno_r', cmap=colormaps.inferno_r)
    plt.register_cmap(name='magma', cmap=colormaps.magma)
    plt.register_cmap(name='magma_r', cmap=colormaps.magma_r)

# Register custom colourmaps
if Version(matplotlib.__version__) < Version("3.5.0"):
    plt.register_cmap(name='parula', cmap=colormaps.parula)
    plt.register_cmap(name='parula_r', cmap=colormaps.parula_r)
    plt.register_cmap(name='hot_desaturated', cmap=colormaps.hot_desaturated_colormap)
    plt.register_cmap(name='hot_desaturated_r', cmap=colormaps.hot_desaturated_colormap_r) # Listed colormap requires making reversed version at earlier step
    plt.register_cmap(name='pale_desaturated', cmap=colormaps.pale_desaturated_colormap)
    plt.register_cmap(name='pale_desaturated_r', cmap=colormaps.pale_desaturated_colormap_r) # Listed colormap requires making reversed version at earlier step
    plt.register_cmap(name='warhol', cmap=colormaps.warhol_colormap)
else:
    matplotlib.colormaps.register(name='parula', cmap=colormaps.parula)
    matplotlib.colormaps.register(name='parula_r', cmap=colormaps.parula_r)
    matplotlib.colormaps.register(name='hot_desaturated', cmap=colormaps.hot_desaturated_colormap)
    matplotlib.colormaps.register(name='hot_desaturated_r', cmap=colormaps.hot_desaturated_colormap_r) # Listed colormap requires making reversed version at earlier step
    matplotlib.colormaps.register(name='pale_desaturated', cmap=colormaps.pale_desaturated_colormap)
    matplotlib.colormaps.register(name='pale_desaturated_r', cmap=colormaps.pale_desaturated_colormap_r) # Listed colormap requires making reversed version at earlier step
    matplotlib.colormaps.register(name='warhol', cmap=colormaps.warhol_colormap)


decimalprecision_ax = 0
decimalprecision_cblin = 0
cb_linear = False


if matplotlib.__version__=="0.99.1.1" and np.__version__=="1.4.1":
   logging.info('Warning, according to loaded numpy and matplotlib versions, user appears to be')
   logging.info('either using csc.taito.fi without loading the mayavi2 module, or by invoking')
   logging.info('the system python interpeter by calling "./scriptname.py" instead of "python ./scriptname.py"')

# Run TeX typesetting through the full TeX engine instead of python's own mathtext. Allows
# for changing fonts, bold math symbols etc, but may cause trouble on some systems.
if not os.getenv('PTNOLATEX'):
   matplotlib.rc('text', usetex=True)
   matplotlib.rcParams['text.latex.preamble'] = r'\boldmath'
   # matplotlib.rcParams['mathtext.fontset'] = 'stix'
   # matplotlib.rcParams['font.family'] = 'STIXGeneral'
   # Matplotlib suppressed logging messages came out after enabling logging.INFO: font.family must be one of (serif, sans-serif, cursive, monospace) when text.usetex is True. serif will be used by default.
   matplotlib.rcParams['font.family'] = 'serif'
   logging.info("Using LaTeX formatting")
   # matplotlib.rcParams['text.dvipnghack'] = 'True' # This hack might fix it on some systems

# Set backends
if matplotlib.get_backend()[:9] == 'module://':
   logging.info("Using backend "+matplotlib.get_backend())
   backend_interactive = matplotlib.get_backend()
   backend_noninteractive = matplotlib.get_backend()
elif not os.getenv('PTBACKEND'):
   backend_interactive = 'TkAgg'
   backend_noninteractive = 'Agg'
else:
   backend_interactive = os.getenv('PTBACKEND')
   backend_noninteractive = os.getenv('PTBACKEND')






if os.getenv('PTNONINTERACTIVE') != None:
   # Non-interactive plotting mode
   try:
      plt.switch_backend(backend_noninteractive)
   except:
      logging.info("Note: Unable to switch to "+backend_noninteractive+" backend")
else:
   # Interactive plotting mode
   plt.ion()
   try:
      plt.switch_backend(backend_interactive)
   except:
      logging.info("Note: Unable to switch to "+backend_interactive+" backend")

pt.backend_interactive=backend_interactive
pt.backend_noninteractive=backend_noninteractive


# Output matplotlib version
logging.info("Using matplotlib version "+matplotlib.__version__)

def get_cmap(colormap):
    if Version(matplotlib.__version__) <= Version("3.6"):
        return matplotlib.cm.get_cmap(name=colormap)
    else:
        return matplotlib.colormaps.get_cmap(colormap)


# Default output directory for plots
defaultoutputdir=os.path.expandvars('$HOME/Plots/')
if os.getenv('PTOUTPUTDIR'):
    defaultoutputdir=os.getenv('PTOUTPUTDIR')

# axisfmt replaces minus sign with en-dash to fix bug with latex descender value return
# nb: axis ticks are never plotted with scientific format
def axisfmt(x, pos):
    f = r'{:.'+decimalprecision_ax+r'f}'
    if not os.getenv('PTNOLATEX'):
        a = f.format(abs(x))
        if np.sign(x)<0: a = r'\mbox{\textbf{--}}'+a
        return r'$'+a+'$'
    else:
        return f.format(x)

# cbfmtsci replaces minus sign with en-dash to fix bug with latex descender value return
# Scientific format for colour bar ticks
def cbfmtsci(x, pos):
    if (cb_linear is True):
        # for linear, use more precision
        a, b = ('{:.'+str(int(decimalprecision_cblin))+'e}').format(x).split('e')
        precisionvalue = int(decimalprecision_cblin)
        if int(b) < 0:
            precisionvalue += int(b)
        if abs(precisionvalue)<1:
            precisionvalue=1
        f = '{:.' + str(abs(precisionvalue)) + 'f}'
        number = f.format(abs(float(a)))+r'{\times}'+'10^{{{}}}'.format(int(b))
    else:
        a, b = '{:.1e}'.format(x).split('e')
        number = '{:.1f}'.format(abs(float(a)))+r'{\times}'+'10^{{{}}}'.format(int(b))
    signchar=r''
    # Multiple braces for b take care of negative values in exponent
    # brackets around \times remove extra whitespace
    if not os.getenv('PTNOLATEX'):
        # replaces minus sign with en-dash to fix big with latex descender value return
        if np.sign(x)<0: signchar=r'\mbox{\textbf{--}}'
    else:
        if np.sign(x)<0: signchar=r'-'
    # Final special treatment for zero value
    if x==0:
        number = r'0.0{\times}10^{{{0}}}'
    return r'$'+signchar+number+'$'

# cbfmt replaces minus sign with en-dash to fix bug with latex descender value return, used for colorbar
# nb: regular floating i.e. non-scientific format for colorbar ticks
def cbfmt(x, pos):
    # Special treatment for zero value
    if x==0:
        return r'$0.0$'
    # Set required decimal precision
    a, b = '{:.1e}'.format(x).split('e')
    # e.g. 9.0e-1 means we need precision 1
    if (cb_linear is True):
        # for linear, use more precision
        precision = str(int(decimalprecision_cblin))
        #if int(b)<1: precision = str(1+abs(int(b)))
    else:
        precision = '0'
        if int(b)<1: precision = str(abs(int(b)))
    f = r'{:.'+precision+'f}'
    if not os.getenv('PTNOLATEX'):
        a = f.format(abs(x))
        if np.sign(x)<0: a = r'\mbox{\textbf{--}}'+a
        return r'$'+a+'$'
    else:
        return f.format(x)


# Helper routines for latex output handling
def bfstring(string):
    if not os.getenv('PTNOLATEX'):
        if len(string)==0:
            return ''
        else:
            return r'\mathbf{'+string+'}'
    # LaTeX output off
    return string

def rmstring(string):
    if len(string)==0:
        return ''
    else:
        return r'\mathrm{'+string+'}'

def mathmode(string):
    if len(string)==0:
        return ''
    else:
        # First remove any internal possible dollar signs, then wrap whole string into math block
        result = string.replace('$','')
        if os.getenv('PTNOLATEX'):
            # Get rid of latex spaces
            result = result.replace(r'\,','~').replace(r'\qquad','~~~~~~')
        return r"$"+result+"$"

def textbfstring(string):
    if not os.getenv('PTNOLATEX'):
        if len(string)==0:
            return ''
        else:
            return r'\textbf{'+string+'}'
    # LaTex output off
    return string

#Draws sharp contour lines along cell edges (xmeshpass,ymeshpass), for some dpi+linewidth combinations antialiasing True causes issues
def cell_edgecontours(ax,XmeshPass,YmeshPass,heightmap,threshold=0,linewidth=0.5,linestyle='solid',colors='black',antialiased=False):
    x=np.array(XmeshPass[0])
    y=np.array([li[0] for li in YmeshPass])

    v = np.diff(heightmap > threshold, axis=1)
    h = np.diff(heightmap > threshold, axis=0)

    # From https://stackoverflow.com/questions/63458863/way-to-contour-outer-edge-of-selected-grid-region-in-python
    # Check that at least one spot exceeds threshold or the below will crash.
    if np.max(heightmap) > threshold:
        
        l = np.argwhere(v.T)
        vlines = np.array(list(zip(np.stack((x[l[:, 0]+1], y[l[:, 1]])).T,
                                    np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
        l = np.argwhere(h.T)
        hlines = np.array(list(zip(np.stack((x[l[:, 0]], y[l[:, 1] + 1])).T,
                                    np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
        lines = np.vstack((vlines, hlines))
        if np.shape(lines)[1]!=0:
            ax.add_collection(matplotlib.collections.LineCollection(lines, lw=linewidth, colors=colors, linestyle=linestyle,zorder=2,antialiased=antialiased))
    else:
       return 1
    
    return 0
def output_path(outputfile,outputfile_default,outputdir,nooverwrite):
        check=False
        if not outputfile:
            if not outputfile_default:
                outputfile="plot.png"

            outputfile=outputfile_default
            if not outputdir:
                # Check later if outputfile contains path information and combine with defaultoutputdir
                check=True

        # Separate possible path information from outputfile
        outputprefixind = outputfile.rfind('/')
        if outputprefixind >= 0:            
            outputdir = outputfile[:outputprefixind+1]
            outputfile = outputfile[outputprefixind+1:]

        if not outputdir: # default initial path
            outputdir=defaultoutputdir
            check=False
    
        if check:
            #remove leading '/' on outputdir as this would cause issues with os.path.join
            outputprefixind = outputdir.find('/')
            if outputprefixind == 0:            
                outputdir = outputdir[1:]
            outputdir= os.path.join(defaultoutputdir,outputdir)
            outputfile=os.path.join(outputdir,outputfile)
        else:
            outputfile = os.path.join(outputdir,outputfile)

        # Ensure output directory exists
        if not os.path.exists(outputdir):
            try:
                os.makedirs(outputdir)
            except FileExistsError: 
                #Parallel jobs might try to create dir simultaneously.
                pass
            except Exception as e:
                raise IOError("Could not create output directory "+outputdir+": "+str(e))
        if not os.access(outputdir, os.W_OK):
            raise IOError("No write access for directory "+outputdir+" Exiting.")

        # Check if target file already exists and overwriting is disabled
        if (nooverwrite and os.path.exists(outputfile)):            
            if os.stat(outputfile).st_size > 0: # Also check that file is not empty
                raise IOError("Found existing file "+outputfile+" and nooverwrite=True.")
            else:
                logging.warning("Found existing file "+outputfile+" of size zero. Re-rendering.")

        return outputfile

[docs] def get_scaled_units(variable_info,vscale=None, env='EarthSpace', manualDict=None): ''' Return scaling metadata :param variable_info: VariableInfo object used for getting the units :param env: A string to choose the scaling dictionary [default: EarthSpace] :param manualDict: a dictionary of {units : {scalingparams}}; used to update the included dictionary :param vscale: float, factor to scale the variable with :returns: (norming factor, scaledUnits, scaledLatexUnits) ''' # if env=='EarthSpace': variable_info.scaleDict = { 's' : {'defaultScale':1, 1e6: {'scaledUnits':'us', 'scaledLatexUnit':r'$\mu\mathrm{s}$'}, 1e3: {'scaledUnits':'ms', 'scaledLatexUnit':r'$\mathrm{ms}$'} }, 'T' : {'defaultScale':1e9, 1e9: {'scaledUnits':'nT', 'scaledLatexUnit':r'$\mathrm{nT}$'} }, 'K' : {'defaultScale':1e-6, 1e-6:{'scaledUnits':'MK', 'scaledLatexUnit':r'$\mathrm{MK}$'} }, 'Pa' : {'defaultScale':1e9, 1e9:{'scaledUnits':'nPa', 'scaledLatexUnit':r'$\mathrm{nPa}$'} }, '1/m' : {'defaultScale':1e3, 1e3:{'scaledUnits':'1/km', 'scaledLatexUnit':r'$\mathrm{km}^{-1}$'} }, '1/m^3' : {'defaultScale':1e-6, 1e-6:{'scaledUnits':'1/cm^3', 'scaledLatexUnit':r'$\mathrm{cm}^{-3}$'} }, '1/m3' : {'defaultScale':1e-6, 1e-6:{'scaledUnits':'1/cm^3', 'scaledLatexUnit':r'$\mathrm{cm}^{-3}$'} }, 'm/s' : {'defaultScale':1e-3, 1e-3:{'scaledUnits':'km/s', 'scaledLatexUnit':r'$\mathrm{km}\,\mathrm{s}^{-1}$'} }, 'V/m' : {'defaultScale':1e3, 1e3:{'scaledUnits':'mV/m', 'scaledLatexUnit':r'$\mathrm{mV}\,\mathrm{m}^{-1}$'} }, 'eV/cm3' : {'defaultScale':1e-3, 1e-3:{'scaledUnits':'keV/cm^3', 'scaledLatexUnit':r'$\mathrm{keV}\,\mathrm{cm}^{-3}$'} }, 'eV/cm^3': {'defaultScale':1e-3, 1e-3:{'scaledUnits':'keV/cm^3', 'scaledLatexUnit':r'$\mathrm{keV}\,\mathrm{cm}^{-3}$'} }, 'T/m': {'defaultScale':1e-12, 1e-9:{'scaledUnits':'nT/m', 'scaledLatexUnit':r'$\mathrm{nT}\,\mathrm{m}^{-1}$'}, 1e-12:{'scaledUnits':'nT/km', 'scaledLatexUnit':r'$\mathrm{nT}\,\mathrm{km}^{-1}$'} }, 'kg/m3' : {'defaultScale':5.97863741e26, 5.97863741e26:{'scaledUnits':'amu/m^3', 'scaledLatexUnit':r'$\mathrm{amu}\,\mathrm{m}^{-3}$'}, 5.97863741e20:{'scaledUnits':'amu/cm^3', 'scaledLatexUnit':r'$\mathrm{amu}\,\mathrm{cm}^{-3}$'} }, 'A/m^2' : {'defaultScale':1, 1e-3:{'scaledUnits':'mA/m^2', 'scaledLatexUnit':r'$\mathrm{mA}\,\mathrm{m}^{-2}$'} }, 'W/m^2' : {'defaultScale':1, 1e-3:{'scaledUnits':'mW/m^2', 'scaledLatexUnit':r'$\mathrm{mW}\,\mathrm{m}^{-2}$'} }, 'm^2' : {'defaultScale':1, 1e-4:{'scaledUnits':'km^2', 'scaledLatexUnit':r'$\mathrm{km}^{2}$'} }, 'V' : {'defaultScale':1, 1e-3:{'scaledUnits':'mV', 'scaledLatexUnit':r'$\mathrm{mV}$'} }, 'Degrees' : {'defaultScale':1, }, 'mho' : {'defaultScale':1, }, } else: variable_info.scaleDict = {} if manualDict is not None: variable_info.scaleDict.update(manualDict) unitScale = 1.0 scaledUnits = variable_info.units scaledLatexUnits = variable_info.latexunits if variable_info.units != '': dictKey = variable_info.units try: udict = variable_info.scaleDict[dictKey] except: if vscale is None: return 1.0, variable_info.units, variable_info.latexunits else: return vscale, variable_info.units, variable_info.latexunits if vscale is None: try: unitScale = udict['defaultScale'] except: return 1.0, variable_info.units, variable_info.latexunits elif np.isclose(vscale, 1.0): return 1.0, variable_info.units, variable_info.latexunits else: unitScale = vscale if not any([np.isclose(unitScale, tryScale) for tryScale in udict.keys() if isinstance(tryScale, Number)]): # return vscale, variable_info.units+" x{vscale:e}".format(vscale=vscale), variable_info.latexunits+r"{\times}"+cbfmtsci(vscale,None) try: #above guarantees the list comprehension does not give an empty list unitScale = [scale for scale in udict.keys() if isinstance(scale, Number) and np.isclose(scale,unitScale)][0] scaledUnits = udict[unitScale]['scaledUnits'] except KeyError: # logging.info('Missing scaledUnits in specialist dict for' + variable_info.units + ' for unitScale='+str(unitScale)) return 1.0, variable_info.units, variable_info.latexunits try: scaledLatexUnits = udict[unitScale]['scaledLatexUnit'] except: # logging.info('Missing scaledLatexUnits in specialist dict for ' + variable_info.units+ ' for unitScale='+str(unitScale)) return 1.0, variable_info.units, variable_info.latexunits else: if vscale is None or np.isclose(vscale, 1.0): return 1.0, variable_info.units, variable_info.latexunits else: return vscale, variable_info.units+"x{vscale:e}".format(vscale=vscale), variable_info.latexunits+r"{\times}"+cbfmtsci(vscale,None) return unitScale, scaledUnits, scaledLatexUnits
# A utility to get variableinfo with corresponding units for simple plotting. Add "canonical" scalings as # necessary, for default/other environments.
[docs] def get_scaled_var(variable_info,vscale=None, data=None, env='EarthSpace', manualDict=None): ''' Automatically scales the variableinfo data and adjusts the units correspondingly with the default dictionaries. :param variable_info: VariableInfo object used for getting the units :param data: in case you wish to provide new data array (why, though?) :param env: A string to choose the scaling dictionary [default: EarthSpace] :param manualDict: a dictionary of {units : {scalingparams}}; used to update the included dictionary :returns: variable_info, with scaled units with pre-formatted units included in the varinfo. ''' if data is None: data = variable_info.data else: variable_info.data = data unitScale, scaledUnits, scaledLatexUnits = get_scaled_units(vscale=vscale, env=env, manualDict=manualDict,variable_info=variable_info) if unitScale == 1: # no change, optimize out the calculation return variable_info variable_info.data = data*unitScale variable_info.units = scaledUnits variable_info.latexunits = scaledLatexUnits return variable_info