18149032752

Committed 01 Oct 2025 02:10AM UTC coverage: 89.455% (+2.0%) from 87.427%

Build # 18149032752

Build Type

push

github

Committed by

jameswilburlewis

Commit Message

Update lock file after changing pyproject.toml

Run Details

40549 of 45329 relevant lines covered (89.45%)

0.89 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

74.53

/pyspedas/tplot_tools/MPLPlotter/tplot.py

import copy
import logging
import numpy as np
import matplotlib as mpl
from datetime import date, datetime, timezone
from matplotlib import pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import pyspedas
from fnmatch import filter as tname_filter
from time import sleep
from pyspedas.tplot_tools import tplot_wildcard_expand, tname_byindex, get_data, var_label_panel
from pyspedas.tplot_tools import lineplot, count_traces, makegap
from pyspedas.tplot_tools import specplot, specplot_make_1d_ybins, reduce_spec_dataset
from pyspedas.tplot_tools import get_var_label_ticks

# the following improves the x-axis ticks labels
import matplotlib.units as munits
import matplotlib.dates as mdates
converter = mdates.ConciseDateConverter()
munits.registry[np.datetime64] = converter
munits.registry[date] = converter
munits.registry[datetime] = converter

def pseudovar_component_props(varname: str):
    """ Return or calculate the plot properties for a single normal tplot variable


    Parameters
    ----------
    varname: str
        Name of the tplot variable to inspect

    Returns
    -------
    dict
        Dictionary of plot properties for this variable::
            is_spec: True if variable is a spectrogram
            ymin, ymax, ylog: Y axis limits
            zmin, zmax, zlog: Z axis limits (if variable is a spectrogram)

    """

    output_dict = {'ymin':np.nan, 'ymax':np.nan, 'ylog':False,
                   'zmin':np.nan, 'zmax': np.nan, 'zlog':False,
                   'is_spec': False}

    attrs = pyspedas.tplot_tools.data_quants[varname].attrs
    plot_opts = attrs.get('plot_options')
    if plot_opts is not None:
        yaxis_opt = plot_opts.get('yaxis_opt')
        if yaxis_opt is not None:
            ylog = yaxis_opt['y_axis_type']
            yrange = yaxis_opt.get('y_range')
    else:
        ylog = False
        yrange=[None, None]

    plot_extras = attrs['plot_options']['extras']
    dat = get_data(varname)

    if ylog is None or ylog is False or ylog == '' or ylog.lower() == 'linear':
        output_dict['ylog'] = False
        ylog = False
    else:
        output_dict['ylog'] = True
        ylog = True

    # ylog is now guaranteed to be boolean

    if plot_extras.get('spec') is not None:
        spec = plot_extras['spec']
        if spec:
            # Inspect a specplot variable
            output_dict['is_spec']=True
            if yrange is not None:
                output_dict['ymin'] = yrange[0]
                output_dict['ymax'] = yrange[1]
            else:
                # Figure out which attribute to use for bin centers (copied from specplot, maybe
                # should be pulled out into a separate routine?)
                var_data=dat
                if len(var_data) == 3:
                    if hasattr(var_data, 'v'):
                        input_bin_centers = var_data.v
                    elif hasattr(var_data, 'v1'):
                        input_bin_centers = var_data.v1
                    else:
                        logging.warning("Multidimensional variable %s has no v or v1 attribute", varname)
                elif len(var_data) == 4:
                    if hasattr(var_data, 'v1'):
                        if 'spec_dim_to_plot' in plot_extras:
                            if plot_extras['spec_dim_to_plot'] == 'v1':
                                input_bin_centers = var_data.v1
                    if hasattr(var_data, 'v2'):
                        if 'spec_dim_to_plot' in plot_extras:
                            if plot_extras['spec_dim_to_plot'] == 'v2':
                                input_bin_centers = var_data.v2

                output_bin_boundaries = specplot_make_1d_ybins(None, input_bin_centers, ylog, no_regrid=True)
                output_dict['ymin']=np.nanmin(output_bin_boundaries)
                output_dict['ymax']=np.nanmax(output_bin_boundaries)

            zaxis_options = attrs['plot_options']['zaxis_opt']
            zrange = zaxis_options.get('zrange')
            if zrange is not None:
                output_dict['zmin'] = zrange[0]
                output_dict['zmax'] = zrange[1]
            else:
                # Determine Z range from data array
                output_dict['zmin'] = np.nanmin(dat.y)
                output_dict['zmax'] = np.nanmax(dat.y)
            zlog = zaxis_options.get('z_axis_type')
            if zlog is None or zlog.lower() == 'linear':
                output_dict['zlog'] = False
            else:
                output_dict['zlog'] = True


    else:
        # Inspect a line plot variable
        output_dict['is_spec']=False
        if yrange is not None:
            output_dict['ymin'] = yrange[0]
            output_dict['ymax'] = yrange[1]
        else:
            output_dict['ymin'] = np.nanmin(dat.y)
            output_dict['ymax'] = np.nanmax(dat.y)
    return output_dict

