interactive.py

import numpy as np
import datetime
import xarray as xr
import pandas as pd

from matplotlib.dates import num2date

from pyxlma.plot.xlma_plot_feature import color_by_time, plot_points, setup_hist, plot_3d_grid, subset
from pyxlma.plot.xlma_base_plot import subplot_labels, inset_view, BlankPlot
from pyxlma.xarray_util import generic_subset

from ipywidgets import Output
output = Output()


class AccumulatorAxesManager(object):
    def __init__(self, a, ax):
        self.accumulator = a
        self.ax = ax
        f = ax.figure
        f.canvas.mpl_connect('draw_event', a.draw_event)
        f.canvas.mpl_connect('button_release_event', a.mouse_up_event)
        f.canvas.mpl_connect('button_press_event', a.mouse_down_event)
        ax.callbacks.connect('xlim_changed', a.axis_limit_changed)
        ax.callbacks.connect('ylim_changed', a.axis_limit_changed)


class Accumulator(object):
    """ Provides for event callbacks for matplotlib drag/release events and
        axis limit changes by accumulating a series of event occurrences.
        Produces a single call to func after a user interacts with the plot.
        Also stores the axes that got the event, and passes them to func.
        Sample usage:
        from pylab import figure, show
        def simple(axes):
            print "update ", axes
        a = Accumulator(simple)
        f=figure()
        ax=f.add_subplot(111)
        plt=ax.plot(range(10))
        f.canvas.mpl_connect('draw_event', a.draw_event)
        f.canvas.mpl_connect('button_release_event', a.mouse_up_event)
        f.canvas.mpl_connect('button_press_event', a.mouse_down_event)
        ax.callbacks.connect('xlim_changed', a.axis_limit_changed)
        ax.callbacks.connect('ylim_changed', a.axis_limit_changed)
        show()
        """

    def __init__(self, func):
        self.func=func
        self.reset()
        self.mouse_up = True
        # print('did init')

    def reset(self):
        """ Reset flags after the update function is called.
            Mouse is tracked separately.
            """
        # print('reset')
        self.limits_changed = 0
        self.got_draw = False
        self.axes = None

    def axis_limit_changed(self, ax):
        # print('ax limits')
        self.limits_changed += 1
        self.axes = ax
        self.check_status()

    def draw_event(self, event):
        # print('draw event')
        if self.limits_changed > 0:
            # only care about draw events if one of the axis limits has changed
            self.got_draw=True
        self.check_status()

    def mouse_up_event(self, event):
        # print('mouse up')
        self.mouse_up = True
        self.check_status()

    def mouse_down_event(self, event):
        # print('mouse down')
        self.mouse_up = False

    def both_limits_changed(self):
        """ Both x and y limits changed and the mouse is up (not dragging)
            This condition takes care of the limits being reset outside of a
            dragging context, such as the view-reset (home) button on the
            Matplotlib standard toolbar.
            """
        # print("both_lim_chg")
        return (self.limits_changed >= 2) & self.mouse_up

    def interaction_complete(self):
        """ x, y, or both limits changed, and the mouse is up (not dragging).
            Also checks if matplotlib has done its final redraw of the screen,
            which comes after the call to *both* set_xlim and set_ylim
            have been triggered. The check for the draw event is the crucial
            step in not producing two calls to self.func.

            New problem: with zoom, after a reset, get a draw event. on next axes change, the draw
            event  combines with an axis change to trigger interaction_complete. Then get another
            reset and ax limit change, and draw, and this passes again.
            Fixed this by adding a check on self.limits_changed > 0  in draw_event.
        """
        # print("interaction_complete")
        return (self.limits_changed>0) & self.got_draw & self.mouse_up

    def check_status(self):
        if self.both_limits_changed() | self.interaction_complete():
            # print('both limits:', self.both_limits_changed(), ', interaction:', self.interaction_complete())
            self.func(self.axes)
            self.reset()