def gather_pseudovar_props(pseudovars:list[str]):
    """ Inspect the components of a pseudovariable to determine plot limits and other settings

    If any component has y or z ranges/scales set, they will be used, otherwise limits will be determined from
    the data values.

    If any component specifies log scaling, it will be applied to all components, otherwise linear scaling will
    be used for all

    Parameters
    ----------
    pseudovars: list[str]
        List of tplot composite variable components to inspect

    Returns
    -------
    dict
        A dictionary containing the plot properties to apply to each component variable::
            line_ymin, line_ymax, line_ylog : Y axis limits for line plots
            spec_ymin, spec_ymax, spec_ylog : Y axis limits for spectrograms
            spec_zmin, spec_zmax, spec_zlog : Z axis limits for spectrograms
            has_line_plots: True if any component is a line plot
            has_spec_plots: True if any component is a spectrogram plot

    """
    output_dict = { 'has_line_plots': False,
                    'has_spec_plots': False,
                    'line_ymin': np.nan,
                    'line_ymax': np.nan,
                    'line_ylog': False,
                    'spec_ymin': np.nan,
                    'spec_ymax': np.nan,
                    'spec_ylog': False,
                    'spec_zmin': np.nan,
                    'spec_zmax': np.nan,
                    'spec_zlog': False,
                    'first_spec_var': None}


    for var in pseudovars:
        props = pseudovar_component_props(var)
        #print(f"{var}: is_spec: {props['is_spec']} ymin: {props['ymin']} ymax: {props['ymax']} ylog: {props['ylog']} zmin: {props['zmin']} zmax: {props['zmax']} zlog: {props['zlog']}")
        if props['is_spec']:
            output_dict['has_spec_plots'] = True
            if output_dict['first_spec_var'] is None:
                output_dict['first_spec_var'] = var
            if props['zlog']:
                output_dict['spec_zlog'] = True
            if props['ylog']:
                output_dict['spec_ylog'] = True
            output_dict['spec_zmin']= np.nanmin([output_dict['spec_zmin'], props['zmin']])
            output_dict['spec_zmax']= np.nanmax([output_dict['spec_zmax'], props['zmax']])
            output_dict['spec_ymin']= np.nanmin([output_dict['spec_ymin'], props['ymin']])
            output_dict['spec_ymax']= np.nanmax([output_dict['spec_ymax'], props['ymax']])
        else:
            output_dict['has_line_plots'] = True
            if props['ylog']:
                output_dict['line_ylog'] = True
            output_dict['line_ymin'] = np.nanmin([output_dict['line_ymin'], props['ymin']])
            output_dict['line_ymax'] = np.nanmax([output_dict['line_ymax'], props['ymax']])
    return output_dict


def tplot(variables,
          trange=None,
          var_label=None,
          xsize=None,
          ysize=None,
          save_png='',
          save_eps='',
          save_svg='',
          save_pdf='',
          save_jpeg='',
          dpi=None,
          display=True,
          fig=None,
          axis=None,
          running_trace_count=None,
          pseudo_idx=None,
          pseudo_right_axis=False,
          pseudo_xaxis_options=None,
          pseudo_yaxis_options=None,
          pseudo_zaxis_options=None,
          pseudo_line_options=None,
          pseudo_extra_options=None,
          show_colorbar=True,
          slice=False,
          return_plot_objects=False):
    """
    Plot tplot variables to the display, or as saved files, using Matplotlib

    Parameters
    ----------
        variables: str or list of str, required
            List of tplot variables to be plotted.  Space-delimited strings may be used instead f
            lists.  Wildcards will be expanded.
        trange: list of string or float, optional
            If set, this time range will be used, temporarily overriding any previous xlim or timespan calls
        var_label : str or list of str, optional
            A list of variables to be displayed as values underneath the X axis major tick marks
        xsize: float, optional
            Plot size in the horizontal direction (in inches)
        ysize: float, optional
            Plot size in the vertical direction (in inches)
        dpi: float, optional
            The resolution of the plot in dots per inch
        save_png : str, optional
            A full file name and path.
            If this option is set, the plot will be automatically saved to the file name provided in PNG format.
        save_eps : str, optional
            A full file name and path.
            If this option is set, the plot will be automatically saved to the file name provided in EPS format.
        save_jpeg : str, optional
            A full file name and path.
            If this option is set, the plot will be automatically saved to the file name provided in JPEG format.
        save_pdf : str, optional
            A full file name and path.
            If this option is set, the plot will be automatically saved to the file name provided in PDF format.
        save_svg : str, optional
            A full file name and path.
            If this option is set, the plot will be automatically saved to the file name provided in SVG format.
        dpi: float, optional
            The resolution of the plot in dots per inch
        display: bool, optional
            If True, then this function will display the plotted tplot variables. Use False to suppress display (for example, if
            saving to a file, or returning plot objects to be displayed later). Default: True
        fig: Matplotlib figure object
            Use an existing figure to plot in (mainly for recursive calls to render composite variables)
        axis: Matplotlib axes object
            Use an existing set of axes to plot on (mainly for recursive calls to render composite variables)
        running_trace_count: int, optional
            In recursive calls for rendering composite variables, the index of the trace currently being rendered.
        pseudo_idx: int, optional
            In recursive calls for rendering composite variables, the index of the variable component currently being rendered.
        pseudo_right_axis: bool, optional
            In recursive calls for rendering composite variables, a flag to indicate the Y scale should be placed on
            the right Y axis.
        pseudo_xaxis_options: dict, optional
            In recursice calls for rendering composite variables, X axis options inherited from the parent variable
        pseudo_yaxis_options: dict, optional
            In recursive calls for rendering composite variables, Y axis options inherited from the parent variable
        pseudo_zaxis_options: dict, optional
            In recursive calls for rendering composite variables, Z axis options inherited from the parent variable
        pseudo_line_options: dict, optional
            In recursive calls for rendering composite variables, line options inherited from the parent variable
        pseudo_extra_options: dict, optional
            In recursive calls for rendering composite variables, extra options inherited from the parent variable
        show_colorbar: bool, optional
            Show a colorbar showing the Z scale for spectrogram plots.
        slice: bool, optional
            If True, show an interactive window with a plot of Z versus Y values for the X axis (time) value under the cursor. Default: False
        return_plot_objects: bool, optional
            If true, returns the matplotlib fig and axes objects for further manipulation. Default: False

    Returns
    -------
        Any
            Returns matplotlib fig and axes objects, if return_plot_objects==True

    Examples
    --------
        >>> # Plot a single line plot
        >>> import pyspedas
        >>> x_data = [2,3,4,5,6]
        >>> y_data = [1,2,3,4,5]
        >>> pyspedas.store_data("Variable1", data={'x':x_data, 'y':y_data})
        >>> pyspedas.tplot("Variable1")

        >>> # Plot two variables
        >>> x_data = [1,2,3,4,5]
        >>> y_data = [[1,5],[2,4],[3,3],[4,2],[5,1]]
        >>> pyspedas.store_data("Variable2", data={'x':x_data, 'y':y_data})
        >>> pyspedas.tplot(["Variable1", "Variable2"])

        >>> # Plot two plots, adding a third variable's values as annotations at X axis major tick marks
        >>> x_data = [1,2,3]
        >>> y_data = [ [1,2,3] , [4,5,6], [7,8,9] ]
        >>> v_data = [1,2,3]
        >>> pyspedas.store_data("Variable3", data={'x':x_data, 'y':y_data, 'v':v_data})
        >>> pyspedas.options("Variable3", 'spec', 1)
        >>> pyspedas.tplot(["Variable2", "Variable3"], var_label='Variable1')

    """
    # This call resolves wildcard patterns and converts integers to variable names
    variables = tplot_wildcard_expand(variables)
    if len(variables) == 0:
        logging.warning("tplot: No matching tplot names were found")
        return


    varlabel_style = pyspedas.tplot_tools.tplot_opt_glob.get('varlabel_style')
    if varlabel_style is None or varlabel_style.lower() == 'extra_axes':
        num_panels = len(variables)
        panel_sizes = [1]*num_panels
    else: # varlabel_style 'extra_panel'
        num_panels = len(variables) + 1
        panel_sizes = [1] * len(variables) + [0.1 * (len(var_label) + 2)]

    # support for the panel_size option
    for var_idx, variable in enumerate(variables):
        if pyspedas.tplot_tools.data_quants.get(variable) is None:
            continue
        panel_size = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras'].get('panel_size')
        if panel_size is not None:
            panel_sizes[var_idx] = panel_size

    if xsize is None:
        xsize = pyspedas.tplot_tools.tplot_opt_glob.get('xsize')
        if xsize is None:
            xsize = 12

    if ysize is None:
        ysize = pyspedas.tplot_tools.tplot_opt_glob.get('ysize')
        if ysize is None:
            if num_panels > 4:
                ysize = 8
            else:
                # This was previously set to 5, which resulted in a non-monotonic progression of panel heights
                default_y_sizes = [5, 5, 6, 7, 8]
                ysize = default_y_sizes[num_panels]

    # The logic here for handling the 'right_axis' option is pretty convoluted, and makes a number of assumptions
    # that may not be warranted: mainly that right_axis will only be set on pseudovariables, and that if
    # set, the first sub-variable gets the left axis and all other sub-variables get a newly twinx-ed right axis.
    # "right axis" implies there will only be at most two Y axes. What if more scales are needed?
    # There also seems to be some conflation of the set of axes for the whole stack of plots, versus
    # the single (left or possibly right) axis for the variable currently being rendered.
    #
    # This whole concept is kind of a mess at the moment.  For now, we'll make it work for the
    # most likely use case, plotting a single spectrum variable followed by a single line variable
    # (for example, an energy spectrum plus spacecraft potential).  JWL 2024-03-26

    if fig is None and axis is None:
        fig, axes = plt.subplots(nrows=num_panels, sharex=True, gridspec_kw={'height_ratios': panel_sizes}, layout='constrained')
        fig.set_size_inches(xsize, ysize)
        plot_title = pyspedas.tplot_tools.tplot_opt_glob['title_text']
        fig.suptitle(plot_title)
        if (plot_title is not None) and plot_title != '':
            if 'title_size' in pyspedas.tplot_tools.tplot_opt_glob:
                title_size = pyspedas.tplot_tools.tplot_opt_glob['title_size']
                fig.suptitle(plot_title, fontsize=title_size)
            else:
                fig.suptitle(plot_title)
        # support for matplotlib styles
        style = pyspedas.tplot_tools.tplot_opt_glob.get('style')
        if style is not None:
            plt.style.use(style)
    else:
        # fig and axis have been passed as parameters, most likely a recursive tplot call to render
        # a pseudovariable
        if pseudo_idx == 0 or pseudo_right_axis == False:
            # setting up first axis
            axes = axis
        elif pseudo_idx > 0 and pseudo_right_axis:
            # generate and use the right axis?  probably still wrong...
            axes = axis.twinx()
        # support for matplotlib styles -- shouldn't be needed if processing pseudovar components?
        style = pyspedas.tplot_tools.tplot_opt_glob.get('style')
        if style is not None:
            plt.style.use(style)


    axis_font_size = pyspedas.tplot_tools.tplot_opt_glob.get('axis_font_size')

    colorbars = {}

    for idx, variable in enumerate(variables):
        var_data_org = get_data(variable, dt=True)
        var_metadata = get_data(variable, metadata=True)

        #Check for a 3d variable, call reduce_spec_dataset
        if hasattr(var_data_org, 'v1') and hasattr(var_data_org, 'v2'):
            temp_dq = reduce_spec_dataset(name=variable)
            var_data_org = get_data(variable, dt=True, data_quant_in=temp_dq)
        
        if var_data_org is None:
            logging.info('Variable not found: ' + variable)
            continue

        var_data = copy.deepcopy(var_data_org)

        # plt.subplots returns a list of axes for multiple panels
        # but only a single axis for a single panel
        if num_panels == 1:
            this_axis = axes
        else:
            this_axis = axes[idx]

        # we need to track the variable name in the axis object
        # for spectrogram slices
        this_axis.var_name = variable

        pseudo_var = False
        overplots = None
        spec = False

        var_quants = pyspedas.tplot_tools.data_quants[variable]

        if not isinstance(var_quants, dict):
            overplots = var_quants.attrs['plot_options'].get('overplots_mpl')
            if overplots is not None and len(overplots) > 0:
                pseudo_var = True

        # deal with pseudo-variables first
        if isinstance(var_data, list) or isinstance(var_data, str) or pseudo_var:
            # this is a pseudo variable
            if isinstance(var_data, str):
                var_data = var_data.split(' ')

            if pseudo_var:
                pseudo_vars = overplots
            else:
                pseudo_vars = var_data

            # pseudo variable metadata should override the metadata
            # for individual variables
            xaxis_options = None
            yaxis_options = None
            zaxis_options = None
            line_opts = None
            plot_extras = None
            if pseudo_var:
                plot_extras = var_quants.attrs['plot_options']['extras']
                if plot_extras.get('spec') is not None:
                    spec = True

                if plot_extras.get('right_axis') is not None:
                    if plot_extras.get('right_axis'):
                        pseudo_right_axis = True

                if pseudo_right_axis or spec:
                    plot_extras = None
                else:
                    yaxis_options = var_quants.attrs['plot_options']['yaxis_opt']
                    zaxis_options = var_quants.attrs['plot_options']['zaxis_opt']
                    line_opts = var_quants.attrs['plot_options']['line_opt']

            traces_processed = 0
            pseudovar_props = gather_pseudovar_props(pseudo_vars)
            #
            # Determine which Y axis options need to be changed to accommodate multiple component variables
            # There might be line variables and spectra, each with their own overall yscale and yrange.
            # If both are present, the spectra will take priority.
            # If explicit y or z scales or ranges are set on the pseudovariable, do not override.
            if pseudovar_props['has_spec_plots']:
                if pseudovar_props['spec_ylog']:
                    new_yscale = 'log'
                else:
                    new_yscale = 'linear'
                new_yr = [pseudovar_props['spec_ymin'], pseudovar_props['spec_ymax']]
            else:
                if pseudovar_props['line_ylog']:
                    new_yscale = 'log'
                else:
                    new_yscale = 'linear'
                new_yr = [pseudovar_props['line_ymin'], pseudovar_props['line_ymax']]
            override_yopts = {}
            if pseudo_right_axis:
                # Who knows?  Let the component variable Y-axis fight it out...
                override_yopts = {}
            elif yaxis_options is None:
                override_yopts = {'y_range':new_yr, 'y_range_user':True,'y_axis_style':new_yscale}
            else:
                if yaxis_options.get('y_range') is None:
                    override_yopts['y_range'] = new_yr
                    override_yopts['y_range_user'] = True
                if yaxis_options.get('y_axis_style') is None:
                    override_yopts['y_axis_style'] = new_yscale
            if yaxis_options is not None:
                yaxis_options = yaxis_options | override_yopts
            else:
                yaxis_options = override_yopts

            # Determine which Z axis options need to be changed to accommodate multiple component variables
            #
            # We only care about specplots here.
            new_zr = [ None, pseudovar_props['spec_zmax']]
            if pseudovar_props['spec_zlog']:
                new_zscale='log'
            else:
                new_zscale='linear'
            override_zopts = {}
            if zaxis_options is None:
                override_zopts = {'z_range':new_zr, 'z_axis_style':new_zscale}
            else:
                if zaxis_options.get('z_range') is None:
                    override_zopts['z_range'] = new_zr
                if zaxis_options.get('z_axis_style') is None:
                    override_zopts['z_axis_style'] = new_zscale
            if zaxis_options is not None:
                zaxis_options = zaxis_options | override_zopts
            else:
                zaxis_options = override_zopts

            for pseudo_idx, var in enumerate(pseudo_vars):
                # We're plotting a pseudovariable.  Iterate over the sub-variables, keeping track of how many
                # traces have been plotted so far, so we can correctly match option values to traces. The pseudovariable
                # y_axis, z_axis, line and extra options are passed as parameters so they can be merged with the
                # sub-variable options, with any pseudovar options overriding the sub-variable options.
                trace_count_thisvar = count_traces(var)
                if var == pseudovar_props['first_spec_var']:
                    pseudo_show_colorbar=True
                else:
                    pseudo_show_colorbar=False
                tplot(var,
                      trange=trange,
                      return_plot_objects=return_plot_objects,
                      xsize=xsize, ysize=ysize,
                      fig=fig, axis=this_axis, display=False,
                      running_trace_count=traces_processed,
                      pseudo_idx=pseudo_idx,
                      pseudo_xaxis_options=xaxis_options, pseudo_yaxis_options=yaxis_options, pseudo_zaxis_options=zaxis_options,
                      pseudo_line_options=line_opts, pseudo_extra_options=plot_extras,
                      pseudo_right_axis=pseudo_right_axis,
                      show_colorbar=pseudo_show_colorbar)
                traces_processed += trace_count_thisvar
            

            continue


        #if data_gap is an option for this variable, or if it's a add
        #gaps here; an individual gap setting should override the
        #global setting
        plot_extras = var_quants.attrs['plot_options']['extras']
        if plot_extras.get('data_gap') is not None and plot_extras.get('data_gap') > 0:
            var_data = makegap(var_data, dt = plot_extras.get('data_gap'))
        else:
            if pyspedas.tplot_tools.tplot_opt_glob['data_gap'] is not None and pyspedas.tplot_tools.tplot_opt_glob['data_gap'] > 0:
                var_data = makegap(var_data, dt = pyspedas.tplot_tools.tplot_opt_glob['data_gap'])

        # set the x-axis range, if it was set with xlim or tlimit or the trange parameter
        if trange is None and pyspedas.tplot_tools.tplot_opt_glob.get('x_range') is None:
            var_data_times = var_data.times
            time_idxs = np.arange(len(var_data_times))
        else:
            if trange is not None:
                if len(trange) != 2:
                    logging.error('Invalid trange setting: must be a 2-element list or array')
                    return
                if isinstance(trange[0], str):
                    x_range = pyspedas.tplot_tools.time_double(trange) # seconds since epoch
                    x_range_start = x_range[0]
                    x_range_stop = x_range[1]
            else:
                x_range = pyspedas.tplot_tools.tplot_opt_glob['x_range']  # Seconds since epoch
                x_range_start = x_range[0]
                x_range_stop = x_range[1]

            # Check for NaN or inf in x_range
            if not np.isfinite(x_range_start):
                logging.warning('tplot: x_range start is not finite, replacing with 0')
                x_range_start = 0

            if not np.isfinite(x_range_stop):
                logging.warning('tplot: x_range end is not finite, replacing with 0')
                x_range_stop = 0

            # Convert to np.datetime64 with nanosecond precision
            x_range = np.array(np.array([x_range_start*1e9, x_range_stop*1e9]),dtype='datetime64[ns]')
            this_axis.set_xlim(x_range)
            time_idxs = np.argwhere((var_data.times >= x_range[0]) & (var_data.times <= x_range[1])).flatten()
            if len(time_idxs) == 0:
                logging.info('No data found in the time range: ' + variable)
                continue
            var_data_times = var_data.times[time_idxs]

        var_times = var_data_times

        # set some more plot options
        xaxis_options = var_quants.attrs['plot_options']['xaxis_opt']
        if pseudo_xaxis_options is not None and len(pseudo_xaxis_options) > 0:
            merged_xaxis_options = xaxis_options | pseudo_xaxis_options
            xaxis_options = merged_xaxis_options

        yaxis_options = var_quants.attrs['plot_options']['yaxis_opt']
        if pseudo_yaxis_options is not None and len(pseudo_yaxis_options) > 0:
            merged_yaxis_options = yaxis_options | pseudo_yaxis_options
            yaxis_options = merged_yaxis_options

        zaxis_options = var_quants.attrs['plot_options']['zaxis_opt']
        if pseudo_zaxis_options is not None and len(pseudo_zaxis_options) > 0:
            merged_zaxis_options = zaxis_options | pseudo_zaxis_options
            zaxis_options = merged_zaxis_options

        line_opts = var_quants.attrs['plot_options']['line_opt']
        if pseudo_line_options is not None and len(pseudo_line_options) > 0:
            merged_line_opts = line_opts | pseudo_line_options
            line_opts = merged_line_opts

        if line_opts is not None:
            if 'name' in line_opts:
                this_axis.set_title(line_opts['name'])
            elif 'title' in line_opts:
                this_axis.set_title(line_opts['title'])

        plot_extras = var_quants.attrs['plot_options']['extras']
        if pseudo_extra_options is not None and len(pseudo_extra_options) > 0:
            merged_plot_extras = plot_extras | pseudo_extra_options
            plot_extras = merged_plot_extras

        xtitle = None
        if xaxis_options.get('axis_label') is not None:
            xtitle = xaxis_options['axis_label']


        xsubtitle = ''
        if xaxis_options.get('axis_subtitle') is not None:
            xsubtitle = xaxis_options['axis_subtitle']

        if style is None:
            xtitle_color = 'black'
        else:
            xtitle_color = None

        if xaxis_options.get('axis_color') is not None:
            xtitle_color = xaxis_options['axis_color']

        ylog = yaxis_options['y_axis_type']

        if ylog == 'log':
            this_axis.set_yscale('log')
        else:
            this_axis.set_yscale('linear')

        ytitle = yaxis_options['axis_label']
        if ytitle == '':
            ytitle = variable

        ysubtitle = ''
        if yaxis_options.get('axis_subtitle') is not None:
            ysubtitle = yaxis_options['axis_subtitle']

        # replace some common superscripts
        ysubtitle = replace_common_exp(ysubtitle)

        if axis_font_size is not None:
            this_axis.tick_params(axis='x', labelsize=axis_font_size)
            this_axis.tick_params(axis='y', labelsize=axis_font_size)

        char_size = pyspedas.tplot_tools.tplot_opt_glob.get('charsize')
        if char_size is None:
            char_size = 12

        if plot_extras.get('char_size') is not None:
            char_size = plot_extras['char_size']

        user_set_yrange = yaxis_options.get('y_range_user')
        if user_set_yrange is not None:
            # the user has set the yrange manually
            yrange = yaxis_options['y_range']
            if not np.isfinite(yrange[0]):
                yrange[0] = None
            if not np.isfinite(yrange[1]):
                yrange[1] = None
            this_axis.set_ylim(yrange)

        ymajor_ticks = yaxis_options.get('y_major_ticks')
        if ymajor_ticks is not None:
            this_axis.set_yticks(ymajor_ticks)

        yminor_tick_interval = yaxis_options.get('y_minor_tick_interval')
        if yminor_tick_interval is not None and ylog != 'log':
            this_axis.yaxis.set_minor_locator(plt.MultipleLocator(yminor_tick_interval))

        if style is None:
            ytitle_color = 'black'
        else:
            ytitle_color = None

        if yaxis_options.get('axis_color') is not None:
            ytitle_color = yaxis_options['axis_color']

        if xtitle is not None and xtitle != '':
            if xtitle_color is not None:
                this_axis.set_xlabel(xtitle + '\n' + xsubtitle, fontsize=char_size, color=xtitle_color)
            else:
                this_axis.set_xlabel(xtitle + '\n' + xsubtitle, fontsize=char_size)

        if ytitle_color is not None:
            this_axis.set_ylabel(ytitle + '\n' + ysubtitle, fontsize=char_size, color=ytitle_color)
        else:
            this_axis.set_ylabel(ytitle + '\n' + ysubtitle, fontsize=char_size)

        border = True
        if plot_extras.get('border') is not None:
            border = plot_extras['border']

        if border == False:
            this_axis.axis('off')

        # axis tick options
        if plot_extras.get('xtickcolor') is not None:
            this_axis.tick_params(axis='x', color=plot_extras.get('xtickcolor'))

        if plot_extras.get('ytickcolor') is not None:
            this_axis.tick_params(axis='y', color=plot_extras.get('ytickcolor'))

        if plot_extras.get('xtick_direction') is not None:
            this_axis.tick_params(axis='x', direction=plot_extras.get('xtick_direction'))

        if plot_extras.get('ytick_direction') is not None:
            this_axis.tick_params(axis='y', direction=plot_extras.get('ytick_direction'))

        if plot_extras.get('xtick_length') is not None:
            this_axis.tick_params(axis='x', length=plot_extras.get('xtick_length'))

        if plot_extras.get('ytick_length') is not None:
            this_axis.tick_params(axis='y', length=plot_extras.get('ytick_length'))

        if plot_extras.get('xtick_width') is not None:
            this_axis.tick_params(axis='x', width=plot_extras.get('xtick_width'))

        if plot_extras.get('ytick_width') is not None:
            this_axis.tick_params(axis='y', width=plot_extras.get('ytick_width'))

        if plot_extras.get('xtick_labelcolor') is not None:
            this_axis.tick_params(axis='x', labelcolor=plot_extras.get('xtick_labelcolor'))

        if plot_extras.get('ytick_labelcolor') is not None:
            this_axis.tick_params(axis='y', labelcolor=plot_extras.get('ytick_labelcolor'))

        # determine if this is a line plot or a spectrogram
        spec = False
        if plot_extras.get('spec') is not None:
            spec = plot_extras['spec']

        if spec:
            # create spectrogram plots
            plot_created = specplot(var_data, var_times, this_axis, yaxis_options, zaxis_options, plot_extras, colorbars, axis_font_size, fig, variable, time_idxs=time_idxs, style=style)
            if not plot_created:
                continue
        else:
            # create line plots
            plot_created = lineplot(var_data, var_times, this_axis, line_opts, yaxis_options, plot_extras, running_trace_count=running_trace_count, time_idxs=time_idxs, style=style, var_metadata=var_metadata)
            if not plot_created:
                continue

        # apply any vertical/horizontal bars
        if pyspedas.tplot_tools.data_quants[variable].attrs['plot_options'].get('time_bar') is not None:
            time_bars = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['time_bar']

            for time_bar in time_bars:
                # vertical bars
                if time_bar['dimension'] == 'height':
                    this_axis.axvline(x=datetime.fromtimestamp(time_bar['location'], tz=timezone.utc),
                        color=np.array(time_bar.get('line_color'))/256.0, lw=time_bar.get('line_width'),
                                      linestyle=time_bar.get('line_dash'))

                # horizontal bars
                if time_bar['dimension'] == 'width':
                    this_axis.axhline(y=time_bar['location'], color=np.array(time_bar.get('line_color'))/256.0,
                                      lw=time_bar.get('line_width'),
                                      linestyle=time_bar.get('line_dash'))

        # highlight time intervals
        if pyspedas.tplot_tools.data_quants[variable].attrs['plot_options'].get('highlight_intervals') is not None:
            highlight_intervals = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['highlight_intervals']

            for highlight_interval in highlight_intervals:
                hightlight_opts = copy.deepcopy(highlight_interval)
                del hightlight_opts['location']
                if highlight_interval['edgecolor'] is not None or highlight_interval['facecolor'] is not None:
                    del hightlight_opts['color']

                this_axis.axvspan(mdates.date2num(datetime.fromtimestamp(highlight_interval['location'][0], timezone.utc)),
                                  mdates.date2num(datetime.fromtimestamp(highlight_interval['location'][1], timezone.utc)),
                                  **hightlight_opts)

        # add annotations
        if pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras'].get('annotations') is not None:
            annotations = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras']['annotations']
            for annotation in annotations:
                this_axis.annotate(annotation['text'], annotation['position'],
                                   xycoords=annotation['xycoords'],
                                   fontsize=annotation['fontsize'],
                                   alpha=annotation['alpha'],
                                   fontfamily=annotation['fontfamily'],
                                   fontvariant=annotation['fontvariant'],
                                   fontstyle=annotation['fontstyle'],
                                   fontstretch=annotation['fontstretch'],
                                   fontweight=annotation['fontweight'],
                                   rotation=annotation['rotation'],
                                   color=annotation['color'])

    # apply any addition x-axes (or panel) specified by the var_label keyword
    if var_label is not None:
        if varlabel_style is None or varlabel_style.lower() == 'extra_axes':
            varlabels_extra_axes(num_panels, this_axis, var_label)
        else:
            var_label_panel(variables, var_label, axes,  axis_font_size)

    # add the color bars to any spectra
    for idx, variable in enumerate(variables):
        if pyspedas.tplot_tools.data_quants.get(variable) is None:
            continue
        plot_extras = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras']
        zaxis_options = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['zaxis_opt']
        if plot_extras.get('spec') is not None:
            spec = plot_extras['spec']
        else:
            spec = False


        if spec and show_colorbar:
            if colorbars.get(variable) is None:
                continue

            if num_panels == 1:
                this_axis = axes
            else:
                this_axis = axes[idx]

            colorbar = fig.colorbar(colorbars[variable]['im'], ax=this_axis)

            if style is None:
                ztitle_color = 'black'
            else:
                ztitle_color = None

            if zaxis_options is None:
                continue

            if zaxis_options.get('axis_color') is not None:
                ztitle_color = zaxis_options['axis_color']

            ztitle_text = colorbars[variable]['ztitle']
            zsubtitle_text = colorbars[variable]['zsubtitle']

            # replace some common superscripts
            ztitle_text = replace_common_exp(ztitle_text)
            zsubtitle_text = replace_common_exp(zsubtitle_text)

            if ztitle_color is not None:
                colorbar.set_label(ztitle_text + '\n ' + zsubtitle_text,
                                   color=ztitle_color, fontsize=char_size)
            else:
                colorbar.set_label(ztitle_text + '\n ' + zsubtitle_text,
                                   fontsize=char_size)

    # plt.tight_layout()
    fig.canvas.draw()

    if save_png is not None and save_png != '':
        if not save_png.endswith('.png'):
            save_png += '.png'
        plt.savefig(save_png, dpi=dpi)

    if save_eps is not None and save_eps != '':
        if not save_eps.endswith('.eps'):
            save_eps += '.eps'
        plt.savefig(save_eps, dpi=dpi)

    if save_svg is not None and save_svg != '':
        if not save_svg.endswith('.svg'):
            save_svg += '.svg'
        plt.savefig(save_svg, dpi=dpi)

    if save_pdf is not None and save_pdf != '':
        if not save_pdf.endswith('.pdf'):
            save_pdf += '.pdf'
        plt.savefig(save_pdf, dpi=dpi)

    if save_jpeg is not None and save_jpeg != '':
        if not save_jpeg.endswith('.jpeg'):
            save_jpeg += '.jpeg'
        plt.savefig(save_jpeg, dpi=dpi)

    if slice:
        slice_fig, slice_axes = plt.subplots(nrows=1)
        slice_plot, = slice_axes.plot([0], [0])
        mouse_event_func = lambda event: mouse_move_slice(event, slice_axes, slice_plot)
        cid = fig.canvas.mpl_connect('motion_notify_event', mouse_event_func)

    if display:
        plt.show()

    if return_plot_objects:
        return fig, axes