def event_space_time_limits(ds):
    xlim = ds.event_longitude.min().values, ds.event_longitude.max().values
    ylim = ds.event_latitude.min().values, ds.event_latitude.max().values
    zlim = ds.event_altitude.min().values/1000, ds.event_altitude.max().values/1000
    tlim = (pd.to_datetime(ds.event_time.min().values.astype('M8[us]')).to_pydatetime(),
            pd.to_datetime(ds.event_time.max().values.astype('M8[us]')).to_pydatetime())
    return xlim, ylim, zlim, tlim



class InteractiveLMAPlot(object):

    def __init__(self, ds, chimax=1.0, stationmin=6,
            plot_cmap='plasma', point_size=5, clon=-101.5, clat=33.5,
            xlim=None, ylim=None, zlim=None, tlim=None):
        """ lma_plot is as returned by pyxlma's BlankPlot

            xlim, ylim, zlim, and tlim are initial limits to use with the dataset;
            otherwise, the limits are taken from the limits of the data in ds.
        """
        xlim_ds, ylim_ds, zlim_ds, tlim_ds = event_space_time_limits(ds)

        if xlim is None: xlim = xlim_ds
        if ylim is None: ylim = ylim_ds
        if zlim is None: zlim = zlim_ds
        if tlim is None: tlim = tlim_ds
        # tlim = pd.to_datetime('2022-06-04T22:15'), pd.to_datetime('2022-06-04T22:20')

        self.stationmin = stationmin
        self.chimax = chimax
        self.plot_cmap = plot_cmap
        self.point_size = point_size
        self.clon = clon
        self.clat = clat
        self.widget_output = output


        # list of all artists in the plot that change when the view subset changes
        self.data_artists = []

        self._accumulator = Accumulator(self.limits_changed)
        self._managers = {}
        self.ds = ds

        self.bounds = {
            'x': xlim,
            'y': ylim,
            'z': zlim,
            't': tlim,
        }

        self.lma_plot = None
        self.inset_ax = None
        self.make_plot()
        self.make_plot_interactive()

    @output.capture()
    def make_plot_interactive(self):
            # Make the just-created plot interactive
        self._managers['xy'] = AccumulatorAxesManager(self._accumulator, self.lma_plot.ax_plan)
        self._managers['tz'] = AccumulatorAxesManager(self._accumulator, self.lma_plot.ax_th)
        self._managers['xz'] = AccumulatorAxesManager(self._accumulator, self.lma_plot.ax_lon)
        self._managers['zy'] = AccumulatorAxesManager(self._accumulator, self.lma_plot.ax_lat)


    def _axes_spacetime_limits(self):
        """ Get the current spacetime limits of the axes """
        xlim = self._managers['xy'].ax.get_xlim()
        ylim = self._managers['xy'].ax.get_ylim()
        zlim = self._managers['tz'].ax.get_ylim()
        t0, t1 = list(map(num2date, self.lma_plot.ax_th.get_xlim()))
        tlim = (t0.replace(tzinfo=None), t1.replace(tzinfo=None))
        return xlim, ylim, tlim, zlim

    @output.capture()
    def limits_changed(self, axes):
        """ updates self.bounds from the limits of all axes in the plot making sure to
            include changes from axes
        """

        # When we start out, we get the old axis limits, naively from the xy and tz
        # plots, and then make sure to update the limits for the axis that changed.

        xlim, ylim, tlim, zlim = self._axes_spacetime_limits()

        tag_changed = None
        for tag, mgr in self._managers.items():
            if axes == mgr.ax:
                tag_changed = tag

        if tag_changed == 'xy':
            xlim = self._managers['xy'].ax.get_xlim()
            ylim = self._managers['xy'].ax.get_ylim()
        if tag_changed == 'tz':
            t0, t1 = list(map(num2date, self.lma_plot.ax_th.get_xlim()))
            tlim = (t0.replace(tzinfo=None), t1.replace(tzinfo=None))
            zlim = self._managers['tz'].ax.get_ylim()
        if tag_changed == 'xz':
            xlim = self._managers['xz'].ax.get_xlim()
            zlim = self._managers['xz'].ax.get_ylim()
        if tag_changed == 'zy':
            zlim = self._managers['zy'].ax.get_xlim()
            ylim = self._managers['zy'].ax.get_ylim()

        # When a zoom happens, the plan view plot changes the region covered
        # by the axes (think of the frame) to ensure that the zoomed region
        # maintains a 1:1 aspect ratio in the map projection coordinates.
        # So, we also need to change the axis limits of the other axes by the
        # same amount.

        # Make it so that we can do time math and get total_seconds.
        # t0 = pd.to_datetime(tlim[0].astype('M8[us]')).to_pydatetime()
        # t1 = pd.to_datetime(tlim[1].astype('M8[us]')).to_pydatetime()
        # tlim = (t0, t1)

        self.bounds['x'] = xlim
        self.bounds['y'] = ylim
        self.bounds['z'] = zlim
        self.bounds['t'] = tlim

        self.make_plot()
        self.make_plot_interactive()


    @output.capture()
    def make_plot(self):
        # pull out relevant plot params from object attributes
        ds = self.ds
        plot_s = self.point_size
        plot_cmap = self.plot_cmap
        xlim, ylim = self.bounds['x'], self.bounds['y']
        zlim, tlim = self.bounds['z'], self.bounds['t']
        xchi = self.chimax
        stationmin = self.stationmin

        # Regularize time limits to the right datetime type.
        tlim_sub = pd.to_datetime(tlim[0]), pd.to_datetime(tlim[1])

        alt_data = ds.event_altitude.values/1000.0
        lon_data = ds.event_longitude.values
        lat_data = ds.event_latitude.values
        time_data = pd.Series(ds.event_time) # because time comparisons
        chi_data = ds.event_chi2.values
        station_data = ds.event_stations.values.astype(int)

        tstring = 'LMA {}-{}'.format(tlim_sub[0].strftime('%H%M'),
                                     tlim_sub[1].strftime('%H%M UTC %d %B %Y '))

        lon_set, lat_set, alt_set, time_set, selection = subset(
                   lon_data, lat_data, alt_data, time_data, chi_data, station_data,
                   xlim, ylim, zlim, tlim_sub, xchi, stationmin)
        # Retain the current LMA data so that subclasses can access the current LMA
        # data and compare to other data to be plotted. For instance, to calculate
        # the offset between the time of the first LMA point and a ground strike point.
        self.this_lma_lon = lon_set
        self.this_lma_lat = lat_set
        self.this_lma_alt = alt_set
        self.this_lma_time = time_set
        self.this_lma_sel = selection

        # if self.lma_plot is not None:
        #     fig = self.lma_plot.fig
        # else:
        #     fig = None
        if self.lma_plot is None:
            self.lma_plot = BlankPlot(pd.to_datetime(tlim_sub[0]), bkgmap=True, #fig=fig,
                      xlim=xlim, ylim=ylim, zlim=zlim, tlim=tlim, title=tstring)
            # Add some subplot labels
            label_art = subplot_labels(self.lma_plot)
        else:
            # clear all the old plots of data
            for a in self.data_artists:
                # print("removing", a)
                a.remove()
            # reset the list of data artists.
            self.data_artists = []
            self.lma_plot.ax_th.set_title(tstring)


        # Add a view of where the subset is
        # Not plotting this because removing the inset_ax created by this funcion does not
        # clean up the callbacks created by GeoAxes._boundary, causing "None" to be
        # passed as the axes to _trigger_patch_reclip
        # if self.inset_ax is not None:
        #     print(self.inset_ax.callbacks)
        #     self.inset_ax.remove()
        # xdiv = ydiv = 0.1
        # self.inset_ax = inset_view(self.lma_plot, lon_data, lat_data, xlim, ylim, xdiv, ydiv,
        #           buffer=0.5, inset_size=0.15, plot_cmap = 'plasma', bkgmap = True)

        # Add a range ring
        ring_art = self.lma_plot.ax_plan.tissot(rad_km=40.0, lons=self.clon, lats=self.clat, n_samples=80,
                          facecolor='none',edgecolor='k')
        self.data_artists.append(ring_art)
        # Add the station locations
        stn_art = self.lma_plot.ax_plan.plot(ds['station_longitude'],
                                       ds['station_latitude'], 'wD', mec='k', ms=5)
        self.data_artists.extend(stn_art)

        if len(lon_set)==0:
            no_src_art = self.lma_plot.ax_hist.text(0.02,1,'No Sources',fontsize=12)
            self.data_artists.append(no_src_art)
        else:
            plot_vmin, plot_vmax, plot_c = color_by_time(time_set, tlim_sub)
            src_art = plot_points(self.lma_plot, lon_set, lat_set, alt_set, time_set,
                              plot_cmap, plot_s, plot_vmin, plot_vmax, plot_c)
            self.data_artists.extend(src_art)

        # Cartopy's enforcement of a 1:1: aspect ratio results in extent of the
        # axes on the plot changing, causing the vertical projection panes
        # to no longer match. Turn off that feature.
        self.lma_plot.ax_plan.set_aspect('auto')

        # Zoom all the axes to the same limits. Don't emit a callback event,
        # so that we don't have an infinite loop.
        self.lma_plot.ax_lon.axis(xlim+zlim, emit=False)
        self.lma_plot.ax_lat.axis(zlim+ylim, emit=False)
        self.lma_plot.ax_th.axis(tlim+zlim, emit=False)
        # print(self.lma_plot.ax_plan.get_position())
        # print(self.lma_plot.ax_lat.get_position())
        # print(self.lma_plot.ax_lon.get_position())

        # Refresh the ticks from the non-Cartopy plot that knows how to
        # generate ticks. If this goes after the ax_plan.axis() below,
        # it causes the proportion of the Cartopy axes frame
        # to change within the figure.
        self.lma_plot.ax_plan.set_xticks(self.lma_plot.ax_lon.get_xticks())
        self.lma_plot.ax_plan.set_yticks(self.lma_plot.ax_lat.get_yticks())

        self.lma_plot.ax_plan.axis(xlim+ylim, emit=False)




        # Refresh the ticks from the non-Cartopy plot that knows how to
        # generate ticks.
        # self.lma_plot.ax_plan.set_xticks(self.lma_plot.ax_lon.get_xticks())
        # self.lma_plot.ax_plan.set_yticks(self.lma_plot.ax_lat.get_yticks())
        # This line does runs the above two lines, and then a set_extent, and
        # promptly crashes the iPy kernel.
        # self.lma_plot.set_ax_plan_labels()



def get_glm_plot_subset(interactive_plot, glm):
    """ Use the plot limits in the interactive plot to subset GLM data.
    
        glm may be an xarray Dataset or a glmtools GLMDataset.
    
        Returns in xarray Dataset subsetted to match the plot.
    """
    
    from glmtools.io.glm import GLMDataset
    
    # We need the subsetting functionality attached to the GLMDataset class, which 
    # can prune the flash-group-event hierarchy to a self-consistent sub-tree.
    if isinstance(glm, xr.Dataset):
        glm = GLMDataset(glm, check_area_units=False, change_energy_units=False)
    else:
        assert isinstance(glm, GLMDataset)

    xlim = interactive_plot.bounds['x']
    ylim = interactive_plot.bounds['y']    
    tlim = interactive_plot.bounds['t']
    start, end = np.datetime64(tlim[0]), np.datetime64(tlim[1])
    
    
    # Find the groups in the time range.
    # In some GLM datasets, perhaps all, the event times are incorrect.
    # Probably missing some unsigned stuff.
    # That is why above we use the group times only and select events by 
    # parent ID through reduce_to_entities
    # print(glm_sub.event_time_offset.min().data, glm_sub.event_time_offset.max().data)
    # print(glm_sub.group_time_offset.min().data, glm_sub.group_time_offset.max().data)
    # print(glm_sub)
    
    glm_bounds = {'group_time_offset':slice(start,end),}
    glm_sub = generic_subset(glm.dataset, glm.dataset.group_id.dims[0], glm_bounds)
    if glm_sub.group_id.data.shape[0] < 1:
        # No data, so just empty everything
        return glm.dataset[{'number_of_events':[], 'number_of_groups':[], 'number_of_flashes':[]}]
    glm_sub = glm.reduce_to_entities('group_id', glm_sub.group_id.data)

    # Recreate the GLMDataset from the reduced dataset.
    # There's probably some way to do this all in one step, perhaps by using the 
    # common set of group_ids. But it may not be faster in the end anyway.
    glm = GLMDataset(glm_sub, check_area_units=False, change_energy_units=False)
    
    # Find the events that are in the view, and keep only their parent groups.
    glm_bounds = {'event_lat':slice(ylim[0], ylim[1]),
                  'event_lon':slice(xlim[0], xlim[1])}
    glm_sub = generic_subset(glm_sub, glm_sub.event_id.dims[0], glm_bounds)
    glm_sub = glm.reduce_to_entities('group_id', glm_sub.event_parent_group_id.data)
                  
    return glm_sub


def get_2d_network_subset(interactive_plot, netw_data, adjust_alt=True, pad_factor=0.01):
    """
    Subset 2D network point data to the current view of an interactive plot.
    
    Parameters
    ----------
    interactive_plot : `pyxlma.plot.interactive.InteractiveLMAPlot`
        The XLMA plot from which the bounds of the current view are extracted
    netw_data : `pandas.DataFrame`
        2d network point location data of the type expected by 
        `pyxlma.plot.xlma_plot_feature.plot_2d_network_points`
    adjust_alt : bool
        if True, adjust the `icheight` column of netw_sub to fit within the plot range
        if False, reject any points whose `icheight` is outside the plot range.
    pad_factor : float
        if adjust_alt==True, factor of the plot's total altitude height range to use
        as padding on the minimum and maximum altitude.

    Returns
    -------
    netw_sub : `pandas.DataFrame`
        subset netw_data to match the current plot
    """
    xlim = interactive_plot.bounds['x']
    ylim = interactive_plot.bounds['y']    
    zlim = [z*1000.0 for z in interactive_plot.bounds['z']]
    tlim = interactive_plot.bounds['t']
    start, end = np.datetime64(tlim[0]), np.datetime64(tlim[1])
    z_span = zlim[1] - zlim[0]
    
    in_lat = (netw_data['latitude'] >= ylim[0]) & (netw_data['latitude'] < ylim[1])
    in_lon = (netw_data['longitude'] >= xlim[0]) & (netw_data['longitude'] < xlim[1])
    in_time = (netw_data['datetime'] >= start) & (netw_data['datetime'] < end)
    
    # Must copy, or it gives a view of the original dataframe/
    netw_sub = netw_data[in_lat & in_lon & in_time].copy()

    # if altitude greater or less than top or bottom of plot, move to top or bottom.
    if adjust_alt:
        min_z = zlim[0]+pad_factor*z_span
        max_z = zlim[1]-pad_factor*z_span
        below = (netw_sub['icheight'] < min_z)
        above = (netw_sub['icheight'] >= max_z)       
        netw_sub['icheight'] = np.where(below, min_z, netw_sub['icheight'])
        netw_sub['icheight'] = np.where(above, max_z, netw_sub['icheight'])
    else:
        in_alt = (netw_sub['icheight'] >= zlim[0]) & (netw_sub['icheight'] < zlim[1])
        netw_sub = netw_sub[in_alt]
    return netw_sub