def varlabels_extra_axes(num_panels, this_axis, var_label):
    # apply any addition x-axes specified by the var_label keyword
    if var_label is not None:
        if not isinstance(var_label, list):
            var_label = [var_label]

        axis_delta = 0.0

        for label in var_label:
            if isinstance(label, int):
                label = tname_byindex(label)
            label_data = get_data(label, xarray=True, dt=True)

            if label_data is None:
                logging.info('Variable not found: ' + label)
                continue

            if len(label_data.values.shape) != 1:
                logging.info(
                    label + ' specified as a vector; var_label only supports scalars. Try splitting the vector into seperate tplot variables.')
                continue

            # set up the new x-axis
            axis_delta = axis_delta - num_panels * 0.1
            new_xaxis = this_axis.secondary_xaxis(axis_delta)
            xaxis_ticks = this_axis.get_xticks().tolist()
            xaxis_ticks_dt = [np.datetime64(mpl.dates.num2date(tick_val).replace(tzinfo=None).isoformat(), 'ns') for
                              tick_val in xaxis_ticks]
            # xaxis_ticks_unix = [tick_val.timestamp() for tick_val in xaxis_ticks_dt]
            xaxis_labels = get_var_label_ticks(label_data, xaxis_ticks_dt)
            new_xaxis.set_xticks(xaxis_ticks_dt)
            new_xaxis.set_xticklabels(xaxis_labels)
            ytitle = pyspedas.tplot_tools.data_quants[label].attrs['plot_options']['yaxis_opt']['axis_label']
            new_xaxis.set_xlabel(ytitle)

        # fig.subplots_adjust(bottom=0.05+len(var_label)*0.1)


def mouse_move_slice(event, slice_axes, slice_plot):
    """
    This function is called when the mouse moves over an axis
    and the slice keyword is set to True; for spectra figures, it
    updates the slice plot based on the mouse location
    """
    if event.inaxes is None:
        return

    # check for a spectrogram
    try:
        data = get_data(event.inaxes.var_name)
    except AttributeError:
        return

    if data is None:
        return

    if len(data) != 3:
        return

    slice_time = mdates.num2date(event.xdata).timestamp()
    idx = np.abs(data.times-slice_time).argmin()

    if len(data.v.shape) > 1:
        # time varying y-axis
        vdata = data.v[idx, :]
    else:
        vdata = data.v

    xaxis_options = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options']['xaxis_opt']
    yaxis_options = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options']['yaxis_opt']
    zaxis_options = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options']['zaxis_opt']

    yrange = yaxis_options.get('y_range')
    if yrange is None:
        yrange = [np.nanmin(vdata), np.nanmax(vdata)]

    zrange = zaxis_options.get('z_range')
    if zrange is None:
        zrange = [np.nanmin(data.y), np.nanmax(data.y)]

    y_label = zaxis_options.get('axis_label')
    if y_label is not None:
        slice_axes.set_ylabel(y_label)

    title = datetime.fromtimestamp(data.times[idx], timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f')

    x_label = yaxis_options.get('axis_label')
    if x_label is not None:
        title = x_label + ' (' + title + ')'

    slice_axes.set_title(title)

    x_subtitle = yaxis_options.get('axis_subtitle')
    if x_subtitle is not None:
        slice_axes.set_xlabel(x_subtitle)

    slice_yaxis_opt = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options'].get('slice_yaxis_opt')

    xscale = None
    yscale = None

    if slice_yaxis_opt is not None:
        xscale = slice_yaxis_opt.get('xi_axis_type')
        yscale = slice_yaxis_opt.get('yi_axis_type')

    if yscale is None:
        # if the user didn't explicitly set the ylog_slice option,
        # use the option from the plot
        yscale = zaxis_options.get('z_axis_type')
        if yscale is None:
            yscale = 'linear'

    if xscale is None:
        # if the user didn't explicitly set the xlog_slice option,
        # use the option from the plot
        xscale = yaxis_options.get('y_axis_type')
        if xscale is None:
            xscale = 'linear'

    if yscale == 'log' and zrange[0] == 0.0:
        zrange[0] = np.nanmin(data.y[idx, :])

    slice_plot.set_data(vdata, data.y[idx, :])
    slice_axes.set_ylim(zrange)
    slice_axes.set_xlim(yrange)
    slice_axes.set_xscale(xscale)
    slice_axes.set_yscale(yscale)

    try:
        plt.draw()
    except ValueError:
        return
    sleep(0.01)

def replace_common_exp(title):
    if hasattr(title, 'decode'):
        title = title.decode('utf-8')
    if '$' in title:
        return title
    if '^' not in title:
        return title
    exp = False
    title_out = ''
    for char in title:
        if char == '^':
            exp = True
            title_out += '$^{'
            continue
        else:
            if exp:
                if not char.isalnum():
                    title_out += '}$' + char
                    exp = False
                    continue
        title_out += char
    if exp:
        title_out += '}$'
    return title_out



spedas / pyspedas / 18149032752

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous